text:
"""
    amp_to_int(field)

converts a complex-valued amplitude field to intensity via `abs2.` and summing over the 4th dimension.
"""
amp_to_int(field) = sum(abs2.(field), dims=4)

"""
    has_z_symmetry(pp::PSFParams)

checks whether the point spread function is expected to be symmetric along the z-direction. Currently this is defined by the aberration list being empty.
"""
function has_z_symmetry(pp::PSFParams)
    return isnothing(pp.aberrations) || isempty(pp.aberrations.indices); # will be changed later, when assymetric aberrations are allowed
end

"""
    get_Abbe_limit(pp::PSFParams)

returns the Abbe limit of incoherent imaging in real space as a Tuple with 3 entries. Note that the coherent limit needs only a factor of
two less sampling, as long as no intensity is calculated. This allows upsampling right before calculating intensities.

See also:
+ get_required_amp_sampling()

Example:
```jdoctest
julia> using PSFs

julia> pp = PSFParams(580.0, 1.4, 1.518)
PSFParams(580.0, 1.4, 1.518, Float32, ModeWidefield, PSFs.pol_scalar, PSFs.var"#42#43"(), PSFs.MethodPropagateIterative, nothing, Aberrations(Any[], Any[], :OSA), nothing)

julia> PSFs.get_Abbe_limit(pp)
(191.04085f0, 191.04085f0, 234.38591f0)
```
"""
function get_Abbe_limit(pp::PSFParams)
    d_xy = pp.λ ./ (2 .* pp.n)
    d_z = (1 - cos(asin(pp.NA/ pp.n))) * pp.λ / pp.n
    pp.dtype.((d_xy,d_xy,d_z))
end

"""
    get_required_amp_sampling(sz::NTuple, pp::PSFParams)

returns the necessary sampling (in real space) for sampling amplitudes. This is almost corresponding to the Abbe limit. Factor of two less because of amplitude but twice because of the Nyquist theorem.
However, this is slight modified by requiring slightly higher sampling (one empty pixel on each side of Fourier space) to stay clear of ambiguities.

See also:
+ get_Abbe_limit()

Example:
```jdoctest
julia> using PSFs

julia> pp = PSFParams(580.0, 1.4, 1.518)
PSFParams(580.0, 1.4, 1.518, Float32, ModeWidefield, PSFs.pol_scalar, PSFs.var"#42#43"(), PSFs.MethodPropagateIterative, nothing, Aberrations(Any[], Any[], :OSA), nothing)

julia> sz = (256,256,64)
(256, 256, 64)

julia> PSFs.get_required_amp_sampling(sz,pp)
(188.05583f0, 188.05583f0, 219.73679f0)

```
"""
function get_required_amp_sampling(sz::NTuple, pp::PSFParams)
    abbe = get_Abbe_limit(pp)[1:length(sz)]
    sz2 = sz .÷ 2
    abbe .* (sz2.-2) ./ sz2 # provide a minimum amount of oversampling to avoid problems with the border pixesl.
end

"""
    get_Ewald_sampling(sz::NTuple, pp::PSFParams)

returns the required minimum sampling for the calculation of a full Ewald sphere.
"""
function get_Ewald_sampling(sz::NTuple, pp::PSFParams)
    s_xyz = pp.λ ./ (2 .* pp.n)
    sz2 = sz .÷ 2
    s_xyz .* (sz2.-2) ./ sz2 # provide a minimum amount of oversampling to avoid problems with the border pixesl.
end

"""
    get_McCutchen_kz_center(ft_shell, pp::PSFParams, sampling)

calculates the (rounded) pixels position half way between both, the kz-borders of the McCutchen pupil to extract from the full sized Ewald sphere.
The pixel z position is returned together with the corresponding kz position.
"""
function get_McCutchen_kz_center(sz, pp::PSFParams, sampling)
    k_z_scale = k_scale(sz, pp, sampling)[3]
    dkz = k_dz(pp) ./ k_z_scale
    pk0 = k_0(pp) ./ k_z_scale
    old_center = sz .÷ 2 .+ 1
    new_center = (old_center[1], old_center[2], round(eltype(old_center), old_center[3] .+ pk0 .- dkz /2))
    # kz_center = (new_center[3] - old_center[3]) * k_z_scale
    return new_center, new_center[3]-old_center[3]
end

"""
    limit_k_z(ft_shell, pp::PSFParams, sampling)

limits the k_z range of the ewald-sphere.
returns the extracted region and the new sampling
"""
function limit_kz(ft_shell, pp::PSFParams, sampling)
    sz = size(ft_shell)
    get_kz_center(sz, pp, sampling)
    dkz = k_dz(pp) ./ k_z_scale
    new_size = (sz[1], sz[2], ceil(eltype(sz), dkz).+1)
    cut_shell = FourierTools.select_region_ft(ft_shell, center = new_center, new_size = new_size)
    sampling = (sampling[1],sampling[2],sampling[2] * sz[3] / new_size[3])
    return cut_shell, sampling
end

"""
    sinc_r(sz::NTuple, pp::PSFParams; sampling=nothing)

calculates the 3-dimensional `sinc(abs(position))` scaled such that the Fourier transformation yields the k-sphere.
Note that for this to work the sampling needs to be sufficient, which may be problematic especially along the z-direction.
This can be checked with the help of the `get_Ewald_sampling()` function.

See also:
+ get_Ewald_sampling()
+ jinc_r_2d

"""
function sinc_r(sz::NTuple, pp::PSFParams; sampling=nothing)
    if isnothing(sampling)
        sampling=get_Ewald_sampling(sz, pp)
    end
    sinc.(rr(pp.dtype, sz, scale=2 .*sampling ./ (pp.λ./pp.n)))
end

"""
    jinc_r_2d(sz::NTuple, pp::PSFParams; sampling=nothing)

calculates a jinc(abs(position)) function in 2D such that its Fourier transform corresponds to the disk-shaped pupil (indcluding the effect ot the numerical aperture).

See also:
+ sinc_r()

"""
function jinc_r_2d(sz::NTuple, pp::PSFParams; sampling=nothing)
    if isnothing(sampling)
        sampling=get_Ewald_sampling(sz, pp)
    end
    jinc.(rr(pp.dtype, sz[1:2], scale=2 .*sampling[1:2] ./ (pp.λ./pp.NA)))
end

# my_disc(sz; rel_border=4/3, eps=0.05) = window_hanning(sz, border_in=rel_border.-eps, border_out=rel_border.+eps)
"""
    my_disc(sz, pp)

creates a disc such that there is no overalp upon wrap-around, when convolving it with itself on this grid.
This is the radius where there is no overlap. However, in the cornes of the calculation, the results are not
100% accurate as something is missing.
"""
my_disc(sz, pp) = disc(pp.dtype, sz, sz .* (1/3))  # This is the radius where there is no overlap and the diagonals (sqrt(2)) are still covered
# my_disc(sz, pp) = disc(pp.dtype, sz, sz .* (sqrt(2)/3))  # This is the radius where there is no overlap and the diagonals (sqrt(2)) are still covered

"""
    iftz(arr)

inverse Fourier transform only along the 3rd dimension (kz).

Arguments:
+ arr:  array to transform along kz
"""
iftz(arr) = ift(arr,(3,))

"""
    theta_z(sz)

a θ function along the z direction, being one for z positon > 0 and zero elsewhere.

Arguments:
+ sz:  size of the array
"""
theta_z(sz) = (zz((1,1,sz[3])) .> 0) # The direction is important due to the highest frequency position at even-sized FFTs

"""
    k_0(pp::PSFParams)

k in the medium as n/lambda.   (1/space units

Arguments:
+ `pp`:  PSF parameter structure
"""
function k_0(pp::PSFParams)
    pp.dtype(pp.n / pp.λ)
end

"""
    k_pupil(pp::PSFParams)

maxim radial k-coordinate (1/space units) where the pupil ends

Arguments:
+ `pp`:  PSF parameter structure
"""
function k_pupil(pp::PSFParams)
    pp.dtype(pp.NA / pp.λ)
end

"""
    k_dz(pp::PSFParams)

relative kz range from pupil boarder to top of Ewald sphere

Arguments:
+ `pp`:  PSF parameter structure
"""
function k_dz(pp::PSFParams)
    pp.dtype((1 - cos(asin(pp.NA/ pp.n))) * k_0(pp))
end

"""
    k_scale(sz,pp,sampling)
    pixelpitch (as NTuple) in k-space
+ `sz`:  size of the real-space array
+ `pp`:  PSF parameter structure
+ `sampling`: pixelpitch in real space as NTuple
"""
function k_scale(sz, pp::PSFParams, sampling)
    pp.dtype.(1 ./ (sz .* sampling))
end

"""
    k_pupil_pos(sz, pp::PSFParams, sampling)

returns the X and Y position of the pupil border in reciprocal space pixels.

Arguments:
+ `sz`:  size of the real-space array
+ `pp`:  PSF parameter structure
+ `sampling`: pixelpitch in real space as NTuple
"""
function k_pupil_pos(sz, pp::PSFParams, sampling)
    k_pupil(pp) ./ k_scale(sz, pp, sampling)
end

"""
    k_0_pos(sz, pp::PSFParams, sampling)

returns the X and Y position of the Ewald-sphere border in reciprocal space pixels.

Arguments:
+ `sz`:  size of the real-space array
+ `pp`:  PSF parameter structure
+ `sampling`: pixelpitch in real space as NTuple
"""
function k_0_pos(sz, pp::PSFParams, sampling)
    k_0(pp) ./ k_scale(sz, pp, sampling)
end

"""
    k_r(sz, pp::PSFParams, sampling)

returns an array of radial k coordinates, |k_xy|
"""
function k_r(sz, pp::PSFParams, sampling)
    min.(k_0(pp), rr(pp.dtype, sz[1:2],scale = k_scale(sz[1:2], pp, sampling[1:2])))
end

"""
    k_xy(sz,pp,sampling)

yields a 2D array with each entry being a 2D Tuple.
"""
function k_xy(sz,pp,sampling)
    idx(pp.dtype, sz[1:2],scale = k_scale(sz[1:2], pp, sampling[1:2]))
end

"""
    k_xy_rel_pupil(sz,pp,sampling)

returns an array of relative distance to the pupil border
"""
function k_xy_rel_pupil(sz,pp,sampling)
    idx(pp.dtype, sz[1:2],scale = k_scale(sz[1:2], pp, sampling[1:2]) ./ k_pupil(pp))
end

"""
    check_amp_sampling_xy(sz, pp,sampling)

issues a warning if the amplitude sampling along X and Y is not within the required limits.

See also:
+ get_Abbe_limit()
+ get_required_amp_sampling()
+ check_amp_sampling()
+ check_amp_sampling_z()

Arguments:
+ `sz`:  size of the real-space array
+ `pp`:  PSF parameter structure
+ `sampling`: pixelpitch in real space as NTuple
"""
function check_amp_sampling_xy(sz, pp,sampling)
    sample_factor = k_pupil(pp) ./ ((sz[1:2] .÷2) .* k_scale(sz, pp, sampling)[1:2])
    if any(sample_factor .> 1.0)
        @warn "Your calculation is undersampled along XY by factors of $sample_factor. The PSF calculation will be incorrect.)"
    end
end

"""
    check_amp_sampling_z(sz, pp,sampling)

issues a warning if the amplitude sampling along Z is not within the required limits.

See also:
+ get_Abbe_limit()
+ get_required_amp_sampling()
+ check_amp_sampling_xy()
+ check_amp_sampling()

Arguments:
+ `sz`:  size of the real-space array
+ `pp`:  PSF parameter structure
+ `sampling`: pixelpitch in real space as NTuple
"""
function check_amp_sampling_z(sz, pp,sampling)
    sample_factor = k_dz(pp) ./ ((sz[3] .÷2) .* k_scale(sz, pp, sampling)[3])
    if (sample_factor > 1.0)
        @warn "Your calculation is undersampled along Z by factors of $sample_factor. The PSF calculation will be incorrect.)"
    end
end

"""
    check_amp_sampling(sz, pp,sampling)

issues a warning if the amplitude sampling (X,Y and Z) is not within the required limits.

See also:
+ get_Abbe_limit()
+ get_required_amp_sampling()
+ check_amp_sampling_xy()
+ check_amp_sampling_z()

Arguments:
+ `sz`:  size of the real-space array
+ `pp`:  PSF parameter structure
+ `sampling`: pixelpitch in real space as NTuple
"""
function check_amp_sampling(sz, pp,sampling)
    check_amp_sampling_xy(sz, pp, sampling)
    check_amp_sampling_z(sz, pp, sampling)
end

"""
    check_amp_sampling_sincr(sz, pp,sampling)

issues a warning if the amplitude sampling (X,Y and Z) is not within the required limits for the `SincR` method, as this requires sampling
the Ewald-sphere according to Nyquist.

See also:
+ get_Abbe_limit()
+ get_required_amp_sampling()
+ check_amp_sampling()
+ get_Ewald_sampling()

Arguments:
+ `sz`:  size of the real-space array
+ `pp`:  PSF parameter structure
+ `sampling`: pixelpitch in real space as NTuple
"""
function check_amp_sampling_sincr(sz, pp,sampling) # The sinc-r method needs (for now without aliasing) to be sampled extremely high along Z
    check_amp_sampling_xy(sz, pp, sampling)
    @show sample_factor = k_0(pp) ./ ((sz[3] .÷2) .* k_scale(sz, pp, sampling)[3])
    if (sample_factor > 1.0)
        @warn "Your calculation is undersampled along Z by factors of $sample_factor. The PSF calculation will be incorrect.)"
    end
end


text:
module BakSneppenSimulation

export applymutations!,
    mutate,
    test,
    runsim,
    plotsim,
    plotzero,
    simlinear,
    Simulation,
    SimCircular,
    SimLinear,
    SimConfig,
    SimState

import Base: show

ENV["GKS_ENCODING"] = "utf8"
using Random, Plots

struct DataPoint
    time::Int
    zeros::Int
    diameter::Int
    death::Int
end

abstract type Simulation end

show(io::IO, sim::Simulation) = print(
    io,
    "Simulation for p=$(sim.config.p) and $(sim.config.mutnum) mutations",
)

struct SimConfig
    n::Int
    p::Float64
    mutnum::Int
    stop::Int
    pointnum::Int
    step::Int
    shift::Int
end

function SimConfig(
    n::Int,
    p::Float64,
    mutnum::Int,
    stop::Int = 10^6,
    pointnum::Int = 10^3,
)
    step = max(1, floor(stop / pointnum))
    shift = convert(Int, ceil(mutnum / 2))
    SimConfig(n, p, mutnum, stop, pointnum, step, shift)
end

struct SimState
    xlist::Array{Int,1}
    x0list::Array{Int,1}
end

function SimState(xlist::Array{Int,1})
    x0list = findall(x -> x == 0, xlist)
    SimState(xlist, x0list)
end

struct SimCircular <: Simulation
    config::SimConfig
    state::SimState
end

function SimCircular(config::SimConfig)
    xlist = zeros(Int, config.n)
    state = SimState(xlist)
    SimCircular(config, state)
end

struct SimLinear <: Simulation
    config::SimConfig
    state::SimState
    bounded::Bool
end

function SimLinear(config::SimConfig, bounded = false)
    xlist = ones(Int, config.n)
    xlist[1] = 0
    state = SimState(xlist)
    SimLinear(config, state, bounded)
end

n(sim::Simulation) = length(sim.state.xlist)
zeros(sim::Simulation) = length(sim.state.x0list)

function circularind(ind, len)
    if ind < 1
        len + ind
    elseif len < ind
        ind - len
    else
        ind
    end
end

function circulardist(dist, len)
    if dist < 0
        dist + len
    else
        dist
    end
end

function diameter(sim::SimCircular)
    num0 = zeros(sim)
    if num0 == 0
        return 0
    end

    if num0 == 1
        return 1
    end

    x0list = sim.state.x0list
    sort!(x0list)
    wlist = similar(x0list)
    for i = 1:length(wlist)
        prev = circularind(i - 1, num0)
        next = circularind(i + 1, num0)

        w1 = circulardist(x0list[next] - x0list[i] - 1, n(sim))
        w2 = circulardist(x0list[i] - x0list[prev] - 1, n(sim))
        wlist[i] = n(sim) - max(w1, w2)
    end
    findmin(wlist)[1]
end

function diameter(sim::SimLinear)
    l = findfirst(x -> x == 0, sim.state.xlist)
    r = findlast(x -> x == 0, sim.state.xlist)
    r - l + 1
end

function runsim(sim::Simulation)
    data = Array{DataPoint,1}()
    death = 0
    for t = 1:sim.config.stop
        mutate(sim)
        if zeros(sim) == 0
            death += 1
        end
        if (t - 1) % sim.config.step == 0
            z = zeros(sim)
            d = 0
            if z != 0
                d = diameter(sim)
            end
            push!(data, DataPoint(t, z, d, death))
        end
    end
    data
end

function simlinear(config::SimConfig, bounded = false)
    sim = SimLinear(config, bounded)
    data = runsim(sim)
    (data, sim)
end

function mutate(sim::Simulation)
    x0len = zeros(sim)
    x0ind = 0
    ind = 0
    if x0len > 0
        x0ind = rand(1:x0len)
        ind = sim.state.x0list[x0ind]
    else
        ind = 1
    end

    mutations = map(x -> (x < sim.config.p ? 1 : 0), rand(sim.config.mutnum))

    applymut!(sim, mutations, ind)
end

function applymut!(sim::SimLinear, mutations::Vector{Int}, ind::Int)
    @debug "applymut!" mutations, ind
    for i = 1:sim.config.mutnum
        mutationind = ind + i - sim.config.shift
        if mutationind < 1
            continue
        elseif mutationind == n(sim) + 1
            # extend the array
            if !sim.bounded
                push!(sim.state.xlist, 1)
            else
                continue
            end
        end

        update!(sim.state, mutations[i], mutationind)
    end
end

function applymut!(sim::SimCircular, mutations::Array{Int}, ind::Int)
    @debug "applymut!" mutations, ind
    for i = 1:sim.config.mutnum
        mutationind = circularind(ind + i - sim.config.shift, n(sim))

        update!(sim.state, mutations[i], mutationind)
    end
end

function update!(state::SimState, new, mutationind)
    old = state.xlist[mutationind]

    if new == old
        return
    end

    state.xlist[mutationind] = new
    if new == 1 && old == 0
        filter!(a -> a != mutationind, state.x0list)
    else
        push!(state.x0list, mutationind)
    end
end

function plotsim(data::Array{DataPoint,1}, sim)
    gr()
    tl = map(p -> p.time, data)
    zl = map(p -> p.zeros, data)
    dl = map(p -> p.diameter, data)
    ylim = max(findmax(dl)[1], sim.config.n)
    p = plot(tl, zl, ylims = (-1, ylim), label = "zeros")
    plot!(tl, dl, label = "diameter", title = string(sim))
end
plotsim((data, sim)) = plotsim(data, sim)

function plotzero(data, sim)
    gr()
    tl = map(p -> p.time, data)
    pl = map(p -> p.zeros / p.diameter, data)
    p = plot(
        tl,
        pl,
        ylims = (0, 1),
        label = "percentage of zeros",
        title = string(sim),
    )
end
plotzero((data, sim)) = plotzero(data, sim)

end # module

ext: jl


text:
### Kernel functions
abstract type SmoothingKernel{T} end

struct GaussianKernel{T} <: SmoothingKernel{T}
  h::T
end

struct EpanechnikovKernel{T} <: SmoothingKernel{T}
  h::T
end

createKernel(::Type{GaussianKernel{T}}, h::T) where {T <: AbstractFloat} = GaussianKernel{T}(h)
createKernel(::Type{EpanechnikovKernel{T}}, h::T) where {T <: AbstractFloat} = EpanechnikovKernel{T}(h)

evaluate(k::GaussianKernel{T}, x::T, y::T) where {T <: AbstractFloat} = exp(-(x-y)^2. / k.h) / k.h
function evaluate(k::EpanechnikovKernel{T}, x::T, y::T) where {T <: AbstractFloat}
    u = (x - y) / k.h
    abs(u) >= 1. ? zero(T) : 0.75 * (1. - u^2.) / k.h
end

############################################################
#
# local polynomial regression with lasso
#
############################################################

function locpolyl1(
  X::Matrix{T}, z::Vector{T}, y::Vector{T},
  zgrid::Vector{T},
  degree::Int64,
  kernel::SmoothingKernel{T},
  λ0::T,
  refit::Bool,
  options::CDOptions=CDOptions()) where {T <: AbstractFloat}

  # construct inner options because we do not want warmStart = false
  # we want to start from the previous iteration since the points
  # on the grid should be close to each other
  opt = CDOptions(options.maxIter, options.optTol, options.randomize, true, options.numSteps)

  n, p = size(X)
  ep = p * (degree + 1)
  out = spzeros(T, ep, length(zgrid))
  outR = spzeros(T, ep, length(zgrid))

  # temporary storage
  w = Array{T}(undef, n)
  expandX = Array{T}(undef, n, ep)
  S = BitArray(undef, ep)
  stdX = Array{T}(undef, ep)
  f = CDWeightedLSLoss(y, expandX, w)        # inner parts of f will be modified in a loop
  g = ProxL1(λ0, stdX)
  β = SparseIterate(ep)

  ind = 0
  for z0 in zgrid
    ind += 1

    # the following two should update f
    w .= evaluate.(Ref(kernel), z, Ref(z0))
    _expand_X!(expandX, X, z, z0, degree)
    _stdX!(stdX, w, expandX)

    # solve for β
    coordinateDescent!(β, f, g, opt)
    out[:, ind] = β

    if refit
        get_nonzero_coordinates!(S, β, p, degree, true)
        Xs = view(expandX, :, S)
        tmp = Xs' * Diagonal(w)
        outR[S, ind] = (tmp * Xs) \ (tmp * y)
    end
  end
  out, outR
end

# leave one out for h selection
function lvocv_locpolyl1(
    X::Matrix{T}, z::Vector{T}, y::Vector{T},
    degree::Int64,                                     # degree of the polynomial
    hArr::Vector{T},
    kernelType::Type{<:SmoothingKernel},
    λ0::T,
    options::CDOptions=CDOptions()) where {T <: AbstractFloat}

    n, p = size(X)
    numH = length(hArr)
    MSE = zeros(numH)

    opt = CDOptions(options.maxIter, options.optTol, options.randomize, true, options.numSteps)

    ep = p * (degree + 1)
    w = Array{T}(undef, n)
    wX = Array{T}(undef, n, ep)
    stdX = Array{T}(undef, ep)
    S = BitArray(undef, ep)

    f = CDWeightedLSLoss(y, wX, w)
    β = SparseIterate(ep)

    for indH = 1:numH
        kernel = createKernel(kernelType, hArr[indH])
        for i = 1:n
            # update variables
            z0 = z[i]
            w .= evaluate.(Ref(kernel), z, Ref(z0))
            w[i] = zero(T)
            _expand_X!(wX, X, z, z0, degree)
            _stdX!(stdX, w, wX)
            _findInitResiduals!(w, wX, y, min(10, ep), f.r)

            # compute sigma
            σ = _getSigma(w, f.r)
            g = ProxL1(λ0 * σ, stdX)
            for iter=1:10
              coordinateDescent!(β, f, g, opt)
              σnew = _getSigma(w, f.r)
              if abs(σnew - σ) / σ < 1e-2
                break
              end
              σ = σnew
              g = ProxL1(λ0 * σ, stdX)
            end

            # refit and make prediction
            get_nonzero_coordinates!(S, β, p, degree, true)
            Xs = view(wX, :, S)
            Yh = dot(wX[i, S], (Xs' * Diagonal(w) * Xs) \ (Xs' * Diagonal(w) * y))
            MSE[indH] += (Yh - y[i])^2.
        end
    end
    MSE
end

function refit_locpolyl1(
    X::Matrix{T}, z::Vector{T}, y::Vector{T},
    z0::T,
    degree::Int64,
    kernel::SmoothingKernel{T},
    β::Union{SparseIterate{T}, SparseVector{T}}
    ) where {T <: AbstractFloat}

    n, p = size(X)
    ep = p * (degree + 1)

    S = BitArray(undef, p)
    get_nonzero_coordinates!(S, β, p, degree, false)
    βr = locpoly(view(X, :, S), z, y, z0, degree, kernel)
    (βr, S)
end

# function refit_locpolyl1(
#     X::Matrix{T}, z::Vector{T}, y::Vector{T},
#     zgrid::Vector{T},
#     degree::Int64,
#     kernel::SmoothingKernel{T},
#     β::SparseMatrixCSC{T}
#     ) where {T <: AbstractFloat}
#
#     n, p = size(X)
#     ep = p * (degree + 1)
#     S = BitArray(undef, p, length(zgrid))
#     fill!(S, false)
#     w = Array{T}(undef, n)
#     expandedX = Array{T}(undef, n, ep)
#     βr = Array{Array{T}}(undef, length(zgrid))
#     tS = BitArray(undef, p)
#
#     for indZ = 1:length(zgrid)
#         o_coordinates!(tS, β[:, indZ], p, degree, false)
#         @show β[:, indZ]
#         @show S[:, indZ] = tS
#         z0 = zgrid[indZ]
#         num_col_refit = sum(S[:, indZ])*(degree+1)
#         Xs = view(X, :, S[:, indZ])
#         βr[indZ] = _locpoly!(view(expandedX, :, 1:num_col_refit), w, Xs, z, y, z0, degree, kernel)
#     end
#     (βr, S)
# end

############################################################
#
# local polynomial regression low dimensions
#
############################################################

function _locpoly!(
  wX::Union{SubArray{T, 2}, Matrix{T}}, w::Vector{T},
  X::Union{SubArray{T, 2}, Matrix{T}},
  z::Union{SubArray{T, 1}, Vector{T}},
  y::Union{SubArray{T, 1}, Vector{T}},
  z0::T,
  degree::Int64,
  kernel::SmoothingKernel{T}) where {T <: AbstractFloat}

  w .= sqrt.(evaluate.(Ref(kernel), z, Ref(z0)))    # square root of kernel weights
  _expand_wX!(wX, w, X, z, z0, degree)              # √w ⋅ x ⊗ [1 (zi - z0) ... (zi-z0)^q]
  @. w *= y                                         # √w ⋅ y
  qr!(wX) \ w
end

locpoly(
  X::Union{SubArray{T, 2}, Matrix{T}}, z::Vector{T}, y::Vector{T},
  z0::T, degree::Int64, kernel::SmoothingKernel=GaussianKernel(one(T))) where {T <: AbstractFloat} =
    _locpoly!(Array{T}(undef, length(y), size(X, 2) * (degree+1)), similar(y), X, z, y, z0, degree, kernel)

function locpoly(
  X::Union{SubArray{T, 2}, Matrix{T}}, z::Vector{T}, y::Vector{T},
  zgrid::Vector{T},
  degree::Int64,                                     # degree of the polynomial
  kernel::SmoothingKernel{T}=GaussianKernel(one(T))) where {T <: AbstractFloat}

  n, p = size(X)
  ep = p * (degree + 1)
  out = Array{T}(undef, ep, length(zgrid))
  w = Array{T}(undef, n)
  wX = Array{T}(undef, n, ep)

  ind = 0
  for z0 in zgrid
    ind += 1
    out[:, ind] = _locpoly!(wX, w, X, z, y, z0, degree, kernel)
  end
  out
end

function getResiduals!(
     ϵhat::Vector{T},
     X::Union{SubArray{T, 2}, Matrix{T}}, z::Vector{T}, y::Vector{T},
     zgrid::Vector{T},
     βhat::Union{Matrix{T}, SparseMatrixCSC{T}},
     degree::Int64,
     kernel::SmoothingKernel{T}=GaussianKernel(one(T))
    ) where {T <: AbstractFloat}

    n, p = size(X)
    ep = p * (degree + 1)
    βi = spzeros(ep)
    for i=1:n
        get_beta!(βi, zgrid, βhat, z[i])
        ϵhat[i] = y[i] - dot(X[i, :], βi[1:(degree+1):ep])
    end

    ϵhat
end

function getStandardError(
    X::Union{SubArray{T, 2}, Matrix{T}}, z::Vector{T},
    σ2::T,
    z0::T,
    degree::Int64,                                     # degree of the polynomial
    kernel::SmoothingKernel{T}) where {T <: AbstractFloat}

    n, p = size(X)
    ep = p * (degree + 1)
    out = Array{T}(undef, p)
    w = Array{T}(undef, n)
    w1 = Array{T}(undef, n)
    XtwX = Array{T}(undef, ep, ep)
    XtwwX = Array{T}(undef, ep, ep)

    w .= evaluate.(Ref(kernel), z, Ref(z0))
    _expand_Xt_w_X!(XtwX, w, X, z, z0, degree)

    w .= w .* w
    _expand_Xt_w_X!(XtwwX, w, X, z, z0, degree)

    A = inv(XtwX)
    varMat = A * XtwwX * A

    for j=1:p
        out[j] = varMat[(j-1)*(degree+1)+1, (j-1)*(degree+1)+1]
    end

    out
end

function getStandardErrorHEW(
    X::Union{SubArray{T, 2}, Matrix{T}}, z::Vector{T},
    ϵ_sqr::Vector{T},
    z0::T,
    degree::Int64,                                     # degree of the polynomial
    kernel::SmoothingKernel{T}) where {T <: AbstractFloat}

    n, p = size(X)
    ep = p * (degree + 1)
    out = Array{T}(undef, p)
    w = Array{T}(undef, n)
    w1 = Array{T}(undef, n)
    XtwX = Array{T}(undef, ep, ep)
    XtwΨwX = Array{T}(undef, ep, ep)

    w .= evaluate.(Ref(kernel), z, Ref(z0))
    _expand_Xt_w_X!(XtwX, w, X, z, z0, degree)

    @. w = w * w * ϵ_sqr
    _expand_Xt_w_X!(XtwΨwX, w, X, z, z0, degree)

    A = inv(XtwX)
    varMat = A * XtwΨwX * A

    for j=1:p
        out[j] = varMat[(j-1)*(degree+1)+1, (j-1)*(degree+1)+1]
    end

    out
end

# function locpoly_alt(
#   X::Matrix{T}, z::Vector{T}, y::Vector{T},
#   zgrid::Vector{T},
#   degree::Int64,                                     # degree of the polynomial
#   kernel::SmoothingKernel{T}=GaussianKernel(one(T))) where {T <: AbstractFloat}
#
#   n, p = size(X)
#   ep = p * (degree + 1)
#   out = Array{T}(undef, ep, length(zgrid))
#   w = Array{T}(undef, n)
#   Xt_w_Y = Array{T}(undef, ep)
#   Xt_w_X = Array{T}(undef, ep, ep)
#
#   ind = 0
#   for z0 in zgrid
#       w .= evaluate.(Ref(kernel), z, Ref(z0))
#       _expand_Xt_w_Y!(Xt_w_Y, w, X, z, y, z0, degree)
#       _expand_Xt_w_X!(Xt_w_X, w, X, z, z0, degree)
#       ind += 1
#       out[:, ind] = Xt_w_X \ Xt_w_Y
#   end
#   out
# end

# leave one out for h selection
function lvocv_locpoly(
    X::Matrix{T}, z::Vector{T}, y::Vector{T},
    degree::Int64,                                     # degree of the polynomial
    hArr::Vector{T},
    kernelType::Type{<:SmoothingKernel}) where {T <: AbstractFloat}

    n, p = size(X)
    numH = length(hArr)
    MSE = zeros(numH)

    ep = p * (degree + 1)
    w = Array{T}(undef, n-1)
    wX = Array{T}(undef, n-1, ep)
    indOut = BitArray(undef, n)

    for indH = 1:numH
        kernel = createKernel(kernelType, hArr[indH])
        for i = 1:n
            fill!(indOut, true)
            indOut[i] = false
            Xview = view(X, indOut, :)
            Yview = view(y, indOut)
            Zview = view(z, indOut)

            hbeta = _locpoly!(wX, w, Xview, Zview, Yview, z[i], degree, kernel)

            # make prediction
            Yh = dot(X[i, :], hbeta[1:(degree+1):ep])
            MSE[indH] += (Yh - y[i])^2.
        end
    end
    MSE
end

# data split for h selection
function split_locpoly(
    X::Matrix{T}, z::Vector{T}, y::Vector{T},
    Xtest::Matrix{T}, ztest::Vector{T}, ytest::Vector{T},
    zgrid::Vector{T},
    degree::Int64,                                     # degree of the polynomial
    hArr::Vector{T},
    kernelType::Type{<:SmoothingKernel}) where {T <: AbstractFloat}

    n, p = size(X)
    numH = length(hArr)
    MSE = zeros(numH)
    ep = p * (degree + 1)
    βhati = zeros(ep)

    for indH = 1:numH
        kernel = createKernel(kernelType, hArr[indH])
        βhat = locpoly(X, z, y, zgrid, degree, kernel)
        for i = 1:n
            get_beta!(βhati, zgrid, βhat, ztest[i])

            # make prediction
            Yh = dot(Xtest[i, :], βhati[1:(degree+1):ep])
            MSE[indH] += (ytest[i] - Yh)^2.
        end
    end
    MSE
end

# # leave one out for h selection
# function lvocv_locpoly(
#     X::Matrix{T}, z::Vector{T}, y::Vector{T},
#     degree::Int64,                                     # degree of the polynomial
#     hArr::Vector{T},
#     kernelType::Type{<:SmoothingKernel}) where {T <: AbstractFloat}
#
#     n, p = size(X)
#     numH = length(hArr)
#     MSE = zeros(numH)
#
#     ep = p * (degree + 1)
#     w = Array{T}(undef, n)
#     Xt_w_Y = Array{T}(undef, ep)
#     Xt_w_X = Array{T}(undef, ep, ep)
#
#     for indH = 1:numH
#         kernel = createKernel(kernelType, hArr[indH])
#         for i = 1:n
#             z0 = z[i]
#             w .= evaluate.(Ref(kernel), z, Ref(z0))
#             w[i] = zero(T)
#             _expand_Xt_w_Y!(Xt_w_Y, w, X, z, y, z0, degree)
#             _expand_Xt_w_X!(Xt_w_X, w, X, z, z0, degree)
#             hbeta = Xt_w_X \ Xt_w_Y
#             # make prediction
#             Yh = dot(X[i, :], hbeta[1:(degree+1):ep])
#             MSE[indH] += (Yh - y[i])^2.
#         end
#     end
#     MSE
# end

############################################################
#
# utils
#
############################################################

"""
For a given z0 finds two closest points in zgrid
and corresponding values of beta. The output is obtained
by interpolating the beta values.
"""
function get_beta!(
       out::Union{SparseVector{T}, Vector{T}},
       zgrid::Vector{T},
       beta_grid::Union{Matrix{T}, SparseMatrixCSC{T}},
       z0::T) where {T <: AbstractFloat}
    id1 = searchsortedlast(zgrid, z0)
    id2 = searchsortedfirst(zgrid, z0)

    if id1 == id2
        out .= beta_grid[:, id1]
    else
        z1 = zgrid[id1]
        z2 = zgrid[id2]
        α = (z0 - z1) / (z2 - z1)
        out .= α * beta_grid[:, id1] .+ (1-α) * beta_grid[:, id2]
    end
    out
end

get_nonzero_coordinates(
    β::Union{SparseVector{T}, SparseIterate{T}},
    p::Int,
    degree::Int,
    expanded::Bool
    ) where {T <: AbstractFloat}  =
       expanded ? get_nonzero_coordinates!(BitArray(undef, p*(degree+1)), β, p, degree, expanded) : get_nonzero_coordinates!(BitArray(undef, p), β, p, degree, expanded)

function get_nonzero_coordinates!(
    S::BitArray,
    β::Union{SparseVector{T}, SparseIterate{T}},
    p::Int,
    degree::Int,
    expanded::Bool) where {T <: AbstractFloat}

    fill!(S, false)
    for j = 1:p
        nonzero = false
        for k=((j-1)*(degree+1)+1):(j*(degree+1))
            nonzero = nonzero | !iszero(β[k])
        end
        if nonzero
            if expanded
                for k=((j-1)*(degree+1)+1):(j*(degree+1))
                    S[k] = true
                end
            else
                S[j] = true
            end
        end
    end
    S
end

"""
Computes matrix whose each row is equal to
w[i] ⋅ (X[i, :] ⊗ [1, (z - z0), ..., (z-z0)^q])
where q is the degree of the polynomial.

The output matrix is preallocated.
"""
function _expand_wX!(
  wX::Union{SubArray{T, 2}, Matrix{T}},
  w::Vector{T},
  X::Union{SubArray{T, 2}, Matrix{T}},
  z::Union{SubArray{T, 1}, Vector{T}}, z0::T, degree::Int64) where {T <: AbstractFloat}

  n, p = size(X)
  # wX = zeros(n, p*(degree+1))
  for j=1:p
    @inbounds for i=1:n
      v = X[i, j] * w[i]
      df = z[i] - z0
      col = (j-1)*(degree+1) + 1
      @inbounds wX[i, col] = v
      for l=1:degree
        v *= df
        @inbounds wX[i, col + l] = v
      end
    end
  end
  # return qrfact!(tX)
  wX
end

# expands the data matrix X
# each row becomes
# (X[i, :] ⊗ [1, (z - z0), ..., (z-z0)^q])
function _expand_X!(
  tX::Matrix{T},
  X::Union{SubArray{T, 2}, Matrix{T}},
  z::Union{SubArray{T, 1}, Vector{T}}, z0::T, degree::Int64) where {T <: AbstractFloat}

  n, p = size(X)
  for j=1:p
    for i=1:n
      v = X[i, j]
      df = z[i] - z0
      col = (j-1)*(degree+1) + 1
      @inbounds tX[i, col] = v
      for l=1:degree
        v *= df
        @inbounds tX[i, col + l] = v
      end
    end
  end
  tX
end

function _expand_Xt_w_X!(
  Xt_w_X::Matrix{T},
  w::Vector{T}, X::Matrix{T},
  z::Vector{T}, z0::T, degree::Int64) where {T <: AbstractFloat}

  n, p = size(X)
  fill!(Xt_w_X, zero(T))
  @inbounds for j=1:p
    @inbounds for k=j:p
      @inbounds for i=1:n
        # if iszero(w[i])
        #   continue
        # end
        v1 = X[i, j] * w[i]
        df1 = z[i] - z0

        col=(j-1)*(degree+1)+1
        for jj=0:degree
          v2 = X[i, k]
          df2 = z[i] - z0
          if k != j
            krange = 0:degree
          else
            krange = jj:degree
            v2 = v2 * df2^jj
          end

          row = (k-1)*(degree+1)+1
          for kk=krange
            # need X[i, j] * (z[i] - z0)^(jj+kk) * X[i, k] * w[i]
            # v2 = (z[i] - z0)^kk * X[i, k]
            # v1 = (z[i] - z0)^jj * X[i, k] * w[i]
            Xt_w_X[row+kk, col+jj] += v2 * v1
            v2 *= df2
          end
          v1 *= df1
        end
      end
    end
  end

  @inbounds for c=1:size(Xt_w_X, 2)
    for r=c+1:size(Xt_w_X, 1)
      Xt_w_X[c, r] = Xt_w_X[r, c]
    end
  end

  Xt_w_X
end

function _expand_Xt_w_Y!(
  Xt_w_Y::Vector{T},
  w::Vector{T}, X::Matrix{T}, z::Vector{T}, y::Vector{T},
  z0::T, degree::Int64) where {T <: AbstractFloat}

  n, p = size(X)
  fill!(Xt_w_Y, zero(T))
  @inbounds for j=1:p
    @inbounds for i=1:n
      # if iszero(w[i])
      #   continue
      # end
      v = X[i, j] * w[i] * y[i]
      df = z[i] - z0

      col=(j-1)*(degree+1)+1
      @inbounds for jj=0:degree
        # need X[i, j] * (z[i] - z0)^jj * Y[i, k] * w[i]
        Xt_w_Y[col+jj] += v
        v *= df
      end
    end
  end

  Xt_w_Y
end


text:
using Documenter
using SwapStreams

makedocs(
    modules=[SwapStreams],
    format=Documenter.HTML(),
    repo="https://github.com/Tokazama/SwapStreams.jl/blob/{commit}{path}#L{line}",
    sitename="SwapStreams.jl",
    authors="Zachary P. Christensen",
)


text:
const ALPHABET = "abcdefghijklmnopqrstuvwxyz" |> collect
const WEIGHTS = [ # from https://www.wikipedia.org/wiki/Letter_frequency
 8.167e-2,
 1.492e-2,
 2.202e-2,
 4.253e-2,
 12.702e-2,
 2.228e-2,
 2.015e-2,
 6.094e-2,
 6.966e-2,
 0.153e-2,
 1.292e-2,
 4.025e-2,
 2.406e-2,
 6.749e-2,
 7.507e-2,
 1.929e-2,
 0.095e-2,
 5.987e-2,
 6.327e-2,
 9.356e-2,
 2.758e-2,
 0.978e-2,
 2.560e-2,
 0.150e-2,
 1.994e-2,
 0.077e-2,
]
s = [ALPHABET[StatsBase.sample(StatsBase.ProbabilityWeights(WEIGHTS))] for i in 1:1000] |> join

# NB: a pangram is technically a _sentence_ which contains all the letters of the alphabet. I have decided to ignore this as detecting the end of sentences (e.g. ignoring e.g.) seems more like a machine-learning problem.

function ispangram(input::AbstractString)
    # `lowercase` is expensive in time and memory for large strings, so we avoid it by checking against lower and uppercase characters
    all(occursin(c,input) || occursin(uppercase(c),input) for c in ALPHABET)
end

# In most cases where there is a pangram, this is slower than the above code.
# If there isn't a pangram, and the string is big, it's ~nproc faster.
#
# Since this exercise is about detecting pangrams, it's fair to say that there is a reasonable likelihood that the strings being tested are a) pangrams and b) short, as they are sentences. So the previous code is better.
#
# It's a nice demonstration that the real-world performance of code depends on the type of input you are expecting.
import ThreadsX
function ispangramthreads(input::AbstractString)
    ThreadsX.mapreduce(c->occursin(c,input),&,ALPHABET)
#    ThreadsX.all(c->occursin(c,input),ALPHABET) # this appears to be single threaded
end

# Variant of above with early stopping; it's still slower than the initial code for most pangrams. Probably due to the overhead of threads:
# ```
# julia> @btime ispangramthreads_earlystop($(""))
#   10.229 μs (179 allocations: 20.00 KiB)
# false
#
# julia> @btime ispangramthreads_earlystop($s) # 10 million character random alphabetic string
#   10.624 μs (213 allocations: 20.53 KiB)
#  true
# ```
#
# ```
# julia> [begin; r = repeat("1",n)*(ALPHABET|>join); (n, (@elapsed ispangram(r)) > (@elapsed ispangramthreads_earlystop(r)) ? "multi-threaded is faster" : "single threaded i
# s faster"); end for n in 10 .^ (1:8)]
# 8-element Array{Tuple{Int64,String},1}:
#  (10, "single threaded is faster")
#  (100, "single threaded is faster")
#  (1000, "single threaded is faster")
#  (10000, "single threaded is faster")
#  (100000, "single threaded is faster")
#  (1000000, "single threaded is faster")
#  (10000000, "multi-threaded is faster")
#  (100000000, "multi-threaded is faster")
# ```
# So, it's faster if you expect 10_000_000 characters before you have a pangram, which is a little longer than most sentences.
function ispangramthreads_earlystop(input::AbstractString)
    tasks = Array{Task,1}()
    foreach(c->push!(tasks,Threads.@spawn occursin(c,input) || occursin(uppercase(c),input)),ALPHABET)

    # This feels like it should be a common pattern: operate on data as soon as it comes through with an early exit if we can.
    # Perhaps channels are the way forward. `try while isready || isopen; take; end; catch; isa(InvalidStateException)`?
    while any(t->!istaskdone(t),tasks)
        any(t->istaskdone(t) && !t.result, tasks) && return false
    end
    # Check the last task
    any(t->istaskdone(t) && !t.result, tasks) && return false
    return true
end

using StaticArrays: @MVector
#
function ispangram3(input)
    # This creates a 26 byte static array. The compiler is able to reason that
    # it does not escape this function, so it gets stack allocated despite
    # being mutable, which is neat.
    present = @MVector zeros(Bool, 26)
    @inbounds for i in 1:ncodeunits(input)
        c = codeunit(input, i)
        if 0x61 <= c <=  0x7a
            present[c - 0x60] = true
        elseif 0x41 <= c <= 0x5a
            present[c - 0x40] = true
        end
    end
    all(present)
end


text:
# Copyright (C) 2020 Argonne National Laboratory
# Written by Alinson Santos Xavier <axavier@anl.gov>

using RELOG

@testset "build_graph" begin
    basedir = dirname(@__FILE__)
    instance = RELOG.parsefile("$basedir/../../instances/s1.json")
    graph = RELOG.build_graph(instance)
    process_node_by_location_name =
        Dict(n.location.location_name => n for n in graph.process_nodes)

    @test length(graph.plant_shipping_nodes) == 8
    @test length(graph.collection_shipping_nodes) == 10
    @test length(graph.process_nodes) == 6

    node = graph.collection_shipping_nodes[1]
    @test node.location.name == "C1"
    @test length(node.incoming_arcs) == 0
    @test length(node.outgoing_arcs) == 2
    @test node.outgoing_arcs[1].source.location.name == "C1"
    @test node.outgoing_arcs[1].dest.location.plant_name == "F1"
    @test node.outgoing_arcs[1].dest.location.location_name == "L1"
    @test node.outgoing_arcs[1].values["distance"] == 1095.62

    node = process_node_by_location_name["L1"]
    @test node.location.plant_name == "F1"
    @test node.location.location_name == "L1"
    @test length(node.incoming_arcs) == 10
    @test length(node.outgoing_arcs) == 2

    node = process_node_by_location_name["L3"]
    @test node.location.plant_name == "F2"
    @test node.location.location_name == "L3"
    @test length(node.incoming_arcs) == 2
    @test length(node.outgoing_arcs) == 2

    @test length(graph.arcs) == 38
end

max_stars_repo_name: ANL-CEEESA/RELOG

max_issues_repo_name: ANL-CEEESA/RELOG

max_forks_repo_name: ANL-CEEESA/RELOG


text:
P = [ 0.5000000   0.0000000   0.0000000   0.0000000   0.0000000   0.0000000 0.0000000;
      0.0000000   1.0593543  -0.2507533   0.0000000   0.0000000   0.0000000 0.0000000;
      0.0000000  -0.2507533   1.0593543   0.0000000   0.0000000   0.0000000 0.0000000;
      0.0000000   0.0000000   0.0000000   0.6666667   0.0000000   0.0000000 0.0000000;
      0.0000000   0.0000000   0.0000000   0.0000000   0.6666667   0.0000000 0.0000000;
      0.0000000   0.0000000   0.0000000   0.0000000   0.0000000   0.6666667 0.0000000;
      0.0000000   0.0000000   0.0000000   0.0000000   0.0000000   0.0000000 0.5000000 ]

basSet = readBasisSetTX93("STO-3G.tx93")
geo = readGeometryXYZ("h2o.xyz")
bas = computeBasis(basSet,geo)

function computeMatrixCoulomb(basis::GaussianBasis, density::Matrix)
  N = length(basis.contractedBFs)

  J = zeros(N,N)
  μIndex = 0
  for (μ in basis.contractedBFs)
    μIndex+=1
    for (νIndex in 1:μIndex)
      ν = basis.contractedBFs[νIndex]
      ERImatrix = [computeElectronRepulsionIntegral(μ,ν,λ,σ) for λ in basis.contractedBFs, σ in basis.contractedBFs]
      J[μIndex,νIndex] = trace(ERImatrix*density)
      J[νIndex,μIndex] = J[μIndex,νIndex]
    end
  end
  return J
end

function computeMatrixExchange(basis::GaussianBasis, density::Matrix)
  N = length(basis.contractedBFs)

  K = zeros(N,N)
  μIndex = 0
  for (μ in basis.contractedBFs)
    μIndex+=1
    for (νIndex in 1:μIndex)
      ν = basis.contractedBFs[νIndex]
      ERImatrix = [computeElectronRepulsionIntegral(μ,λ,ν,σ) for λ in basis.contractedBFs, σ in basis.contractedBFs]
      K[μIndex,νIndex] = trace(ERImatrix*density)
      K[νIndex,μIndex] = K[μIndex,νIndex]
    end
  end
  return K
end

function computeEnergyInteratomicRepulsion(
  atom1::Atom,
  atom2::Atom)
  q1 = atom1.element.atomicNumber
  q2 = atom2.element.atomicNumber
  A = atom1.position
  B = atom2.position
  return q1*q2/(distance(A,B)^2)
end

function computeEnergyInteratomicRepulsion(geo::Geometry)
  result = 0.::Float64
  for (A in 1:length(geo.atoms))
    for (B in 1:A-1)
      result += computeEnergyInteratomicRepulsion(geo.atoms[A],geo.atoms[B])
    end
  end
  return result
end

function computeEnergyHartreeFock(
  basis::GaussianBasis,
  geometry::Geometry,
  density::Matrix)
  P = density
  T = computeMatrixKinetic(basis)
  V = computeMatrixNuclearAttraction(basis,geometry)
  J = computeMatrixCoulomb(basis,density)
  K = computeMatrixExchange(basis,density)
  G = 2J-K
  return trace((T+V+1/2*G)*P)
end

function computeMatrixFock(
  basis::GaussianBasis,
  geometry::Geometry,
  density::Matrix)
  T = computeMatrixKinetic(basis)
  V = computeMatrixNuclearAttraction(basis,geometry)
  J = computeMatrixCoulomb(basis,density)
  K = computeMatrixExchange(basis,density)
  return T+V+2J-K
end

function evaluateSCFStep(
  basis::GaussianBasis,
  geometry::Geometry,
  initialGuessDensity::Matrix,
  overlap::Matrix,
  electronNumber::Integer)

  N = electronNumber
  S = overlap
  F = computeMatrixFock(basis,geometry,initialGuessDensity)
  eig = eigfact(F,S)
  C = eig.vectors
  moEnergies = eig.values
  Cocc = eig.vectors[:,1:N]
  P = Cocc * Cocc.'

  return (moEnergies,P)
end

function evaluateSCF(
  basis::GaussianBasis,
  geometry::Geometry,
  initialGuessDensity::Matrix,
  electronNumber::Integer)

  N = electronNumber
  normalize!(basis)
  S = computeMatrixOverlap(basis)
  P = initialGuessDensity

  for (i in 1:30)
    energies,P = evaluateSCFStep(basis,geometry,P,S,N)
    println("E[$i]: $(computeEnergyHartreeFock(basis,geometry,P))")
  end
end

h2o = Geometry([Atom(Element("O"),Position(0.0, 0.0, 0.04851804)),
		Atom(Element("H"),Position(0.75300223,0.00000000,-0.51923377)),
		Atom(Element("H"),Position(-0.75300223,0.00000000,-0.51923377))])
angstrom2bohr!(h2o)

bfs = computeBasis(basSet,h2o)
normalize!(bfs)
bfs.contractedBFs[3],bfs.contractedBFs[5] = bfs.contractedBFs[5],bfs.contractedBFs[3]

function prettyprint(cgbf::ContractedGaussianBasisFunction,indent="")
  print(indent); dump(typeof(cgbf))
  for idx in 1:length(cgbf.coefficients)
    print(indent * "  $(cgbf.coefficients[idx]) × "); dump(typeof(cgbf.primitiveBFs[idx]))
    print(indent * "    center:   "); println(cgbf.primitiveBFs[idx].center)
    print(indent * "    exponent: "); println(cgbf.primitiveBFs[idx].exponent)
    print(indent * "    mqn:      "); println(cgbf.primitiveBFs[idx].mqn)
  end
end

function prettyprint(basis::GaussianBasis)
  dump(typeof(basis))
  indent = "  "

  println(indent * "contractedBFs: ")
  indent = indent * "  "

  for cgbf in basis.contractedBFs
    prettyprint(cgbf,indent)
  end

end

max_stars_repo_name: tkelman/QuantumLab

max_issues_repo_name: tkelman/QuantumLab

max_forks_repo_name: tkelman/QuantumLab


text:
using MLStyle.Render
module S_test_render
        b = 5
        node = :(function func_name(a)
            a + b
        end)
end
@testset "ast interpolation" begin
    x = 1
    y = 2

    node = @format [b = x + y] quote
        b + 2
    end
    @test eval(node) == 5

    S = S_test_render
    S.eval(render(S.node, Dict{Symbol, Any}(:b => 2, :func_name => :f1)))
    @test S.f1(1) == 3

end


text:
struct Cubic{T} <: SPHKernel
    n_neighbours::Int64
    norm_1D::T
    norm_2D::T
    norm_3D::T
end

"""
    Cubic(T::DataType=Float64, n_neighbours::Integer=64)

Set up a `Cubic` kernel for a given DataType `T`.
"""
Cubic(T::DataType=Float64, n_neighbours::Integer=64) = Cubic{T}(n_neighbours, 4.0/3.0, 40.0/7π, 8.0/π)

"""
    kernel_value_1D(kernel::Cubic{T}, u::Real, h_inv::Real) where T

Evaluate cubic spline at position ``u = \\frac{x}{h}``.
"""
function kernel_value_1D(kernel::Cubic{T}, u::Real, h_inv::Real) where T

    n = kernel.norm_1D * h_inv

    if u < 0.5
        return ( 1 + 6 * ( u - 1 ) * u^2) * n |> T
    elseif u < 1
        u_m1 = 1 - u
        return  2u_m1^3 * n |> T
    else
        return 0.0 |> T
    end

end

"""
    kernel_deriv_1D(kernel::Cubic{T}, u::Real, h_inv::Real) where T

Evaluate the derivative of the Cubic spline at position ``u = \\frac{x}{h}``.
"""
function kernel_deriv_1D(kernel::Cubic{T}, u::Real, h_inv::Real) where T

    n = kernel.norm_1D * h_inv^2

    if u < 0.5
        return ( u * (18u - 12 )) * n |> T
    elseif u < 1
        u_m1 = 1 - u
        return -6u_m1^2 * n |> T
    else
        return 0.0 |> T
    end

end

"""
    bias_correction_1D(kernel::Cubic{T}, density::Real, m::Real, h_inv::Real) where T

Does not do anything for the BSplines. Implemented for stability.
"""
function bias_correction_1D(kernel::Cubic{T}, density::Real, m::Real, h_inv::Real) where T
    return density
end

"""
    kernel_value_2D(kernel::Cubic{T}, u::Real, h_inv::Real) where T

Evaluate cubic spline at position ``u = \\frac{x}{h}``.
"""
function kernel_value_2D(kernel::Cubic{T}, u::Real, h_inv::Real) where T

    n = kernel.norm_2D * h_inv^2

    if u < 0.5
        return ( 1 + 6 * ( u - 1 ) * u^2) * n |> T
    elseif u < 1.0
        u_m1 = 1 - u
        return  2u_m1^3 * n |> T
    else
        return 0.0 |> T
    end

end

"""
    kernel_deriv_2D(kernel::Cubic{T}, u::Real, h_inv::Real) where T

Evaluate the derivative of the Cubic spline at position ``u = \\frac{x}{h}``.
"""
function kernel_deriv_2D(kernel::Cubic{T}, u::Real, h_inv::Real) where T

    n = kernel.norm_2D * h_inv^3

    if u < 0.5
        return ( u * (18u - 12 )) * n |> T
    elseif u < 1.0
        u_m1 = 1 - u
        return  -6u_m1^2 * n |> T
    else
        return 0.0 |> T
    end

end

"""
    bias_correction_2D(kernel::Cubic, density::Real, m::Real, h_inv::Real)

Does not do anything for the BSplines. Implemented for stability.
"""
function bias_correction_2D(kernel::Cubic{T}, density::Real, m::Real, h_inv::Real) where T
    return density
end

"""
    kernel_value_3D(kernel::Cubic{T}, u::Real, h_inv::Real) where T

Evaluate cubic spline at position ``u = \\frac{x}{h}``.
"""
function kernel_value_3D(kernel::Cubic{T}, u::Real, h_inv::Real) where T

    n = kernel.norm_3D * h_inv^3

    if u < 0.5
        return ( 1 + 6 * ( u - 1 ) * u^2) * n |> T
    elseif u < 1.0
        u_m1 = 1 - u
        return 2u_m1^3 * n |> T
    else
        return 0.0 |> T
    end

end

"""
    kernel_deriv_3D(kernel::Cubic{T}, u::Real, h_inv::Real) where T

Evaluate the derivative of the Cubic spline at position ``u = \\frac{x}{h}``.
"""
function kernel_deriv_3D(kernel::Cubic{T}, u::Real, h_inv::Real) where T

    n = kernel.norm_3D * h_inv^2 * h_inv^2

    if u < 0.5
        return ( u * (18u - 12 )) * n |> T
    elseif u < 1.0
        u_m1 =  1 - u
        return -6u_m1^2 * n |> T
    else
        return 0.0 |> T
    end

end

"""
    bias_correction_3D(kernel::Cubic{T}, density::Real, m::Real, h_inv::Real) where T

Does not do anything for the BSplines. Implemented for stability.
"""
function bias_correction_3D(kernel::Cubic{T}, density::Real, m::Real, h_inv::Real)  where T
    return density |> T
end


text:
export get_features

"""
    features(jlt::JLBoostTree)
    features(jlt::JLBoostTreeModel)

Return all the features used in the tree
"""
get_features(jlt::JLBoostTreeModel) = get_features(trees(jlt))

get_features(jlt::JLBoostTree) = begin
    d = get_features!(jlt::JLBoostTree, Dict{Symbol, Bool}())
    keys(d) |> collect |> sort
end

get_features(jlt::AbstractVector{<:JLBoostTree}) = begin
    d = get_features!(jlt[1])
    for jlt1 in @view(jlt[2:end])
        get_features!(jlt1, d)
    end
    keys(d) |> collect |> sort
end

get_features!(jlt::JLBoostTree, d = Dict{Symbol, Bool}()) = begin

    if isequal(jlt.splitfeature, missing)
        return d
    else
        d[jlt.splitfeature] = true
    end

    for c in jlt.children
        get_features!(c, d)
    end

    d
end


text:
@testset "Computation graph tests" begin

    using Functors: fmap

    @testset "Scalar computation graphs" begin

        # Setup a simple scalar graph which sums two numbers
        ins = InputVertex.(1:3)
        sumvert = CompVertex(+, ins[1], ins[2])
        scalevert = CompVertex(x -> 2x, sumvert)
        graph = CompGraph(inputs(sumvert), [scalevert])
        sumvert2 = CompVertex((x,y) -> x+y+2, ins[1], ins[3])
        graph2out = CompGraph(ins, [scalevert, sumvert2])

        @testset "Structural tests" begin
            @test nvertices(graph) == 4
            @test nvertices(graph2out) == 6

            @test vertices(graph) == [ins[1], ins[2], sumvert, scalevert]
            @test vertices(graph2out) == [ins[1], ins[2], sumvert, scalevert, ins[3], sumvert2]
        end

        @testset "Computation tests" begin
            @test graph(2,3) == 10
            @test graph([2], [3]) == [10]
            @test graph(0.5, 1.3) ≈ 3.6
            @test graph2out(4,5,8) == (18, 14)
        end
    end

    @testset "Array computation graphs" begin

        # Setup a graph which scales one of the inputs by 3 and then merges is with the other
        ins = InputVertex.(1:2)
        scalevert = CompVertex(x -> 3 .* x, ins[1])
        mergegraph1 = CompGraph(ins, [CompVertex(hcat, ins[2], scalevert)])
        mergegraph2 = CompGraph(ins, [CompVertex(hcat, scalevert, ins[2])])

        @testset "Computation tests" begin
            @test CompGraph([ins[1]], [scalevert])(ones(Int64, 1, 2)) == [3 3]

            @test mergegraph1(ones(Int64, 3,2), ones(Int64, 3,3)) == [1 1 1 3 3; 1 1 1 3 3; 1 1 1 3 3]
            @test mergegraph2(ones(Int64, 3,2), ones(Int64, 3,3)) == [3 3 1 1 1; 3 3 1 1 1; 3 3 1 1 1]
        end

    end

    @testset "Simple graph copy" begin
        ins = InputVertex.(1:3)
        v1 = CompVertex(+, ins[1], ins[2])
        v2 = CompVertex(vcat, v1, ins[3])
        v3 = CompVertex(vcat, ins[1], v1)
        v4 = CompVertex(-, v3, v2)
        v5 = CompVertex(/, ins[1], v1)
        graph = CompGraph(ins, [v5, v4])

        @testset "Copy graph with $label" for (label, cfun) in (
            (deepcopy, deepcopy),
            ("fmap", g -> fmap(identity, g))
        )
            gcopy = cfun(graph)

            @test issame(graph, gcopy)
            @test graph(3,4,10) == gcopy(3,4,10)
        end
    end

    @testset "Indexing" begin
        in1 = inputvertex("in1", 1)
        in2 = inputvertex("in2", 1)
        v1 = absorbvertex(MatMul(nout(in1), nout(in1)), in1; traitdecoration=named("v1"))
        v2 = "v2" >> in2 + v1
        v3 = absorbvertex(MatMul(nout(v2), 2), v2; traitdecoration=named("v3"))
        v4 = conc(v1, v3; dims=2)

        graph = CompGraph([in1, in2], v4)

        @test graph[1:nvertices(graph)] == graph[begin:end] == vertices(graph)
        @test graph[begin] == graph[1] == in1
        @test graph[nvertices(graph)] == graph[end] == v4

        @test findvertices("v1", graph) == [v1]
        @test findvertices(r"^in", graph) == [in1, in2]
    end

    @testset "Mutation graph copy" begin
        ins = inputvertex.(1:3, 1)
        v1 = ins[1] + ins[2]
        v2 = conc(v1, ins[3], dims=1)
        v3 = conc(ins[1], v1, dims=1)
        v4 = v3 - v2
        v5 = ins[1] / v1
        v6 = absorbvertex(identity, v5)
        graph = CompGraph(ins, [v4, v6])
        vs = vertices(graph)

        @testset "Copy graph with $label" for (label, cfun) in (
            (deepcopy, deepcopy),
            ("fmap", g -> fmap(identity, g))
        )
            gcopy = cfun(graph)

            @test issame(graph, gcopy)
            @test graph(3,4,10) == gcopy(3,4,10)

            vscopy = vertices(gcopy)
            @test length(vscopy) == length(vs)
            @test any(v -> v ∈ vs, vscopy) == false

            allgraph = all_in_graph(inputs(gcopy)[1])
            @test length(allgraph) == length(vs)
            @test any(v -> v ∈ vs, allgraph) == false

            @testset "Consistency for vertex $i" for (i, v) in enumerate(vscopy)
                @testset "Is in outputs of input $j" for (j, vi) in enumerate(inputs(v))
                    @test v in outputs(vi)
                end
                @testset "Is in inputs of output $j" for (j, vo) in enumerate(outputs(v))
                    @test v in inputs(vo)
                end
            end
        end
    end

    @testset "Graph add/remove trait" begin
        struct MockTrait <: DecoratingTrait
            t::MutationTrait
        end

        @functor MockTrait

        inver = inputvertex("in", 3)
        v1 = absorbvertex(+, inver)
        v2 = conc(inver, v1, dims=1)
        graph = CompGraph(inver, v2)

        # slightly annoying that Functors treats leaves and non-leaves differently
        # MutationTraits happen to be leaves as they have no children
        addtrait(x) = x
        addtrait(t::MutationTrait) = MockTrait(t)
        graphnewtrait = fmap(addtrait, graph)

        @testset "Check trait for vertex $i" for (i, v) in enumerate(filter(v -> v ∉ inputs(graphnewtrait), vertices(graphnewtrait)))
            @test trait(v) isa MockTrait
        end

        rmtrait(f, x) = Functors._default_walk(f, x)
        rmtrait(f, t::MockTrait) = t.t
        grapholdtrait = fmap(identity, graphnewtrait; walk=rmtrait)

        expected = (SizeAbsorb, SizeStack)
        @testset "Check trait for vertex $i" for (i, v) in enumerate(filter(v -> v ∉ inputs(grapholdtrait), vertices(grapholdtrait)))
            @test trait(v) isa expected[i]
        end
    end

    @testset "Graph rename" begin
        v0 = inputvertex("in", 3)
        v1 = absorbvertex(+, v0, traitdecoration = named("v1"))
        v2 = conc(v0, v1, dims=1, traitdecoration = named("v2"))
        graph = CompGraph(v0, v2)

        rename(x) = x
        rename(s::String) = s * "new"

        graphnew = fmap(rename, graph)

        @test name.(vertices(graphnew)) == ["innew", "v1new", "v2new"]
    end

    @testset "Topological sort" begin
        in1,in2 = InputVertex.(("in1", "in2"))
        v1 = CompVertex(+, in1, in2)
        v2 = CompVertex(-, in2, in1)
        v3 = CompVertex(vcat, v1,v2,v1)
        g = CompGraph([in1, in2], v3)

        @test vertices(g) == [in1, in2, v1, v2, v3]
    end
end


text:
using Base.Test
using Tinyurl

@test Tinyurl.tinyurl("https://rahulkp220.github.io") == "http://tinyurl.com/ydgk73kx"


text:
"""
	AbstractKnotPoint{Nx,Nu,V,T}

Stores the states, controls, time, and time step at a single point along a trajectory of a
forced dynamical system with `n` states and `m` controls.

# Interface
All instances of `AbstractKnotPoint` should support the following methods:

| **Required methods**  |    | **Brief description**  |
|:--------------------- |:-- |:---------------------- |
| `state_dim(z)` | | State vector dimension |
| `control_dim(z)` | | Control vector dimension |
| `get_data(z)::V` | | Get the vector of concatenated states and controls |
| **Optional Methods** | **Default definition** | **Brief description** |
| `time(z)` | `z.t` | Get the time |
| `timestep(z)` | `z.dt` | Get the time step |
| `settime!(z, t)` | `z.t = t` | Set the time |
| `settimestep!(z, dt)` | `z.dt = dt` | Set the time step |
| `is_terminal(z)` | `timestep(z) === zero(datatype(z))` |

# Methods
Given the above interface, the following methods are defined for an `AbstractKnotPoint`:

* `dims(z)`
* [`getstate`](@ref)
* [`getcontrol`](@ref)
* [`state`](@ref)
* [`control`](@ref)
* [`state`](@ref)
* [`setdata!`](@ref)
* [`setstate!`](@ref)
* [`setcontrol!`](@ref)
* [`is_terminal`](@ref)
* [`datatype`](@ref)
* [`vectype`](@ref)
"""
abstract type AbstractKnotPoint{Nx,Nu,V,T} <: AbstractVector{T} end

dims(z::AbstractKnotPoint) = (state_dim(z), control_dim(z))

"""
    getstate(z, v)

Extract the state vector from a vector `v`. Returns a view into the matrix by default,
or an `SVector` if the knot point uses `SVector`s.
"""
getstate(z::AbstractKnotPoint, v) = view(v, 1:state_dim(z))

"""
    getcontrol(z, v)

Extract the state vector from a vector `v`. Returns a view into the matrix by default,
or an `SVector` if the knot point uses `SVector`s. If the knot point is a terminal
knot point, it will return an empty view or an `SVector` of zeros.
"""
function getcontrol(z::AbstractKnotPoint, v)
    if is_terminal(z)
        return view(v, length(v):length(v)-1)
    else
        n,m = dims(z)
        return view(v, n+1:n+m)
    end
end

getstate(z::AbstractKnotPoint{Nx,Nu,<:SVector}, v) where {Nx,Nu} = v[SVector{Nx}(1:Nx)]
getcontrol(z::AbstractKnotPoint{Nx,Nu,<:SVector}, v) where {Nx,Nu} = !is_terminal(z) * v[SVector{Nu}(Nx+1:Nx+Nu)]

"""
    state(z)

Get the state vector. Returns either a view or an `SVector`.
"""
state(z::AbstractKnotPoint) = getstate(z, getdata(z))

"""
    control(z)

Get the control vector. Returns either a view of an `SVector`. For a terminal knot point
it will return an empty view or an `SVector` of zeros.
"""
control(z::AbstractKnotPoint) = getcontrol(z, getdata(z))

"""
    setdata!(z, v)

Set the data vector, or the concatenated vector of states and controls.
"""
setdata!(z::AbstractKnotPoint, v) = z.z .= v
setdata!(z::AbstractKnotPoint{<:Any,<:Any,<:SVector}, v) = z.z = v

"""
    setstate!(z, x)

Set the state vector for an [`AbstractKnotPoint`](@ref).
"""
setstate!(z::AbstractKnotPoint, x) = state(z) .= x
setstate!(z::AbstractKnotPoint{Nx,<:Any,<:SVector}, x) where Nx = setdata!(z, [SVector{Nx}(x); control(z)])

"""
    setcontrol!(z, x)

Set the control vector for an [`AbstractKnotPoint`](@ref).
"""
setcontrol!(z::AbstractKnotPoint, u) = control(z) .= u
setcontrol!(z::AbstractKnotPoint{<:Any,Nu,<:SVector}, u) where Nu = setdata!(z, [state(z); SVector{Nu}(u)])

settime!(z::AbstractKnotPoint, t) = z.t = t
settimestep!(z::AbstractKnotPoint, dt) = z.dt = dt

time(z::AbstractKnotPoint) = getparams(z).t
timestep(z::AbstractKnotPoint) = getparams(z).dt

"""
    is_terminal(z)

Determines if the knot point `z` is a terminal knot point with no
controls, only state information. By default a knot point is assumed to be a
terminal knot point if the time step is zero. By convention, if the last
knot point has control values assigned (for example, when doing first-order
hold on the controls), the final time step is set to infinity instead of zero.
"""
is_terminal(z::AbstractKnotPoint) = timestep(z) === zero(datatype(z))

"""
    datatype(x)

Get the numeric data type used by the object `x`. Typically a floating point data type.
"""
datatype(z::AbstractKnotPoint{<:Any,<:Any,<:Any,T}) where T = T

"""
    vectype(x)

Get the vector type used by the object `x`. Used to allow either static or
dynamic arrays in structs such as [`AbstractKnotPoint`](@ref).
"""
vectype(z::AbstractKnotPoint{<:Any,<:Any,V}) where V = V

# Array interface
Base.size(z::AbstractKnotPoint) = (state_dim(z) + control_dim(z), )
Base.getindex(z::AbstractKnotPoint, i::Int) = Base.getindex(getdata(z), i)
Base.setindex!(z::AbstractKnotPoint, v, i::Integer) = Base.setindex!(getdata(z), v, i)
Base.IndexStyle(z::AbstractKnotPoint) = IndexLinear()

for (name,mutable) in [(:KnotPoint, true), (:StaticKnotPoint, false)]
    expr = quote
        struct $name{Nx,Nu,V,T} <: AbstractKnotPoint{Nx,Nu,V,T}
            z::V
            t::Float64
            dt::Float64
            n::Int
            m::Int
            function $name{Nx,Nu}(z::V,t,dt) where {Nx,Nu,V}
                @assert Nx > 0
                @assert Nu > 0
                @assert Nx + Nu == length(z)
                new{Nx,Nu,V,eltype(V)}(z, t, dt, Nx, Nu)
            end
            function $name{Nx,Nu}(x::V, u::V, t,dt) where {Nx,Nu,V}
                Nx != Any && @assert Nx == length(x)
                Nu != Any && @assert Nu == length(u)
                new{Nx,Nu,V,eltype(V)}([x;u], t, dt, Nx, Nu)
            end
            function $name(n::Integer, m::Integer, z::V, t, dt) where V
                @assert n > 0
                @assert m > 0
                @assert n + m == length(z)
                new{Any,Any,V,eltype(V)}(z, t, dt, n, m)
            end
            function $name{Nx,Nu}(n::Integer, m::Integer, z::V, t, dt) where {Nx,Nu,V}
                Nx != Any && @assert Nx == n
                Nu != Any && @assert Nu == m
                @assert n + m == length(z)
                new{Nx,Nu,V,eltype(V)}(z, t, dt, n, m)
            end
        end
        function $name(x::StaticVector{Nx}, u::StaticVector{Nu}, t, dt) where {Nx, Nu}
            $name{Nx,Nu}([x;u], t, dt)
        end
        function $name(x::AbstractVector, u::AbstractVector, t, dt)
            $name(length(x), length(u), [x;u], t, dt)
        end
        function Base.copy(z::$name{Nx,Nu}) where {Nx,Nu}
            $name{Nx,Nu}(z.n, z.m, z.z, z.t, z.dt)
        end

        state_dim(z::$name{Any}) = z.n
        control_dim(z::$name{<:Any,Any}) = z.m
        state_dim(z::$name{Nx}) where Nx = Nx
        control_dim(z::$name{<:Any,Nu}) where Nu = Nu
        vectype(::$name{<:Any,<:Any,V}) where V = V
        vectype(::Type{<:$name{<:Any,<:Any,V}}) where V = V

        @inline getparams(z::$name) = (t=z.t, dt=z.dt)
        @inline getdata(z::$name) = z.z
        function Base.copyto!(dest::$name, src::$name)
            if ismutabletype($name)
                dest.t = src.t
                dest.dt = src.dt
            end
            copyto!(dest.z, src.z)
            dest
        end

        # Array interface
        Base.similar(z::$name{Nx,Nu,V,T}, ::Type{S}, dims::Dims) where {Nx,Nu,V,T,S} =
            KnotPoint{Nx,Nu}(z.n, z.m, similar(z.z, S, dims), z.t, z.dt)
        Base.similar(z::$name{Nx,Nu,<:SVector,T}, ::Type{S}, dims::Dims) where {Nx,Nu,T,S} =
            KnotPoint{Nx,Nu}(z.n, z.m, similar(MVector{Nx+Nu,S}), z.t, z.dt)
    end
    # Set struct mutability
    if mutable
        struct_expr = expr.args[2]
        @assert struct_expr.head == :struct
        struct_expr.args[1] = mutable
    end
    eval(expr)
end

function (::Type{<:AbstractKnotPoint{Nx,Nu}})(z, t, dt) where {Nx,Nu}
    KnotPoint{Nx,Nu}(z, t, dt)
end

function StaticKnotPoint(z::AbstractKnotPoint{Nx,Nu}, v::AbstractVector) where {Nx,Nu}
    StaticKnotPoint{Nx,Nu}(state_dim(z), control_dim(z), v, time(z), timestep(z))
end
@inline StaticKnotPoint(z::AbstractKnotPoint) = StaticKnotPoint(z, getdata(z))
@inline setdata(z::StaticKnotPoint, v) = StaticKnotPoint(z, v)
setstate(z::StaticKnotPoint, x) = setdata(z, [x; control(z)])
setcontrol(z::StaticKnotPoint, u) = setdata(z, [state(z); u])

function Base.:*(c::Real, z::KP) where {KP<:AbstractKnotPoint}
    KP(getdata(z)*c, getparams(z)...)
end

"""
    KnotPoint{Nx,Nu,V,T}

A mutable [`AbstractKnotPoint`](@ref) with `Nx` states, `Nu` controls, stored using
a vector type `V` with data type `T`. Since the struct is mutable, the time, timestep,
and data can all be changed, which can be very efficient when the data being stored
as an `SVector`.

# Constructors

    KnotPoint{n,m}(v, t, dt)
    KnotPoint{n,m}(x, u, t, dt)
    KnotPoint{n,m}(n, m, v, t, dt)
    KnotPoint(n, m, v, t, dt)       # create a KnotPoint{Any,Any}
    KnotPoint(x, u, t, dt)

The last method will create a `KnotPoint{Any,Any}` if `x` and `u` are not `StaticVector`s.

The vector type `V` can be queried using `vectype(z)`.
"""
KnotPoint

"""
    StaticKnotPoint

A static version of [`KnotPoint`](@ref). Uses all of the same methods and constructors,
but also adds the following methods:

    StaticKnotPoint(z, v)

which creates a new `StaticKnotPoint` using the information from the `AbstractKnotPoint` z,
but using data from `v`. Useful for creating temporary knot points from existing ones
without any runtime allocations.

The following methods are similar to their mutable versions, but create a new
`StaticKnotPoint`:

    setdata
    setstate
    setcontrol
"""
StaticKnotPoint


text:
module Pretend

using Base: Callable
using ExprTools: splitdef, combinedef
using MacroTools: postwalk, rmlines

export @mockable
export apply, mocked, spy
export Fallback

export
    called_at_least_once,
    called_exactly_once,
    was_not_called

include("types.jl")
include("activation.jl")
include("patchstore.jl")
include("mock.jl")
include("spy.jl")

end


text:
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# Description
# ==============================================================================
#
#   Structures related to the orbit propagators.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

export OrbitPropagator

"""
    OrbitPropagator{T}

Abstract type of the orbit propagator. Every propagator structure must be a
subtype of this type and must implement the following API functions:

    propagate!(orbp, t::Number)
    propagate!(orbp, t::AbstractVector)
    propagate_to_epoch!(orbp, JD::Number)
    propagate_to_epoch!(orbp, JD::AbstractVector)
    step!(orbp, Δt::Number)

"""
abstract type OrbitPropagator{T} end

#                          Two Body Orbit Propagator
# ==============================================================================

export TwoBody_Structure, OrbitPropagatorTwoBody

"""
    TwoBody_Structure{T}

Low level Two Body orbit propagator structure.

"""
@with_kw mutable struct TwoBody_Structure{T}
    # Initial mean orbital elements
    # ==========================================================================
    epoch::T # ......................... Epoch of the initial mean elements [JD]
    a_0::T   # ..................... Initial mean normalized semi-major axis [m]
    n_0::T   # ..................................... Initial mean motion [rad/s]
    e_0::T   # .................................. Initial mean eccentricity [  ]
    i_0::T   # .................................. Initial mean inclination [rad]
    Ω_0::T   # ......................................... Initial mean RAAN [rad]
    ω_0::T   # .......................... Initial mean argument of perigee [rad]
    f_0::T   # ................................. Initial mean true anomaly [rad]
    M_0::T   # ................................. Initial mean mean anomaly [rad]
    μ::T     # .. Standard gravitational parameter of the central body [m^3/s^2]

    # Current mean orbital elements
    # ==========================================================================
    Δt::T   # .............. Timespan from the epoch of the initial elements [s]
    f_k::T  # .................................. Current mean true anomaly [rad]
    M_k::T  # .................................. Current mean mean anomaly [rad]
end

"""
    OrbitPropagatorTwoBody{T} <: OrbitPropagator{T}

Structure that holds the information related to the Two Body orbit propagator.

# Fields

* `orb`: Mean orbital elements (see `Orbit`).
* `tbd`: Structure that stores the Two Body orbit propagator data (see
        `TwoBody_Structure`).

"""
struct OrbitPropagatorTwoBody{T} <: OrbitPropagator{T}
    tbd::TwoBody_Structure{T}
end

#                             J2 Orbit Propagator
# ==============================================================================

export J2_GravCte, J2_Structure, OrbitPropagatorJ2

"""
    J2_GravCte{T}

Gravitational constants for J2 orbit propagator.

# Fields

* `R0`: Earth equatorial radius [m].
* `μm`: √GM [er/s]^(3/2).
* `J2`: The second gravitational zonal harmonic of the Earth.

"""
@with_kw struct J2_GravCte{T}
    R0::T
    μm::T
    J2::T
end

"""
    J2_Structure{T}

Low level J2 orbit propagator structure.

"""
@with_kw mutable struct J2_Structure{T}
    # Initial mean orbital elements
    # ==========================================================================
    epoch::T              # ............ Epoch of the initial mean elements [JD]
    al_0::T               # ....... Initial mean normalized semi-major axis [er]
    n_0::T                # ........................ Initial mean motion [rad/s]
    e_0::T                # ..................... Initial mean eccentricity [  ]
    i_0::T                # ..................... Initial mean inclination [rad]
    Ω_0::T                # ............................ Initial mean RAAN [rad]
    ω_0::T                # ............. Initial mean argument of perigee [rad]
    f_0::T                # .................... Initial mean true anomaly [rad]
    M_0::T                # .................... Initial mean mean anomaly [rad]
    dn_o2::T              # ...... First time derivative of mean motion [rad/s²]
    ddn_o6::T             # ..... Second time derivative of mean motion [rad/s³]
    j2_gc::J2_GravCte{T}  # ........ J2 orbit propagator gravitational constants

    # Current mean orbital elements
    # ==========================================================================
    Δt::T   # .............. Timespan from the epoch of the initial elements [s]
    al_k::T # ..................... Current mean normalized semi-major axis [er]
    e_k::T  # ................................... Current mean eccentricity [  ]
    i_k::T  # ................................... Current mean inclination [rad]
    Ω_k::T  # .......................................... Current mean RAAN [rad]
    ω_k::T  # ........................... Current mean argument of perigee [rad]
    f_k::T  # .................................. Current mean true anomaly [rad]
    M_k::T  # .................................. Current mean mean anomaly [rad]

    # Auxiliary variables
    # ==========================================================================
    δa::T   # ........................... Semi-major axis time derivative [er/s]
    δe::T   # ............................... Eccentricity time derivative [1/s]
    δΩ::T   # ..................................... RAAN time derivative [rad/s]
    δω::T   # ...................... Argument of perigee time derivative [rad/s]
    δM_0::T # ............................. Mean anomaly time derivative [rad/s]
end

"""
    OrbitPropagatorJ2{T} <: OrbitPropagator{T}

Structure that holds the information related to the J2 orbit propagator.

# Fields

* `j2d`: Structure that stores the J2 orbit propagator data (see
         `J2_Structure`).

"""
struct OrbitPropagatorJ2{T} <: OrbitPropagator{T}
    j2d::J2_Structure{T}
end

#                        J2 osculating orbit propagator
# ==============================================================================

export J2osc_Strutcture, OrbitPropagatorJ2osc

"""
    J2osc_Structure{T}

Low level J2 osculating orbit propagator structure.

"""
@with_kw mutable struct J2osc_Structure{T}
    # J2 orbit propagator to propagate the mean elements.
    j2d::J2_Structure{T}

    # Propagation time from epoch.
    Δt::T

    # Current osculating Keplerian elements
    # ==========================================================================
    a_k::T # ................................... Osculating semi-major axis [er]
    e_k::T # ....................................... Osculating eccentricity [ ]
    i_k::T # ...................................... Osculating inclination [rad]
    Ω_k::T # ............................................. Osculating RAAN [rad]
    ω_k::T # .............................. Osculating argument of perigee [rad]
    f_k::T # ..................................... Osculating true anomaly [rad]
    M_k::T # ..................................... Osculating mean anomaly [rad]
end

"""
    OrbitPropagatorJ2osc{T} <: OrbitPropagator{T}

Structure that holds the information related to the J2 osculating orbit
propagator.

# Fields

* `j2oscd`: Structure that stores the J2 osculating orbit propagator data (see
            `J2osc_Structure`).

"""
struct OrbitPropagatorJ2osc{T} <: OrbitPropagator{T}
    j2oscd::J2osc_Structure{T}
end

#                             J4 orbit propagator
# ==============================================================================

export J4_GravCte, J4_Structure, OrbitPropagatorJ4

"""
    J4_GravCte{T}

Gravitational constants for J4 orbit propagator.

# Fields

* `R0`: Earth equatorial radius [m].
* `μm`: √GM [er/s]^(3/2).
* `J2`: The second gravitational zonal harmonic of the Earth.
* `J4`: The fourth gravitational zonal harmonic of the Earth.

"""
@with_kw struct J4_GravCte{T}
    R0::T
    μm::T
    J2::T
    J4::T
end

"""
    J4_Structure{T}

Low level J4 orbit propagator structure.

"""
@with_kw mutable struct J4_Structure{T}
    # Initial mean orbital elements
    # ==========================================================================
    epoch::T              # ............ Epoch of the initial mean elements [JD]
    al_0::T               # ....... Initial mean normalized semi-major axis [er]
    n_0::T                # ........................ Initial mean motion [rad/s]
    e_0::T                # ..................... Initial mean eccentricity [  ]
    i_0::T                # ..................... Initial mean inclination [rad]
    Ω_0::T                # ............................ Initial mean RAAN [rad]
    ω_0::T                # ............. Initial mean argument of perigee [rad]
    f_0::T                # .................... Initial mean true anomaly [rad]
    M_0::T                # .................... Initial mean mean anomaly [rad]
    dn_o2::T              # ...... First time derivative of mean motion [rad/s²]
    ddn_o6::T             # ..... Second time derivative of mean motion [rad/s³]
    j4_gc::J4_GravCte{T}  # ........ J4 orbit propagator gravitational constants

    # Current mean orbital elements
    # ==========================================================================
    Δt::T   # .............. Timespan from the epoch of the initial elements [s]
    al_k::T # ..................... Current mean normalized semi-major axis [er]
    e_k::T  # ................................... Current mean eccentricity [  ]
    i_k::T  # ................................... Current mean inclination [rad]
    Ω_k::T  # .......................................... Current mean RAAN [rad]
    ω_k::T  # ........................... Current mean argument of perigee [rad]
    f_k::T  # .................................. Current mean true anomaly [rad]
    M_k::T  # .................................. Current mean mean anomaly [rad]

    # Auxiliary variables
    # ==========================================================================
    δa::T   # ........................... Semi-major axis time derivative [er/s]
    δe::T   # ............................... Eccentricity time derivative [1/s]
    δΩ::T   # ..................................... RAAN time derivative [rad/s]
    δω::T   # ...................... Argument of perigee time derivative [rad/s]
    δM_0::T # ............................. Mean anomaly time derivative [rad/s]
end

"""
    OrbitPropagatorJ4{T} <: OrbitPropagator{T}

Structure that holds the information related to the J4 orbit propagator.

# Fields

* `j4d`: Structure that stores the J4 orbit propagator data (see
         `J4_Structure`).

"""
struct OrbitPropagatorJ4{T} <: OrbitPropagator{T}
    j4d::J4_Structure{T}
end

#                                     SGP4
# ==============================================================================

export OrbitPropagatorSGP4

"""
    OrbitPropagatorSGP4{T} <: OrbitPropagator{T}

Structure that holds the information related to the SGP4 propagator.

# Fields

* `sgp4d`: Structure that stores the SGP4 data (see `SGP4_Structure`).

"""
struct OrbitPropagatorSGP4{T} <: OrbitPropagator{T}
    sgp4d::SGP4_Structure{T}
end


text:
__precompile__()

module NLPModelsJuMP

include("nlsmpb_model.jl")
include("nlp_to_mpb.jl")
include("jump_model.jl")

end


text:
"""
Abstract problem.
"""
abstract type AbstractProblem{T} end

"""
    GeneralBTPDE(;
        model,
        matrices,
        abstol = 1e-6,
        reltol = 1e-4,
        odesolver = QNDF(autodiff = false),
    )

General BTPDE problem.
"""
Base.@kwdef struct GeneralBTPDE{T,S} <: AbstractProblem{T}
    model::Model{T}
    matrices::NamedTuple
    abstol::T = 1e-6
    reltol::T = 1e-4
    odesolver::S = QNDF(autodiff = false)
end

"""
    IntervalConstantBTPDE(; model, matrices, θ = 0.5, timestep)

BTPDE problem specialized on intervalwise constant `ScalarGradient`s, e.g [`PGSE`](@ref),
[`DoublePGSE`](@ref).
"""
Base.@kwdef struct IntervalConstantBTPDE{T} <: AbstractProblem{T}
    model::Model{T}
    matrices::NamedTuple
    θ::T = 0.5
    timestep::T
end

"""
    HADC(;
        model,
        matrices,
        abstol = 1e-6,
        reltol = 1e-4,
        odesolver = QNDF(autodiff = false),
    )

HADC problem.
"""
Base.@kwdef struct HADC{T,S} <: AbstractProblem{T}
    model::Model{T}
    matrices::NamedTuple
    abstol::T = 1e-6
    reltol::T = 1e-4
    odesolver::S = QNDF(autodiff = false)
end

"""
    Karger(; model, difftensors, odesolver, timestep)

Karger problem.
"""
Base.@kwdef struct Karger{T,S} <: AbstractProblem{T}
    model::Model{T}
    difftensors::Vector{SMatrix{3,3,T,9}}
    odesolver::S = MagnusGL6()
    timestep::T
end

"""
    Laplace(; model, matrices, neig_max)

Laplace eigenvalue problem.
"""
Base.@kwdef struct Laplace{T} <: AbstractProblem{T}
    model::Model{T}
    matrices::NamedTuple
    neig_max::Int
end

"""
    MatrixFormalism(; model, matrices, lap_eig, ninterval)

Matrix formalism problem. Given a Laplace eigendecomposition `lap_eig`, this problem
consists of computing the MF magnetization.
"""
Base.@kwdef struct MatrixFormalism{T} <: AbstractProblem{T}
    model::Model{T}
    matrices::NamedTuple
    lap_eig::NamedTuple
    ninterval::Int
end

"""
    AnalyticalLaplace(; ρ, r, D, W, T₂, γ, dim, eiglim, eigstep)

Analytical radial Laplace eigenvalue problem.
"""
Base.@kwdef struct AnalyticalLaplace{T} <: AbstractProblem{T}
    ρ::Vector{Complex{T}}
    r::Vector{T}
    D::Vector{T}
    W::Vector{T}
    T₂::Vector{T}
    γ::T
    dim::Int
    eiglim::T
    eigstep::T
end

"""
    AnalyticalMatrixFormalism(; analytical_laplace, lap_mat, volumes)

Analytical Matrix formalism problem. Given a radial Laplace eigendecomposition
`analytical_laplace`, this problem consists of computing the MF compartment signals.
"""
Base.@kwdef struct AnalyticalMatrixFormalism{T} <: AbstractProblem{T}
    analytical_laplace::AnalyticalLaplace{T}
    lap_mat::NamedTuple
    volumes::Vector{T}
end


text:
if get(ENV, "CI", "false") == "true"
    using Distributed
    #! format: off
    VERSION < v"1.8-" && # workaround the hang in CI
    if nprocs() < 4
        addprocs(4 - nprocs())
    end
    #! format: on
end

using TestFunctionRunner
TestFunctionRunner.@run


text:
using Test
import Forecast: nextodd, drop, compact, insert_row, insert_column

@testset "utils" begin

    # nextodd
    @test nextodd(5) == 5
    @test nextodd(5.) isa Integer
    @test nextodd(2) == 3

    # drop
    M = [1 2 3
         4 5 6
         7 8 9]
    @test drop(M;r=2,c=2) == [1 3
                              7 9]
    @test drop(M;r=1,c=2) == [4 6
                              7 9]
    M = [1 2 3.
         4 5 6
         7 8 9]
    @test drop(M;r=1,c=3) == [4. 5.
                              7. 8.]

    @test drop(M) == drop(M;r=[],c=[]) == M
    @test drop(M;r=[1,2],c=3) == [7 8]
    @test drop(M;r=[1,2],c=[2,3]) == reshape([7],1,1)

    # compact
    compact(5) ==  5
    compact(5) isa Integer
    compact(5.) ==  5.
    compact(5.) isa AbstractFloat

    x = reshape([1],1,1,1,1,1)
    compact(x) == 1

    x = reshape([1 2],2,1,1,1,1)
    compact(x) == [1,2]
    x = reshape([1 2],1,1,2,1,1)
    compact(x) == [1,2]
    x = reshape([1 2; 3 4],1,1,2,1,2)
    compact(x) == [1 2
                   3 4]

    #insert row / column
    M = zeros(2,2)
    @test insert_row(M, 2, 1) == [0 0
                                  1 1
                                  0 0]
    @test insert_column(M, 2, 1) == [0 1 0
                                     0 1 0]

    @test insert_column(insert_row(M,2,1),2,1) == [0 1 0
                                                   1 1 1
                                                   0 1 0]
end


text:
## SESSION PREPARATION
pwd()

using CSV, DataFrames, DataFramesMeta, MixedModels, PooledArrays
using StatsBase: mean, std, zscore

include("spr_fns.jl")

## DATA PREPARATION
# Load data and scale continuous covariates
dt = read_spr_data("data/CAPSPR.csv");

dt.srp = dt.srp .* -1
dt.rcnoun = dt.rcnoun .* -1

# exclude data
dt = exclude_trial(dt[((dt.Region .!== "critical -2") .& (dt.Duplicated .!== "multi")),:], 50, 2500, 50, 6000)

#### TODO: zscore rcnoun / rcverb. Update formula. Run models anew.

## MODEL PREPARATION
contr = Dict(:Subject => Grouping(), :Item => Grouping())

f1 = @formula(logRT ~ 1 + Cloze + (1 + Cloze | Item) + (1 + Cloze | Subject))
f2 = @formula(logRT ~ 1 + srp + (1 + srp | Item) + (1 + srp | Subject))
f3 = @formula(logRT ~ 1 + rcnoun + (1 + rcnoun | Item) + (1 + rcnoun | Subject))
f3a = @formula(logRT ~ 1 + logrcnoun + (1 + logrcnoun | Item) + (1 + logrcnoun | Subject))
f4 = @formula(logRT ~ 1 + rcverb + (1 + rcverb | Item) + (1 + rcverb | Subject))
f5 = @formula(logRT ~ 1 + rcnoun + rcverb + (1 + rcnoun + rcverb | Item) + (1 + rcnoun + rcverb | Subject))
f6 = @formula(logRT ~ 1 + rcnoun + srp + (1 + rcnoun + srp | Item) + (1 + rcnoun + srp | Subject))
f7 = @formula(logRT ~ 1 + srp + rcnoun + inter + (1 + srp + rcnoun + inter | Item) + (1 + srp + rcnoun + inter | Subject))
f8 = @formula(logRT ~ 1 + Cloze + rcnoun + (1 + Cloze + rcnoun | Item) + (1 + Cloze + rcnoun | Subject))

f = f2

lmerSPR = fit_models(dt, f, contr);

numpred = length(f.rhs) - 2
coefs = [Symbol(x) for x in string.(f.rhs[1:numpred])];

lmerSPR = combine_datasets(dt, lmerSPR, coefs[2:length(coefs)])
lmerSPR = compute_RTs(lmerSPR, coefs);
estimates = names(lmerSPR)[occursin.("est_", names(lmerSPR))]
lmerSPR = compute_residuals(lmerSPR, estimates);
lmerSPR = compute_coefs(lmerSPR, coefs);

CSV.write("data/lmerSPR_Asrp.csv", lmerSPR);


text:
"""
    finite-difference acceleration objective
"""
struct AccelerationObjective <: Objective
    Q
    n
    h
end

acceleration_objective(Q, n; h = 0.0) = AccelerationObjective(Q, n, h)

function objective(Z, obj::AccelerationObjective, idx, T)
    n = obj.n

    J = 0.0

    for t = 2:T-1
        # configurations
        q⁻ = view(Z, idx.x[t - 1][n .+ (1:n)])
        q = view(Z, idx.x[t][n .+ (1:n)])
        q⁺ = view(Z, idx.x[t + 1][n .+ (1:n)])

        # time steps
        h⁻ = obj.h == 0.0 ? view(Z, idx.u[t - 1])[end] : obj.h
        h⁺ = obj.h == 0.0 ? view(Z, idx.u[t])[end] : obj.h

        # velocities
        v⁺ = (q⁺ - q) / h⁺
        v⁻ = (q - q⁻) / h⁻

        # acceleration
        a = (v⁺ - v⁻) / (0.5 * h⁺ + 0.5 * h⁻)

        J += a' * obj.Q[t] * a
    end

    return J
end

function objective_gradient!(∇J, Z, obj::AccelerationObjective, idx, T)

    n = obj.n

    for t = 2:T-1
        # configurations
        q⁻ = view(Z, idx.x[t - 1][n .+ (1:n)])
        q = view(Z, idx.x[t][n .+ (1:n)])
        q⁺ = view(Z, idx.x[t + 1][n .+ (1:n)])

        # time steps
        h⁻ = obj.h == 0.0 ? view(Z, idx.u[t - 1])[end] : obj.h
        h⁺ = obj.h == 0.0 ? view(Z, idx.u[t])[end] : obj.h

        # velocities
        v⁺ = (q⁺ - q) / h⁺
        v⁻ = (q - q⁻) / h⁻

        # acceleration
        a = (v⁺ - v⁻) / (0.5 * h⁺ + 0.5 * h⁻)

        dJda = 2.0 * obj.Q[t] * a
        dadv⁺ = 1.0 / (0.5 * h⁺ + 0.5 * h⁻)
        dadv⁻ = -1.0 / (0.5 * h⁺ + 0.5 * h⁻)
        dadh⁺ = -0.5 * (v⁺ - v⁻) / (0.5 * h⁺ + 0.5 * h⁻)^2.0
        dadh⁻ = -0.5 * (v⁺ - v⁻) / (0.5 * h⁺ + 0.5 * h⁻)^2.0
        dv⁺dq⁺ = 1.0 / h⁺
        dv⁺dq = -1.0 / h⁺
        dv⁺dh⁺ = -1.0 * v⁺ / h⁺
        dv⁻dq = 1.0 / h⁻
        dv⁻dq⁻ = -1.0 / h⁻
        dv⁻dh⁻ = -1.0 * v⁻ / h⁻

        ∇J[idx.x[t - 1][n .+ (1:n)]] += dJda * dadv⁻ * dv⁻dq⁻
        ∇J[idx.x[t][n .+ (1:n)]] += dJda * (dadv⁺ * dv⁺dq + dadv⁻ * dv⁻dq)
        ∇J[idx.x[t + 1][n .+ (1:n)]] += dJda * dadv⁺ * dv⁺dq⁺
        if obj.h == 0.0
            ∇J[idx.u[t - 1][end]] += dJda' * (dadh⁻ + dadv⁻ * dv⁻dh⁻)
            ∇J[idx.u[t][end]] += dJda' * (dadh⁺ + dadv⁺ * dv⁺dh⁺)
        end
    end

    return nothing
end

max_stars_repo_name: zacmanchester/DirectMotionPlanning

max_issues_repo_name: zacmanchester/DirectMotionPlanning

max_forks_repo_name: zacmanchester/DirectMotionPlanning


text:
using DiffEqFlux, OrdinaryDiffEq, Test # , Plots

function lotka_volterra(du,u,p,t)
  x, y = u
  α, β, δ, γ = p
  du[1] = dx = (α - β*y)x
  du[2] = dy = (δ*x - γ)y
end
p = [2.2, 1.0, 2.0, 0.4]
u0 = [1.0,1.0]
prob = ODEProblem(lotka_volterra,u0,(0.0,10.0),p)

# Reverse-mode

function predict_rd(p)
  Array(solve(prob,Tsit5(),p=p,saveat=0.1,reltol=1e-4,sensealg=TrackerAdjoint()))
end
loss_rd(p) = sum(abs2,x-1 for x in predict_rd(p))
loss_rd() = sum(abs2,x-1 for x in predict_rd(p))
loss_rd()

grads = Zygote.gradient(loss_rd, p)
@test !iszero(grads[1])

opt = ADAM(0.1)
cb = function ()
  display(loss_rd())
  # display(plot(solve(remake(prob,p=p),Tsit5(),saveat=0.1),ylim=(0,6)))
end

# Display the ODE with the current parameter values.
loss1 = loss_rd()
Flux.train!(loss_rd, Flux.params(p), Iterators.repeated((), 100), opt, cb = cb)
loss2 = loss_rd()
@test 10loss2 < loss1

# Forward-mode, R^n -> R^m layer

p = [2.2, 1.0, 2.0, 0.4]
function predict_fd()
  vec(Array(solve(prob,Tsit5(),p=p,saveat=0.0:0.1:1.0,reltol=1e-4,sensealg=ForwardDiffSensitivity())))
end
loss_fd() = sum(abs2,x-1 for x in predict_fd())
loss_fd()

ps = Flux.params(p)
grads = Zygote.gradient(loss_fd, ps)
@test !iszero(grads[p])

data = Iterators.repeated((), 100)
opt = ADAM(0.1)
cb = function ()
  display(loss_fd())
  # display(plot(solve(remake(prob,p=p),Tsit5(),saveat=0.1),ylim=(0,6)))
end

# Display the ODE with the current parameter values.
loss1 = loss_fd()
Flux.train!(loss_fd, ps, data, opt, cb = cb)
loss2 = loss_fd()
@test 10loss2 < loss1

# Adjoint sensitivity
p = [2.2, 1.0, 2.0, 0.4]
ps = Flux.params(p)
function predict_adjoint()
    solve(remake(prob,p=p),Tsit5(),saveat=0.1,reltol=1e-4)
end
loss_reduction(sol) = sum(abs2,x-1 for x in vec(sol))
loss_adjoint() = loss_reduction(predict_adjoint())
loss_adjoint()

grads = Zygote.gradient(loss_adjoint, ps)
@test !iszero(grads[p])

data = Iterators.repeated((), 100)
opt = ADAM(0.1)
cb = function ()
  display(loss_adjoint())
  # display(plot(solve(remake(prob,p=p),Tsit5(),saveat=0.1),ylim=(0,6)))
end

# Display the ODE with the current parameter values.
loss1 = loss_adjoint()
Flux.train!(loss_adjoint, ps, data, opt, cb = cb)
loss2 = loss_adjoint()
@test 10loss2 < loss1


text:
export bissect_sec
function bissect_sec(h :: AbstractNLPModel;
                     t1 :: Float64 = h.meta.x0[1],
                     t2 :: Float64 = 100.0,
                     tol :: Float64 = 1e-7,
                     maxiter :: Int = 50,
                     verbose :: Bool = false)

    (length(h.meta.x0) > 1) && warn("Not a 1-D problem ")
  (tₐ, tᵦ) = trouve_intervalle(h, t1, 2.5)
  γ = 0.8; t = tᵦ; tₚ = tₐ; tqnp = tₐ
  hₖ = 0; hkm1 = 0; gkm1 = 0; hplus = 0
  iter = 0

  gₖ = grad(h, [t])[1]
  gkm1 = grad(h, [tqnp])[1]

  verbose &&
    @printf(" iter        tₚ        t         dN         gₖ      ")
  verbose &&
    @printf("gplus        \n")
  verbose &&
    @printf(" %7.2e %7.2e %7.2e  %7.2e  %7.2e  %7.2e\n",
                iter, tₚ, t, 0.0, gₖ,  0.0)

  while ((abs(gₖ) > tol) && (iter < maxiter)) || (iter == 0)
    s = t - tqnp; y = gₖ - gkm1
    dN = -gₖ * s / y

    if ((tₚ - t) * dN > 0.0) && (dN / (tₚ - t) < γ)
      tplus = t + dN
      #hplus = obj(h, tplus)
      gplus = grad(h, [tplus])[1]
      verbose && print_with_color(:green, "N")
    else
      tplus = (t + tₚ) / 2
      #hplus = obj(h,  tplus)
      gplus = grad(h, [tplus])[1]
      verbose && print_with_color(:green, "B")
    end

    if t > tₚ
      if gplus < 0.0
        tₚ = t; tqnp = t; t = tplus
      else
        tqnp = t; t = tplus
      end
    else
      if gplus > 0.0
        tₚ = t; tqnp = t; t = tplus
      else
        tqnp = t; t = tplus
      end
    end

    #mise à jour des valeurs
    hkm1 = hₖ; gkm1 = gₖ; hₖ = hplus; gₖ = gplus
    iter = iter + 1
    verbose &&
        @printf(" %7.2e %7.2e %7.2e  %7.2e  %7.2e  %7.2e\n",
                iter, tₚ, t, dN, gₖ, gplus)
  end
    topt = t
    status = :NotSolved
    (abs(gₖ) < tol) && (status = :Optimal)
    (iter >= maxiter) && (status = :Tired)
    tired = iter > maxiter
    optimal = abs(gₖ) < tol
    return (topt, obj(h, [topt])[1], norm(gₖ, Inf), iter, optimal, tired, status, h.counters.neval_obj, h.counters.neval_grad, h.counters.neval_hess)
end


text:
include("PRM.jl")
pyplot()

################
# Setup

connectRadius = 5

startstate = Point(1.,1)
goalstate = Point(20.,20)

# create obstacles (fixed map for now)
obs1 = HyperRectangle(Vec(8,3.), Vec(2,2.)) #Todo
obs2 = HyperRectangle(Vec(4,4.), Vec(2,10.)) #Todo

obstacles = Vector{HyperRectangle}()
push!(obstacles, obs1, obs2)

# create room
walls = Vector{LineSegment}()
w,h  = 20,20
perimeter = HyperRectangle(Vec(0.,0), Vec(w,h))
roomPerimeter = algfxn.decompRect(perimeter)
for l in roomPerimeter
    push!(walls, l)
end

r = algT.Room(w,h,walls,obstacles)

###################
# Run iterations

nTrials = 30 #average multiple runs for the same numSamples
cost = 0
listCosts = Vector{Float32}()
listpSucc= Vector{Float32}()
nSamples_list = [10 20 30 80 150 200 300]

tic()
for nSamples in nSamples_list
    totalcost = 0.
    idx = 0.
    nSuccess = 0.
    pathcost = 0.
    param = algT.AlgParameters(nSamples, connectRadius)
    while idx < nTrials
        nodeslist, edgeslist = preprocessPRM(r, param)
        pathcost, isPathFound, solPath = queryPRM(startstate, goalstate, nodeslist, edgeslist, obstacles)
        if isPathFound
            nSuccess += 1
        end
        if !(pathcost == Void)
            totalcost += pathcost
        end
        idx += 1
    end

    avgCost = totalcost /  nSuccess
    pSucc = nSuccess / nTrials
    #@printf("For the iter of %d the avg cost was %d across %d trials", maxIter, avgCost, nTrials)
    push!(listCosts, avgCost)
    push!(listpSucc, pSucc)
end

timeExperiment = toc();
print("Time to: $blah\n")
timestamp = Base.Dates.now()

@show listCosts
@show listpSucc
sizeplot = (800, 800)

# sanity check
plot()

supTitle="\nPRM with maxDist=$connectRadius, nTrials=$nTrials.
        (time to run:$timeExperiment, timestamp=$timestamp)\n\n"
costTitle= "numsamples vs pathcost\n"

pPRMcost = scatter(nSamples_list, listCosts',
    color=:black,
    title = supTitle * costTitle, ylabel = "euclidean path cost", xlabel = "numsamples",
    yaxis=((0,100), 0:20:100))

successTitle = "\nnumsamples vs pSuccess\n"
pPRMsuccess = scatter(nSamples_list, listpSucc',
    color = :orange, markersize= 6,
    title = successTitle, ylabel = ("P(success)=numSucc/$nTrials trials"), xlabel = "numsamples",
    yaxis=((0, 1.2), 0:0.1:1))

plot(pPRMcost, pPRMsuccess, layout=(2,1), legend=false,
    xaxis=((0, 320), 0:50:300),
    size = sizeplot, sizeplot = (800, 800))

# todo: save one "representative" figure from each run , to plot
# todo: truncate timestamp and elapsed time, so that it doesn't give silly amounts of precision


text:
import Base: +, -, *, /, ==, inv
import Base: size
import Base: print, show

abstract type AbstractFOTransferFunction end

"""
FOTF object
"""
mutable struct FOTF <: AbstractFOTransferFunction
    num
    nn
    den
    nd
    ioDelay
end

"""
    fotf(numerator, order_of_numerator, denumerator, order_of_denumerator, delay)

Construct the transfer function in fractional order systems.

### Example

```julia-repl
julia> G = fotf([1, 2], [0.3, 0.4], [1, 2], [0.5, 0.6], 2)
FOTF

s^{0.3} + 2s^{0.4}
------------------
s^{0.5} + 2s^{0.6}
```
"""
function fotf(a, na, b, nb, T)
    a = a
    na = na
    b = b
    nb = nb
    return FOTF(a, na, b, nb, T)
end

fotf(a, na, b, nb) = FOTF(a, na, b, nb, 0)

fotf(s::String) = s == "s" ? (return fotf([1], [1], [1], [0], 0)) : (nothing)
fotf(c::Char) = c == 's' ? (return fotf([1], [1], [1], [0], 0)) : (nothing)
fotf(x::Number) = x

#FIXME: size()
function size(tf::FOTF)
    return size(tf.num)
end

function simplify(G::FOTF)
    (b, nb) = polyuniq(G.num, G.nn, eps())
    (a, na) = polyuniq(G.den, G.nd, eps())

    if length(a) == length(b)
        da = a[1]
        db = b[1]
        if a./da == b./db && na == nb
            a = 1
            b = db/da
            na = 0
            nb = 0
        end
    end

    if length(nb) == 0
        nb = 0
        b = 0
        na = 0
        a = 0
    end

    nn = min(na[end], nb[end])
    nb = nb.-nn
    na = na.-nn
    return fotf( b, nb, a, na, G.ioDelay)
end

function polyuniq(a, an, tol)
    ii = sortperm(an, rev=true)
    an = sort(an, rev=true)

    a = a[ii]
    ax = diff(an)
    key = 1
    for i = 1:length(ax)
        if abs(ax[i]) <= tol
            a[key] = a[key] + a[key+1]
            deleteat!(a, key+1)
            deleteat!(an, key+1)
        else
            key = key+1
        end
    end
    ii = findall(x -> abs(x) > tol, a)
    a = a[ii]
    an = an[ii]
    # Here  a and an become one-element vector
    return a, an
end

"""
    kronsum(A, B)

Return the Kronecker sum of two matrices
"""
function kronsum(A::T, B::T) where {T <: AbstractArray}
    ma, na = size(A, 2), size(A, 1)
    mb, nb = size(B, 2), size(B, 1)
    A = reshape(A, 1, ma, 1, na)
    B= reshape(B, mb, 1, nb, 1)
    C = reshape(broadcast(+, A, B), (mb*mb, na*nb))
    return C
end

kronsum(A, B) = A .+ B

"""
    fotf2cotf(tf)

Convert an FOTF object to a commensurate order object.

### Example

```julia-repl
julia> fotf2cotf(G)
TransferFunction{Discrete{Int64}, ControlSystems.SisoRational{Float64}}
2.0z^4 + 1.0z^3
---------------
2.0z^6 + 1.0z^5

Sample Time: 2 (seconds)
Discrete-time transfer function model
```
"""
function fotf2cotf(G::FOTF)
    α = base_order(G)

    if α == 0
        a = G.den
        b = G.num
        d = b[1]/a[1]
        return tf(d), α
    else
        n0 = Int64.(round.(G.nd/α))
        a = zeros(maximum(n0)+1)
        a[n0.+1] = G.den
        a = a[end:-1:1]

        m0 = Int64.(round.(G.nn/α))
        b = zeros(maximum(m0)+1)
        b[m0.+1] = G.num
        b = b[end:-1:1]
        return G.ioDelay == 0 ? tf(b, a) : tf(b, a, G.ioDelay)
    end
end

"""
    base_order(G)

Find the base order of an FOTF.

### Example

```julia-repl
julia> G = fotf([1, 2, 3], [0.1, 0.2, 0.34], [4, 5, 6], [0.65, 0.67, 0.78]);
julia> base_order(G)
0.01
```
"""
function base_order(G::FOTF)
    a=[]
    a=[a; G.nn; G.nd]

    # Julia doesn't has `rationalize` method for Integer😟
    nume = Int64[]
    denume = Int64[]
    for (_, i) in enumerate(a)
        isa(i, Int64) ? continue : i = rationalize(i)
        push!(nume, numerator(i))
        push!(denume, denominator(i))
    end
    return gcd(nume)/lcm(denume)
end

"""
    fotfdata(G)

Get the data from an FOTF object

### Example

```julia-repl
julia> G = fotf([1, 2], [0.3, 0.4], [1, 2], [0.5, 0.6], 2)
julia> fotfdata(G)
5-element Vector{Any}:
  [1, 2]
  [0.5, 0.6]
  [1, 2]
  [0.3, 0.4]
 2
```
"""
function fotfdata(G::FOTF)
    b = G.num
    a = G.den
    nb = G.nn
    na = G.nd
    L = G.ioDelay

    if typeof(a) <: Number && typeof(b) <: Number
        a=[a]
        b=[b]
    end

    if typeof(na) <: Number && typeof(nb) <: Number
        na=[na]
        nb=[nb]
    end

    return a, na, b, nb, L
end

function Base.iszero(G::FOTF)
    (a, na, b, nb) = fotfdata(G)

    length(nb) == 1 & abs(b[1]) < eps() ? true : false
end

function isstable(G::FOTF)
    a0=0.001
    g=G
    a=g.nd
    a1=floor.(a./a0)
    g1= fotf2cotf(G)
    α = base_order(G)
    c = denvec(g1)[1]
    p0 = roots(c)# Need to verify
    kk=[]

    for k=1:length(p0)
        a=g.den
        na=g.nd
        pa=p0[k]^(1/α)
        if norm(a*(pa.^na)) < 1e-6
            push!(kk,k)
        end
    end
    p=p0[kk]

end

"""
Addition operation of two fractional order model.

!!! warning
    You can't add two systems with different I/O delay.

"""
function +(G1::FOTF, G2::FOTF)
    if iszero(G2)
        return G1
    elseif G1.ioDelay == G2.ioDelay
        key = 0
        G1 = simplify(G1)
        G2 = simplify(G2)
        (a1, na1, b1, nb1) = fotfdata(G1)
        (a2, na2, b2, nb2) = fotfdata(G2)
        if length(a1) == length(a2) #eps()
            if a1 == a2 && na1 == na2
                a=a1
                na=na1
                b=[b1 b2]
                nb=[nb1 nb2]
                key=1
            end
        end

        if key == 0
            a = kron(a1, a2)
            na = kronsum(na1, na2)
            b = [kron(a1, b2); kron(b1, a2)]
            nb = [kronsum(na1, nb2); kronsum(nb1, na2)]
        end# Here the a, na, b, nb should be vector, instead of matrix here.
        a=a[:]
        na=na[:]
        b=b[:]
        nb=nb[:]
        G = fotf(b, nb, a, na)
        return simplify(G)
    end
end

"""
Minus operation of FOTF
"""
function -(G1::FOTF, G2::FOTF)
    G = G1+(-G2)
    if typeof(G) <: Nothing
        return 0
    end
    return G
end

-(G::FOTF) = fotf(-G.num, G.nn, G.den, G.nd)

function *(G1::FOTF, G2::FOTF)
    G = sisotimes(G1, G2)
    return G
end

*(x::Number, G::FOTF) = fotf(x*G.num, G.nn, G.den, G.nd, G.ioDelay)

*(G::FOTF, x::Number) = *(x::Number, G::FOTF)

function sisotimes(G1::FOTF, G2::FOTF)
    (a1, na1, b1, nb1) = fotfdata(G1)
    (a2, na2, b2, nb2) = fotfdata(G2)
    a = kron(a1, a2)
    na = kronsum(na1, na2)
    b = kron(b1, b2)
    nb = kronsum(nb1, nb2)
    a=a[:]
    na=na[:]
    b=b[:]
    nb=nb[:]
    G = simplify(fotf(b, nb, a, na, G1.ioDelay+G2.ioDelay))
    return G
end

# Equality judgement
"""
Test two FOTF are  equal or not.
"""
function ==(G1::FOTF, G2::FOTF)
    key = 0
    G = G1-G2
    if typeof(G) == Nothing
        return true
    end
    b = G.num
    key = key + (length(b)==1 && b[1]==0)
    key == 0 ? false : true
end

"""
The inversion of an FOTF object
"""
function inv(G::FOTF)
    (a, na, b, nb, L) = fotfdata(G)
    L > 0 ? error("Delay term is not allowed in inversion operation.") : nothing

    G1 = fotf(a, na, b, nb, L)
    dd = b[1]
    G1.num = a/dd
    G1.den = b/dd
    G1.nd = nb
    G1.nn = na

    return simplify(G1)
end

"""
The right division of two FOTF objects
"""
/(G1::FOTF, G2::FOTF) = G1.ioDelay == 0 && G2.ioDelay == 0 ? G1*inv(G2) : nothing
/(x::Number, G2::FOTF) = x*inv(G2)
/(G::FOTF, x::Number) = G*inv(x)

"""
The left division of two FOTF objects
"""
\(G1::FOTF, G2::FOTF) = G1.ioDelay == 0 && G2.ioDelay == 0 ? inv(G1)*G2 : nothing