using DynamicHMCModels

ProjDir = @__DIR__

df = DataFrame(CSV.read(joinpath(ProjDir, "..", "..", "data", "chimpanzees.csv"), delim=';'))

Base.@kwdef struct Chimpanzees_02
    "Number of actors"
    N_actors::Int
    pulled_left::Vector{Int}
    prosoc_left::Vector{Int}
    condition::Vector{Int}
    actor::Vector{Int}
end

function make_transformation(model::Chimpanzees_02)
    as((a = as(Vector, model.N_actors), bp = asℝ, bpC = asℝ))
end

model = Chimpanzees_02(; N_actors = maximum(df.actor), pulled_left = df.pulled_left,
                    prosoc_left = df.prosoc_left, condition = df.condition,
                    actor = df.actor)

function (model::Chimpanzees_02)(θ)
    @unpack pulled_left, prosoc_left, condition, actor = model
    @unpack a, bp, bpC = θ
    ℓ_likelihood = mapreduce(+, actor, condition, prosoc_left,
       pulled_left) do actor, condition, prosoc_left, pulled_left
           p = logistic(a[actor] + (bp + bpC * condition) * prosoc_left)
           logpdf(Bernoulli(p), pulled_left)
       end
    P = Normal(0, 10)
    ℓ_prior = logpdf(P, bpC) + logpdf(P, bp) + sum(a -> logpdf(P, a), a)
    ℓ_prior + ℓ_likelihood
end

P = TransformedLogDensity(make_transformation(model), model)
∇P = ADgradient(:ForwardDiff, P)
results = mcmc_with_warmup(Random.GLOBAL_RNG, ∇P, 1000)
posterior = P.transformation.(results.chain)

DynamicHMC.Diagnostics.EBFMI(results.tree_statistics)

DynamicHMC.Diagnostics.summarize_tree_statistics(results.tree_statistics)

a3d = Array{Float64, 3}(undef, 1000, 9, 1);
for j in 1:1
  for i in 1:1000
    a3d[i, 1, j] = values(posterior[i].bp)
    a3d[i, 2, j] = values(posterior[i].bpC)
    a3d[i, 3:9, j] = values(posterior[i].a)
  end
end

# Create MCMCChains object

parameter_names = ["bp", "bpC"]
pooled_parameter_names = ["a[$i]" for i in 1:7]
sections =   Dict(
  :parameters => parameter_names,
  :pooled => pooled_parameter_names
)
cnames = vcat(parameter_names, pooled_parameter_names)
chns = create_mcmcchains(a3d, cnames, sections, start=1)

# Result rethinking

rethinking = "
Iterations = 1:1000
Thinning interval = 1
Chains = 1,2,3,4
Samples per chain = 1000

Empirical Posterior Estimates:
        Mean        SD       Naive SE       MCSE      ESS
a.1 -0.74503184 0.26613979 0.0042080396 0.0060183398 1000
a.2 10.77955494 5.32538998 0.0842018089 0.1269148045 1000
a.3 -1.04982353 0.28535997 0.0045119373 0.0049074219 1000
a.4 -1.04898135 0.28129307 0.0044476339 0.0056325117 1000
a.5 -0.74390933 0.26949936 0.0042611590 0.0052178124 1000
a.6  0.21599365 0.26307574 0.0041595927 0.0045153523 1000
a.7  1.81090866 0.39318577 0.0062168129 0.0071483527 1000
 bp  0.83979926 0.26284676 0.0041559722 0.0059795826 1000
bpC -0.12913322 0.29935741 0.0047332562 0.0049519863 1000
";

# Describe draws

describe(chns) |> display

# Describe pooled draws

describe(chns, sections=[:pooled])


# brent-tests.jl -
#
# Tests for Brent's algorithms.
#

module BrentTests

using Test, Printf
using OptimPackNextGen.Brent

# Counter of function evaluations.
const cnt = Ref{Int}(0)

# Yields `y` converted to a floating-point type according to the type of `x`.
fp(x::Number, y::Real) = convert(float(typeof(x)), y)

fzero_test0(x::Real) = (cnt[] += 1; 1/(x - 3) - 6)
params(::typeof(fzero_test0)) = (3, 4, 19/6)

# Test functions for root finding from Vakkalagadda Satya Sai Prakash,
# "Implementation of Brent-Dekker and A Better Root Finding Method and
# Brent-Dekker Method's Parallelization".

fzero_test1(x::Real) = (cnt[] += 1; exp(x)/2 - 5x + 2)
params(::typeof(fzero_test1)) = (1, 4, NaN)

fzero_test2(x::Real) = (cnt[] += 1;
                                 -2x^4 + 2x^3 - 16x^2 - 60x + 100)
params(::typeof(fzero_test2)) = (0, 4, NaN)

fzero_test3(x::Real) = (cnt[] += 1; exp(x)*cos(x) - x*sin(x))
params(::typeof(fzero_test3)) = (0, 3, NaN)

fzero_test4(x::Real) = (cnt[] += 1; x^5 - 5x + 3)
params(::typeof(fzero_test4)) = (0, 1, NaN)

fzero_test5(x::Real) = (cnt[] += 1;
                        (x^3 - fp(x,0.926)*x^2 +
                         fp(x,0.0371)*x + fp(x,0.043)))
params(::typeof(fzero_test5)) = (0, 0.8, NaN)

fzero_test6(x::Real) = (cnt[] += 1;
                        (-9 + sqrt(99 + 2x - x^2) + cos(2x)))
params(::typeof(fzero_test6)) = (-2, 6, NaN)

fzero_test7(x::Real) = (cnt[] += 1; sin(cosh(x)))
params(::typeof(fzero_test7)) = (0, 2, NaN)

fzero_test8(x::Real) = (cnt[] += 1; exp(-exp(-x)) - x)
params(::typeof(fzero_test8)) = (0, 1, NaN)

# Test functions from the GNU Scientific Library (GSL).

fzero_gsl1(x::Real) = (cnt[] += 1; x^20 - 1)
params(::typeof(fzero_gsl1)) = (0.1, 2.0, 1.0)

fzero_gsl2(x::Real) = (cnt[] += 1; sqrt(abs(x))*sign(x))
params(::typeof(fzero_gsl2)) = (-1.0/3.0, 1.0, 0.0)

fzero_gsl3(x::Real) = (cnt[] += 1; x^2 - fp(x, 1e-8))
params(::typeof(fzero_gsl3)) = (0, 1, sqrt(1e-8))

fzero_gsl4(x::Real) = (cnt[] += 1; x*exp(-x))
params(::typeof(fzero_gsl4)) = (-1.0/3.0, 2.0, 0.0)

fzero_gsl5(x::Real) = (cnt[] += 1; (x - 1)^7)
params(::typeof(fzero_gsl5)) = (0, 3, 1) # (0.9995, 1.0002, 1)

# Problem 4 from AMPGO, xm = 2.868034, fm = -3.85045:
problem04(x) = (cnt[] += 1; -(16x^2 - 24x + 5)*exp(-x))
bounds(::typeof(problem04)) = [1.9, 3.9]
solution(::typeof(problem04)) = [2.868034, -3.85045]

# Problem 13 from AMPGO, xm = 1/sqrt(2), fm = -1.5874:
problem13(x) = (cnt[] += 1; -x^(2/3) - (1 - x^2)^(1/3))
bounds(::typeof(problem13)) = [0.001, 0.99]
solution(::typeof(problem13)) = [1/sqrt(2), -1.5874]

# Problem 18 from AMPGO, xm = 2, fm = 0:
problem18(x) = (cnt[] += 1; x ≤ 3 ? (x - 2)^2 : 2log(x - 2) + 1)
bounds(::typeof(problem18)) = [0, 6]
solution(::typeof(problem18)) = [2, 0]

verb = true
@testset "Brent fzero" begin
    let f = fzero_test0
        (a, b, x0) = params(f)
        @test_throws ArgumentError fzero(f, a, b; rtol=0)
        @test_throws ArgumentError fzero(f, a, b; rtol=1)
        @test_throws ArgumentError fzero(f, a, b; atol=NaN)
    end
    k = 0
    for f in (fzero_test0,
              fzero_test1, fzero_test2, fzero_test3, fzero_test4,
              fzero_test5, fzero_test6, fzero_test7, fzero_test8,
              fzero_gsl1, fzero_gsl2, fzero_gsl3, fzero_gsl4, fzero_gsl5)
        for T in (Float32, Float64)
            (a, b, x0) = params(f)
            rtol = Brent.fzero_rtol(T)
            atol = (x0 == 0 ? eps(T) : Brent.fzero_atol(T))
            prec = sqrt(eps(T))
            cnt[] = 0
            (x, fx) = if T === Float64
                fzero(f, a, b; rtol = rtol, atol = atol)
            else
                fzero(T, f, a, b; rtol = rtol, atol = atol)
            end
            if verb
                if k ≤ 8
                    name = "fzero_test$k"
                else
                    name = "fzero_gsl$(k - 8)"
                end
                @printf("%-12s T = %s, n = %3d, x = % .15f, f(x) = % .15f\n",
                        name, repr(T), cnt[], x, fx)
            end
            @test abs(fx) ≈ zero(T) rtol=zero(T) atol=prec
            if !isnan(x0)
                @test x ≈ T(x0) rtol=rtol atol=atol
            end
            fa = f(a)
            (x1, f1) =  if T === Float64
                fzero(f, a, fa, b; rtol = rtol, atol = atol)
            else
                fzero(T, f, a, fa, b; rtol = rtol, atol = atol)
            end
            @test x1 == x
            fb = f(b)
            (x2, f2) =  if T === Float64
                fzero(f, a, fa, b, fb; rtol = rtol, atol = atol)
            else
                fzero(T, f, a, fa, b, fb; rtol = rtol, atol = atol)
            end
            @test x2 == x
        end
        k += 1
    end
end

@testset "Brent fmin" begin
    for (f, id) in ((problem04,  4),
                    (problem13, 13),
                    (problem18, 18),)
        for (T, prec) in ((Float16, 3e-2),
                          (Float32, 3e-6),
                          (Float64, 6e-15))
            atol = Brent.fmin_atol(T)
            rtol = Brent.fmin_rtol(T)

            cnt[] = 0
            (xm, fm, lo, hi) = if T === Float64
                fmin(f, bounds(f)...)
            else
                fmin(T, f, bounds(f)...)
            end
            (xm_true, fm_true) = solution(f)
            if verb
                println("Problem $id (T=$T):")
                @printf(" ├─ nevals = %d\n", cnt[])
                @printf(" ├─ x_approx = %.15f ± %.3e\n", xm, (hi - lo)/2)
                @printf(" ├─ x_true   = %.15f\n", xm_true)
                @printf(" ├─ x_error  = %.3e\n", abs(xm - xm_true))
                @printf(" ├─ f(x_approx) = %.15f\n", fm)
                @printf(" ├─ f(x_true)   = %.15f\n", fm_true)
                @printf(" └─ f_error     = %.3e\n", abs(fm - fm_true))
            end
            @test abs(xm - xm_true) ≤ atol + rtol*abs(xm_true)
            #@test abs(fm - solution(f)[2]) ≤ prec
        end
        verb && println()
    end
end

end # module


"""
     IndexedMarkovChain{T}

Representation of a Markov Chain type with "arbitrary" support (`T`). Assumes
that **rows** add up to one.

"""
abstract type IndexedMarkovChain{T} <: AbstractArray{T, 1} end
Base.getindex(mc::IndexedMarkovChain, i)=supp(mc)[i]
Base.length(mc::IndexedMarkovChain)=length(supp(mc))
Base.size(mc::IndexedMarkovChain)=(length(mc), )
Base.IndexStyle(::Type{<:IndexedMarkovChain}) = IndexLinear()
supp(mc::IndexedMarkovChain, i)=error("Method `supp` must be implemented for $(typeof(MC)) ")
transition(mc::IndexedMarkovChain, i, j)=transition(mc)[i,j]

" Unconditional mean of a MarkovChain "
Statistics.mean(mc::IndexedMarkovChain)=dot(markov_invariant(mc), supp(mc))
Statistics.mean(mc::IndexedMarkovChain, v::AbstractVector)=dot(markov_invariant(mc), v)

"Markov Chain with pre-stored cumulative transition matrix"
struct MarkovChain{T, S} <: IndexedMarkovChain{T}
    support::Vector{T}
    transition::Matrix{S}
    transition_cumu::Matrix{S}
end
supp(mc::MarkovChain)=mc.support
transition(MC::MarkovChain)=MC.transition # Make rows add up to one
transition_cumu(MC::MarkovChain)=MC.transition_cumu
Base.eltype(mc::MarkovChain{T, S}) where {T, S}=T

" Simples constructor for any iterable `v`. Note: no sanity check for dimensions."
MarkovChain(v, M::Matrix)=MarkovChain([v[i] for i in 1:length(v)], M,
                                      cumsum(M, dims=2))

" Same as `MarkovChain`, but pre-store invariant distribution "
struct MarkovChainExtended{T, S} <: IndexedMarkovChain{T}
    mc::MarkovChain{T, S}
    invariant::Array{Float64, 1}
end

MarkovChainExtended(mc::MarkovChain{T, S}) where {T,S}=
    MarkovChainExtended{T, S}(mc, markov_invariant(mc))
MarkovChainExtended(v, M::Matrix, tol=1e-12)=
    MarkovChainExtended(MarkovChain(v, M))

supp(mc::MarkovChainExtended)=supp(mc.mc)
transition(mc::MarkovChainExtended)=transition(mc.mc)
transition_cumu(mc::MarkovChainExtended)=transition_cumu(mc.mc)

# Note: will also work with `Aggregate` (& anything for which `transition` is implemented)
function expect_markov(v::AbstractArray, current_state, MC)
    e = 0.

    for next_state in 1:length(MC)
        e += v[next_state] * transition(MC, current_state, next_state)
    end

    e
end

" TODO: unrolled version for arbitrary tuple of arrays "
function expect_markov(v1::AbstractArray, v2::AbstractArray, current_state, MC)
    e1 = e2 = 0.

    for next_state in 1:length(MC)
        e1 += v1[next_state] * transition(MC, current_state, next_state)
        e2 += v2[next_state] * transition(MC, current_state, next_state)
    end

    e1, e2
end

markov_invariant(mc::MarkovChainExtended)=mc.invariant
markov_invariant(mc::MarkovChain, tol=1e-12)=markov_invariant(mc.transition,tol)

" Returns eigenvector associated with A's last eigenvalue "
markov_invariant(A, tol)=begin
    ed = eigen(transpose(A))

    !within_tolerance(ed.values[end], 1, tol) &&
        error("No invariant distribution found")

    evec = Float64.(ed.vectors[:, end])

    makeprob!(evec)

    return evec
end

makeprob!(u)=(u ./= sum(u))

"""

Simulate a Markov Chain  for `len` periods, starting with index `i0`

The `rand_tmp` vector should contain `len` draws from a uniform distribution.

Overwrites vector `path_prealloc`.

"""
function markov_genpath!(path_prealloc, mc, i0, len, rand_tmp)
    for t in 1:len
        i0 = draw_next(mc, i0, rand_tmp[t])
        path_prealloc[t] = i0
    end
end

function markov_genpath_seeded(mc, i0, len, seed=1234)
    rand_tmp=rand(len)
    path_prealloc = fill(0, len)
    markov_genpath!(path_prealloc, mc, i0, len, rand_tmp)
    return path_prealloc
end

function markov_genpath(mc, i0, len)
    rand_tmp=rand(len)
    path_prealloc = fill(0, len)
    markov_genpath!(path_prealloc, mc, i0, len, rand_tmp)
    return path_prealloc
end

" Guaranteed to not be `nothing`; use with EXTREME caution "
function findfirstg(bvec)
    idx = 1

    for i in 1:length(bvec)
        if bvec[i]
            idx = i
            break
        end
    end

    return idx
end

" Guaranteed to not be `nothing`; use with EXTREME caution "
function findfirstpositive(nvec)
    j = 0
    while (nvec[j+1] < 0) && (j+1 < length(nvec))
        j += 1
    end
    j+1
end

draw_next(mc::IndexedMarkovChain, i0, r)=begin
    @views findfirstpositive(transition_cumu(mc)[i0, :] .- r)
end

draw_next(mc::IndexedMarkovChain, i0)=draw_next(mc, i0, rand())

indeptensor(mc0::IndexedMarkovChain, mc1::IndexedMarkovChain)=MarkovChainExtended(vec(collect(Base.product(mc0, mc1))),
                                                                                Base.kron(transition(mc1), transition(mc0)))

function makemean1!(mc::IndexedMarkovChain{T}) where {T <: Number}
    mm = Statistics.mean(mc)
    foreach(1:length(mc)) do i
        supp(mc)[i] /= mm
    end
end


mutable struct AdjacencyMatrix <: AbstractGraph
    cntedge::Int
    cntvertex::Int
    arr::AbstractMatrix # Should eventually be of type Vertex
    mark::Vector # 0...unvisited ; 1...visited
    function AdjacencyMatrix(n::Int)
        cntedge = 0
        cntvertex = n
        arr = zeros(Int, cntvertex, cntvertex)
        mark = zeros(Int, cntvertex)
        new(cntedge, cntvertex, arr, mark)
    end
end

countedges(g::AdjacencyMatrix) = g.cntedge
countvertices(g::AdjacencyMatrix) = g.cntvertex

function first(g::AdjacencyMatrix, v::Vertex)
    for i = 1:g.cntvertex
        (g.arr[v, i] != 0) && return i
    end
    return g.cntvertex + 1
end

function next(g::AdjacencyMatrix, v::Vertex, w::Vertex)
    for i = w+1:g.cntvertex
        (g.arr[v, i] != 0) && return i
    end
    return g.cntvertex + 1
end

function setedge!(g::AdjacencyMatrix, v::Vertex, w::Vertex, weight)
    (weight > 0) || throw(ArgumentError("Illegal weight value."))
    (g.arr[v, w] == 0) && (g.cntedge += 1)
    g.arr[v, w] = weight
end

function remove!(g::AdjacencyMatrix, v::Vertex, w::Vertex)
    (g.arr[v, w] != 0) && (g.cntedge -= 1)
    g.arr[v, w] = 0
end

isedge(g::AdjacencyMatrix, v::Vertex, w::Vertex) = g.arr[v, w] != 0
getweight(g::AdjacencyMatrix, v::Vertex, w::Vertex) = g.arr[v, w]
getmark(g::AdjacencyMatrix, v::Vertex) = g.mark[v]
setmark!(g::AdjacencyMatrix, v::Vertex, value) = (g.mark[v] = value)

==(a::AdjacencyMatrix, b::AdjacencyMatrix) = a.arr == b.arr


function _prepare(nodes::Matrix{T},
                  kernel::RK
                 ) where {T <: AbstractFloat, RK <: ReproducingKernel_0}
    n = size(nodes, 1)
    n_1 = size(nodes, 2)
    min_bound = Vector{T}(undef, n)
    compression::T = 0
    @inbounds for i = 1:n
        min_bound[i] = nodes[i,1]
        maxx::T = nodes[i,1]
        for j = 2:n_1
            min_bound[i] = min(min_bound[i], nodes[i,j])
            maxx = max(maxx, nodes[i,j])
        end
        compression = max(compression, maxx - min_bound[i])
    end
    if compression <= eps(T(1.0))
        error("Cannot prepare the spline: `nodes` data are not correct.")
    end

    t_nodes = similar(nodes)
    @inbounds for j = 1:n_1
        for i = 1:n
            t_nodes[i,j] = (nodes[i,j] - min_bound[i]) / compression
        end
    end

    if T(kernel.ε) == T(0.0)
        ε = _estimate_ε(t_nodes)
        if isa(kernel, RK_H0)
            kernel = RK_H0(ε)
        elseif isa(kernel, RK_H1)
            ε *= T(1.5)
            kernel = RK_H1(ε)
        elseif isa(kernel, RK_H2)
            ε *= T(2.0)
            kernel = RK_H2(ε)
        else
            error("incorrect `kernel` type.")
        end
    end

    gram = _gram(t_nodes, kernel)
    chol = nothing
    try
        chol = cholesky(gram)
    catch
        error("Cannot prepare the spline: Gram matrix is degenerate.")
    end

    cond = _estimate_cond(gram, chol)

    spline = NormalSpline(kernel,
                          compression,
                          t_nodes,
                          nothing,
                          nothing,
                          nothing,
                          nothing,
                          min_bound,
                          gram,
                          chol,
                          nothing,
                          cond
                         )
    return spline
end

function _construct(spline::NormalSpline{T, RK},
                    values::Vector{T},
                    cleanup::Bool = false
                   ) where {T <: AbstractFloat, RK <: ReproducingKernel_0}
    if(length(values) != size(spline._nodes, 2))
        error("Number of data values does not correspond to the number of nodes.")
    end
    if isnothing(spline._chol)
        error("Gram matrix was not factorized.")
    end

    mu = Vector{T}(undef, size(spline._gram, 1))
    ldiv!(mu, spline._chol, values)

    spline = NormalSpline(spline._kernel,
                          spline._compression,
                          spline._nodes,
                          values,
                          nothing,
                          nothing,
                          nothing,
                          spline._min_bound,
                          cleanup ? nothing : spline._gram,
                          cleanup ? nothing : spline._chol,
                          mu,
                          spline._cond
                         )
    return spline
end

###################

function _prepare(nodes::Matrix{T},
                  d_nodes::Matrix{T},
                  es::Matrix{T},
                  kernel::RK
                 ) where {T <: AbstractFloat, RK <: ReproducingKernel_1}
    n = size(nodes, 1)
    n_1 = size(nodes, 2)
    n_2 = size(d_nodes, 2)

    if(size(es, 2) != n_2)
        error("Number of derivative directions does not correspond to the number of derivative nodes.")
    end

    t_es = similar(es)
    try
        @inbounds for i = 1:n_2
            t_es[:,i] = es[:,i] ./ norm(es[:,i])
        end
    catch
        error("Cannot normalize derivative direction: zero direction vector.")
    end

    min_bound = Vector{T}(undef, n)
    compression::T = 0
    @inbounds for i = 1:n
        min_bound[i] = nodes[i,1]
        maxx::T = nodes[i,1]
        for j = 2:n_1
            min_bound[i] = min(min_bound[i], nodes[i,j])
            maxx = max(maxx, nodes[i,j])
        end
        for j = 1:n_2
            min_bound[i] = min(min_bound[i], d_nodes[i,j])
            maxx = max(maxx, d_nodes[i,j])
        end
        compression = max(compression, maxx - min_bound[i])
    end

    if compression <= eps(T(1.0))
        error("Cannot prepare the spline: `nodes` data are not correct.")
    end

    t_nodes = similar(nodes)
    t_d_nodes = similar(d_nodes)

    @inbounds for j = 1:n_1
        for i = 1:n
            t_nodes[i,j] = (nodes[i,j] - min_bound[i]) / compression
        end
    end

    @inbounds for j = 1:n_2
        for i = 1:n
            t_d_nodes[i,j] = (d_nodes[i,j] - min_bound[i]) / compression
        end
    end

    if T(kernel.ε) == T(0.0)
        ε = _estimate_ε(t_nodes, t_d_nodes)
        if isa(kernel, RK_H1)
            ε *= T(2.0)
            kernel = RK_H1(ε)
        elseif isa(kernel, RK_H2)
            ε *= T(2.5)
            kernel = RK_H2(ε)
        else
            error("incorrect `kernel` type.")
        end
    end

    gram = _gram(t_nodes, t_d_nodes, t_es, kernel)
    chol = nothing
    try
        chol = cholesky(gram)
    catch
        error("Cannot prepare the spline: Gram matrix is degenerate.")
    end

    cond = _estimate_cond(gram, chol)

    spline = NormalSpline(kernel,
                          compression,
                          t_nodes,
                          nothing,
                          t_d_nodes,
                          t_es,
                          nothing,
                          min_bound,
                          gram,
                          chol,
                          nothing,
                          cond
                         )
    return spline
end

function _construct(spline::NormalSpline{T, RK},
                    values::Vector{T},
                    d_values::Vector{T},
                    cleanup::Bool = false
                   ) where {T <: AbstractFloat, RK <: ReproducingKernel_0}
    if(length(values) != size(spline._nodes, 2))
        error("Number of data values does not correspond to the number of nodes.")
    end
    if(length(d_values) != size(spline._d_nodes, 2))
        error("Number of derivative values does not correspond to the number of derivative nodes.")
    end
    if isnothing(spline._chol)
        error("Gram matrix was not factorized.")
    end

    mu = Vector{T}(undef, size(spline._gram, 1))
    ldiv!(mu, spline._chol, [values; spline._compression .* d_values])

    spline = NormalSpline(spline._kernel,
                          spline._compression,
                          spline._nodes,
                          values,
                          spline._d_nodes,
                          spline._es,
                          d_values,
                          spline._min_bound,
                          cleanup ? nothing : spline._gram,
                          cleanup ? nothing : spline._chol,
                          mu,
                          spline._cond
                         )
    return spline
end

function _evaluate(spline::NormalSpline{T, RK},
                   points::Matrix{T},
                   do_parallel::Bool = false
                  ) where {T <: AbstractFloat, RK <: ReproducingKernel_0}
    if isnothing(spline)
        error("Spline was not prepared.")
    end
    if isnothing(spline._mu)
        error("Spline coefficients were not calculated.")
    end

    if size(points, 1) != size(spline._nodes, 1)
        if size(points, 1) == 1 && size(points, 2) > 1
            error("Incorrect first dimension of the `points` parameter (use 'evaluate_one' function for evaluating the spline at one point).")
        else
            error("Incorrect first dimension of the `points` parameter (the spline was built in the space of different dimension).")
        end
    end

    n = size(spline._nodes, 1)
    n_1 = size(spline._nodes, 2)
    m = size(points, 2)

    pts = similar(points)
    @inbounds for j = 1:m
        for i = 1:n
            pts[i,j] = (points[i,j] - spline._min_bound[i]) / spline._compression
        end
    end

    spline_values = Vector{T}(undef, m)
    iend1 = iend2 = iend3 = 0
    istart2 = istart3 = istart4 = 0
    mu = spline._mu[1:n_1]
    d_mu = spline._mu[(n_1 + 1):end]
    if do_parallel && m >= 1000 && Threads.nthreads() >= 4
        step = m ÷ 4
        iend1 = 1 + step
        istart2 = iend1 + 1
        iend2 = istart2 + step
        istart3 = iend2 + 1
        iend3 = (istart3 + step) < m ? (istart3 + step) : m
        istart4 = iend3 + 1
        @inbounds Threads.@threads for t = 1:4
            if t == 1
                _do_work(1, iend1, pts, spline._nodes, mu, spline._kernel, spline_values)
                if !isnothing(spline._d_nodes)
                    _do_work_d(1, iend1, pts, spline._d_nodes, spline._es, d_mu, spline._kernel, spline_values)
                end
            elseif t == 2
                _do_work(istart2, iend2, pts, spline._nodes, mu, spline._kernel, spline_values)
                if !isnothing(spline._d_nodes)
                    _do_work_d(istart2, iend2, pts, spline._d_nodes, spline._es, d_mu, spline._kernel, spline_values)
                end
            elseif t == 3
                _do_work(istart3, iend3, pts, spline._nodes, mu, spline._kernel, spline_values)
                if !isnothing(spline._d_nodes)
                    _do_work_d(istart3, iend3, pts, spline._d_nodes, spline._es, d_mu, spline._kernel, spline_values)
                end
            elseif t == 4
                if istart4 <= m
                    _do_work(istart4, m, pts, spline._nodes, mu, spline._kernel, spline_values)
                    if !isnothing(spline._d_nodes)
                        if istart4 <= m
                            _do_work_d(istart4, m, pts, spline._d_nodes, spline._es, d_mu, spline._kernel, spline_values)
                        end
                    end
                end
            end
        end
    else
        _do_work(1, m, pts, spline._nodes, mu, spline._kernel, spline_values)
        if !isnothing(spline._d_nodes)
            _do_work_d(1, m, pts, spline._d_nodes, spline._es, d_mu, spline._kernel, spline_values)
        end
    end

    return spline_values
end

function _do_work(istart::Int,
                  iend::Int,
                  points::Matrix{T},
                  nodes::Matrix{T},
                  mu::Vector{T},
                  kernel::RK,
                  spline_values::Vector{T}
                ) where {T <: AbstractFloat, RK <: ReproducingKernel_0}
    n_1 = size(nodes, 2)
    h_values = Vector{T}(undef, n_1)
    @inbounds for p = istart:iend
        for i = 1:n_1
            h_values[i] = _rk(kernel, points[:,p], nodes[:,i])
        end
        spline_values[p] = sum(mu .* h_values)
    end
end

function _do_work_d(istart::Int,
                    iend::Int,
                    pts::Matrix{T},
                    d_nodes::Matrix{T},
                    es::Matrix{T},
                    d_mu::Vector{T},
                    kernel::RK,
                    spline_values::Vector{T}
                   ) where {T <: AbstractFloat, RK <: ReproducingKernel_1}
   n_2 = size(d_nodes, 2)
   d_h_values = Vector{T}(undef, n_2)
   @inbounds for p = istart:iend
       for i = 1:n_2
           d_h_values[i] = _∂rk_∂e(kernel, pts[:,p], d_nodes[:,i], es[:,i])
       end
       spline_values[p] += sum(d_mu .* d_h_values)
    end
end

function _evaluate_gradient(spline::NormalSpline{T, RK},
                            point::Vector{T}
                           ) where {T <: AbstractFloat, RK <: ReproducingKernel_0}
    if isnothing(spline._mu)
        error("Spline coefficients were not calculated.")
    end

    n = size(spline._nodes, 1)
    n_1 = size(spline._nodes, 2)

    pt = Vector{T}(point)
    @inbounds for i = 1:n
        pt[i] = (point[i] - spline._min_bound[i]) / spline._compression
    end

    d_h_values = Vector{T}(undef, n_1)
    grad = Vector{T}(undef, n)
    mu = spline._mu[1:n_1]
    @inbounds for k = 1:n
        for i = 1:n_1
            d_h_values[i] = _∂rk_∂η_k(spline._kernel, pt, spline._nodes[:,i], k)
        end
        grad[k] = sum(mu .* d_h_values)
    end

    if !isnothing(spline._d_nodes)
        n_2 = size(spline._d_nodes, 2)
        d_h_values = Vector{T}(undef, n_2)
        d_mu = spline._mu[n_1+1:end]
        @inbounds for k = 1:n
            for i = 1:n_2
                d_h_values[i] = T(0.0)
                for l = 1:n
                    d_h_values[i] += (_∂²rk_∂η_r_∂ξ_k(spline._kernel, pt, spline._d_nodes[:,i], k, l) * spline._es[l,i])
                end
            end
            grad[k] += sum(d_mu .* d_h_values)
        end
    end
    @inbounds for k = 1:n
        grad[k] /= spline._compression
    end
    return grad
end


import DCELs._pseudoangle, DCELs._pseudoangle_unsigned, DCELs._is_sorted_ccw

@testset "Pseudoangle" begin
    @test _pseudoangle(1, 0) == 0
    @test _pseudoangle(1, 1) == 0.5
    @test _pseudoangle(0, 1) == 1
    @test _pseudoangle(-1, 1) == 1.5
    @test _pseudoangle(-1, 0) == 2
    @test _pseudoangle(-1, -1) == -1.5
    @test _pseudoangle(0, -1) == -1
    @test _pseudoangle(1, -1) == -0.5
end

@testset "Pseudoangle unsigned" begin
    @test _pseudoangle_unsigned(1, 0) == 0
    @test _pseudoangle_unsigned(1, 1) == 0.5
    @test _pseudoangle_unsigned(0, 1) == 1
    @test _pseudoangle_unsigned(-1, 1) == 1.5
    @test _pseudoangle_unsigned(-1, 0) == 2
    @test _pseudoangle_unsigned(-1, -1) == 2.5
    @test _pseudoangle_unsigned(0, -1) == 3
    @test _pseudoangle_unsigned(1, -1) == 3.5
end

@testset "Sorting" begin
    u = Vertex(1, 0)
    v = Vertex(0, 1)
    w = Vertex(-1, 0)
    center = Vertex(0, 0)

    @test _is_sorted_ccw(u, v, w, center)
    @test !_is_sorted_ccw(u, w, v, center)
end


        # isleft
        l = LineSegment2D(Point2D{T}(0, 0), Point2D{T}(1, 0))
        @test isleft(Point2D{T}(0, 1) , l)
        @test !isleft(Point2D{T}(0, -1) , l)
        @test !isleft(Point2D{T}(0, -1e-6) , l)
        @test isleft(Point2D{T}(0, 1e-6) , l)
        @test isleft(Point2D{T}(0.5, 0) , l)

        # nearest_point
        𝘅₁ = Point2D{T}(0, 0)
        𝘅₂ = Point2D{T}(2, 0)
        𝘅₃ = Point2D{T}(1, 1)
        q = QuadraticSegment2D(𝘅₁, 𝘅₂, 𝘅₃)
        p = Point2D{T}(1, 1.1)
        r, p_c = nearest_point(p, q)
        @test r ≈ 0.5
        @test 𝘅₃ ≈ p_c

        # isleft
        𝘅₁ = Point2D{T}(0, 0)
        𝘅₂ = Point2D{T}(2, 0)
        𝘅₃ = Point2D{T}(1, 1)
        q = QuadraticSegment2D(𝘅₁, 𝘅₂, 𝘅₃)
        @test !isleft(Point2D{T}(1, 0), q)
        @test isleft(Point2D{T}(1, 2), q)
        @test !isleft(Point2D{T}(1, 0.9), q)

        # in
        @test Point2D{F}(0,0) ∈ plane
        @test Point2D{F}(4,4) ∈ plane
        @test Point2D{F}(1,2) ∉ plane

        # in_halfspace
        @test in_halfspace(Point2D{F}(0,0), plane)
        @test in_halfspace(Point2D{F}(0,1), plane)
        @test !in_halfspace(Point2D{F}(0,-1), plane)

        # in
        @test Point3D{F}(1,0,2) ∈ plane
        @test Point3D{F}(2,2,2) ∈ plane
        @test Point3D{F}(1,0,0) ∉ plane

        # in_halfspace
        @test in_halfspace(Point3D{F}(0,0,2), plane)
        @test in_halfspace(Point3D{F}(0,0,3), plane)
        @test !in_halfspace(Point3D{F}(0,0,-1), plane)

        # in
        @test Point2D{T}(2, 1) ∈ aab
        @test Point2D{T}(3, 1) ∈ aab
        @test !(Point2D{T}(4, 1) ∈ aab)
        @test !(Point2D{T}(2, 5) ∈ aab)

        # in
        @test Point3D{T}(2, 1, 0) ∈ aab
        @test Point3D{T}(3, 1, 0) ∈ aab
        @test !(Point3D{T}(4, 1, 0) ∈ aab)
        @test !(Point3D{T}(2, 5, 0) ∈ aab)
        @test !(Point3D{T}(2, 1, 2) ∈ aab)

        # in
        p = Point2D{T}(1//2, 1//10)
        @test p ∈ tri
        p = Point2D{T}(1//2, -1//10)
        @test p ∉ tri
        # in
        p = Point3D{T}(0, 1//2, 1//10)
        @test p ∈ tri
        p = Point3D{T}(0, 1//2, -1//10)
        @test p ∉ tri
        p = Point3D{T}(1//100, 1//2, 1//10)
        @test p ∉ tri

        # in
        p = Point2D{T}(1//2, 1//10)
        @test p ∈  quad
        p = Point2D{T}(1//2, -1//10)
        @test p ∉ quad

                    # in
            p₁ = Point2D{T}(0, 0)
            p₂ = Point2D{T}(2, 0)
            p₃ = Point2D{T}(2, 2)
            p₄ = Point2D{T}(3//2, 1//4)
            p₅ = Point2D{T}(3, 1)
            p₆ = Point2D{T}(1, 1)
            tri6 = QuadraticTriangle(p₁, p₂, p₃, p₄, p₅, p₆)
            @test Point2D{T}(1, 1//2) ∈  tri6
            @test Point2D{T}(1, 0) ∉  tri6

        # in
        p₁ = Point2D{T}(0, 0)
        p₂ = Point2D{T}(2, 0)
        p₃ = Point2D{T}(2, 3)
        p₄ = Point2D{T}(0, 3)
        p₅ = Point2D{T}(3//2, 1//2)
        p₆ = Point2D{T}(5//2, 3//2)
        p₇ = Point2D{T}(3//2, 5//2)
        p₈ = Point2D{T}(0, 1)
        quad8 = QuadraticQuadrilateral(p₁, p₂, p₃, p₄, p₅, p₆, p₇, p₈)
        @test Point2D{T}(1, 1) ∈  quad8
        @test Point2D{T}(1, 0) ∉  quad8


@testset "Relative fractional ideals" begin
  Qx, x = FlintQQ["x"]
  f = x^2 + 12*x - 92
  K, a = number_field(f, "a")
  OK = maximal_order(K)
  Ky, y = K["y"]
  g = y^2 - 54*y - 73
  L, b = number_field(g, "b")
  OL = maximal_order(L)

  I = L(fmpq(1, 2))*OL
  @test denominator(I) == fmpz(2)
  @test Hecke.is_integral(I.den*I)

  PM = basis_pmatrix(OL)
  PM.matrix[1, 1] = K(fmpq(1, 2))
  PM.matrix[2, 1] = K()
  PM.matrix[2, 2] = K(fmpq(1, 3))
  PM = pseudo_hnf(PM, :lowerleft)
  J = fractional_ideal(OL, PM)
  @test denominator(J) == fmpz(6)
  @test Hecke.is_integral(J.den*J)

  @testset "Weird modulus" begin
    K, a = Hecke.rationals_as_number_field()
    Kt, t = K["t"]
    E, z = NumberField(t^2 + 1, "z")
    OE = Order(E, pseudo_matrix(matrix(K, 2, 2, [1, 0, 0, 1]), [1 * maximal_order(K), 2 * maximal_order(K)]))
    I = E(1) * OE
    @test I * I == I
    @test I + I == I
    @test isone(I//I)
  end
end


@testset "Selections" begin

    rng = StableRNG(42)

    @testset "Rank" begin
        fitness = [5, 1, 3, 4, 2.0]

        s = ranklinear(1.0)
        Random.seed!(rng, 2)
        @test s(fitness, 2, rng=rng) == [5,1]

        s = ranklinear(2.0)
        Random.seed!(rng, 2)
        @test s([1,2], 2, rng=rng) == [2,2]
    end

    @testset "Uniform" begin
        s = uniformranking(2)
        Random.seed!(rng, 2);
        @test sort(unique(s([1.0,2.0,3.0], 10, rng=rng))) == [1,2]
        Random.seed!(rng, 2);
        @test sort(unique(s([5,2,3], 5, rng=rng))) == [2,3]
        @test_throws AssertionError s([1.], 4)
    end

    @testset "Roulette" begin
        Random.seed!(rng, 2);
        @test roulette([30.0, -0.1,-0.2, -30.0], 2, rng=rng) == [4,1]
        Random.seed!(rng, 2);
        @test roulette([0.0,0.0,3.0], 2, rng=rng) == [3, 3]
        Random.seed!(rng, 2);
        @test roulette([0,2,0], 2, rng=rng) == [2,2]
        Random.seed!(rng, 2);
        @test rouletteinv([0.0,0.0,3.0], 2, rng=rng) == [1, 1]
        Random.seed!(rng, 2);
        @test rouletteinv([1,2,0], 2, rng=rng) == [3,3]
    end

    @testset "Truncation" begin
        @test truncation([0.0,2.0,3.0], 2) == [1,2]
        @test truncation([0,2,1], 1) == [1]
        @test_throws AssertionError truncation([1.,2.], 5)
    end

    @testset "Tournament" begin
        @test_throws AssertionError tournament(0)
        t = tournament(3, select=argmax)
        Random.seed!(rng, 2);
        @test all(t([0,2,0],100,rng=rng) .== 2)
        t = tournament(3)
        Random.seed!(rng, 2);
        @test all(i->i∈[1,3], t([0,2,0],100,rng=rng))
        t = tournament(2)
        Random.seed!(rng, 2);
        @test all(t([0.0,0.0,1.0],100,rng=rng) .< 3.0)
        fitness = [0 0 1 1; 0 1 0 1]
        t = tournament(2, select=Evolutionary.twowaycomp)
        Random.seed!(rng, 2);
        @test all(t(fitness,100,rng=rng) .< 4)
        @test mean(t(fitness,100,rng=rng)) < 2
    end

    @testset "SUS" begin
        Random.seed!(rng, 2);
        @test sort(unique(sus([1.0,0.0,2.0], 5, rng=rng))) == [1,3]
        Random.seed!(rng, 2);
        @test sort(unique(sus([0,1,2], 5, rng=rng))) == [2,3]
    end

    @testset "DE" begin
        N = 5
        col = collect(1:10)

        Random.seed!(rng, 1)
        idxs = random(col, N, rng=rng)
        @test length(idxs) == N
        @test all(i-> i ∈ col, idxs)
        @test !all(length(unique(random(col, N, rng=rng))) == N for i in 1:100)

        Random.seed!(rng, 1)
        idxs = randomoffset(col, N, rng=rng)
        @test length(idxs) == N
        @test all(i-> i ∈ col, idxs)
        @test all( let idxs = randomoffset(col, N, rng=rng); let l = (idxs[1] - 1 + N)%10; idxs[N] == (l == 0 ? 10 : l) end end for i in 1:100 )

        idxs = best(reverse(col), N)
        @test length(idxs) == N
        @test all(i-> i ∈ col, idxs)
        @test idxs == fill(10, N)

        Random.seed!(rng, 1)
        idxs = permutation(col, N, rng=rng)
        @test length(idxs) == N
        @test all(i-> i ∈ col, idxs)
        @test all(length(unique(permutation(col, N, rng=rng))) == N for i in 1:100)

        Random.seed!(rng, 1)
        @testset "Mutually Exclusive Indices" for i in idxs
            targets = Evolutionary.randexcl(rng, col, [i], 4)
            @test length(unique(push!(targets, i))) == N
        end
    end

end


using RDatasets, KernelDensity
mlmf = dataset("mlmRev", "Gcsemv");

df = mlmf[complete_cases(mlmf[[:Written, :Course]]), :];
dc = array(df[:Course]);  kdc = kde(dc);
dw = array(df[:Written]); kdw = kde(dw);

using Winston
kdc = kde(dc); kdw = kde(dw);
plot(kdc.x, kdc.density, "r--", kdw.x, kdw.density, "b;")


# Portions translated from SLICOT-Reference distribution
# Copyright (c) 2002-2020 NICONET e.V.
function run_mb03wd(datfile, io=stdout)
    NIN = 5
    NOUT = 6
    NMAX = 20
    PMAX = 20
    LDA1 = NMAX
    LDA2 = NMAX
    LDTAU = NMAX-1
    LDZ1 = NMAX
    LDZ2 = NMAX
    LDZTA = NMAX
    LDWORK = max( NMAX, NMAX + PMAX - 2 )
    ZERO = 0.0e0
    ONE = 1.0e0
    A = Array{Float64,3}(undef, LDA1,LDA2,PMAX)
    TAU = Array{Float64,2}(undef, LDTAU,PMAX)
    Z = Array{Float64,3}(undef, LDZ1,LDZ2,PMAX)
    WR = Array{Float64,1}(undef, NMAX)
    WI = Array{Float64,1}(undef, NMAX)
    AS = Array{Float64,3}(undef, LDA1,LDA2,PMAX)
    DWORK = Array{Float64,1}(undef, LDWORK)
    ZTA = Array{Float64,2}(undef, LDZTA,NMAX)
    f = open(datfile,"r")
    readline(f)
    vs = split(readline(f))
    N = parse(BlasInt, vs[1])
    P = parse(BlasInt, vs[2])
    ILO = parse(BlasInt, vs[3])
    IHI = parse(BlasInt, vs[4])
    ILOZ = parse(BlasInt, vs[5])
    IHIZ = parse(BlasInt, vs[6])
    JOB =  vs[7][1]
    COMPZ =  vs[8][1]
    if ( N<0 || N>min( LDA1, LDA2 ) )
        @error "Illegal N=$N"
    end
    if ( P<=0 || P>PMAX )
        @error "Illegal P=$P"
    end

    _isz,_jsz,_ksz = (N,N,P)
    for k in 1:_ksz
      vs = String[]
      while length(vs) < _isz*_jsz
        append!(vs, replace.(split(readline(f)),'D'=>'E'))
      end
      for i in 1:_isz
        _i0 = (i-1)*_jsz
        A[i,1:_jsz,k] .= parsex.(Float64, vs[_i0+1:_i0+_jsz])
      end
    end
    close(f)
# interp call 1

    #FOREIGN.dlacpy!( 'F', N, N, A(1,1,K), LDA1, AS(1,1,K), LDA1 )
    AS[1:N,1:N,1:P] .= A[1:N,1:N,1:P]

# interp call 2

    INFO = SLICOT.mb03vd!(N, P, ILO, IHI, A, TAU)

    if ( INFO!=0 )
        @warn "mb03vd returns info=$INFO"
        return
    end
    if ( LSAME( COMPZ, 'V' ) )
# interp call 3

    #FOREIGN.dlacpy!( 'L', N, N, A(1,1,K), LDA1, Z(1,1,K), LDZ1 )
        for k in 1:P
            Z[1:N,1:N,k] .= tril(A[1:N,1:N,k])
        end

# interp call 4

        INFO = SLICOT.mb03vy!(N, P, ILO, IHI, Z, TAU)
        @test INFO == 0
    INFO == 0 || return

        if ( INFO > 0 )
            @error "mb03vy returns info=$INFO"
        end
        # deviate from source here; continue even if COMPZ != 'V'
    end
# interp call 5

    INFO = SLICOT.mb03wd!(JOB, COMPZ, N, P, ILO, IHI, ILOZ, IHIZ, A, Z, WR, WI, LDWORK)

    if ( INFO > 0 )
        @warn "mb03wd returns info=$INFO"
        println(io, "valid eigvals:")
        imin = max(ILO, INFO+1)
        show(io, "text/plain", WR[imin:N] + im * WI[imin:N])
        println(io)
    end # if

# interp call 6

    INFO = SLICOT.mb03wx!(ILO-1, P, A, WR, WI)

    if  IHI<N
# interp call 7

        INFO = SLICOT.mb03wx!(N-IHI, P, view(A,IHI+1:N, IHI+1:N, 1:P))

    end
# unable to translate write statement:
#   in do block [('40', 'I', 'N')]
#   write WR(I), WI(I)
# interp output 1
    println(io, "A:")
    _nc = N
    _nr = N
    _nk = P
    show(io, "text/plain", A[1:_nr,1:_nc,1:_nk])
    println(io)

# interp output 2
    println(io, "Z:")
    _nc = N
    _nr = N
    _nk = P
    show(io, "text/plain", Z[1:_nr,1:_nc,1:_nk])
    println(io)

    SSQ = ZERO
    for K in 1:P
        KP1 = mod(K,P) + 1
# interp call 8

    # FOREIGN.dgemm!( 'T', 'N', N, N, N, ONE, Z(1,1,K), LDZ1, AS(1,1,K), LDA1, ZERO, ZTA, LDZTA )
        zta = Z[1:N,1:N,K]' * AS[1:N,1:N,K]
# interp call 9

    # FOREIGN.dgemm!( 'N', 'N', N, N, N, ONE, ZTA, LDZTA, Z(1,1,KP1), LDZ1, -ONE, A(1,1,K), LDA1 )
        A[1:N,1:N,K] .-= zta * Z[1:N,1:N,KP1]
        # SSQ = DLAPY2( SSQ, DLANGE( 'Frobenius', N, N, A(1,1,K), LDA1, DWORK ) )
        SSQ = hypot(SSQ, norm(A[1:N,1:N,K]))
    end
    println(io, "decomposition residual: ",SSQ)
end # run_mb03wd()


cd(@__DIR__)
using Pkg; Pkg.activate("."); Pkg.instantiate()

using OrdinaryDiffEq
using ModelingToolkit
using DataDrivenDiffEq
using LinearAlgebra, DiffEqSensitivity, Optim
using DiffEqFlux, Flux

using JLD2, FileIO
using Statistics
# Set a random seed for reproduceable behaviour
using Random
Random.seed!(1234)

## Data generation
function lotka!(du, u, p, t)
    α, β, γ, δ = p
    du[1] = α*u[1] - β*u[2]*u[1]
    du[2] = γ*u[1]*u[2]  - δ*u[2]
end

@info "Generate data "

# Define the experimental parameter
tspan = (0.0f0,3.0f0)
u0 = Float32[0.44249296,4.6280594]
p_ = Float32[1.3, 0.9, 0.8, 1.8]
prob = ODEProblem(lotka!, u0,tspan, p_)
solution = solve(prob, Vern7(), abstol=1e-12, reltol=1e-12, saveat = 0.1)

# Ideal data
X = Array(solution)
t = solution.t
Xₙ = similar(X)
# Add noise in terms of the mean
x̄ = mean(X, dims = 2)

# Create a function to adapt the noise magnitude
function noisy_magnitude(iteration)
    iteration <= 100 && return Float32(1e-3)
    iteration <= 200 && return Float32(5e-3)
    iteration <= 300 && return Float32(1e-2)
    iteration <= 400 && return Float32(2.5e-2)
    return Float32(5e-2)
end

## Define the network
# Gaussian RBF as activation
rbf(x) = exp.(-(x.^2))

include("./loop_recoveries.jl")

@info "Start evaluation loops"
for i in 1:1:500
    Xₙ .= X .+ (noisy_magnitude(i)*x̄) .* randn(eltype(X), size(X)...)
    try
        recover_dynamics(Xₙ, t, "Scenario_1", "$i")
    catch
        @info "Failed on run $i"
    end
end
@info "Finished evaluation loops"


"""
    array_geometry(lon, lat, alt) -> easting, northing, elevation, mean_lon, mean_lat
    array_geometry(traces::AbstractArray{Seis.Trace}) -> easting, northing, elevation, mean_lon, mean_lat
    array_geometry(stations::AbstractArray{Seis.Station}) -> easting, northing, elevation, missing, missing

Return the local `easting` and `northing` in km of each station in the array `s` in
a transverse Mercator transform, relative to the centre of the array.  The longitude
and latitude of the array centre are given by `mean_lon` and `mean_lat` (°) respectively.

`lon` and `lat` should be in ° and `alt` in m.  Or provide an array of `Seis.Trace`s
or `Seis.Station`s.  If `traces` or `stations` is an array of `CartTrace` or
`CartStation`s, then no mean longitude or latitude are calculated and instead
they are returned as `nothing`.
"""
function array_geometry(lon, lat, alt)
    T = eltype(lon)
    length(lon) == length(lat) == length(alt) ||
        throw(DimensionMismatch("`lon` and `lat` must be the same length"))
    # Find spherical mean of the array, ignoring altitude
    x, y, z = geog2cart(lon, lat, true)
    X, Y, Z = Float64.(sum.((x, y, z)))
    mean_lon, mean_lat, r = cart2geog(X, Y, Z, true)
    # Create transform into a local TM projection
    origin = Geodesy.LLA(mean_lat, mean_lon, zero(eltype(lat)))
    transform = Geodesy.ENUfromLLA(origin, Geodesy.wgs84)
    # Do the transform.  x and y are in metres
    x = Array{T}(undef, length(lon))
    y = similar(x)
    z = similar(x)
    for (i, (lo, la, al)) in enumerate(zip(lon, lat, alt))
        p = transform(Geodesy.LLA(la, lo, al))
        x[i] = p.e/1e3
        y[i] = p.n/1e3
        z[i] = p.u/1e3
    end
    x, y, z, mean_lon, mean_lat
end

array_geometry(s::TraceArray) = array_geometry(s.sta)

array_geometry(stas::AbstractArray{<:Seis.GeogStation{T}}) where T =
    array_geometry(stas.lon, stas.lat, coalesce.(stas.elev, zero(T)))

function array_geometry(stas::AbstractArray{<:Seis.CartStation{T}}) where T
    to_km = T(0.001)
    xx = stas.x
    yy = stas.y
    zz = stas.z
    any(x -> x[1] === missing || x[2] === missing, zip(xx, yy)) &&
        throw(ArgumentError("all stations must have both x and y coordinates"))
    x = to_km*(xx .- mean(xx))
    y = to_km*(yy .- mean(yy))
    z = to_km*coalesce.((zz .- mean(zz)), zero(T))
    x, y, z, nothing, nothing
end

"""
    array_centre(lon, lat) -> midlon, midlat

Return the mean location of the array in spherical coordinates.
`lon` and `lat` are vectors in °.
"""
function array_centre(lon, lat)
    length(lon) == length(lat) ||
        throw(DimensionMismatch("`lon` and `lat` must be the same length"))
    x, y, z = geog2cart(lon, lat, true)
    X, Y, Z = sum.((x, y, z))
    mlon, mlat, r = cart2geog(X, Y, Z, true)
    mlon, mlat
end
array_centre(s::TraceArray) = array_centre(s.sta.lon, s.sta.lat)

"""
    delay_times(x, y, sx, sy) -> delays
    delay_times(s::AbstractArray{Seis.Trace}, sx, sy) -> delays

Return the delay times `delays` relative to the array centre for an incoming
plane wave with horizontal slownesses `sx` and `sy` s/km.  Stations have
coordinates `x` and `y` in the east and north directions respectively (km).
Or give an array of `Seis.Trace`s.
"""
function delay_times(x, y, sx, sy)
    length(x) == length(y) ||
        throw(ArgumentError("Length of `x` and `y` must be the same"))
    delays = Array{float(eltype(x))}(undef, length(x), length(sx), length(sy))
    @inbounds for (j, ssyy) in enumerate(sy), (i, ssxx) in enumerate(sx)
        @simd for k in eachindex(x)
            delays[k,i,j] = ssxx*x[k] + ssyy*y[k]
        end
    end
    delays
end
delay_times(s::TraceArray, sx, sy) = delay_times(array_geometry(s)[1:2]..., sx, sy)

"""
    find_maxima(sx, sy, power) -> sx_max, sy_max, slowness, backazimuth

Find `n` maxima in the beam power grid `power`, and return vectors
of `sx_max`, `sy_max`, `slowness` and `backazimuth` for each.
"""
function find_maxima(sx, sy, power::AbstractArray{T,2}, n=1) where T
    inds = sortperm(vec(power), rev=true)
    sx_max, sy_max, p, β = T[], T[], T[], T[]
    for i in 1:n
        ix, iy = Tuple(CartesianIndices(power)[inds[i]])
        push!(sx_max, sx[ix])
        push!(sy_max, sy[iy])
        push!(p, sqrt(sx[ix]^2 + sy[iy]^2))
        push!(β, rad2deg(atan(sx[ix], sy[iy])))
    end
    sx_max, sy_max, p, β
end

"List of available stacking routines"
const AVAILABLE_STACK_METHODS = (:linear, :nthroot, :phaseweight)

"""
    stack(S::AbstractArray{Seis.Trace}, time_range, align=zeros(length(S)), weight; method=:linear) -> s::Seis.Trace

Return the linear stack of all traces in `S`, aligned in time on the
value in `align`, between `times[1]` and `times[end]`.

`times` can be a Tuple, Range or Array; in all cases, only the first and last values
are used.  **N.B. If using a non-integer Range, this may not include the end value
you expect.  E.g., `0.1:0.9` by default contains only the value `0.1` because the
default step is `1`.**

`align` can be a header symbol, or an array of equal length to `S`.  If it is a Symbol,
then the values in the headers are used.

All traces must have the same sampling frequency, and an error is thrown
if not all traces can be included in the stack.

If provided, traces are weighted by the values in `weight`, which are normalised
to sum to unity.

Available stacking methods are: $(AVAILABLE_STACK_METHODS).
"""
function stack(S::AbstractArray{<:Seis.Trace{T}}, time_range,
               align::Union{Symbol,AbstractArray}=zeros(T, length(S));
               weights=ones(T, length(S))./length(S),
               method=:linear,
               n=nothing) where T
    # Convert to a vector of times if given as a symbol
    if typeof(align) == Symbol
        picks = getproperty(S.picks, align)
        any(ismissing, picks) && error("not all traces have a pick ")
        align = picks.time
    end
    all(x -> x == S[1].delta, S.delta) || error("all traces must have same delta")
    length(S) == length(weights) ||
        throw(ArgumentError("weights vector must be same length as traces"))
    delay = -align
    t1, t2 = time_range[1], time_range[end]
    weights ./= sum(weights) # Normalise weights
    if method == :linear
        stack_trace = _stack_linear([Seis.trace(s) for s in S],
            S[1].delta, S.b, delay, t1, t2, weights)
    elseif method == :nthroot
        n === nothing &&
            throw(ArgumentError("n must be provided for nthroot stacking"))
        stack_trace = _stack_nthroot([Seis.trace(s) for s in S],
            n, S[1].delta, S.b, delay, t1, t2, weights)
    elseif method == :phaseweight
        n === nothing &&
            throw(ArgumentError("n must be provided for phaseweighted stacking"))
        stack_trace = _stack_phaseweight([Seis.trace(s) for s in S],
            n, S[1].delta, S.b, delay, t1, t2, weights)
    else
        throw(ArgumentError("unrecognised stacking method"))
    end
    Seis.Trace{T}(t1, S[1].delta, stack_trace)
end

"""
    _stack_linear(a, delta, b, delay, t1, t2) -> stacked_trace::Vector{eltype(a)}

Return an array containing the linear stack of the vector of vectors `a`, whose
start times are given in `b`, and where the stack is returned between times `t1` and `t2`
(which can be negative if delay is also negative).  `delta` is the sampling interval
for all the traces.

The trace starts at time `t1` and has sample spacing `delta`.
"""
function _stack_linear(a::Vector{<:AbstractArray}, delta::Real, b::Vector,
        delay::Vector, t1::Real, t2::Real, weights)
    N, npts, b_shift, stack_trace = _stack_array(a, delta, b, delay, t1, t2)
    # stack_trace checks that we are always in array bounds
    @inbounds for j in 1:N
        ip1 = round(Int, (t1 - b_shift[j])/delta) + 1
        w = weights[j]
        @simd for i in ip1:(ip1+npts-1)
            stack_trace[i-ip1+1] += a[j][i]*w
        end
    end
    stack_trace
end

"""
    _stack_nthroot(a::Vector{<:AbstractArray}, n, delta, b, delay, t1, t2, weights)

Return an array containing the `n`throot stack of the vector of vectors `a`, whose
start times are given in `b`, and where the stack is returned between times `t1` and `t2`
(which can be negative if delay is also negative).  `delta` is the sampling interval
for all the traces.

The trace starts at time `t1` and has sample spacing `delta`.
"""
function _stack_nthroot(a::Vector{<:AbstractArray}, n, delta, b, delay, t1, t2, weights)
    N, npts, b_shift, stack_trace = _stack_array(a, delta, b, delay, t1, t2)
    @inbounds for j in 1:N
        ip1 = round(Int, (t1 - b_shift[j])/delta) + 1
        w = weights[j]
        @simd for i in ip1:(ip1+npts-1)
            stack_trace[i-ip1+1] += sign(a[j][i])*abs(a[j][i])^(1/n)*w
        end
    end
    stack_trace .= sign.(stack_trace).*abs.(stack_trace).^n
    stack_trace
end

"""
    _stack_phaseweight(a::Vector{<:AbstractArray}, delta, b, delay, t1, t2, weights)

Return an array containing the phaseweighted stack of the vector of vectors `a`, whose
start times are given in `b`, and where the stack is returned between times `t1` and `t2`
(which can be negative if delay is also negative).  `delta` is the sampling interval
for all the traces.  The weighting power is given by `n`.

The trace starts at time `t1` and has sample spacing `delta`.
"""
function _stack_phaseweight(a::Vector{<:AbstractArray{T}}, n, delta, b, delay, t1, t2, weights) where T
    N, npts, b_shift, stack_trace = _stack_array(a, delta, b, delay, t1, t2)
    # Collect shifted traces together
    traces = Array{T}(undef, npts, N)
    @inbounds for j in 1:N
        ip1 = round(Int, (t1 - b_shift[j])/delta) + 1
        w = weights[j]
        @simd for i in ip1:(ip1+npts-1)
            traces[i-ip1+1,j] = a[j][i]*w
        end
    end
    # Create instantaneous phases
    phases = DSP.hilbert(traces)
    phases .= phases./abs.(phases)
    # Stack
    @inbounds for i in 1:npts
        weight = abs(sum(phases[i,:])/N)^n
        @simd for j in 1:N
            stack_trace[i] += traces[i,j]*weight/N
        end
    end
    stack_trace
end

"""
    _stack_array(a::Vector{Vector{<:Real}}, delta, b::Vector, delay::Vector, t1, t2) -> N, npts, b_shifted, stack_trace

Return the number of traces `N`, the number of points in the output stack trace,
`npts`, an array of shifted trace start times `b_shifted` and a zeroed array
for the stack itself, `stack_trace`.  This routine also checks that the input
arrays are the required length and no points are missing from the desired stack.
"""
function _stack_array(a::Vector{A}, delta, b, delay, t1, t2) where {A<:AbstractArray}
    N = length(a)
    length(b) == length(delay) == N ||
        throw(ArgumentError("Lengths of `a`, `b` and `delay` must all be the same"))
    t1 < t2 || throw(ArgumentError("`t2` is before or the same as `t1`"))
    b_shift = b .+ delay
    e_shift = b_shift .+ (length.(a) .- 1).*delta
    all(b_shift .<= t1) && all(t2 .<= e_shift) ||
        error("Not all traces have data within time range $(t1):$(t2) s
        range in b: $(extrema(b_shift))
        range in e: $(extrema(e_shift))")
    npts = round(Int, (t2 - t1)/delta) + 1
    stack_trace = zeros(eltype(eltype(a)), npts)
    N, npts, b_shift, stack_trace
end


"""
	violates_bapat_sunder(A,B, tol = ATOL)

Checks if matrices `A` and `B` violate the Bapat-Sunder conjecture, see
[Boson bunching is not maximized by indistinguishable particles](https://arxiv.org/abs/2203.01306)
"""
function violates_bapat_sunder(A,B, tol = ATOL)

	diagonal_el_product = prod([B[i,i] for i in 1:size(B)[1]])

	# @test (imag(diagonal_el_product) ≈ 0.)

	if !(imag(diagonal_el_product) ≈ 0.)
		throw(Exception("product of diagonal elements of B is not real"))
	elseif !(is_positive_semidefinite(A) && is_positive_semidefinite(B))
		throw(Exception("A or B not semidefinitepositive"))
	else
		return abs(ryser(A .* B)) / (abs(ryser(A)) * real(prod([B[i,i] for i in 1:size(B)[1]]))) > 1+tol
	end

	return nothing

end


using SparseArrays, DiffEqOperators, LinearAlgebra, Random,
      Test, BandedMatrices, FillArrays, LazyArrays, BlockBandedMatrices

# This test file tests for the correctness of higher dimensional concretization.
# The tests verify that multiplication in the concretized case agrees with the matrix-free
# multiplication

@testset "First Dimension" begin

      # Test that even when we have a vector, the concretizations using the higher dimension dispatch still function
      # correctly
      M = rand(22)

      L1 = CenteredDifference(1,2,1.0,20)
      L2 = CenteredDifference(1,2,1.0,20)
      L3 = CenteredDifference(3,3,1.0,20)

      @test L1*M ≈ Array(L1, size(M))*vec(M) ≈ sparse(L1,size(M))*M ≈ BandedMatrix(L1, size(M))*M ≈ BandedBlockBandedMatrix(L1,size(M))*M
      @test L2*M ≈ Array(L2, size(M))*vec(M) ≈ sparse(L2,size(M))*M ≈ BandedMatrix(L2, size(M))*M ≈ BandedBlockBandedMatrix(L2,size(M))*M
      @test L3*M ≈ Array(L3, size(M))*vec(M) ≈ sparse(L3,size(M))*M ≈ BandedMatrix(L3, size(M))*M ≈ BandedBlockBandedMatrix(L3,size(M))*M

      M = rand(22,2,2,2,2)

      @test vec(L1*M) ≈ Array(L1, size(M))*vec(M) ≈ sparse(L1,size(M))*vec(M) ≈ BandedMatrix(L1, size(M))*vec(M) ≈ BandedBlockBandedMatrix(L1,size(M))*vec(M)
      @test vec(L2*M) ≈ Array(L2, size(M))*vec(M) ≈ sparse(L2,size(M))*vec(M) ≈ BandedMatrix(L2, size(M))*vec(M) ≈ BandedBlockBandedMatrix(L2,size(M))*vec(M)
      @test vec(L3*M) ≈ Array(L3, size(M))*vec(M) ≈ sparse(L3,size(M))*vec(M) ≈ BandedMatrix(L3, size(M))*vec(M) ≈ BandedBlockBandedMatrix(L3,size(M))*vec(M)

end

@testset "Second Dimension" begin

      M = rand(2,22,2)

      L1 = CenteredDifference{2}(1,2,1.0,20)
      L2 = CenteredDifference{2}(1,2,1.0,20)
      L3 = CenteredDifference{2}(3,3,1.0,20)

      @test vec(L1*M) ≈ Array(L1, size(M))*vec(M) ≈ sparse(L1,size(M))*vec(M) ≈ BandedMatrix(L1, size(M))*vec(M) ≈ BandedBlockBandedMatrix(L1,size(M))*vec(M)
      @test vec(L2*M) ≈ Array(L2, size(M))*vec(M) ≈ sparse(L2,size(M))*vec(M) ≈ BandedMatrix(L2, size(M))*vec(M) ≈ BandedBlockBandedMatrix(L2,size(M))*vec(M)
      @test vec(L3*M) ≈ Array(L3, size(M))*vec(M) ≈ sparse(L3,size(M))*vec(M) ≈ BandedMatrix(L3, size(M))*vec(M) ≈ BandedBlockBandedMatrix(L3,size(M))*vec(M)

      M = rand(2,22,2,2)

      @test vec(L1*M) ≈ Array(L1, size(M))*vec(M) ≈ sparse(L1,size(M))*vec(M) ≈ BandedMatrix(L1, size(M))*vec(M) ≈ BandedBlockBandedMatrix(L1,size(M))*vec(M)
      @test vec(L2*M) ≈ Array(L2, size(M))*vec(M) ≈ sparse(L2,size(M))*vec(M) ≈ BandedMatrix(L2, size(M))*vec(M) ≈ BandedBlockBandedMatrix(L2,size(M))*vec(M)
      @test vec(L3*M) ≈ Array(L3, size(M))*vec(M) ≈ sparse(L3,size(M))*vec(M) ≈ BandedMatrix(L3, size(M))*vec(M) ≈ BandedBlockBandedMatrix(L3,size(M))*vec(M)

      M = rand(2,22,2,2,3)

      @test vec(L1*M) ≈ Array(L1, size(M))*vec(M) ≈ sparse(L1,size(M))*vec(M) ≈ BandedMatrix(L1, size(M))*vec(M) ≈ BandedBlockBandedMatrix(L1,size(M))*vec(M)
      @test vec(L2*M) ≈ Array(L2, size(M))*vec(M) ≈ sparse(L2,size(M))*vec(M) ≈ BandedMatrix(L2, size(M))*vec(M) ≈ BandedBlockBandedMatrix(L2,size(M))*vec(M)
      @test vec(L3*M) ≈ Array(L3, size(M))*vec(M) ≈ sparse(L3,size(M))*vec(M) ≈ BandedMatrix(L3, size(M))*vec(M) ≈ BandedBlockBandedMatrix(L3,size(M))*vec(M)

end

@testset "Third Dimension" begin

      M = rand(3,2,22)

      L1 = CenteredDifference{3}(1,2,1.0,20)
      L2 = CenteredDifference{3}(1,2,1.0,20)
      L3 = CenteredDifference{3}(3,3,1.0,20)

      @test vec(L1*M) ≈ Array(L1, size(M))*vec(M) ≈ sparse(L1,size(M))*vec(M) ≈ BandedMatrix(L1, size(M))*vec(M) ≈ BandedBlockBandedMatrix(L1,size(M))*vec(M)
      @test vec(L2*M) ≈ Array(L2, size(M))*vec(M) ≈ sparse(L2,size(M))*vec(M) ≈ BandedMatrix(L2, size(M))*vec(M) ≈ BandedBlockBandedMatrix(L2,size(M))*vec(M)
      @test vec(L3*M) ≈ Array(L3, size(M))*vec(M) ≈ sparse(L3,size(M))*vec(M) ≈ BandedMatrix(L3, size(M))*vec(M) ≈ BandedBlockBandedMatrix(L3,size(M))*vec(M)

      M = rand(3,2,22,2)

      @test vec(L1*M) ≈ Array(L1, size(M))*vec(M) ≈ sparse(L1,size(M))*vec(M) ≈ BandedMatrix(L1, size(M))*vec(M) ≈ BandedBlockBandedMatrix(L1,size(M))*vec(M)
      @test vec(L2*M) ≈ Array(L2, size(M))*vec(M) ≈ sparse(L2,size(M))*vec(M) ≈ BandedMatrix(L2, size(M))*vec(M) ≈ BandedBlockBandedMatrix(L2,size(M))*vec(M)
      @test vec(L3*M) ≈ Array(L3, size(M))*vec(M) ≈ sparse(L3,size(M))*vec(M) ≈ BandedMatrix(L3, size(M))*vec(M) ≈ BandedBlockBandedMatrix(L3,size(M))*vec(M)

      M = rand(3,2,22,2,3)

      @test vec(L1*M) ≈ Array(L1, size(M))*vec(M) ≈ sparse(L1,size(M))*vec(M) ≈ BandedMatrix(L1, size(M))*vec(M) ≈ BandedBlockBandedMatrix(L1,size(M))*vec(M)
      @test vec(L2*M) ≈ Array(L2, size(M))*vec(M) ≈ sparse(L2,size(M))*vec(M) ≈ BandedMatrix(L2, size(M))*vec(M) ≈ BandedBlockBandedMatrix(L2,size(M))*vec(M)
      @test vec(L3*M) ≈ Array(L3, size(M))*vec(M) ≈ sparse(L3,size(M))*vec(M) ≈ BandedMatrix(L3, size(M))*vec(M) ≈ BandedBlockBandedMatrix(L3,size(M))*vec(M)

end

@testset "Fifth Dimension" begin

      M = rand(3,2,3,2,22)

      L1 = CenteredDifference{5}(1,2,1.0,20)
      L2 = CenteredDifference{5}(1,2,1.0,20)
      L3 = CenteredDifference{5}(3,3,1.0,20)

      @test vec(L1*M) ≈ Array(L1, size(M))*vec(M) ≈ sparse(L1,size(M))*vec(M) ≈ BandedMatrix(L1, size(M))*vec(M) ≈ BandedBlockBandedMatrix(L1,size(M))*vec(M)
      @test vec(L2*M) ≈ Array(L2, size(M))*vec(M) ≈ sparse(L2,size(M))*vec(M) ≈ BandedMatrix(L2, size(M))*vec(M) ≈ BandedBlockBandedMatrix(L2,size(M))*vec(M)
      @test vec(L3*M) ≈ Array(L3, size(M))*vec(M) ≈ sparse(L3,size(M))*vec(M) ≈ BandedMatrix(L3, size(M))*vec(M) ≈ BandedBlockBandedMatrix(L3,size(M))*vec(M)

      M = rand(3,2,3,2,22,3)

      @test vec(L1*M) ≈ Array(L1, size(M))*vec(M) ≈ sparse(L1,size(M))*vec(M) ≈ BandedMatrix(L1, size(M))*vec(M) ≈ BandedBlockBandedMatrix(L1,size(M))*vec(M)
      @test vec(L2*M) ≈ Array(L2, size(M))*vec(M) ≈ sparse(L2,size(M))*vec(M) ≈ BandedMatrix(L2, size(M))*vec(M) ≈ BandedBlockBandedMatrix(L2,size(M))*vec(M)
      @test vec(L3*M) ≈ Array(L3, size(M))*vec(M) ≈ sparse(L3,size(M))*vec(M) ≈ BandedMatrix(L3, size(M))*vec(M) ≈ BandedBlockBandedMatrix(L3,size(M))*vec(M)

      M = rand(2,2,2,2,22,3,2)

      @test vec(L1*M) ≈ Array(L1, size(M))*vec(M) ≈ sparse(L1,size(M))*vec(M) ≈ BandedMatrix(L1, size(M))*vec(M) ≈ BandedBlockBandedMatrix(L1,size(M))*vec(M)
      @test vec(L2*M) ≈ Array(L2, size(M))*vec(M) ≈ sparse(L2,size(M))*vec(M) ≈ BandedMatrix(L2, size(M))*vec(M) ≈ BandedBlockBandedMatrix(L2,size(M))*vec(M)
      @test vec(L3*M) ≈ Array(L3, size(M))*vec(M) ≈ sparse(L3,size(M))*vec(M) ≈ BandedMatrix(L3, size(M))*vec(M) ≈ BandedBlockBandedMatrix(L3,size(M))*vec(M)

end


using Statistics, LinearAlgebra, PyPlot

mutable struct Lattice
    name::String
    state::Array{Array{Float64,3},1}
    probability::Array{Array{Float64,3},1}
    Lattice(name::String, state::Array{Array{Float64,3},1}, probability::Array{Array{Float64,3},1}) = new(name, state, probability)
end

"""
	lattice_approximation(bstructure::Array{Int64,1}, path::Function, nIterations::Int64, r::Int64 = 2, dim::Int64 = 1)
Returns a valuated approximated lattice for the stochastic process provided in any dimension. The default dimension is dim = 1 for a lattice in 1 dimension.

Args:
- bstructure - Branching structure of the scenario lattice e.g., bstructure = [1,2,3,4,5] represents a 5-staged lattice
- path - Function generating samples from a known distribution with length equal to the length of bstructure of the lattice.
- nIterations - Number of iterations to be performed.
- r - Parameter for the transportation distance.
- dim - dimension of the lattice to be generated. This depends entirely on the dimension of the samples generated.
"""
function lattice_approximation(bstructure::Array{Int64,1}, path::Function, nIterations::Int64, r::Int64 = 2, dim::Int64 = 1)
    tdist = zeros(dim)         # multistage distance
    T = length(bstructure)     # number of stages in the lattice
    states = [zeros(bstructure[j], 1, dim) for j = 1 : T] # States of the lattice at each time t
    probabilities = vcat([zeros(bstructure[1], 1, dim)],[zeros(bstructure[j-1], bstructure[j], dim) for j = 2 : T]) # Probabilities of the lattice at each time t
    init_path = path()
    #Check the dimension of the sample path if it is equal to the dimension of the lattice.
    if dim != size(init_path)[2]
        @error("Dimension of lattice ($dim) is not equal to the dimension of the input array ($(size(init_path)[2]))")
    end
    # Initialize the states of the nodes of the lattice
    for t = 1 : T
        states[t] .= init_path[t]
    end
    Z = Array{Float64}(undef, T, dim) # Array to hold the new samples generated
    #Stochastic approximation step starts here
    for n = 1 : nIterations
        Z .= path() # Draw a new sample Gaussian path
        last_index = Int64.(ones(dim))
        dist = zeros(dim)
        for t = 1 : T
            for i = 1 : dim
                # corrective action to include lost nodes
                states[t][:, :, i][findall(sum(probabilities[t][:, :, i], dims = 2) .< 1.3 * sqrt(n) / bstructure[t])] .= Z[t,i]
                min_dist, new_index = findmin(abs.(vec(states[t][:, :, i] .- Z[t,i]))) # find the closest lattice entry
                dist[i] += min_dist^2  # Euclidean distance for the paths
                probabilities[t][last_index[i], new_index, i] += 1.0  # increase the counter for the nodes
                # use stochastic approximation algorithm to calculate for the new states of the nodes in the lattice.
                states[t][new_index, :, i] = states[t][new_index, :, i] - r * (min_dist)^(r-1) * (states[t][new_index, :, i] .- Z[t,i]) / ((Z[1] + 30 + n))
                last_index[i] = new_index
            end
        end
        # calculate the multistage distance
        dist = dist .^ (1/2); tdist = (tdist .* (n - 1) + dist .^r )/ n
        #tdist = (tdist*(n-1) + dist)/n
    end
    # calculate the probabilities
    probabilities = probabilities ./ nIterations
    # functions to print the results of the lattices
    # these functions depends on the dimension of the lattice
    if dim == 1
        # Round the results to make sure that they are neat.
        # Remember it is a 3D array, so a loop will do.
        return Lattice("Approximated Lattice with $bstructure branching structure and distance=$(tdist .^ (1 / r)) at $(nIterations) iterations",
        [round.(states[i], digits = 6) for i = 1 : length(states)], [round.(probabilities[j], digits = 6) for j = 1 : length(probabilities)])
    else
        lattices = Lattice[] # create an empty array of type Lattice to hold the resulting lattices
        for i = 1 : length(states[1])
            st = [zeros(bstructure[j], 1, 1) for j = 1 : length(states)]
            pp = vcat([zeros(bstructure[1], 1, 1)],[zeros(bstructure[j-1] , bstructure[j], 1) for j = 2 : length(states)])

            for j = 1 : length(states)
                st[j] .= states[j][:, :, i]
                pp[j] .= probabilities[j][:, :, i]
            end
            sublattice = Lattice("Lattice $i with distance=$((tdist.^(1/r))[i])",
            [round.(st[i], digits = 6) for i = 1 : length(st)], [round.(pp[j], digits = 6) for j = 1 : length(pp)])
            push!(lattices,sublattice)
        end
        return lattices
    end
end

"""
	PlotLattice(lt::Lattice,fig=1)

Returns a plot of a lattice.
"""
function plot_lattice(lt::Lattice,fig = 1)
    if !isempty(fig)
        figure(figsize=(6,4))
    end
    lts = subplot2grid((1,4),(0,0),colspan = 3)
    title("states")
    xlabel("stage,time",fontsize=11)
    #ylabel("states")
    xticks(1:length(lt.state))
    lts.spines["top"].set_visible(false)                                                         # remove the box at the top
    lts.spines["right"].set_visible(false)                                                       # remove the box at the right
    for t = 2:length(lt.state)
        for i=1:length(lt.state[t-1])
            for j=1:length(lt.state[t])
                lts.plot([t-1,t],[lt.state[t-1][i],lt.state[t][j]])
            end
        end
    end

    prs = subplot2grid((1,4), (0,3))
    title("probabilities")
    prs.spines["top"].set_visible(false)
    prs.spines["left"].set_visible(false)
    prs.spines["right"].set_visible(false)
# Use the states and probabilites at the last stage to plot the marginal distribution
    stts = lt.state[end]
    n = length(stts)                                    # length of terminal nodes of the lattice.
    h = 1.05 * std(stts) / (n^0.2) + eps()                  #Silverman rule of thumb
    #lts.set_ylim(minimum(stts)-1.5*h, maximum(stts)+1.5*h)
    #prs.set_ylim(minimum(stts)-3.0*h, maximum(stts)+3.0*h)
    proba = sum(lt.probability[end], dims=1)
    yticks(())                                          #remove the ticks on probability plot
    t = LinRange(minimum(stts) - h, maximum(stts) + h, 100) #100 points on probability plot
    density = zero(t)
    for (i, ti) in enumerate(t)
        for (j, xj) in enumerate(stts)
            tmp = (xj - ti) / h
            density[i] += proba[j]* 35/32 * max(1.0 -tmp^2, 0.0)^3 / h #triweight kernel
        end
    end
    plot(density, t)
    prs.fill_betweenx(t, 0 , density)
end


## --- Load required pacages, install Chron if required

    try
        using Chron
    catch
        using Pkg
        Pkg.add("Chron")
        using Chron
    end

    using Statistics, StatsBase, SpecialFunctions
    using Plots; gr();

## --- Define sample properties

    # # # # # # # # # # # Enter sample information here! # # # # # # # # # # # #
    # Input the number of samples we wish to model (must match below)
    nSamples = 4
    # Make an instance of a ChronSection object for nSamples
    smpl = NewChronAgeData(nSamples)
    smpl.Name           = ("Sample 1", "Sample 2", "Sample 3", "Sample 4") # Et cetera
    smpl.Age_14C       .= [ 6991,  7088,  7230,  7540,] # Measured ages
    smpl.Age_14C_sigma .= [   30,    70,    50,    50,] # Measured 1-σ uncertainties
    smpl.Height        .= [ -355,  -380,-397.0,-411.5,] # Depths below surface should be negative
    smpl.Height_sigma  .= fill(0.01, nSamples) # Usually assume little or no sample height uncertainty
    smpl.Age_Sidedness .= zeros(nSamples) # Sidedness (zeros by default: geochron constraints are two-sided). Use -1 for a maximum age and +1 for a minimum age, 0 for two-sided
    smpl.Age_Unit = "Years BP" # Unit of measurement for ages
    smpl.Height_Unit = "m" # Unit of measurement for Height and Height_sigma

    # IMPORTANT: smpl.Height must increase with increasing stratigraphic height
    # -- i.e., stratigraphically younger samples must be more positive. For this
    # reason, it is convenient to represent depths below surface as negative
    # numbers.
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

## --- Calculate and plot calendar age PDFs for each sample

    smpl.Params = fill(NaN, length(intcal13["Age_Calendar"]), nSamples)
    for i = 1:nSamples
        # The likelihood that a measured 14C age could result from a sample of
        # a given calendar age is proportional to the intergral of the product
        # of the two respective distributions
        likelihood = normproduct.(smpl.Age_14C[i], smpl.Age_14C_sigma[i], intcal13["Age_14C"], intcal13["Age_sigma"])
        likelihood ./= sum(likelihood) * intcal13["dt"] # Normalize

        # Estimate mean and standard deviation by drawing samples from distribution
        samples = draw_from_distribution(likelihood, 10^6) .* maximum(intcal13["Age_Calendar"])
        smpl.Age[i] = mean(samples)
        smpl.Age_sigma[i] = std(samples)
        smpl.Age_025CI[i] = percentile(samples,2.5)
        smpl.Age_975CI[i] = percentile(samples,97.5)

        # Populate smpl.Params with log likelihood for each sample
        smpl.Params[:,i] = normproduct_ll.(smpl.Age_14C[i], smpl.Age_14C_sigma[i], intcal13["Age_14C"], intcal13["Age_sigma"])

        # Plot likelihood vector for each sample
        t = (intcal13["Age_Calendar"] .> smpl.Age[i] - 5*smpl.Age_sigma[i]) .& (intcal13["Age_Calendar"] .< smpl.Age[i] + 5*smpl.Age_sigma[i])
        plot(intcal13["Age_Calendar"][t], likelihood[t], label=smpl.Name[i], xlabel="Calendar Age", ylabel="Likelihood")
        savefig("$(smpl.Name[i])_CalendarAge.pdf")
    end

## --- Run stratigraphic model

    # # # # # # # # # # Configure stratigraphic model here! # # # # # # # # # #
    # Configure the stratigraphic Monte Carlo model
    config = NewStratAgeModelConfiguration()
    # If you in doubt, you can probably leave these parameters as-is
    config.resolution = 0.2 # Same units as sample height. Smaller is slower!
    config.bounding = 0.5 # how far away do we place runaway bounds, as a fraction of total section height
    (bottom, top) = extrema(smpl.Height)
    npoints_approx = round(Int,length(bottom:config.resolution:top) * (1 + 2*config.bounding))
    config.nsteps = 15000 # Number of steps to run in distribution MCMC
    config.burnin = 10000*npoints_approx # Number to discard
    config.sieve = round(Int,npoints_approx) # Record one out of every nsieve steps
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

    # Run the stratigraphic MCMC model
    @time (mdl, agedist, lldist) = StratMetropolis14C(smpl, config)

    # Plot results (mean and 95% confidence interval for both model and data)
    hdl = plot([mdl.Age_025CI; reverse(mdl.Age_975CI)],[mdl.Height; reverse(mdl.Height)], fill=(minimum(mdl.Height),0.5,:blue), label="model")
    plot!(hdl, mdl.Age, mdl.Height, linecolor=:blue, label="")
    plot!(hdl, smpl.Age, smpl.Height, xerror=(smpl.Age-smpl.Age_025CI,smpl.Age_975CI-smpl.Age),label="data",seriestype=:scatter,color=:black)
    plot!(hdl, xlabel="Age ($(smpl.Age_Unit))", ylabel="Height ($(smpl.Height_Unit))")
    savefig(hdl,"AgeDepthModel.pdf");
    display(hdl)

## --- Interpolate results at a specific height

    # Stratigraphic height at which to interpolate
    height = -404

    age_at_height = linterp1s(mdl.Height,mdl.Age,height)
    age_at_height_min = linterp1s(mdl.Height,mdl.Age_025CI,height)
    age_at_height_max = linterp1s(mdl.Height,mdl.Age_975CI,height)
    print("Interpolated age at height=$height: $age_at_height +$(age_at_height_max-age_at_height)/-$(age_at_height-age_at_height_min) $(smpl.Age_Unit)")

    # Optional: interpolate full age distribution
    interpolated_distribution = Array{Float64}(undef,size(agedist,2))
    for i=1:size(agedist,2)
        interpolated_distribution[i] = linterp1s(mdl.Height,agedist[:,i],height)
    end
    hdl = histogram(interpolated_distribution, nbins=50, label="")
    plot!(hdl, xlabel="Age ($(smpl.Age_Unit)) at height=$height", ylabel="Likelihood (unnormalized)")
    savefig(hdl, "Interpolated age distribution.pdf")
    display(hdl)

## --- Calculate deposition rate binned by age  - - - - - - - - - - - - - - - -

    # Set bin width and spacing
    binwidth = round(nanrange(mdl.Age)/10,sigdigits=1) # Can also set manually, commented out below
    # binwidth = 100 # Same units as smpl.Age
    binoverlap = 10
    ages = collect(minimum(mdl.Age):binwidth/binoverlap:maximum(mdl.Age))
    bincenters = ages[1+Int(binoverlap/2):end-Int(binoverlap/2)]
    spacing = binoverlap

    # Calculate rates for the stratigraphy of each markov chain step
    dhdt_dist = Array{Float64}(undef, length(ages)-binoverlap, config.nsteps)
    @time for i=1:config.nsteps
        heights = linterp1(reverse(agedist[:,i]), reverse(mdl.Height), ages)
        dhdt_dist[:,i] .= abs.(heights[1:end-spacing] - heights[spacing+1:end]) ./ binwidth
    end

    # Find mean and 1-sigma (68%) CI
    dhdt = nanmean(dhdt_dist,dim=2)
    dhdt_50p = nanmedian(dhdt_dist,dim=2)
    dhdt_16p = nanpctile(dhdt_dist,15.865,dim=2) # Minus 1-sigma (15.865th percentile)
    dhdt_84p = nanpctile(dhdt_dist,84.135,dim=2) # Plus 1-sigma (84.135th percentile)
    # Other confidence intervals (10:10:50)
    dhdt_20p = nanpctile(dhdt_dist,20,dim=2)
    dhdt_80p = nanpctile(dhdt_dist,80,dim=2)
    dhdt_25p = nanpctile(dhdt_dist,25,dim=2)
    dhdt_75p = nanpctile(dhdt_dist,75,dim=2)
    dhdt_30p = nanpctile(dhdt_dist,30,dim=2)
    dhdt_70p = nanpctile(dhdt_dist,70,dim=2)
    dhdt_35p = nanpctile(dhdt_dist,35,dim=2)
    dhdt_65p = nanpctile(dhdt_dist,65,dim=2)
    dhdt_40p = nanpctile(dhdt_dist,40,dim=2)
    dhdt_60p = nanpctile(dhdt_dist,60,dim=2)
    dhdt_45p = nanpctile(dhdt_dist,45,dim=2)
    dhdt_55p = nanpctile(dhdt_dist,55,dim=2)

    # Plot results
    hdl = plot(bincenters,dhdt, label="Mean", color=:black, linewidth=2)
    plot!(hdl,[bincenters; reverse(bincenters)],[dhdt_16p; reverse(dhdt_84p)], fill=(0,0.2,:darkblue), linealpha=0, label="68% CI")
    plot!(hdl,[bincenters; reverse(bincenters)],[dhdt_20p; reverse(dhdt_80p)], fill=(0,0.2,:darkblue), linealpha=0, label="")
    plot!(hdl,[bincenters; reverse(bincenters)],[dhdt_25p; reverse(dhdt_75p)], fill=(0,0.2,:darkblue), linealpha=0, label="")
    plot!(hdl,[bincenters; reverse(bincenters)],[dhdt_30p; reverse(dhdt_70p)], fill=(0,0.2,:darkblue), linealpha=0, label="")
    plot!(hdl,[bincenters; reverse(bincenters)],[dhdt_35p; reverse(dhdt_65p)], fill=(0,0.2,:darkblue), linealpha=0, label="")
    plot!(hdl,[bincenters; reverse(bincenters)],[dhdt_40p; reverse(dhdt_60p)], fill=(0,0.2,:darkblue), linealpha=0, label="")
    plot!(hdl,[bincenters; reverse(bincenters)],[dhdt_45p; reverse(dhdt_55p)], fill=(0,0.2,:darkblue), linealpha=0, label="")
    plot!(hdl,bincenters,dhdt_50p, label="Median", color=:grey, linewidth=1)
    plot!(hdl, xlabel="Age ($(smpl.Age_Unit))", ylabel="Depositional Rate ($(smpl.Height_Unit) / $(smpl.Age_Unit) over $binwidth $(smpl.Age_Unit))", fg_color_legend=:white)
    # savefig(hdl,"DepositionRateModelCI.pdf")
    display(hdl)

## --- Optional: Stratigraphic model including hiatuses

    # Data about hiatuses
    nHiatuses = 2 # The number of hiatuses you have data for
    hiatus = NewHiatusData(nHiatuses) # Struct to hold data
    hiatus.Height         = [-371.5, -405.0 ]
    hiatus.Height_sigma   = [   0.0,    0.0 ]
    hiatus.Duration       = [ 100.0,   123.0]
    hiatus.Duration_sigma = [  30.5,    20.0]

    # Run the model. Note the additional `hiatus` arguments
    @time (mdl, agedist, hiatusdist, lldist) = StratMetropolis14C(smpl, hiatus, config)

    # Plot results (mean and 95% confidence interval for both model and data)
    hdl = plot([mdl.Age_025CI; reverse(mdl.Age_975CI)],[mdl.Height; reverse(mdl.Height)], fill=(minimum(mdl.Height),0.5,:blue), label="model")
    plot!(hdl, mdl.Age, mdl.Height, linecolor=:blue, label="", fg_color_legend=:white)
    plot!(hdl, smpl.Age, smpl.Height, xerror=(smpl.Age-smpl.Age_025CI,smpl.Age_975CI-smpl.Age),label="data",seriestype=:scatter,color=:black)
    plot!(hdl, xlabel="Age ($(smpl.Age_Unit))", ylabel="Height ($(smpl.Height_Unit))")

## --- End of File


function minimum_area(p::poly_eight{3,Float64}, n̂::SVector{3,Float64}, r::SVector{3,Float64})
    isempty(p) && (return -Inf)
    min_area = Inf
    length_clip = length(p)
    for k = 1:length_clip
        area_k = triangle_area((p[k], p[mod1(k + 1, length_clip)], r), n̂)
        min_area = min(area_k, min_area)
    end
    return min_area
end

@testset "static clip" begin
    tol = 1.0e-13
    n_empty = 0
    n_hits = zeros(MVector{8,Int64})
    for k_tet = 1:5000000
        if n_hits[8] <= 2
            r_orig = make_4_sided()

            n̂ = triangleNormal(r_orig[1], r_orig[2], r_orig[3])
            n̂_check = triangleNormal(r_orig[2], r_orig[3], r_orig[4])
            (n̂ ≈ n̂_check) || error("quad points do not lie on a plane")
            plane = SMatrix{1,4,Float64,4}(n̂[1], n̂[2], n̂[3], -dot(n̂, r_orig[1]))

            tet = roll_non_degenerate_tet()
            A = asMatOnePad(tet[1], tet[2], tet[3], tet[4])
            inv_A = inv(A)
            ζ_orig = one_pad_then_mul(inv_A, r_orig)
            ζ_clip = clip_in_tet_coordinates(ζ_orig)
            r_clip = mul_then_un_pad(A, ζ_clip)

            for k = 1:length(r_clip)  # confirm that all vertices on plane
                # TODO: adjust the test to use the origional normal of the clipped shape so this tolerance
                # can be tightened
                @test abs(dist_from_plane(plane, r_clip[k])) < (2000 * tol)
            end

            for k = 1:30  # test 30 random points
                r_in_plane = project_into_plane(plane, randn(SVector{3,Float64}))
                min_ζ = minimum(inv_A * onePad(r_in_plane))
                min_area_orig = minimum_area(r_orig, n̂, r_in_plane)  # is_inside origional
                min_area_clip = minimum_area(r_clip, n̂, r_in_plane)  # is_inside clip
                if tol < min_area_clip  # inside clipping
                    point_inside_ζ = -tol < min_ζ
                    point_inside_orig = -tol < min_area_orig
                    @test point_inside_ζ && point_inside_orig
                else
                    point_outside_ζ = min_ζ < tol
                    point_outside_orig = min_area_orig < tol
                    @test point_outside_ζ || point_outside_orig
                end
            end

            if isempty(r_clip)
                n_empty += 1
            else
                n_hits[length(r_clip)] += 1
            end
        end
    end
    @test 2 <= n_hits[8]  # sanity check
    @test 1000 < n_empty  # sanity check
end


######################################################################
# circle.jl
######################################################################

export circle

"""
`circle(d::Domain, xc::VF, r::Float64)`

Returns a distribution that evenly weights points in a circle centered at `xc` with radius `r`. The distribution is normalized (so it integrates to 1) before being returned to the user.
"""
function circle(d::Domain, xc::Vector{Float64}, radius::Float64)
	phi = zeros(x_cells(d), y_cells(d))
	R2 = radius * radius

	for xi = 1:x_cells(d)
		x = x_min(d) + (xi-0.5)*x_size(d)
		for yi = 1:y_cells(d)
			y = y_min(d) + (yi-0.5)*y_size(d)
			dx = xc[1] - x
			dy = xc[2] - y
			if (dx*dx + dy*dy) < R2
				phi[xi,yi] = 1.0
			end
		end
	end
	normalize!(phi, d.cell_size)
	return phi
end


## return the right Cholesky factor from the LDLt decomposition of A⁻¹
function Tinvt{T}(p::Pedigree{T})
    n = length(p.sire)
    cpi = Array(Int32,n + 1)            # column pointers
    cpi[1] = one(Int32)
    rvi = sizehint(Int32[],n + n>>1)    # row values of T⁻¹
    for i in 1:n
        r1 = min(p.sire[i],p.dam[i])
        r2 = max(p.sire[i],p.dam[i])
        if r2 == 0
            cpi[i+1] = cpi[i]
        elseif r1 == 0
            cpi[i+1] = cpi[i] + one(Int32)
            push!(rvi,r2)
        else
            cpi[i+1] = cpi[i] + convert(Int32,2)
            push!(rvi,r1)
            push!(rvi,r2)
        end
    end
    Triangular(SparseMatrixCSC(n,n,cpi,rvi,fill(-0.5,(length(rvi,)))),:L,true)
end

## return the right Cholesky factor, Lt, from the LLt factorization of A
function Ltrans{T}(p::Pedigree{T})
    n = length(p.sire)
    tri = Tinvt(p).data
    cpi = tri.colptr
    rvi = tri.rowval
    cpl = sizehint(Int32[1],n+1)             # column pointers for L'
    rvl = sizehint(Int32[], 8n)              # row values for L'
    nvl = sizehint(Float64[], 8n)            # nonzero values in L'
    anc = sizehint(IntSet([]),n)             # set of ancestors
    for j in 1:n
        empty!(anc)
        l2parents = 0.0   # accumulator for sum of squares of L cols of parents
        for k in cpi[j]:(cpi[j+1]-1)    # each parent in the pedigree
            for i in cpl[rvi[k]]:(cpl[rvi[k]+1]-1)
                push!(anc,rvl[i])       # add the parents ancestors to the set
                l2parents += abs2(nvl[i])
            end
        end
        push!(anc,j)                    # the animal counts as one of its ancestors
        push!(cpl,cpl[end] + length(anc))
        rv = collect(anc)
        append!(rvl,rv)
        lrv = length(rv)
        cc = zeros(eltype(tri),(lrv,))
        cc[lrv] = sqrt(1.0 - 0.25*l2parents)
        for k in cpi[j]:(cpi[j+1]-1)    # each parent in the pedigree
            for i in cpl[rvi[k]]:(cpl[rvi[k]+1]-1)  # add 0.5*each parent's column
                cc[searchsortedfirst(rv,rvl[i])] += 0.5*nvl[i]
            end
        end
        append!(nvl,cc)
    end
    SparseMatrixCSC(n,n,cpl,rvl,nvl)
end

function inbreeding(p::Pedigree)
    Lt = Ltrans(p)
    inb = zeros(Lt.n)
    nvl = Lt.nzval
    cpt = Lt.colptr
    for i in 1:Lt.n
        for k in cpt[i]:(cpt[i+1]-1)
            inb[i] += abs2(nvl[k])
        end
        inb[i] -= one(eltype(inb))
    end
    inb,Lt
end


####### SIMULATION OF MARGINAL ###########################################

### +++ INITIALIZATION +++

### import packages
using Random, LinearAlgebra, RDatasets, MixedModels, Distributions, Optim, CSV, Dates

### define functions
vcov2 = function(x, xbar, σ, τ, n, γ)
    m = length(xbar)
    a = γ / τ^2 * m # 1-1 entry
    d = (sum(x.^2) - n * γ * sum(xbar.^2) ) / σ^2 # 2-2 entry
    b = γ / τ^2 / n * sum(x) # 2-1 entry
    [d -b; -b a] / (a * d - b^2) # return inverse (X^t V^-1 X)^-1
end

beta = function(vcov1, x, xbar, y, ybar, σ, τ, n, γ)
    c1 = γ / τ^2 * sum(ybar) # = 1^t V^-1 y
    c2 = (x' * y - n * γ * sum(ybar .* xbar))/σ^2 # = x^t V^-1 y
    vcov1 * [c1 c2]'  # (X^tV^-1X)^-1 X^tV^-1y
end

### fix parameters
p = 2
nsim = 10000
Δ = [0:0.1:2;]
Random.seed!(12345)

### +++ PER SPECIFICATION SETTING +++

### to compile
τ1 = 8^0.5
σ1 = 4^0.5
m = 10
n = 1

### variable parameters
N = m * n

β0 = [-4.9, 0.03] # rand(Uniform(0, 1), p)
x = rand(Uniform(0, 10), N)
xbar = vec(sum(reshape(x, n, :), dims=1) / n)

### +++ KNOWN DELTA +++

###     MARGINAL COVARIANCE MATRIX

### helpful parameters
γ0 = τ1^2 / (τ1^2 + σ1^2 / n)
vcov0 = vcov2(x, xbar, σ1, τ1, n, γ0) # = (X^tV^-1X)^-1
d0 = (1 - γ0) .* hcat(fill(1, m), xbar) # d' = l' - b'X
dd0 = d0' * d0

### assemble covariance matrix parts
g1 = γ0 * σ1^2 / n

### assemble inverse of covariance matrix with Woodbury
Σ0diag = g1
Σ0i = (I - d0 * inv(Σ0diag * inv(vcov0) + dd0) * d0' ) / Σ0diag

###     QUANTILES
qm = cquantile(Chisq(m), 0.05) # quantile of alpha = 0.05

### simulation
Coverage = zeros(2, length(Δ))
for i in 1:nsim

    ### +++ EMPIRICAL +++

    ### marginal evaluation
    v = τ1 * randn(m)
    u = β0[1] .+ β0[2] * x + repeat(v, inner = n)
    μ0 = β0[1] .+ β0[2] .* xbar + v

    ### variable data
    y = u + σ1 * randn(N) # conditional
    ybar = vec(sum(reshape(y, n, :), dims=1) / n)
    data = DataFrame(x = x, v = repeat(categorical(v), inner = n), y = y)

    ### fit model
    model = fit(MixedModel, @formula(y ~ x + (1 | v)), data)

    ### make sure the variance parameters are positive
    τ = maximum([VarCorr(model).σρ.v.σ[1], 0.001])
    σ = maximum([VarCorr(model).s, 0.001])

    ###     MARGINAL

    ### helpful parameters
    γ = τ^2 / (τ^2 + σ^2 / n)
    vcov1 = vcov2(x, xbar, σ, τ, n, γ) # = (X^tV^-1X)^-1
    d = (1-γ) .* hcat(fill(1, m), xbar) # d' = l' - b'X
    dd = d' * d
    a = n * τ^2 + σ^2

    ### assemble diagonal matrix parts
    g1 = γ * σ^2 / n
    g3 = 2 * σ^4 / τ^2 * γ / N/ (n - 1)

    ### assemble inverse of covariance matrix with Woodbury
    Σdiag = g1 + 2 * g3
    #Σ = Σdiag * I + d * vcov1 * d'
    Σi = (I - d * inv(Σdiag * inv(vcov1) + dd) * d' ) / Σdiag

    ###     COMPUTE MEAN
    β1 = beta(vcov0, x, xbar, y, ybar, σ1, τ1, n, γ0)
    μ1 = (1 - γ0) * (β1[1] .+ β1[2] .* xbar) + γ0 * ybar

    β2 = beta(vcov1, x, xbar, y, ybar, σ, τ, n, γ) # = fixef(model)
    μ2 = (1 - γ) * (β2[1] .+ β2[2] .* xbar) + γ * ybar

    ###     COMPUTE STATISTICS
    t1 = (μ1 - μ0) .+ Δ'
    t2 = (μ2 - μ0) .+ Δ'
    Coverage[1,:] += (sum(t1' .* (t1' * Σ0i), dims = 2) .< qm) / nsim
    Coverage[2,:] += (sum(t2' .* (t2' * Σi), dims = 2) .< qm) / nsim
end


module LogExpFunctions

using DocStringExtensions: SIGNATURES
using Base: Math.@horner

import ChainRulesCore
import ChangesOfVariables
import InverseFunctions
import IrrationalConstants
import LinearAlgebra

export xlogx, xlogy, xlog1py, xexpx, xexpy, logistic, logit, log1psq, log1pexp, log1mexp, log2mexp, logexpm1,
    softplus, invsoftplus, log1pmx, logmxp1, logaddexp, logsubexp, logsumexp, softmax,
    softmax!, logcosh

include("basicfuns.jl")
include("logsumexp.jl")
include("chainrules.jl")
include("inverse.jl")
include("with_logabsdet_jacobian.jl")

end # module


using LinearAlgebra
using Random
using Distributions
using Cubature
using Roots
using SpecialFunctions

include("kernels/kernel.jl")
include("kernels/mean_basis.jl")
include("transform/nonlinfun.jl")
include("quadrature/quadgrid.jl")
include("param/param_prior.jl")
include("computation/param_grid_comp.jl")
include("computation/post_prob_comp.jl")
include("computation/weight_comp.jl")
include("computation/distribution_comp.jl")

# Inputs:
    # x0: where to predict, testing data
    # training_data: including features/index/x and labels/values/z
    # kernel: a specific RBF kernel
    # nonlinfun: a specific nonlinear transformation, parameterized
    # p: the degree of freedom using in polynomial basis for the prior mean
    #    p = 1 by default (constant prior mean function)

# Outputs:
    # the distribution of function values z0, at a given index x0
    # including mean, median, standard deviation and condidence interval

function model(x0, training_data, kernel, nonlinfun, p = 1)
    # training_data
    x = training_data[1]
    z = training_data[2]

    # define nonlinear transformation
    # both function values and derivatives
    g =  (x, lambda) -> nonlinfun(x, lambda)[1]
    dg = (x, lambda) -> nonlinfun(x, lambda)[2]

    # define polynomial basis
    X = meanbasisMat(x, p)
    X0 = meanbasisMat(x0, p)

    # precompute the constant coefficient in p(z0|z)
    n = size(x, 1) # #training points
    k = size(x0, 1) # #predicting points
    Gamma = gamma((n-p+k)/2)/(gamma((n-p)/2) * pi^(k/2))

    # wrap up basic training info
    traindata = (idx = x, val = z)
    nonlintrans = (fun = g, deriv = dg)
    trainBasicInfo = (data = training_data, kernel = kernel,
                    polydof = p, nonlintrans = nonlintrans)

    ## compute weights
    # compute quadrature grids and weights
    n_quad = [8, 8, 8] # number of quad nodes for each param, could increase if needed
    n_theta = [n_quad[1], n_quad[2]]
    n_lambda = n_quad[3]
    n_param = (theta = n_theta, lambda = n_lambda)
    param_gridInfo = param_grid_comp(n_param, param_priorInfo)

    # compute alpha_ij = w_i w_j h_ij for theta_i and lambda_j
    alpha, sideSetInfo = weight_comp(param_gridInfo, param_priorInfo, trainBasicInfo)

    # compute distribution of z0 at x0
    distribution_z0 = distribution_comp(x0, trainBasicInfo, sideSetInfo, param_gridInfo, alpha, Gamma)

    return distribution_z0, alpha
end


using LinearAlgebra
using StaticArrays
using Plots; pyplot()
using FFTW
using FastClosures
using NPZ
using LaTeXStrings

include("/Users/tfiguei/Work/Magnus/MagnusExpansion.jl")
include("AuxFunctions.jl")
include("HarmonicCorrectionSnap.jl")
include("snap.jl")

σ = pauli_matrices()
χ_GHz = (-8281.0, 0.0, 0.0)#(-8281.0, 48.8, 0.5)
χ = χ_GHz./abs(χ_GHz[1])
N = 10
levels = SA[0, 4]#collect(0:1:9)
α = (pi/2, pi/3)#pi*(rand(10) .- 0.5)
tf = 50.0
Ncorr = 10

system_params = Dict("chi"=>χ, "N"=>N, "levels"=>levels, "alpha"=>α,
					 "tf"=>tf, "Ncorr"=>Ncorr)

UI = U(system_params, σ; ideal_dynamics=false, frame="interaction", matrix=true)
Identity = Matrix{Complex{Float64}}(I, 2*N, 2*N)
println(1 - fidelity(UI, Identity, 2*N))

Uideal = U(system_params, σ; ideal_dynamics=true, matrix=true)
Ur = U(system_params, σ; ideal_dynamics=false, matrix=true)
println(1 - fidelity(Ur, Uideal, 2*N))

# "First" order correction
magnus_order = 1
max_order = 2
coeffs = get_coeffs(magnus_order, max_order, Ncorr, system_params)
println("---")
Ucorr = U(system_params, σ, coeffs, max_order; frame="qubit", matrix=true)
println(1 - fidelity(Ucorr, Uideal, 2*N))
# "Second" order correction
magnus_order = 2
coeffs2 = get_coeffs(coeffs, magnus_order, max_order, Ncorr, system_params)
println("---")
coeffs2 = coeffs + coeffs2
Ucorr = U(system_params, σ, coeffs2, max_order; frame="qubit", matrix=true)
println(1 - fidelity(Ucorr, Uideal, 2*N))

# In the lab frame
println("Results in the lab frame")
Uideal = U(system_params, σ; frame="lab", ideal_dynamics=true, matrix=true)
Unideal = U(system_params, σ; frame="lab", ideal_dynamics=false, matrix=true)
Ucorr = U(system_params, σ, coeffs2, max_order; frame="lab", matrix=true)
println(1 - fidelity(Unideal, Uideal, 2*N))
println(1 - fidelity(Ucorr, Uideal, 2*N))

#FFT results
Npoints = 1000
data = zeros(Npoints, 5)
for i = 1: Npoints
	t = (i-1)*tf/(Npoints-1)
	data[i,1] = t
	data[i, 2: 3] .= pulse_envelope(t, tf, χ, levels, α)
	data[i, 4: 5] .= data[i, 2: 3]
	data[i, 4: 5] .+= correction_pulse_envelope(t, tf, χ, coeffs2, max_order, Ncorr)
end

x = data[:, 1]
y = data[:, [2, 4]]
p1 = plot(x, y,
	 linewidth = 2,
	 labels = ["Initial\npulse" "Corrected\npulse"],
	 layout = (2, 1),
	 title = ["x quadrature" " "])
plot!(xlim=(0, tf),
	  tickfont = font(12, "Serif"),
	  xlabel = "time "*L"t/t_\mathrm{f}",
	  ylabel = "pulse "*L"\times |\chi|^{-1}",
      guidefont = font(14, "Serif"),
	  legendfont = font(12, "Serif"),
	  margin = 5Plots.mm)

x = data[:, 1]
y = data[:, [3, 5]]
p2 = plot(x, y,
	 linewidth = 2,
	 labels = ["Initial\npulse" "Corrected\npulse"],
	 layout = (2, 1),
	 title = ["y quadrature" " "])
plot!(xlim=(0, tf),
	  tickfont = font(12, "Serif"),
	  xlabel = "time "*L"t/t_\mathrm{f}",
	  ylabel = "pulse "*L"\times |\chi|^{-1}",
	  guidefont = font(14, "Serif"),
	  legendfont = font(12, "Serif"),
	  background_color_legend = false,
	  margin = 5Plots.mm)
plot(p1, p2, layout=(1,2), size = (800, 800))
savefig("pulse.pdf")
# Add padding to the pulse
data = vcat(zeros(Npoints, 5), data, zeros(Npoints, 5))

fft_data = zeros(size(data))
fft_data[:, 1] = fftfreq(size(data)[1], tf/(Npoints-1))
for i = 2: 5
	fft_pulse = fft(data[:, i])
	fft_data[:, i] = abs.(fft_pulse)
end

a = div(size(data)[1], 2)
x = fft_data[1: a, 1]
y = fft_data[1: a, [2, 4]]
p1 = plot(x, y,
		  linewidth = 2,
		  labels = ["Initial\npulse" "Corrected\npulse"],
		  linestyle = [:solid :dash])
plot!(xlim=(0, 10),
	  tickfont = font(12, "Serif"),
	  xlabel = "Frequency "*L"\omega/|\chi|",
	  ylabel = L"|S_x(x)|",
	  guidefont = font(14, "Serif"),
	  legendfont = font(12, "Serif"),
	  background_color_legend = false,
	  margin = 5Plots.mm)

x = fft_data[1: a, 1]
y = fft_data[1: a, [3, 5]]
p2 = plot(x, y,
	 linewidth = 2,
	 labels = ["Initial\npulse" "Corrected\npulse"],
	 linestyle = [:solid :dash])
plot!(xlim=(0, 10),
	  tickfont = font(12, "Serif"),
	  xlabel = "Frequency "*L"\omega/|\chi|",
	  ylabel = L"|S_y(x)|",
	  guidefont = font(14, "Serif"),
	  legendfont = font(12, "Serif"),
	  background_color_legend = false,
	  margin = 5Plots.mm)

plot(p1, p2, layout=(1,2), size = (800, 400))

npzwrite("fft.npy", fft_data[1: div(Npoints,2), :])
savefig("fft.pdf")


using Test
using FeatureEng

using Dates
using DataFrames

println("Running tests...")
println("Meep. Morp.")

@testset "Numeric" begin

    @testset "Scaling" begin
        data = [ 1:100; ]
        scale = fit_transform!(ScaleMinMax(),data)
        transformed = apply_transform(scale,data)
        @test minimum(transformed) == 0.0
        @test maximum(transformed) == 1.0

        data = [ 0, 5, 10 ]
        scale = fit_transform!(ScaleVariance(),data)
        transformed = apply_transform(scale,data)
        @test transformed == [-1.,0.,1.]

        data = [ 0, 0, 2^.5, 2^.5 ]
        scale = fit_transform!(ScaleL2(),data)
        transformed = apply_transform(scale,data)
        @test transformed == data ./ sum(data .^ 2) ^ .5

    end

    @testset "Transforming" begin
        @testset "Log Transform" begin
            data = ℯ .^ [ 1:10; ]
            @test isapprox(transformLog(data), [1:10;])

            for base = [ℯ,2,10]
                data = base .^ [ 1:10; ]
                @test isapprox(transformLog(data,base), [1:10;])
            end
        end

        @testset "Root Transform" begin
            for base = [ℯ,2,10]
                data = [1:10;] .^ base
                @test isapprox(transformRoot(data,base), [1:10;])
            end
        end

        @testset "BoxCox Transform" begin
            data = [ 1:10; ]

            @test isapprox(transformBoxCox(data,0.0), log.(data))
            @test isapprox(transformBoxCox(data,1.0), [ 0:9; ])

        end
    end

    @testset "Binning" begin
    end

    @testset "Interaction" begin
        df = DataFrame(
            x = 1:10,
            y = repeat([0,1],5)
        )
        col_names = ["x","x_x","x_y","y","y_y"]
        res = polynomial(df)
    end
end

@testset "Categorical" begin
    @testset "One-Hot Encoding" begin
        data = [3,1,2,4]
        @test encode_onehot(data) == DataFrame(Dict(
            "1" => [0,1,0,0],
            "2" => [0,0,1,0],
            "3" => [1,0,0,0],
            "4" => [0,0,0,1]
        ))
        @test encode_onehot(data,"col_") == DataFrame(Dict(
            "col_1" => [0,1,0,0],
            "col_2" => [0,0,1,0],
            "col_3" => [1,0,0,0],
            "col_4" => [0,0,0,1]
        ))
        @test encode_onehot(data,[1:6;]) == DataFrame(Dict(
            "1" => [0,1,0,0],
            "2" => [0,0,1,0],
            "3" => [1,0,0,0],
            "4" => [0,0,0,1],
            "5" => [0,0,0,0],
            "6" => [0,0,0,0]
        ))
        @test encode_onehot(data,[1:6;],"col_") == DataFrame(Dict(
            "col_1" => [0,1,0,0],
            "col_2" => [0,0,1,0],
            "col_3" => [1,0,0,0],
            "col_4" => [0,0,0,1],
            "col_5" => [0,0,0,0],
            "col_6" => [0,0,0,0],
        ))
        @test encode_onehot(["dog","cat","dog"]) == DataFrame(Dict(
            "cat" => [0,1,0],
            "dog" => [1,0,1]
        ))
    end

    @testset "Dummy Encoding" begin
    end

    @testset "Hash Encoding" begin
        data = [1:100:1000;]
        @test size(encode_hash(data,4)) == (10,4)
        @test encode_hash(data,4) == encode_hash(data,4)
    end
end

@testset "Datetime" begin
    @testset "Parse DateTime Strings" begin
        date_strings = [
            "2021-01-27 14:03:25",
            "1999-10-05 01:13:43",
            "abcdefg"
        ]
        datetimes = [
            DateTime(2021,1,27,14,03,25),
            DateTime(1999,10,5,1,13,43),
            missing
        ]
        @test all(a === b for (a,b) = zip(
            strp_datetimes(date_strings), datetimes))
    end
    @testset "Get Date Features" begin
        data = [
            DateTime(2021,1,27,14,3,25),
            DateTime(1999,10,5,1,13,43),
            DateTime(2010,6,12,11,0,0)
        ]
        @test extract_date_features(data) == DataFrame(
            year=[2021,1999,2010],
            month=["January","October","June"],
            dayofmonth=[27,5,12],
            dayofweek=["Wednesday","Tuesday","Saturday"],
            isweekend=[0,0,1],
            quarter=[1,4,2]
        )
    end
    @testset "Get Time Features" begin
        data = [
            DateTime(2021,1,27,14,3,25),
            DateTime(1999,10,5,1,13,43),
            DateTime(2010,6,12,11,0,0)
        ]
        @test extract_time_features(data) == DataFrame(
            hour=[14,1,11],
            minute=[3,13,0],
            second=[25.,43.,0.],
            isAM=[0,1,1]
        )
    end
    @testset "Get Date Features" begin
        data = [
            DateTime(2021,1,27,14,3,25),
            DateTime(1999,10,5,1,13,43),
            DateTime(2010,6,12,11,0,0)
        ]
        @test extract_datetime_features(data) == DataFrame(
            year=[2021,1999,2010],
            month=["January","October","June"],
            dayofmonth=[27,5,12],
            dayofweek=["Wednesday","Tuesday","Saturday"],
            isweekend=[0,0,1],
            quarter=[1,4,2],
            hour=[14,1,11],
            minute=[3,13,0],
            second=[25.,43.,0.],
            isAM=[0,1,1]
        )
    end
end


type HiddenMarkovModel <: DynamicDiscreteModel
	#HMM specific field

	#DynamicDiscreteModel fields
	m::Array{Float64,4}			  	#the transition matrix given as m[x,y,x',y']
	mu::Array{Float64,2}  			#initial distribution (dx,dy)
	rho::Array{Float64,1}
	phi::Array{Float64,1}
	psi::Array{Float64,1}
end

function hmm(ab::Tuple,mu)
	a,b=ab
	dx,dy=size(mu)
	model=HiddenMarkovModel(Array(Float64,dx,dy,dx,dy),mu,Array(Float64,1),Array(Float64,dx),Array(Float64,dx))
	hmm2ddm!(model,a,b)
	model
end

#default initial distribution uniform
function hmm(ab::Tuple)
	dx,dy=size(ab[2])
	mu=fill(1/(dx*dy),dx,dy)
	hmm(ab,mu)
end

#ab=(a,b) parametrization
coef!(model::HiddenMarkovModel,ab::Tuple)=hmm2ddm!(model,ab[1],ab[2])

#canonical parametrization in case someone wants to apply mle() to it
coef!(model::HiddenMarkovModel,theta::Array{Float64,1})=coef!(model,theta2ab(theta,size(model.mu,1)))

theta2ab(theta::Array{Float64,1},dx::Int)=(z2q(theta[1:dx*(dx-1)]),z2q(theta[dx*(dx-1)+1:end],dx))

function dim(model::HiddenMarkovModel)
	dx,dy=size(model.mu)
	dx*(dx-1)+dx*(dy-1)
end

function em(model::HiddenMarkovModel,data)
	dx,dy=size(model.mu)
	w=zeros(dx,dy,dx,dy)
	estep(model,data,w)
	a=nsm(reshape(sum(w,(2,4)),dx,dx))
	b=nsm(reshape(sum(w,(1,2)),dx,dy))
	llk=loglikelihood(model,data,(a,b))
	println()
	println(" estimating hidden Markov model via Baum-Welch algorithm...")
	println("  log-likelihood: $(round(llk,4))")
	(a,b)
end

baumwelch=em