heat_eq_dir = joinpath(output_root, "HeatEquation")

struct DifussionImplicit end
function TridiagonalMatrix(
    q::StateVec,
    aux::StateVec,
    Δt::T,
    ρ::S,
    a::S,
    K::S,
    grid::Grid{T},
    ::DifussionImplicit,
    i::Int,
) where {S <: Symbol, T}
    #! format: off
    a_ = [   -Δt * grid.Δzi2 *      (aux[ρ, Dual(k)] .* q[a, Dual(k), i] .* aux[K, Dual(k), i])[2] for k in over_elems_real(grid)[1:(end - 1)]]
    b_ = [1 + Δt * grid.Δzi2 * (sum((aux[ρ, Dual(k)] .* q[a, Dual(k), i] .* aux[K, Dual(k), i]))) for k in over_elems_real(grid)]
    c_ = [   -Δt * grid.Δzi2 *      (aux[ρ, Dual(k)] .* q[a, Dual(k), i] .* aux[K, Dual(k), i])[1] for k in over_elems_real(grid)[2:end]]
    #! format: on
    return Tridiagonal(a_, b_, c_)
end

"""
    solve_tdma!(q::StateVec, tendencies::StateVec, aux::StateVec, x::S, ρ::S, a_τp1::S, a_τ::S, K::S, grid::Grid{T}, Δt::T, i::Int=1)

Solve tri-diagonal system using Diffusion-implicit time-marching.
"""
function solve_tdma!(
    q::StateVec,
    tendencies::StateVec,
    aux::StateVec,
    x::S,
    ρ::S,
    a_τp1::S,
    a_τ::S,
    K::S,
    grid::Grid{T},
    Δt::T,
    i::Int = 1,
) where {S <: Symbol, T}
    B = [
        aux[a_τ, k, i] / q[a_τp1, k, i] * q[x, k, i] +
        Δt * tendencies[x, k, i] / (aux[ρ, k] * q[a_τp1, k, i])
        for k in over_elems_real(grid)
    ]
    A = TridiagonalMatrix(q, aux, Δt, ρ, a_τp1, K, grid, DifussionImplicit(), i)
    X = inv(A) * B
    assign_real!(q, x, grid, X, i)
end

@testset "∂_t T = K ΔT + 1, T = 0 ∈ ∂Ω, diffusion-explicit" begin
    domain_set = DomainSet(gm = 1)
    domain_subset = DomainSubSet(gm = true)
    K, maxiter, Δt = 1.0, 1000, 0.005
    grid = UniformGrid(0.0, 1.0, 10)
    q = StateVec(((:T, domain_subset),), grid, domain_set)
    aux = StateVec(((:ΔT, domain_subset),), grid, domain_set)
    rhs = deepcopy(q)
    for i in 1:maxiter
        for k in over_elems_real(grid)
            aux[:ΔT, k] = Δ_z(q[:T, Cut(k)], grid)
            rhs[:T, k] = K * aux[:ΔT, k] + 1
        end
        for k in over_elems(grid)
            q[:T, k] += Δt * rhs[:T, k]
        end
        apply_Dirichlet!(q, :T, grid, 0.0, Zmax())
        apply_Dirichlet!(q, :T, grid, 0.0, Zmin())
    end
    sol_analtyic = 1 / 2 * grid.zc - grid.zc .^ 2 / 2
    sol_error = [abs(q[:T, k] - sol_analtyic[k]) for k in over_elems_real(grid)]
    @test all(sol_error .< grid.Δz)
    # Skip plotting...
    # mkpath(heat_eq_dir)
    # nc = NetCDFWriter(joinpath(heat_eq_dir, "T"))
    # export_state(nc, grid, q)
    # rm(full_name(nc))
end

@testset "∂_t T = K ΔT, T={1, 0} ∈ {z_{min}, z_{max}}, diffusion-explicit" begin
    domain_set = DomainSet(gm = 1)
    domain_subset = DomainSubSet(gm = true)
    K, maxiter, Δt = 1.0, 1000, 0.002
    grid = UniformGrid(0.0, 1.0, 10)
    q = StateVec(((:T, domain_subset),), grid, domain_set)
    aux = StateVec(((:ΔT, domain_subset),), grid, domain_set)
    rhs = deepcopy(q)
    for i in 1:maxiter
        for k in over_elems_real(grid)
            aux[:ΔT, k] = Δ_z(q[:T, Cut(k)], grid)
            rhs[:T, k] = K * aux[:ΔT, k]
        end
        for k in over_elems(grid)
            q[:T, k] += Δt * rhs[:T, k]
        end
        apply_Dirichlet!(q, :T, grid, 1.0, Zmin())
        apply_Dirichlet!(q, :T, grid, 0.0, Zmax())
    end
    sol_analtyic = 1 .- grid.zc
    sol_error = [abs(q[:T, k] - sol_analtyic[k]) for k in over_elems_real(grid)]
    @test all(sol_error .< grid.Δz)
end

@testset "∂_t T = K ΔT, T = {1, 0} ∈ {z_{min}, z_{max}}, diffusion implicit" begin
    domain_set = DomainSet(gm = 1)
    domain_subset = DomainSubSet(gm = true)
    K, maxiter, Δt = 1.0, 1000, 0.1
    grid = UniformGrid(0.0, 1.0, 10)
    k_star1 = first_interior(grid, Zmin())
    k_star2 = first_interior(grid, Zmax())
    vars = ((:T, domain_subset), (:a, domain_subset))
    q = StateVec(vars, grid, domain_set)
    vars = (
        (:ΔT, domain_subset),
        (:a_tau, domain_subset),
        (:ρ_0, domain_subset),
        (:K, domain_subset),
        (:a, domain_subset),
    )
    aux = StateVec(vars, grid, domain_set)
    rhs = deepcopy(q)
    for k in over_elems(grid)
        q[:a, k] = 1.0
        aux[:ρ_0, k] = 1.0
        aux[:a_tau, k] = 1.0
        aux[:K, k] = 1.0
    end
    for i in 1:maxiter
        for k in over_elems_real(grid)
            rhs[:T, k] = 0.0
            if k == k_star1
                gv = 2 * 1 - q[:T, k]
                grad_T = ∇_z_flux([gv, 0], grid)
                bc_surf = ∇_z_flux([grad_T, 0], grid)
                rhs[:T, k] += bc_surf
            end
            if k == k_star2
                gv = -q[:T, k]
                grad_T = ∇_z_flux([0, gv], grid)
                bc_surf = ∇_z_flux([0, grad_T], grid)
                rhs[:T, k] += bc_surf
            end
        end
        solve_tdma!(q, rhs, aux, :T, :ρ_0, :a, :a_tau, :K, grid, Δt)
        extrap!(q, :T, grid)
    end
    sol_analtyic = 1 .- grid.zc
    sol_numerical = [q[:T, k] for k in over_elems_real(grid)]
    sol_error = [abs(q[:T, k] - sol_analtyic[k]) for k in over_elems_real(grid)]
    @test all(sol_error .< grid.Δz)
end

@testset "∂_t T = K ΔT, -K∇T = F ∈ z_{min}, T = 0 ∈ z_{max}, explicit Euler" begin
    domain_set = DomainSet(gm = 1)
    domain_subset = DomainSubSet(gm = true)
    K, maxiter, Δt = 1.0, 10, 0.001, 10
    grid = UniformGrid(0.0, 1.0, 10)
    k_star1 = first_interior(grid, Zmin())
    k_star2 = first_interior(grid, Zmax())
    q = StateVec(((:T1, domain_subset), (:T2, domain_subset)), grid, domain_set)
    aux = StateVec(
        ((:ΔT1, domain_subset), (:ΔT2, domain_subset)),
        grid,
        domain_set,
    )
    rhs = deepcopy(q)
    for i in 1:maxiter
        q_out = -K * 1
        q_in = -q_out
        for k in over_elems_real(grid)
            aux[:ΔT1, k] = Δ_z(q[:T1, Cut(k)], grid)
            aux[:ΔT2, k] = Δ_z(q[:T2, Cut(k)], grid)
            rhs[:T1, k] = K * aux[:ΔT1, k]
            rhs[:T2, k] = K * aux[:ΔT2, k]
            if k == k_star1 && i > 1
                rhs[:T2, k] += -∇_z_flux([q_in, 0.0], grid)
            end
        end
        for k in over_elems(grid)
            q[:T1, k] += Δt * rhs[:T1, k]
            q[:T2, k] += Δt * rhs[:T2, k]
        end
        apply_Neumann!(q, :T1, grid, q_out, Zmin())
        apply_Dirichlet!(q, :T1, grid, 0.0, Zmax())
        apply_Neumann!(q, :T2, grid, 0.0, Zmin())
        apply_Dirichlet!(q, :T2, grid, 0.0, Zmax())
        @test all([q[:T1, k] ≈ q[:T2, k] for k in over_elems_real(grid)])
    end
end

@testset "∂_t T = ∇ • (K(z) ∇T) + 1, T = 0 ∈ ∂Ω, explicit Euler" begin
    domain_set = DomainSet(gm = 1)
    domain_subset = DomainSubSet(gm = true)
    K, maxiter, Δt = 1.0, 1000, 0.001
    grid = UniformGrid(0.0, 1.0, 10)
    unknowns = ((:T, domain_subset),)
    vars = ((:ΔT, domain_subset), (:K_thermal, domain_subset))
    q = StateVec(unknowns, grid, domain_set)
    aux = StateVec(vars, grid, domain_set)
    rhs = deepcopy(q)
    cond_thermal(z) = z > 0.5 ? 1 : 0.1
    for k in over_elems(grid)
        aux[:K_thermal, k] = cond_thermal(grid.zc[k])
    end
    for i in 1:maxiter
        for k in over_elems_real(grid)
            aux[:ΔT, k] = Δ_z(q[:T, Cut(k)], grid, aux[:K_thermal, Cut(k)])
            rhs[:T, k] = K * aux[:ΔT, k] + 1
        end
        for k in over_elems(grid)
            q[:T, k] += Δt * rhs[:T, k]
        end
        apply_Dirichlet!(q, :T, grid, 0.0, Zmax())
        apply_Dirichlet!(q, :T, grid, 0.0, Zmin())
    end
end

@testset "∂_t T = ∇ • (K(z) ∇T) + 1, T = {0, 0} ∈ {z_{min}, z_{max}}, comparison" begin
    domain_set = DomainSet(gm = 1)
    domain_subset = DomainSubSet(gm = true)
    VS, TS, F, K = 1.0, 1.0, 1.0, 1.0
    maxiter, Δt = 5000, 0.005
    grid = UniformGrid(0.0, 1.0, 10)
    k_star1 = first_interior(grid, Zmin())
    k_star2 = first_interior(grid, Zmax())
    vars = (
        (:T_explicit_surf, domain_subset),
        (:T_explicit_vol, domain_subset),
        (:T_implicit, domain_subset),
        (:a, domain_subset),
    )
    q = StateVec(vars, grid, domain_set)
    vars = (
        (:ΔT, domain_subset),
        (:a_tau, domain_subset),
        (:ρ_0, domain_subset),
        (:K_thermal, domain_subset),
        (:a, domain_subset),
    )
    aux = StateVec(vars, grid, domain_set)
    rhs = deepcopy(q)
    for k in over_elems(grid)
        q[:a, k] = 1.0
        aux[:ρ_0, k] = 1.0
        aux[:a_tau, k] = 1.0
    end
    cond_thermal(z) = z > 0.5 ? 1 : 0.1
    for k in over_elems(grid)
        aux[:K_thermal, k] = cond_thermal(grid.zc[k])
    end
    for i in 1:maxiter
        for k in over_elems_real(grid)
            aux[:ΔT, k] =
                Δ_z(q[:T_explicit_surf, Cut(k)], grid, aux[:K_thermal, Cut(k)])
            rhs[:T_explicit_surf, k] = aux[:ΔT, k] + VS
            aux[:ΔT, k] =
                Δ_z(q[:T_explicit_vol, Cut(k)], grid, aux[:K_thermal, Cut(k)])
            rhs[:T_explicit_vol, k] = aux[:ΔT, k] + VS
            rhs[:T_implicit, k] = VS
            k == k_star1 && (
                rhs[:T_implicit, k] += bc_source(
                    q,
                    grid,
                    aux,
                    :T_implicit,
                    :ρ_0,
                    :a,
                    :K_thermal,
                    Zmin(),
                    TS,
                    Dirichlet(),
                    DiffusionAbsorbed(),
                )
            )
            k == k_star2 && (
                rhs[:T_implicit, k] += bc_source(
                    q,
                    grid,
                    aux,
                    :T_implicit,
                    :ρ_0,
                    :a,
                    :K_thermal,
                    Zmax(),
                    TS,
                    Dirichlet(),
                    DiffusionAbsorbed(),
                )
            )
            k == k_star1 && (
                rhs[:T_explicit_vol, k] += bc_source(
                    q,
                    grid,
                    aux,
                    :T_explicit_vol,
                    :ρ_0,
                    :a,
                    :K_thermal,
                    Zmin(),
                    TS,
                    Dirichlet(),
                    DiffusionAbsorbed(),
                )
            )
            k == k_star2 && (
                rhs[:T_explicit_vol, k] += bc_source(
                    q,
                    grid,
                    aux,
                    :T_explicit_vol,
                    :ρ_0,
                    :a,
                    :K_thermal,
                    Zmax(),
                    TS,
                    Dirichlet(),
                    DiffusionAbsorbed(),
                )
            )
        end

        for k in over_elems_real(grid)
            q[:T_explicit_surf, k] += Δt * rhs[:T_explicit_surf, k]
            q[:T_explicit_vol, k] += Δt * rhs[:T_explicit_vol, k]
        end
        solve_tdma!(
            q,
            rhs,
            aux,
            :T_implicit,
            :ρ_0,
            :a,
            :a_tau,
            :K_thermal,
            grid,
            Δt,
        )

        apply_Dirichlet!(q, :T_explicit_surf, grid, TS, Zmax())
        apply_Dirichlet!(q, :T_explicit_surf, grid, TS, Zmin())

        assign_ghost!(q, :T_explicit_vol, grid, 0.0, Zmin())
        assign_ghost!(q, :T_explicit_vol, grid, 0.0, Zmax())

        assign_ghost!(q, :T_implicit, grid, 0.0, Zmin())
        assign_ghost!(q, :T_implicit, grid, 0.0, Zmax())
    end
    extrap!(q, :T_explicit_vol, grid, TS, Zmin())
    extrap!(q, :T_explicit_vol, grid, TS, Zmax())
    extrap!(q, :T_implicit, grid, TS, Zmin())
    extrap!(q, :T_implicit, grid, TS, Zmax())

    @test q[:T_explicit_vol, Dual(k_star1)][1] ≈ TS
    @test q[:T_explicit_vol, Dual(k_star2)][2] ≈ TS
    @test all([
        q[:T_explicit_vol, k] ≈ q[:T_explicit_surf, k]
        for k in over_elems_real(grid)
    ])
    @test all([
        q[:T_implicit, k] ≈ q[:T_explicit_surf, k]
        for k in over_elems_real(grid)
    ])

end

@testset "∂_t T = ∇ • (K(z) ∇T) + VS, -K(z) ∇T = 1 ∈ z_{min}, T = 0 ∈ z_{max}, comparison" begin
    domain_set = DomainSet(gm = 1)
    domain_subset = DomainSubSet(gm = true)
    VS, TS, F = 1.0, 0.0, 2.0
    IC = 0.0
    maxiter, Δt = 10000, 0.0001
    grid = UniformGrid(0.0, 1.0, 20)
    k_star1 = first_interior(grid, Zmin())
    k_star2 = first_interior(grid, Zmax())
    vars = (
        (:T_explicit_surf, domain_subset),
        (:T_explicit_vol, domain_subset),
        (:T_implicit, domain_subset),
        (:a, domain_subset),
    )
    q = StateVec(vars, grid, domain_set)
    vars = (
        (:ΔT, domain_subset),
        (:a_tau, domain_subset),
        (:ρ_0, domain_subset),
        (:K_thermal, domain_subset),
    )
    aux = StateVec(vars, grid, domain_set)
    rhs = deepcopy(q)
    for k in over_elems(grid)
        q[:a, k] = 1.0
        aux[:ρ_0, k] = 1.0
        aux[:a_tau, k] = 1.0
    end
    cond_thermal(z) = 5 + 3 * cos(6 * π * z)
    for k in over_elems(grid)
        aux[:K_thermal, k] = cond_thermal(grid.zc[k])
        q[:T_explicit_surf, k] = IC
        q[:T_explicit_vol, k] = IC
        q[:T_implicit, k] = IC
    end
    K_surf = aux[:K_thermal, Dual(k_star1)][1]
    apply_Neumann!(q, :T_explicit_surf, grid, F / K_surf, Zmin())
    apply_Dirichlet!(q, :T_explicit_surf, grid, TS, Zmax())
    for i in 1:maxiter
        assign_ghost!(q, :T_explicit_vol, grid, 0.0, Zmin())
        assign_ghost!(q, :T_explicit_vol, grid, 0.0, Zmax())
        for k in over_elems_real(grid)
            aux[:ΔT, k] = Δ_z_dual(
                q[:T_explicit_surf, Cut(k)],
                grid,
                aux[:K_thermal, Dual(k)],
            )
            rhs[:T_explicit_surf, k] = aux[:ΔT, k] + VS

            aux[:ΔT, k] = Δ_z_dual(
                q[:T_explicit_vol, Cut(k)],
                grid,
                aux[:K_thermal, Dual(k)],
            )
            rhs[:T_explicit_vol, k] = aux[:ΔT, k] + VS
            k == k_star1 && (
                rhs[:T_explicit_vol, k] += bc_source(
                    q,
                    grid,
                    aux,
                    :T_explicit_vol,
                    :ρ_0,
                    :a,
                    :K_thermal,
                    Zmin(),
                    F,
                    Neumann(),
                    DiffusionAbsorbed(),
                )
            )
            k == k_star2 && (
                rhs[:T_explicit_vol, k] += bc_source(
                    q,
                    grid,
                    aux,
                    :T_explicit_vol,
                    :ρ_0,
                    :a,
                    :K_thermal,
                    Zmax(),
                    TS,
                    Dirichlet(),
                    DiffusionAbsorbed(),
                )
            )

            rhs[:T_implicit, k] = VS

            k == k_star1 && (
                rhs[:T_implicit, k] += bc_source(
                    q,
                    grid,
                    aux,
                    :T_implicit,
                    :ρ_0,
                    :a,
                    :K_thermal,
                    Zmin(),
                    F,
                    Neumann(),
                    DiffusionAbsorbed(),
                )
            )
            k == k_star2 && (
                rhs[:T_implicit, k] += bc_source(
                    q,
                    grid,
                    aux,
                    :T_implicit,
                    :ρ_0,
                    :a,
                    :K_thermal,
                    Zmax(),
                    TS,
                    Dirichlet(),
                    DiffusionAbsorbed(),
                )
            )
        end

        for k in over_elems_real(grid)
            q[:T_explicit_surf, k] += Δt * rhs[:T_explicit_surf, k]
            q[:T_explicit_vol, k] += Δt * rhs[:T_explicit_vol, k]
        end
        solve_tdma!(
            q,
            rhs,
            aux,
            :T_implicit,
            :ρ_0,
            :a,
            :a_tau,
            :K_thermal,
            grid,
            Δt,
        )

        apply_Neumann!(q, :T_explicit_surf, grid, F / K_surf, Zmin())
        apply_Dirichlet!(q, :T_explicit_surf, grid, TS, Zmax())

        apply_Neumann!(q, :T_explicit_vol, grid, F / K_surf, Zmin())
        apply_Dirichlet!(q, :T_explicit_vol, grid, TS, Zmax())

        apply_Neumann!(q, :T_implicit, grid, F / K_surf, Zmin())
        apply_Dirichlet!(q, :T_implicit, grid, TS, Zmax())
    end

    apply_Neumann!(q, :T_explicit_vol, grid, F / K_surf, Zmin())
    extrap!(q, :T_explicit_vol, grid, TS, Zmax())

    apply_Neumann!(q, :T_implicit, grid, F / K_surf, Zmin())
    extrap!(q, :T_implicit, grid, TS, Zmax())

    @test q[:T_explicit_vol, Dual(k_star2)][2] ≈ TS
    @test q[:T_implicit, Dual(k_star2)][2] ≈ TS

    @test all([
        abs(q[:T_explicit_vol, k] - q[:T_explicit_surf, k]) < 10 * Δt
        for k in over_elems(grid)
    ])
    @test all([
        abs(q[:T_explicit_vol, k] - q[:T_explicit_surf, k]) < 10 * Δt
        for k in over_elems(grid)
    ])
    @test all([
        abs(q[:T_implicit, k] - q[:T_explicit_surf, k]) < 10 * Δt
        for k in over_elems(grid)
    ])
end


using CairoMakie, Random, Dates, TimeSeries, MarketData, StatsBase
tempo = string.(timestamp(ohlc))
lentime = length(tempo)
slice_dates = range(1, lentime, step = 3)
colors = values(ohlc.Close) .> values(ohlc.Open)
lowV = values(ohlc.Low)
highV = values(ohlc.High)
linesegs = []
for i in 1:lentime
    push!(linesegs, Point2f0(i, lowV[i]))
    push!(linesegs, Point2f0(i, highV[i]))
end
linesegs = Point2f0.(linesegs)
cmap = ((:orangered, 0.85), (:dodgerblue, 0.85))

set_theme!(theme_dark())
fig = Figure(figure_padding = (25,15,60,15), resolution = (600, 400))
ax = Axis(fig[1,1]; xgridstyle=:dash, ygridstyle=:dash)
barplot!(ax, 1:lentime, values(ohlc.Open), fillto = values(ohlc.Close),
    width = 0.65, color = colors, strokewidth = 0.5, strokecolor = colors,
    colormap = cmap)
linesegments!(ax, linesegs; linewidth = 3, color = colors, colormap = cmap)
lines!(5:500, values(moving(mean, ohlc.Close, 5)); color = :yellow)
xlims!(ax, 194.5, 225.5)
ylims!(ax, 17,24)
ax.xticks = (slice_dates, tempo[slice_dates])
ax.xticklabelrotation = π/4
save("condleStick.png", fig)
fig


#
# This file just illustrates the nuts-and-bolts of the approximate inference API provided
# by Stheno. It's worth mentally contrasting this with what you would have written to
# perform exact inference (hint: it's remarkably similar).
#

# Set up the environment to run this example. Make sure you're within the folder that this
# file lives in.
using Pkg
Pkg.activate(@__DIR__)
Pkg.instantiate()

# The time-to-first-plot issue means that this might take a little while.
using AbstractGPs, BenchmarkTools, LinearAlgebra, Plots, Random, Stheno
rng = MersenneTwister(123456);

# Construct a Matern-5/2 kernel with lengthscale 0.5 and variance 1.2.
k = 1.2 * transform(Matern52Kernel(), 2.0);

# Construct a zero-mean GP with a simple kernel. Don't worry about the GPC object.
f = Stheno.wrap(GP(k), GPC());

# Specify some locations at which to consider the GP.
N = 5_000
x = vcat(rand(rng, div(N, 2)) * 3, rand(rng, div(N, 2)) * 4 .+ 6)

# Specify the variance of the noise under which we'll make observations of the GP.
# We could also have made this a `Real` to specify isotropic noise.
Σ = Diagonal(rand(rng, N) .+ 0.1)

# Construct marginal distribution over `f` at `x` added to some independent zero-mean
# Gaussian noise with covariance matrix `Σ`.
fx = f(x, Σ)

# Generate a sample from the prior.
y = rand(rng, fx)

# Now pretend that we can't perform inference because `N` is too large.

# Specify some pseudo-point locations.
z = range(-3.0, 13.0; length=150)

# Specify inducing-points. Add a small amount of jitter for numerical stability.
u = f(z, 1e-6)

# Compute the Evidence Lower BOund.
println("The elbo is reasonable tight.")
@show elbo(fx, y, u), logpdf(fx, y)

println("Benchmark logpdf")
display(@benchmark logpdf($fx, $y))
println()

println("Benchmark elbo")
display(@benchmark elbo($fx, $y, $u))
println()

# Compute the approximate posterior process.
f_post = approx_posterior(VFE(), fx, y, u);

# Specify some points at which to plot the approximate posterior.
Npr = 1000
xpr = range(-3.0, 13.0; length=Npr)

# It's possible to work efficently with the approximate posterior marginals.
# Unfortunately it's not possible to work efficiently with the entire posterior process --
# this is limitation of pseudo-point approximations generally, as opposed to a limitation
# of Stheno.
posterior_marginals = marginals(f_post(xpr))

# Visualise the posterior. At the time of writing, it remains important to do this manually
# for the sake of efficiency. If you would like to have a high-level interface similar to
# the one available in the exact inference setting, please feel free to raise an issue, or
# implement it yourself and open a PR!

# Specify our plotting backend.
gr();

# Construct a new plot object.
posterior_plot = plot(legend=nothing);

# Generate several samples from the posterior predictive distribution.
plot!(posterior_plot, f_post(xpr); color=:blue);

# Plot the observations.
scatter!(posterior_plot, x, y;
    markercolor=:blue,
    markershape=:circle,
    markerstrokewidth=0.0,
    markersize=1,
    markeralpha=0.2,
);

# Show the plot.
display(posterior_plot);


abstract type ConvexApproximator <: ParametrisedConvexApproximator end

function _construct_convex_approximator(α_is::AbstractVector, β_is::AbstractVector)
    l = length(α_is[1])
    i_max = length(α_is)
    @assert length(β_is) == i_max
    _α_is = hcat(α_is...)'
    _β_is = hcat(β_is...)'
    l, i_max, _α_is, _β_is
end

function affine_map(nn::ConvexApproximator, z::AbstractArray)
    @unpack _α_is, _β_is = nn
    _α_is * z .+ _β_is
end

function affine_map(nn::ConvexApproximator, z::Convex.AbstractExpr)
    @unpack _α_is, _β_is = nn
    _α_is * z + (_α_is*zeros(size(z)) .+ _β_is)
end

include("MA.jl")
include("LSE.jl")


#=

The "Logarithmic Method" in [wdht].

=#

export WdhtCoordinator, insert_node!, remove_node!, choose_nodes, num_nodes, node_weight, reseed!

# Make it mutable for reseeding
mutable struct WdhtNode
    id::String
    pos::Float64
    weight::Float64
end

"Weighted Distributed Hash Tables"
struct WdhtCoordinator <: RvCoordinator
    nodes::Dict{String, WdhtNode}

    WdhtCoordinator() = new(Dict{String, WdhtNode}())
end

function hash1f(x::String)::Float64
    h = hash(x)
    (h & MASK) / MAX_HASH
end

function insert_node!(coordinator::WdhtCoordinator, id::String, weight::Float64)
    node = WdhtNode(id, hash1f(id), weight)
    coordinator.nodes[id] = node
end

function remove_node!(coordinator::WdhtCoordinator, id::String)
    delete!(coordinator.nodes, id)
end

function choose_nodes(
    coordinator::WdhtCoordinator,
    data::String,
    n::Union{Int, Nothing}=nothing,
)::Vector{String}
    h = hash1f(data)
    nodes = [
        begin
            s = -log(mod(1.0 - (h - node.pos), 1.0)) / node.weight
            s => id
        end
        for (id, node) in coordinator.nodes
    ]

    n::Int = isnothing(n) ? length(nodes) : n
    last.(partialsort!(nodes, 1:n))
end

function num_nodes(coordinator::WdhtCoordinator)::Int
    length(coordinator.nodes)
end

function node_weight(coordinator::WdhtCoordinator, id::String)::Float64
    coordinator.nodes[id].weight
end

function reseed!(coordinator::WdhtCoordinator)
    for node in values(coordinator.nodes)
        node.pos = rand()
    end
end


using PkgSRA, Test, Random
Random.seed!(13)

# include("../utils/regression_utils.jl")
# include("../utils/sparse_regression_functions.jl")

# Generate test data: Lorenz
include("../examples/example_lorenz.jl")
dat = Array(solve_lorenz_system())
numerical_grad = numerical_derivative(dat, ts)

sindy_library = Dict("cross_terms"=>2,"constant"=>nothing);

make_model(x) = sindyc(slstNumber(2, x), dat, numerical_grad)

#####
##### Test Sparse Regression function
#####
thresh = 0.1
m1 = slstHard(2, thresh)
A = sparse_regression(m1, dat, numerical_grad)
@testset "Hard threshold" begin
    @test all(abs.(A).<thresh)
end

#####
##### Test sparse models: all equations same terms
#####
num_terms5 = [5]
terms5 = get_nonzero_terms(make_model(num_terms5))
num_terms3 = [3]
terms3 = get_nonzero_terms(make_model(num_terms3))
@testset "Number of terms" begin
    @test length(terms5) == size(dat,1)*num_terms5
    @test length(terms3) == size(dat,1)*num_terms5
end

#####
##### Test sparse models: all equations different terms
#####
@testset "Number of terms2" begin
    num_terms = [1, 2, 3]
    terms = get_nonzero_terms(make_model(num_terms))
    @test length(terms) == sum(num_terms)

    num_terms = [2, 3, 4]
    terms = get_nonzero_terms(make_model(num_terms))
    @test length(terms) == sum(num_terms)
end

#####
##### Test ensemble models
#####
make_ensemble(x) = sindyc_ensemble(dat, numerical_grad,
                    sindy_library,
                    x,
                    selection_criterion=my_aicc,
                    sparsification_mode="num_terms")

# Each ensemble has all same
val_list = [1, 2]
(best_model_raw,best_criterion,all_criteria,all_models) =
    make_ensemble(val_list)

@testset "Ensemble models" begin
    i = 1
    @test length(get_nonzero_terms(all_models[i])) == size(dat,1)*val_list[i]
    i = 2
    @test length(get_nonzero_terms(all_models[i])) == size(dat,1)*val_list[i]
end

# Permutations
# val_list = calc_permutations(5,3)
# val_list = combinations(1:5, 3)
val_list = Iterators.product(1:3,1:3,1:3) # DEBUG
(best_model,best_criterion,all_criteria,all_models) =
    make_ensemble(val_list)

@testset "Permutations" begin
    for (i, v) in enumerate(val_list)
        @test length(get_nonzero_terms(all_models[i])) == sum(v)
    end
end

# The best model should be almost the real one
@testset "Accuracy" begin
    for (real_val, model_val) in zip(vec(core_dyn_true.A), vec(best_model.A))
        @test isapprox(Float64(real_val), model_val, atol=1e-1)
    end
end


using Gilbert
using Test

@testset "size $m,$n" for m = 1:20, n = 1:20
    list = gilbertindices((m,n))
    @test length(list) == m*n
    @test list[1] == CartesianIndex(1,1)
    if m >= n
        if !(n == 2 && isodd(m))
            @test list[end] == CartesianIndex(m,1)
        end
    else
        if !(m == 2 && isodd(n))
            @test list[end] == CartesianIndex(1,n)
        end
    end

    ndiag = 0
    for i = 1:m*n-1
        Δ = map(abs, (list[i+1] - list[i]).I)
        @test Δ[1] <= 1
        @test Δ[2] <= 1
        if Δ[1] + Δ[2] > 1
            ndiag += 1
        end
    end
    if m > n && isodd(m) && iseven(n) || m < n && isodd(n) && iseven(m)
        @test ndiag <= 1
    else
        @test ndiag == 0
    end

    L = Gilbert.linearindices(list)
    @test !any(iszero, L)
    @test sum(L) == sum(1:m*n)
end


function fermi(epsilon, beta)
  x = epsilon * beta
  x < 0.0 ? (1.0 / (1.0 + exp(x))) : (1.0 - 1.0 / (1.0 + exp(-x)))
end


#!/usr/bin/env julia

import GR

function mandel(x, y, iters)
    c = complex(x, y)
    z = 0.0im
    ci = 0
    inc = 1

    for i in 0:iters
        z = z^2 + c
        if abs2(z) >= 4
            return ci
        end
        ci += inc
        if ci == 0 || ci == 255
            inc = -inc
        end
    end

    return 255
end

function create_fractal(min_x, max_x, min_y, max_y, image, iters)
    height = size(image, 1)
    width = size(image, 2)

    pixel_size_x = (max_x - min_x) / width
    pixel_size_y = (max_y - min_y) / height
    for x in 1:width
        real = min_x + (x - 1) * pixel_size_x
        for y in 1:height
            imag = min_y + (y - 1) * pixel_size_y
            color = mandel(real, imag, iters)
            image[y, x] = color
        end
    end
end

x = -0.9223327810370947027656057193752719757635
y = 0.3102598350874576432708737495917724836010

f = 0.5
for i in 0:200
    tic()
    image = zeros(Int32, 500, 500)
    create_fractal(x-f, x+f, y-f, y+f, image, 400)
    dt = toq()

    println(@sprintf("Mandelbrot created in %f s", dt))

    GR.clearws()
    GR.setviewport(0, 1, 0, 1)
    GR.setcolormap(13)
    GR.cellarray(0, 1, 0, 1, 500, 500, image + 1000)
    GR.updatews()

    f *= 0.9
end


include("euler/euler.jl")
using .Primes: prime_numbers
using Combinatorics: permutations
using BenchmarkTools
BenchmarkTools.DEFAULT_PARAMETERS.samples = 100

function compute(n::Int)::Int
    function get_permutation_numbers(used::Vector{Int}, prime_index::Int)::Int
        sub_result = 0
        for i ∈ setdiff(digits, used)
            if parse(Int, string(i, join(used[1:2]))) % primes[prime_index] == 0
                if prime_index < 7
                    sub_result += get_permutation_numbers(vcat(i, used), prime_index + 1)
                else
                    used_digits = vcat(i, used)
                    sub_result += parse(Int, join(vcat(setdiff(digits, used_digits), used_digits)))
                end
            end
        end
        return sub_result
    end
    digits = collect(0:n)
    primes = reverse(prime_numbers(17))
    result = 0
    for permutation ∈ permutations(digits, 3)
        if parse(Int, join(permutation)) % primes[10 - n] == 0
            result += get_permutation_numbers(permutation, 11 - n)
        end
    end
    return result
end

compute(9)

@benchmark compute(9)


function  PCLUSTERS = GetPositiveSpanningTrees(PCLUSTERS,Lx,Ly)
clusternames = fieldnames(PCLUSTERS);
numclusters = length(clusternames);
for zz = 1:numclusters
    positions = eval(['PCLUSTERS.cluster' num2str(zz) '.positions']);
    xi = positions(:,1);
    yi = positions(:,2);
    nc = length(xi);

    %Get the spanning tree
    xdiff= repmat(xi,1,nc) - repmat(xi',nc,1);
    ydiff= repmat(yi,1,nc) - repmat(yi',nc,1);
    rij = sqrt(xdiff.^2 + ydiff.^2);
    rij = tril(rij);
    rij = sparse(rij);

    [T,~] = graphminspantree(rij);
    %[I,J] = find(T);

    %Try wrapping all positions, save the spanning tree which is
    %smallest
    T_old = T;
    xi_new = xi;
    yi_new = yi;
    [~,I] = sort(xi);
    %Wrapping in x...
    if ( (sum(xi<-Lx/4)>0) && (sum(xi>Lx/4)>0) )
      for ii =  1:length(xi)
        test_xi = xi;
        test_xi(I(1:ii)) = test_xi(I(1:ii)) + Lx;
        xdiff= repmat(test_xi,1,nc) - repmat(test_xi',nc,1);
        ydiff= repmat(yi,1,nc) - repmat(yi',nc,1);
        rij = sqrt(xdiff.^2 + ydiff.^2);
        rij = tril(rij);
        rij = sparse(rij);
        [T,~] = graphminspantree(rij);
        if sum(T(:)) < sum(T_old(:))
          disp('Found smaller spanning tree by wrapping in x')
          T_old = T;
          xi_new = test_xi;
        end
      end
    end

     %Wrapping in y
     if ((sum(yi<-Ly/4)>0) && (sum(yi>Ly/4)>0))
       [~,I] = sort(yi);
       for ii =  1:length(yi)
         test_yi = yi;
         test_yi(I(1:ii)) = test_yi(I(1:ii)) + Ly;
         ydiff= repmat(test_yi,1,nc) - repmat(test_yi',nc,1);
         xdiff= repmat(xi_new,1,nc) - repmat(xi_new',nc,1);
         rij = sqrt(xdiff.^2 + ydiff.^2);
         rij = tril(rij);
         rij = sparse(rij);
         [T,~] = graphminspantree(rij);
         if sum(T(:)) < sum(T_old(:))
           disp('Found smaller spanning tree by wrapping in y')
           T_old = T;
           yi_new = test_yi;
         end
       end
     end

%Get the radius and centre of mass position
XC = mean(xi_new);
YC = mean(yi_new);
Xrad = mean(abs(xi_new-XC));
Yrad = mean(abs(yi_new-YC));
Radius = sqrt(Xrad.^2 + Yrad.^2); %#ok
XC = XC -Lx*(XC > Lx/2);
YC = YC -Ly*(YC > Ly/2);
XC = XC +Lx*(XC < -Lx/2);
YC = YC +Ly*(YC < -Ly/2);
RC = [XC YC]; %#ok

    clear negindex
    [I,J] = find(T_old); %#ok
   eval(['PCLUSTERS.cluster' num2str(zz) '.spanningtreepointers = [I J];']);
   eval(['PCLUSTERS.cluster' num2str(zz) '.RC = RC;']);
   eval(['PCLUSTERS.cluster' num2str(zz) '.Radius = Radius;']);
end


function yinyang(n::Int=3)
    radii   = (i * n for i in (1, 3, 6))
    ranges  = collect(collect(-r:r) for r in radii)
    squares = collect(collect((x, y) for x in rnge, y in rnge) for rnge in ranges)
    circles = collect(collect((x, y) for (x,y) in sqrpoints if hypot(x, y) ≤ radius)
                      for (sqrpoints, radius) in zip(squares, radii))
    m = Dict((x, y) => ' ' for (x, y) in squares[end])
    for (x, y) in circles[end] m[(x, y)] = x > 0 ? '·' : '*' end
    for (x, y) in circles[end-1]
        m[(x, y + 3n)] = '*'
		m[(x, y - 3n)] = '·'
    end
    for (x, y) in circles[end-2]
        m[(x, y + 3n)] = '·'
		m[(x, y - 3n)] = '*'
    end
    return join((join(m[(x, y)]  for x in reverse(ranges[end])) for y in ranges[end]), '\n')
end

println(yinyang(4))


function calc_grad( pw::PWGrid, Potentials, Focc, psi::Array{ComplexF64,2} )

    Npoints = size(psi)[1]
    Nstates = size(psi)[2]
    Ω = pw.Ω
    Ns = pw.Ns
    #
    grad = zeros( ComplexF64, Npoints, Nstates )

    H_psi = op_H( pw, Potentials, psi )
    for i = 1:Nstates
        grad[:,i] = H_psi[:,i]
        for j = 1:Nstates
            grad[:,i] = grad[:,i] - dot( psi[:,j], H_psi[:,i] ) * psi[:,j]
        end
        grad[:,i] = Focc[i]*grad[:,i]
    end
    return grad

end


module CalculatorController
  using JSON
  using TypedPolynomials
  using StaticArrays
  using LinearAlgebra
  using Optim
  using ReverseDiff
  using DiffResults
  using BenchmarkTools

  #if !(@isdefined coeffs)
      # Configuration constants
      const PATH = "../../electrode/clients/lib/efield/comsol/data/fit_coeffs.json"
      const KEY_ORDER = ["LP", "UP", "LW", "LE", "UE", "UW"]
      const PARAM_ORDER = ["bias", "dipole", "angle", "dEdx", "dEdy", "nux", "nuy", "d2Edx2", "d2Edy2"]

      # Optimization constants
      const weights = SA_F64[1, 0, 1, 1, 1, 0, 0, 1, 1]
      const Vmax = SA_F64[5000, 5000, 5000, 5000, 5000, 5000]
      const Vmin = SA_F64[-5000, -5000, -5000, -5000, -5000, -5000]

      # Physical constants
      const amu = 1.66054e-27 # kg
      const DVcmToJ = 3.33564e-28 # J

      # Physical parameters for KRb
      const D_POLY_FIT = SA_F64[0, 0.0561, -0.00337, 1.03812e-4, -1.42283e-6, 5.34778e-9]
      const m = 127*amu

      # Polynomial orders
      const xorder = 7
      const yorder = 7

      # Location to evaluate polynomials
      const xval = 0.0
      const yval = 0.0

      # Load JSON with polynomial coefficients
      const coeffs = JSON.parsefile(PATH)

      # Construct potential as polynomial
      @polyvar x y
      const monomials = [10^(i + j) * x^i * y^j for i in 0:xorder for j in 0:yorder]
      const Vpoly = [dot(convert(Vector{Float64}, coeffs[k]), monomials) for k in KEY_ORDER]

      # Compute polynomials corresponding to electric field and its derivatives
      const Epoly = SMatrix{6,2}([differentiate(-v, var) for v in Vpoly, var in [x, y]])
      const dEdxpoly = SMatrix{6,2}([differentiate.(e, x) for e in Epoly])
      const dEdypoly = SMatrix{6,2}([differentiate.(e, y) for e in Epoly])
      const d2Edx2poly = SMatrix{6,2}([differentiate.(de, x) for de in dEdxpoly])
      const d2Edy2poly = SMatrix{6,2}([differentiate.(de, y) for de in dEdypoly])

      # Evaluate polynomials corresponding to electric field and its derivatives
      const ExV = subs.(Epoly, (x, y) => (xval, yval))
      const dEdxV = subs.(dEdxpoly, (x, y) => (xval, yval))
      const dEdyV = subs.(dEdypoly, (x, y) => (xval, yval))
      const d2Edx2V = subs.(d2Edx2poly, (x, y) => (xval, yval))
      const d2Edy2V = subs.(d2Edy2poly, (x, y) => (xval, yval))

      # Construct dipole moment and its derivatives as polynomials
      @polyvar E
      const Dpoly = dot(D_POLY_FIT, [(E/1000)^i for i in 0:length(D_POLY_FIT)-1])
      const dDdEpoly = differentiate(Dpoly, E)
      const d2DdE2poly = differentiate(dDdEpoly, E)
 # end

  # Compute the parameters of the field given electrode potentials
  function get_params(V)
      # Evaluate the electric field and its derivatives given electrode potentials
      Ev = sum(ExV .* V, dims=1)
      dEdxv = sum(dEdxV .* V, dims=1)
      dEdyv = sum(dEdyV .* V, dims=1)
      d2Edx2v = sum(d2Edx2V .* V, dims=1)
      d2Edy2v = sum(d2Edy2V .* V, dims=1)

      # Compute the electric field magnitude, direction and gradient
      bias = norm(Ev)
      angle = atan(Ev[1], Ev[2]) * 180/π
      dEdx = dot(Ev, dEdxv)/bias
      dEdy = dot(Ev, dEdyv)/bias
      d2Edx2 = (dot(dEdxv, dEdxv) + dot(d2Edx2v, Ev) - dEdx^2)/bias
      d2Edy2 = (dot(dEdyv, dEdyv) + dot(d2Edy2v, Ev) - dEdy^2)/bias

      # Evaluate the dipole moment at the bias field and its derivatives
      dipole = subs(Dpoly, E => bias)
      dDdEv = subs(dDdEpoly, E => bias)
      d2DdE2v = subs(d2DdE2poly, E => bias)

      # Compute the trap frequency
      biaspoly = (-dipole/bias*0 + 2*dDdEv + bias*d2DdE2v)
      kx = (dipole/bias + dDdEv) * dot(Ev, d2Edx2v) + dEdx^2 * biaspoly
      ky = (dipole/bias + dDdEv) * dot(Ev, d2Edy2v) + dEdy^2 * biaspoly
      # TODO: Correct these formulae after testing program
      νx = -sign(kx)*sqrt(abs(100*DVcmToJ*kx) / (2 * π * m))
      νy = -sign(ky)*sqrt(abs(100*DVcmToJ*ky) / (2 * π * m))
      out = SVector(bias, dipole, angle, dEdx, dEdy, νx, νy, d2Edx2, d2Edy2)

      return out
  end

  function barrier(x, xmin, xmax)
      return 1E0 .* (-log.(max.(x .- xmin, nextfloat(0.0))) -log.(max.(xmax .- x, nextfloat(0.0))))
  end
  # Compute the least squares error
  function err(p, V, w)
      return dot(w, (get_params(V) .- p).^2) + sum(barrier(V, Vmin, Vmax))
  end

  # Optimize electrode potentials for parameters p
  function opt(p; V0=nothing, w=weights)
      if V0 === nothing
          V0 = randn(6)
      end
      f(V) = err(p, V, w)
      tape = ReverseDiff.GradientTape(f, V0)
      compiled_tape = ReverseDiff.compile(tape)
      result = DiffResults.GradientResult(V0)
      function fg!(F, G, V)
          ReverseDiff.gradient!(result, compiled_tape, V)
          G .= DiffResults.gradient(result)
          return DiffResults.value(result)
      end
      opts = Optim.Options(x_tol=1E-8, f_tol=1E-8, g_tol=1E-8)
      optimize(Optim.only_fg!(fg!), V0, BFGS(), opts)
      # TODO: Why doesn't this work?
      #optimize(Optim.only_fg!(fg!), Vmin, Vmax, V0, Fminbox(BFGS()), opts)
  end

  function opt_json(p; V0=nothing, w=weights)
    optres = opt(p; V0=V0, w=w)
    iter = optres.iterations
    bestV = max.(min.(optres.minimizer, Vmax), Vmin)
    p2 = get_params(bestV)
    Vdict = Dict(k => bestV[i] for (i, k) in enumerate(KEY_ORDER))
    pdict = Dict(k => p2[i] for (i, k) in enumerate(PARAM_ORDER))
    edict = Dict("err" => err(p, bestV, w), "iter" => iter, "barrier" => sum(barrier(bestV, Vmin, Vmax)))
    return merge(Vdict, pdict, edict)
  end

  function params_json(V)
    p = get_params(V)
    return Dict(k => p[i] for (i, k) in enumerate(PARAM_ORDER))
  end

  # Benchmark the optimization
  # goal = [12446.91, 0, 45, 250, 63.83, 0, 0]
  # println("\nBenchmark:")
  # @benchmark opt(x) setup=(x = randn(length(goal)))
end


using Test
using OpenQuantumSystems
using Random, SparseArrays, LinearAlgebra, StableRNGs

@testset "evolution" begin

    Random.seed!(StableRNG(0), 1)

    D(op1::Array, op2::Array) = abs(norm(op1 - op2))
    D(x1::StateVector, x2::StateVector) = norm(x2 - x1)
    D(op1::AbstractOperator, op2::AbstractOperator) =
        abs(tracedistance_nh(dense(op1), dense(op2)))
    D(op1::AbstractSuperOperator, op2::AbstractSuperOperator) =
        abs(tracedistance_nh(dense(op1), dense(op2)))

    # TODO: change to macro
    mode1 = Mode(0.2, 1.0)
    Energy = [0.0, 200.0]
    mol1 = Molecule([mode1], 2, Energy)
    mol2 = Molecule([mode1], 2, Energy)
    agg = Aggregate([mol1, mol2])
    aggInds = getIndices(agg)
    aggIndsLen = length(aggInds)
    basis = GenericBasis([aggIndsLen])
    FCFact = getFranckCondonFactors(agg, aggInds)
    Ham = getAggHamiltonian(agg, aggInds, FCFact)

    H_lambda, H_S = eigen(Ham.data)
    H_Sinv = inv(H_S)
    H_lambda = diagm(H_lambda)

    t = 0.0
    U_op_ref = exp(-1im * Ham * t)
    U_op = evolutionOperator(Ham, t)
    @test 1e-12 > D(U_op, U_op_ref)
    U_op = evolutionOperatorA(H_lambda, H_S, H_Sinv, t)
    @test 1e-12 > D(U_op, U_op_ref.data)

    U_sop = evolutionSuperOperator(Ham, t)
    U_sop_ref = spre(U_op_ref) * spost(U_op_ref')
    @test 1e-12 > D(U_sop, U_sop_ref)

    t = 1.0
    U_op_ref = exp(-1im * Ham * t)
    U_op = evolutionOperator(Ham, t)
    @test 1e-12 > D(U_op, U_op_ref)
    U_op = evolutionOperatorA(H_lambda, H_S, H_Sinv, t)
    @test 1e-12 > D(U_op, U_op_ref.data)

    U_sop = evolutionSuperOperator(Ham, t)
    U_sop_ref = spre(U_op_ref) * spost(U_op_ref')
    @test 1e-12 > D(U_sop, U_sop_ref)

    tspan = [0.0:0.5:1.0;]
    U_op_array = evolutionOperatorArray(Ham, tspan)
    U_op1 = evolutionOperator(Ham, 0.0)
    U_op2 = evolutionOperator(Ham, 0.5)
    U_op3 = evolutionOperator(Ham, 1.0)
    @test 1e-12 > D(U_op_array[1], U_op1)
    @test 1e-12 > D(U_op_array[2], U_op2)
    @test 1e-12 > D(U_op_array[3], U_op3)

    U_sop_array = evolutionSuperOperatorArray(Ham, tspan)
    U_sop1 = evolutionSuperOperator(Ham, 0.0)
    U_sop2 = evolutionSuperOperator(Ham, 0.5)
    U_sop3 = evolutionSuperOperator(Ham, 1.0)
    @test 1e-11 > D(U_sop_array[1], U_sop1)
    @test 1e-11 > D(U_sop_array[2], U_sop2)
    @test 1e-11 > D(U_sop_array[3], U_sop3)

    t = 0.0
    foreach(evolutionOperatorIterator(Ham, tspan)) do U_op
        U_op_ref = evolutionOperator(Ham, t)
        @test 1e-12 > D(U_op, U_op_ref)
        t += 0.5
    end

    t = 0.0
    foreach(
        evolutionOperatorIterator(Ham, tspan; diagonalize = false, approximate = false),
    ) do U_op
        U_op_ref = evolutionOperator(Ham, t)
        @test 1e-12 > D(U_op, U_op_ref)
        t += 0.5
    end

    t = 0.0
    foreach(
        evolutionOperatorIterator(Ham, tspan; diagonalize = false, approximate = true),
    ) do U_op
        U_op_ref = evolutionOperator(Ham, t)
        @test 1e-12 > D(U_op, U_op_ref)
        t += 0.5
    end

    t = 0.0
    foreach(evolutionSuperOperatorIterator(Ham, tspan)) do U_sop
        U_sop_ref = evolutionSuperOperator(Ham, t)
        @test 1e-11 > D(U_sop, U_sop_ref)
        t += 0.5
    end

    t = 0.0
    foreach(
        evolutionSuperOperatorIterator(
            Ham,
            tspan;
            diagonalize = false,
            approximate = false,
        ),
    ) do U_op
        U_op_ref = evolutionSuperOperator(Ham, t)
        @test 1e-12 > D(U_op, U_op_ref)
        t += 0.5
    end

    t = 0.0
    foreach(
        evolutionSuperOperatorIterator(Ham, tspan; diagonalize = false, approximate = true),
    ) do U_op
        U_op_ref = evolutionSuperOperator(Ham, t)
        @test 1e-12 > D(U_op, U_op_ref)
        t += 0.5
    end

    ket = randstate(basis)
    ket1 = U_op1 * ket
    ket2 = U_op2 * ket
    ket3 = U_op3 * ket
    ket_array = evolutionExact(ket, tspan, Ham)
    @test 1e-12 > D(ket1, ket_array[1])
    @test 1e-12 > D(ket2, ket_array[2])
    @test 1e-12 > D(ket3, ket_array[3])

    ket_array = Array{Array{ComplexF32,1},1}(undef, 0)
    evolutionExact!(ket_array, ket, tspan, Ham)
    @test 1e-7 > D(ket1.data, ket_array[1])
    @test 1e-7 > D(ket2.data, ket_array[2])
    @test 1e-7 > D(ket3.data, ket_array[3])

    ket_array = evolutionApproximate(ket, tspan, Ham)
    @test 1e-12 > D(ket1, ket_array[1])
    @test 1e-12 > D(ket2, ket_array[2])
    @test 1e-12 > D(ket3, ket_array[3])

    ket_array = Array{Array{ComplexF32,1},1}(undef, 0)
    evolutionApproximate!(ket_array, ket, tspan, Ham)
    @test 1e-7 > D(ket1.data, ket_array[1])
    @test 1e-7 > D(ket2.data, ket_array[2])
    @test 1e-7 > D(ket3.data, ket_array[3])

    op = dm(ket)
    op1 = U_sop1 * op
    op2 = U_sop2 * op
    op3 = U_sop3 * op
    op_array = evolutionExact(op, tspan, Ham)
    @test 1e-12 > D(op1, op_array[1])
    @test 1e-12 > D(op2, op_array[2])
    @test 1e-12 > D(op3, op_array[3])

    op_array = Array{Array{ComplexF32,2},1}(undef, 0)
    evolutionExact!(op_array, op, tspan, Ham)
    @test 1e-7 > D(op1.data, op_array[1])
    @test 1e-7 > D(op2.data, op_array[2])
    @test 1e-7 > D(op3.data, op_array[3])

    op_array = evolutionApproximate(op, tspan, Ham)
    @test 1e-12 > D(op1, op_array[1])
    @test 1e-12 > D(op2, op_array[2])
    @test 1e-12 > D(op3, op_array[3])

    op_array = Array{Array{ComplexF32,2},1}(undef, 0)
    evolutionApproximate!(op_array, op, tspan, Ham)
    @test 1e-7 > D(op1.data, op_array[1])
    @test 1e-7 > D(op2.data, op_array[2])
    @test 1e-7 > D(op3.data, op_array[3])

    data = Matrix(Hermitian(rand(ComplexF64, aggIndsLen, aggIndsLen)))
    rho0 = DenseOperator(basis, basis, data)
    normalize!(rho0)

    tspan = [0.0:0.1:1.0;]
    rho_t_ref = evolutionExact(rho0, tspan, Ham)

    T, rho_t = evolution_exact(rho0, tspan, Ham; diagonalize = false)
    U_op_array = evolutionOperatorArray(Ham, tspan)
    for t_i = 1:length(tspan)
        U_op = U_op_array[t_i]
        rho_ref = U_op * rho0 * U_op'
        1e-12 > D(rho_t_ref[t_i], rho_t[t_i])
    end

    T, rho_t = evolution_exact(rho0, tspan, Ham; diagonalize = true)
    for t_i = 1:length(tspan)
        U_op = U_op_array[t_i]
        rho_ref = U_op * rho0 * U_op'
        1e-12 > D(rho_t_ref[t_i], rho_t[t_i])
    end

    T, rho_t = evolution_approximate(rho0, tspan, Ham; diagonalize = false)
    for t_i = 1:length(tspan)
        U_op = U_op_array[t_i]
        rho_ref = U_op * rho0 * U_op'
        1e-12 > D(rho_t_ref[t_i], rho_t[t_i])
    end

    T, rho_t = evolution_approximate(rho0, tspan, Ham; diagonalize = true)
    for t_i = 1:length(tspan)
        U_op = U_op_array[t_i]
        rho_ref = U_op * rho0 * U_op'
        1e-12 > D(rho_t_ref[t_i], rho_t[t_i])
    end

end


@testset "chainrules" begin
    Random.seed!(1)

    @testset "general: single input" begin
        for x in (1.0, -1.0, 0.0, 0.5, 10.0, -17.1, 1.5 + 0.7im)
            test_scalar(erf, x)
            test_scalar(erfc, x)
            test_scalar(erfi, x)

            test_scalar(airyai, x)
            test_scalar(airyaiprime, x)
            test_scalar(airybi, x)
            test_scalar(airybiprime, x)

            test_scalar(erfcx, x)
            test_scalar(dawson, x)

            if x isa Real
                test_scalar(invdigamma, x)
            end

            if x isa Real && 0 < x < 1
                test_scalar(erfinv, x)
                test_scalar(erfcinv, x)
            end

            if x isa Real && x > 0 || x isa Complex
                test_scalar(gamma, x)
                test_scalar(digamma, x)
                test_scalar(trigamma, x)
            end
        end
    end

    @testset "Bessel functions" begin
        for x in (1.5, 2.5, 10.5, -0.6, -2.6, -3.3, 1.6 + 1.6im, 1.6 - 1.6im, -4.6 + 1.6im)
            test_scalar(besselj0, x)
            test_scalar(besselj1, x)

            isreal(x) && x < 0 && continue

            test_scalar(bessely0, x)
            test_scalar(bessely1, x)

            for nu in (-1.5, 2.2, 4.0)
                test_frule(besseli, nu, x)
                test_rrule(besseli, nu, x)

                test_frule(besselj, nu, x)
                test_rrule(besselj, nu, x)

                test_frule(besselk, nu, x)
                test_rrule(besselk, nu, x)

                test_frule(bessely, nu, x)
                test_rrule(bessely, nu, x)

                # use complex numbers in `rrule` for FiniteDifferences
                test_frule(hankelh1, nu, x)
                test_rrule(hankelh1, nu, complex(x))

                # use complex numbers in `rrule` for FiniteDifferences
                test_frule(hankelh2, nu, x)
                test_rrule(hankelh2, nu, complex(x))
            end
        end
    end

    @testset "beta and logbeta" begin
        test_points = (1.5, 2.5, 10.5, 1.6 + 1.6im, 1.6 - 1.6im, 4.6 + 1.6im)
        for _x in test_points, _y in test_points
            # ensure all complex if any complex for FiniteDifferences
            x, y = promote(_x, _y)
            test_frule(beta, x, y)
            test_rrule(beta, x, y)

            test_frule(logbeta, x, y)
            test_rrule(logbeta, x, y)
        end
    end

    @testset "log gamma and co" begin
        # It is important that we have negative numbers with both odd and even integer parts
        test_points = (1.5, 2.5, 10.5, -0.6, -2.6, -3.3, 1.6 + 1.6im, 1.6 - 1.6im, -4.6 + 1.6im)
        for x in test_points
            for m in (0, 1, 2, 3)
                test_frule(polygamma, m, x; rtol=1e-8)
                test_rrule(polygamma, m, x; rtol=1e-8)
            end

            isreal(x) && x < 0 && continue
            test_scalar(loggamma, x)
            for a in test_points
                # ensure all complex if any complex for FiniteDifferences
                _a, _x = promote(a, x)
                test_frule(gamma, _a, _x; rtol=1e-8)
                test_rrule(gamma, _a, _x; rtol=1e-8)

                test_frule(loggamma, _a, _x)
                test_rrule(loggamma, _a, _x)
            end

            isreal(x) || continue
            test_frule(logabsgamma, x)
            test_rrule(logabsgamma, x; output_tangent=(randn(), randn()))
            for a in test_points
                isreal(a) && a > 0 || continue
                test_frule(gamma_inc, a, x, 0)
                test_rrule(gamma_inc, a, x, 0; output_tangent=(randn(), randn()))
            end
        end
    end
end


using Distributed
@everywhere using JuMP, Gurobi
@everywhere function kmer_coverage_ilp(locus, kmer_class, allele_ids, coverages)
  # m = Model(solver=GurobiSolver(Presolve=0))
  # one minute at most per locus, gap=0.1, no need to have optimal per se;
  # One thread per solve, since every ILP is already being run in parallel.
  m = Model(solver=GurobiSolver(OutputFlag=0, MIPGap=0.1, TimeLimit=1200, Threads=1))
  n_kmer = length(kmer_class)
  n_alleles = length(allele_ids)
  # kmer decision variable:
  @variable(m, x[1:n_kmer], Bin)
  @variable(m, c[1:n_alleles], Int)
  # coverage constraints:
  kmer_list = collect(keys(kmer_class))
  kmer_cover = DefaultDict{Int,Vector{Int}}(() -> Int[]) # allele -> list of kmers that cover it
  for (i, kmer) in enumerate(kmer_list)
    for al in kmer_class[kmer]
      push!(kmer_cover[al], i)
    end
  end
  # define coverage:
  for j in 1:n_alleles
    @constraint(m, c[j] == sum(x[i] for i in kmer_cover[j]))
  end
  # min coverage:
  for j in 1:n_alleles
    @constraint(m, c[j] >= coverages[j])
    # @constraint(m, c[j] == coverages[j])
  end
  # Forcing same coverage: (takes much longer, get many more kmers; better to store coverage)
  # Update: consider the cardinality of the kmers, to minimize it also:
  # @variable(m, max_card, Int) # max cardinality of the selected kmers
  # for (i, kmer) in enumerate(kmer_list)
  #   @constraint(m, max_card >= x[i] * length(kmer_class[kmer]))
  # end
  # minimize # of selected kmers:
  @objective(m, Min, sum(x[i] for i in 1:n_kmer))
  # @objective(m, Min, sum(x[i] for i in 1:n_kmer) + 100000*max_card)
  # TODO: e parameters works well, reduces variance but increases kmer number as expected.
  status = solve(m)
  if status == :Infeasible
    return nothing
  end
  # println("Objective value: ", getobjectivevalue(m))
  # for (i,kmer) in enumerate(kmer_list)
  #   if getvalue(x[i]) == 1
  #     println("$kmer : $(kmer_class[kmer])")
  #   end
  # end
  # println(join([Int(round(getvalue(c[j]),0)) for j in 1:n_alleles],", "))
  try
    kmers = [kmer_list[i] for i in 1:n_kmer if getvalue(x[i]) == 1]
    coverages = [Int16(round(getvalue(c[j]))) for j in 1:n_alleles]
    return kmers, coverages
  catch
    println("What happened??? -> LP status: $status LOCUS:$locus.")
    exit()
  end
  # return kmers, coverages, Int16(round(getvalue(max_card)))
end


using BenchmarkTools
using NumberTheoreticTransforms, FFTW

const x = mod.(rand(Int, 4096), 65537);
@btime fnt($x, $169, $65537);
@btime fft($x);

const x2 = mod.(rand(Int, 8192), 65537);
@btime fnt($x2, $225, $65537);
@btime fft($x2);


using VortexDistributions, Distributions, Statistics, BenchmarkTools

include("../util.jl")

function euclid(x, y)
    return sqrt((x[1]-y[1])^2 + (x[2]-y[2])^2)
end

##
Nx = 4000; Ny = 4000
Lx = 200; Ly = Lx
x = LinRange(-Lx/2,Lx/2, Nx+1)[1:end-1]; y = LinRange(-Ly/2,Ly/2, Ny+1)[1:end-1];
psi0 = one.(x*y') |> complex
psi = Torus(psi0,x,y)

##

rand_vorts = rand_pointvortex(2, psi)
# p1 = PointVortex(0., 0., 1)
# p2 = PointVortex(0.001, 0.002, 1)

vortex!(psi, rand_vorts)
# vortex!(psi, [p1, p2])

@time vfound1 = findvortices(psi)
@time vfound2 = naivePlaquette(psi.ψ, [psi.x, psi.y], 0)
@time vfound3 = remove_vortices_edge(findvortices_jumps(psi), psi)
@time vfound4 = remove_vortices_edge(findvortices_grid(psi), psi)

# sort!(rand_vorts, by = x -> x.xv)
sort!(vfound1, by = x -> x.xv)
sort!(vfound2, by = x -> x.xv)
sort!(vfound3, by = x -> x.xv)
sort!(vfound4, by = x -> x.xv)

vf1diff = [euclid([rand_vorts[i].xv, rand_vorts[i].yv], [vfound1[i].xv, vfound1[i].yv]) for i in 1:length(rand_vorts)]
vf2diff = [euclid([rand_vorts[i].xv, rand_vorts[i].yv], [vfound2[i].xv, vfound2[i].yv]) for i in 1:length(rand_vorts)]
vf2_3diff = [euclid([vfound2[i].xv, vfound2[i].yv], [vfound3[i].xv, vfound3[i].yv]) for i in 1:length(rand_vorts)]

vf1avg = mean(vf1diff)
vf2avg = mean(vf2diff)
vf2_3avg = mean(vf2_3diff)

naive = benchmark2Dnaive(2)

naive = @benchmark naivePlaquette(psi.ψ, [psi.x, psi.y], 0)
optimised = @benchmark remove_vortices_edge(findvortices_grid(psi), psi)

benchmarks = benchmark2Dnaive2(2)

using JLD2

@save "data/naiveBenchmarks.jld2" benchmarks;


function toh06a()
#
#  generalization of K.C. Toh and L. N. Trefethen
#
   p = ones(21)
   z = [0.95557280578614 + 0.29475517441090im
        0.95557280578614 - 0.29475517441090im
        0.82623877431600 + 0.56332005806362im
        0.82623877431600 - 0.56332005806362im
        0.62348980185873 + 0.78183148246803im
        0.62348980185873 - 0.78183148246803im
        0.36534102436640 + 0.93087374864420im
        0.36534102436640 - 0.93087374864420im
        0.07473009358642 + 0.99720379718118im
        0.07473009358642 - 0.99720379718118im
        -0.98883082622513 + 0.14904226617617im
        -0.98883082622513 - 0.14904226617617im
        -0.90096886790242 + 0.43388373911756im
        -0.90096886790242 - 0.43388373911756im
        -0.73305187182983 + 0.68017273777092im
        -0.73305187182983 - 0.68017273777092im
        -0.50000000000000 + 0.86602540378444im
        -0.50000000000000 - 0.86602540378444im
        -0.22252093395631 + 0.97492791218182im
        -0.22252093395631 - 0.97492791218182im];
    z = [z ones(20)]

    p, PolyZeros(z)
end


#--- Setup
function alloc_con(con,z)
    ∇c = TO.gen_jacobian(con)
    allocs  = @ballocated TO.evaluate($con, $z) samples=1 evals=1
    allocs += @ballocated TO.jacobian!($∇c, $con, $z) samples=1 evals=1
end

model = Cartpole()
n,m = size(model)
x,u = rand(model)
z = KnotPoint(x,u,0.1)

#--- Goal Constraint
@testset "Goal Constraint" begin
    xf = @SVector rand(n)
    goal = GoalConstraint(xf)
    @test TO.evaluate(goal, x) ≈ x - xf
    C = zeros(n,n)
    @test TO.jacobian!(C, goal, z) == true
    @test C ≈ I(n)

    @test length(goal) == n
    @test TO.upper_bound(goal) ≈ zero(xf)
    @test TO.lower_bound(goal) ≈ zero(xf)
    @test TO.is_bound(goal)
    @test TO.check_dims(goal, n, m)
    @test TO.check_dims(goal, n+1, m) == false
    @test TO.check_dims(goal, n, m+1) == true
    @test TO.widths(goal) == (n,)
    @test TO.widths(goal, n, m) == (n,)
    @test TO.widths(goal, n+1, m) == (n+1,)
    @test state_dim(goal) == n

    @test GoalConstraint(Vector(xf)).xf isa MVector{n}
end

#--- Linear Constraint
@testset "Linear Constraint" begin
    p = 5
    A = @SMatrix rand(p,n+m)
    b = @SVector rand(p)
    ∇c = zeros(p,n+m)
    lin = LinearConstraint(n,m,A,b, Inequality())
    @test TO.evaluate(lin, z) ≈ A*z.z - b
    @test jacobian!(∇c, lin, z) == true
    @test ∇c ≈ A

    lin2 = LinearConstraint(n,m, Matrix(A), Vector(b), Inequality())
    @test lin2.A isa SizedMatrix{p,n+m}
    @test lin2.b isa SVector{p}

    @test length(lin) == p
    @test TO.widths(lin) == (n+m,)
    @test state_dim(lin) == n
    @test control_dim(lin) == m
    @test TO.upper_bound(lin) ≈ zeros(p)
    @test TO.lower_bound(lin) ≈ fill(-Inf,p)
    @test TO.is_bound(lin) == false

    A = @SMatrix rand(p,n)
    lin2 = LinearConstraint(n,m, A,b, Equality(), 1:n)
    @test TO.evaluate(lin2, z) ≈ A*x - b
    ∇c = TO.gen_jacobian(lin2)
    @test TO.jacobian!(∇c, lin2, z) == true
    @test ∇c ≈ [A zeros(p)]

    @test TO.widths(lin2) == (n+m,)
    @test alloc_con(lin2,z) == 0
    @test state_dim(lin2) == n
    @test control_dim(lin2) == m
    @test TO.upper_bound(lin2) ≈ zeros(p)
    @test TO.lower_bound(lin2) ≈ zeros(p)
    @test TO.is_bound(lin2) == false

    @test TO.sense(lin) == Inequality()
    @test TO.sense(lin2) == Equality()

    A = @SMatrix rand(p,m)
    lin3 = LinearConstraint(n,m, A,b, Inequality(), n .+ (1:m))
    @test TO.evaluate(lin3, z) ≈ A*u - b
    ∇c = TO.gen_jacobian(lin3)
    @test TO.jacobian!(∇c, lin3, z) == true
    @test ∇c ≈ [zeros(p,n) A]
end

#--- Circle/Sphere Constraints
@testset "Circle/Sphere Constraints" begin
    xc = SA[1,1,1]
    yc = SA[1,2,3]
    r  = SA[1,1,1]
    cir = CircleConstraint(n, xc, yc, r)
    @test TO.evaluate(cir, z) ≈ -((x[1] .- xc).^2 + (x[2] .- yc).^2 .- r.^2)
    ∇c = TO.gen_jacobian(cir)
    @test TO.jacobian!(∇c, cir, z) == false
    @test ∇c ≈ hcat(-2*(x[1] .- xc), -2*(x[2] .- yc), zeros(3,n-2))
    @test cir isa CircleConstraint{3,Int}
    @test cir isa TO.StateConstraint
    @test length(cir) == 3
    @test state_dim(cir) == n
    @test_throws MethodError control_dim(cir)
    @test TO.widths(cir) == (n,)

    cir_ = CircleConstraint(n, Float64.(xc), yc, r)
    @test cir_ isa CircleConstraint{3,Float64}
    cir_ = CircleConstraint{3,Float64}(n, xc, yc, r)
    @test cir_ isa CircleConstraint{3,Float64}

    cir2 = CircleConstraint(n, Float64.(xc), yc, r, 2, 3)
    @test TO.evaluate(cir2, z) ≈ -((x[2] .- xc).^2 + (x[3] .- yc).^2 .- r.^2)
    ∇c = TO.gen_jacobian(cir2)
    @test TO.jacobian!(∇c, cir2, z) == false
    @test ∇c ≈ hcat(zeros(3), -2*(x[2] .- xc), -2*(x[3] .- yc), zeros(3,n-3))

    @test_throws AssertionError CircleConstraint(n, push(xc,2), yc, r)

    zc = SA[3,3,3]
    sph = SphereConstraint{3,Int}(n, xc, yc, zc, r)
    @test TO.evaluate(sph, z) ≈ -((x[1] .- xc).^2 + (x[2] .- yc).^2  .+ (x[3] .- zc).^2 .- r.^2)
    ∇c = TO.gen_jacobian(sph)
    @test TO.jacobian!(∇c, sph, z) == false
    @test ∇c ≈ hcat(-2*(x[1] .- xc), -2*(x[2] .- yc), -2*(x[3] .- zc), zeros(3,n-3))
    @test sph isa SphereConstraint{3,Int}
    @test sph isa TO.StateConstraint
    @test length(sph) == 3

    sph_ = SphereConstraint(n, Float64.(xc), yc, zc, r)
    @test sph_ isa SphereConstraint{3,Float64}
    sph_ = SphereConstraint{3,Float64}(n, xc, yc, zc, r)
    @test sph_ isa SphereConstraint{3,Float64}

    sph2 = SphereConstraint(n, Float64.(xc), yc, zc, r, 2, 3, 1)
    @test sph2.zi == 1
    @test TO.evaluate(sph2, z) ≈ -((x[2] .- xc).^2 + (x[3] .- yc).^2 + (x[1] .- zc).^2 .- r.^2)
    ∇c = TO.gen_jacobian(sph2)
    @test TO.jacobian!(∇c, sph2, z) == false
    @test ∇c ≈ hcat(-2*(x[1] .- zc), -2*(x[2] .- xc), -2*(x[3] .- yc), zeros(3,n-3))
end

#--- Collision Constraint
@testset "Collision Constraint" begin
    x1 = SA[1,2]
    x2 = SA[3,4]
    col = CollisionConstraint(n, x1, x2, 2.)
    d = x[x1] - x[x2]
    @test TO.evaluate(col, z) ≈ SA[4 - d'd]
    ∇c = TO.gen_jacobian(col)
    TO.jacobian!(∇c, col, z) == false
    @test ∇c ≈ [-2d' 2d']
    @test length(col) == 1

    col_ = CollisionConstraint(n, x1, x2, 1)
    @test col_.radius isa Float64
    col_ = CollisionConstraint(n, 1:2, x2, 1)
    @test col_.x1 isa SVector{2,Int}
    @test_throws AssertionError CollisionConstraint(n, 1:2, 1:3, 1.0)
end

#--- Norm Constraint
@testset "Norm Constraint" begin
    ncon = NormConstraint(n,m, 2.0, Inequality(), 1:n)
    @test TO.evaluate(ncon, z) ≈ [x'x - 2]
    ∇c = TO.gen_jacobian(ncon)
    @test TO.jacobian!(∇c, ncon, z) == false
    @test ∇c ≈ [2x; 0]'

    @test length(ncon) == 1
    @test TO.widths(ncon) == (n+m,)
    @test TO.sense(ncon) == Inequality()

    ncon2 = NormConstraint(n,m, 2.0, Inequality(), :state)
    @test TO.evaluate(ncon, z) ≈ TO.evaluate(ncon2, z)

    ncon2 = NormConstraint(n, m, 3.0, Equality(), :control)
    @test TO.evaluate(ncon2, z) ≈ [u'u - 3]
    ∇c = TO.gen_jacobian(ncon2)
    @test TO.jacobian!(∇c, ncon2, z) == false
    @test ∇c ≈ [zeros(n); 2u]'

    ncon3 = NormConstraint(n, m, 4.0, Inequality(), SA[1,3,5])
    @test TO.evaluate(ncon3, z) ≈ [x[1]^2 + x[3]^2 + u'u - 4]
    ∇c = TO.gen_jacobian(ncon3)
    @test TO.jacobian!(∇c, ncon3, z) == false
    @test ∇c ≈ [2x[1] 0 2x[3] 0 2u[1]]
end

#--- Bound Constraint
@testset "Bound Constraint" begin
    xmin = -@SVector rand(n)
    xmax = +@SVector rand(n)
    umin = -@SVector rand(m)
    umax = +@SVector rand(m)

    bnd = BoundConstraint(n,m, x_min=xmin, x_max=xmax, u_min=umin, u_max=umax)
    @test TO.evaluate(bnd, z) ≈ [x - xmax; u - umax; xmin - x; umin - u]
    ∇c = TO.gen_jacobian(bnd)
    @test TO.jacobian!(∇c, bnd, z) == true
    @test ∇c ≈ [I(n+m); -I(n+m)]
    @test length(bnd) == 2(n+m)
    @test TO.widths(bnd) == (n+m,)
    @test TO.upper_bound(bnd) == [xmax; umax]
    @test TO.lower_bound(bnd) == [xmin; umin]
    @test TO.is_bound(bnd) == true

    xmin = pop(pushfirst(xmin, -Inf))
    umax = popfirst(push(umax, Inf))
    bnd = BoundConstraint(n,m, x_min=xmin, x_max=xmax, u_min=umin, u_max=umax)
    @test TO.evaluate(bnd, z) ≈ [x - xmax; u[1:end-1] - umax[1:end-1];
        xmin[2:end] - x[2:end]; umin - u]
    ∇c = TO.gen_jacobian(bnd)
    @test TO.jacobian!(∇c, bnd, z) == true
    iz = ones(Bool,2(n+m))
    iz[n+1] = 0
    iz[n+m+1] = 0
    @test ∇c ≈ [I(n+m); -I(n+m)][iz, :]
    @test length(bnd) == 2(n+m) - 2
    @test TO.widths(bnd) == (n+m,)
    @test TO.upper_bound(bnd) == [xmax; umax]
    @test TO.lower_bound(bnd) == [xmin; umin]
    @test TO.is_bound(bnd) == true

    bnd_ = BoundConstraint(n,m, x_min=-10, x_max=10, u_min=umin, u_max=umax)
    @test TO.evaluate(bnd_, z) ≈ [x .- 10; u[1:end-1] - umax[1:end-1];
        -10 .- x; umin - u]
    bnd_ = BoundConstraint(n,m, x_min=-10, x_max=10, u_min=Vector(umin), u_max=umax)
    @test TO.evaluate(bnd_, z) ≈ [x .- 10; u[1:end-1] - umax[1:end-1];
        -10 .- x; umin - u]
    bnd_ = BoundConstraint(n,m, x_min=-10, x_max=10, u_min=Vector(umin), u_max=MVector(umax))
    @test TO.evaluate(bnd_, z) ≈ [x .- 10; u[1:end-1] - umax[1:end-1];
        -10 .- x; umin - u]

    xmin = -rand(1:10,n)
    xmax = rand(1:10,n)
    bnd_ = BoundConstraint(n,m, x_min=xmin, x_max=xmax)
    @test TO.evaluate(bnd_, z) ≈ [x .- xmax; xmin .- x]

    @test_throws ArgumentError BoundConstraint(n,m, x_min=10, x_max=-10, u_min=umin, u_max=umax)
end

#--- Indexed Constraint
@testset "Indexed Constraint" begin
    xmin = -@SVector rand(n)
    xmax = +@SVector rand(n)
    umin = -@SVector rand(m)
    umax = +@SVector rand(m)

    bnd = BoundConstraint(n,m, x_min=xmin, x_max=xmax, u_min=umin, u_max=umax)

    n2,m2 = 2n, 2m
    x2,u2 = [x; 2x], [u; 2u]
    z2 = KnotPoint(x2,u2,z.dt,z.t)

    idx = TO.IndexedConstraint(n2,m2, bnd)
    @test TO.evaluate(idx, z2) ≈ TO.evaluate(bnd, z)
    ∇c = TO.gen_jacobian(idx)
    ∇c0 = TO.gen_jacobian(bnd)
    @test TO.jacobian!(∇c, idx, z2) == true
    @test TO.jacobian!(∇c0, bnd, z) == true
    @test ∇c ≈ [∇c0[:,1:n] zeros(length(bnd), n) ∇c0[:,n+1:end] zeros(length(bnd), m)]

    @test TO.length(idx) == TO.length(bnd)
    @test TO.sense(idx) == TO.sense(bnd)
    @test TO.state_dim(idx) == 2n
    @test TO.control_dim(idx) == 2m
    @test TO.upper_bound(idx) == TO.upper_bound(bnd)
    @test TO.lower_bound(idx) == TO.lower_bound(bnd)
    @test TO.is_bound(idx) == true

    xc = SA[1,1,1]
    yc = SA[1,2,3]
    r  = SA[1,1,1]
    cir = CircleConstraint(n, xc, yc, r)

    idx = TO.IndexedConstraint(n2,m2, cir, n+1:2n, m+1:2m)
    @test size(idx.A) == size(idx.∇c)
    @test isempty(idx.B)
    @test TO.evaluate(idx, z2) ≈ TO.evaluate(cir, 2z)
    ∇c  = TO.gen_jacobian(idx)
    ∇c0 = TO.gen_jacobian(cir)
    @test TO.jacobian!(∇c, idx, z2) == false
    @test TO.jacobian!(∇c0, cir, 2z) == false
    @test ∇c ≈ [zeros(length(cir), n) ∇c0[:,1:n] zeros(length(cir), 2)]

    @test TO.length(idx) == TO.length(cir)
    @test TO.sense(idx) == TO.sense(cir)
    @test TO.state_dim(idx) == 2n
    @test TO.control_dim(idx) == 2m
    @test TO.upper_bound(idx) == TO.upper_bound(cir)
    @test TO.lower_bound(idx) == TO.lower_bound(cir)
    @test TO.is_bound(idx) == TO.is_bound(cir)

    # TODO: test IndexedConstraint with a ControlConstraint
end


workspace()

using IntervalArithmetic
using EAGODomainReduction

Y = Interval{Float64}[Interval(1.0,4.0),Interval(1.0,4.0),Interval(1.0,4.0),Interval(1.0,4.0)]
ymult_lo = [50, 1.0, 2.0, 3.0]
ymult_hi = [0, 1.0, 2.0, 3.0]
yLBD = 1.0
yUBD = 3.0
Variable_DR!(Y,ymult_lo,ymult_hi,yLBD,yUBD)
println("Y post full duality-based bound tightening: ", Y)

X = Interval{Float64}[Interval(1.0,4.0),Interval(1.0,4.0),Interval(1.0,4.0),Interval(1.0,4.0)]
mult_lo = [50, 1.0]
mult_hi = [-1.0, 0]
LBD = 1.0
UBD = 2.0
nx = 2
Variable_DR_Imp!(X,mult_lo,mult_hi,LBD,UBD,nx)
println("X post implicit duality-based bound tightening: ", X)


using Zygote, LinearAlgebra
using PyCall
# Make sure pytorch is installed

torch = pyimport("torch")

# From Mike Innes' JuliaCon 2019 talk
# https://youtu.be/OcUXjk7DFvU?t=1106
Zygote.@adjoint function pycall(f, x...; kw...)
    x = map(py, x)
    y = pycall(f, x...; kw...)
    y.detach().numpy(), function (ȳ)
        y.backward(gradient = py(ȳ))
        (nothing, map(x->x.grad.numpy(), x)...)
    end
end

# Non-symmetric version, which gives complex eigenvalues
# function torch_eig(A)
#     tA = torch.from_numpy(A)
#     out = torch.eig(tA)
#     evs = out[1].numpy()
#     return [ evs[j, 1] + im*evs[j,2] for j = 1:size(evs,1) ]
# end

function torch_symeig(A)
    tA = torch.from_numpy(A)
    out = torch.symeig(tA)
    evs = out[1].numpy()
    return evs
end

function test_eigmin(A)
    evals = torch_symeig(A)
    return minimum(evals)
end

A = rand(4, 4); A = A + A';
@show eigmin(A)
@show test_eigmin(A)
@show Zygote.gradient(test_eigmin, A)


using ACE, Test, Printf, LinearAlgebra, StaticArrays, BenchmarkTools

##
@testset "ACE.jl" begin
    # ------------------------------------------
    #   basic polynomial basis building blocks
    include("polynomials/test_ylm.jl")
    include("testing/test_wigner.jl")
    include("polynomials/test_transforms.jl")
    include("polynomials/test_orthpolys.jl")

    # --------------------------------------------
    # core permutation-invariant functionality
    include("test_scal1pbasis.jl")
    include("test_1pbasis.jl")
    include("test_pibasis.jl")

    # ------------------------
    #   rotation_invariance
    include("test_cg.jl")
    include("test_symmbasis.jl")
    include("test_euclvec.jl")

    # Model tests
    include("test_linearmodel.jl")

end

    # -----------------------------------------
    #    old tests to be re-introduced - maybe
    # include("test_real.jl")
    # include("test_orth.jl")
    # include("bonds/test_cylindrical.jl")
    # include("bonds/test_fourier.jl")
    # include("bonds/test_envpairbasis.jl")


module Day3

using DelimitedFiles

data_as_matrix(strings, bitcount) = begin
    m = zeros(Int, (length(strings), bitcount))
    for (row, string) in enumerate(strings)
        for bit in 1:bitcount
            b = parse(Int, string[bit])
            m[row, bit] = b
        end
    end
    m
end

bit_ratings(m) = begin
    hist = sum(m, dims=1) |> vec
    @assert size(hist) == (size(m, 2),)

    avgs = hist ./ size(m, 1)
    most_common = Int.(round.(avgs, RoundNearestTiesUp))
    least_common = 1 .- most_common
    (most_common, least_common)
end

rates_product(strings, bitcount) = begin
    m = data_as_matrix(strings, bitcount)

    γ_bits, ϵ_bits = bit_ratings(m)

    γ = parse(Int, join(γ_bits), base=2)
    ϵ = parse(Int, join(ϵ_bits), base=2)

    γ * ϵ
end

life_support_rating(strings, bitcount) = begin
    m = data_as_matrix(strings, bitcount)
    oxygen_candidates = co2_candidates = m

    for i in 1:bitcount
        if size(oxygen_candidates, 1) > 1
            most, _ = bit_ratings(oxygen_candidates)
            oxygen_mask = view(oxygen_candidates, :, i) .== most[i]
            oxygen_candidates = oxygen_candidates[oxygen_mask, :]
        end

        if size(co2_candidates, 1) > 1
            _, least = bit_ratings(co2_candidates)
            co2_mask = view(co2_candidates, :, i) .== least[i]
            co2_candidates = co2_candidates[co2_mask, :]
        end
    end

    oxygen_rating = parse(Int, join(oxygen_candidates[1, :]), base=2)
    co2_rating = parse(Int, join(co2_candidates[1, :]), base=2)

    oxygen_rating * co2_rating
end

test = ["00100"
        "11110"
        "10110"
        "10111"
        "10101"
        "01111"
        "00111"
        "11100"
        "10000"
        "11001"
        "00010"
        "01010"];

@assert rates_product(test, 5) == 198

@assert life_support_rating(test, 5) == 230

input = readdlm("resources/day3_input.txt", '\t', String, '\n') |> vec;

display(rates_product(input, 12))
display(life_support_rating(input, 12))

end


using SimpleWeightedGraphs
include("bijection_map.jl")

function edges_from_rule(rule_string)
    # rules are like this: light red bags contain 1 bright white bag, 2 muted yellow bags.
    parent_bag_string, other_bags_string = split(rule_string, "contain")
    parent_bag = strip(split(parent_bag_string, "bag")[1])
    other_bags_descriptions = split(other_bags_string, ",")

    edges = []
    for children_bag_description in split(other_bags_string, ",")
        child_bag_with_count = strip(split(children_bag_description, "bag")[1])

        # special case: handle no other
        if occursin("no other", child_bag_with_count)
            return [(parent_bag, nothing, 0)]
        end

        child_bag_count = parse(Int32, child_bag_with_count[1])
        child_bag = strip(child_bag_with_count[2:end])

        push!(edges, (parent_bag, child_bag, child_bag_count))
    end
    edges
end

function add_nodes_and_ids!(nodes_to_ids, node)
    if node === nothing
        return
    end

    node = String(node)
    if !haskey(nodes_to_ids, node)
        node_id = length(nodes_to_ids) + 1
        nodes_to_ids[node] = node_id
    end
end

function get_edges_and_nodes(input_file)
    edges, nodes_to_ids = open(input_file) do file
        file_string = read(file, String)
        rules = split(file_string, "\r\n")

        # note - Julia doesn't have bidicts... annoying.
        # also annoying - Julia graphs require you to use integers
        # for nodes.

        edges = []
        nodes_to_ids = BiDict{String,Int64}()
        for rule in rules
            rule_edges = edges_from_rule(rule)
            for edge in rule_edges
                node_1, node_2, edge_weight = edge
                add_nodes_and_ids!(nodes_to_ids, node_1)
                add_nodes_and_ids!(nodes_to_ids, node_2)
            end
            # println(edge_rules)
            append!(edges, rule_edges)
            # println("edges: $edges")
        end
        edges, nodes_to_ids
    end
    edges, nodes_to_ids
end

function build_graph(edges, nodes_to_ids)
    num_nodes = length(nodes_to_ids)
    g = SimpleWeightedDiGraph(num_nodes)
    for (n1, n2, num_bags) in edges
        if !(n2 === nothing)
            id1, id2 = nodes_to_ids[String(n1)], nodes_to_ids[String(n2)]
            # println("edge from: ($n1 $id1) to ($n2 $id2)")
            add_edge!(g, id1, id2, num_bags)
        end
    end
    println("number of vertex: $(nv(g))")
    println("number of edges: $(ne(g))")
    g
end

function find_parent_bags(g, nodes_to_ids, target_node)

    target_node_id = nodes_to_ids[String(target_node)]
    println("target_node: $target_node node_id: $target_node_id")
    all_parent_nodes = []
    for (node_name, node_id) in nodes_to_ids.key_to_value
        if (node_id != target_node_id) && (has_path(g, node_id, target_node_id))
            push!(all_parent_nodes, node_name)
        end
    end
    println("all bags: $all_parent_nodes")
    println("num bags: $(length(all_parent_nodes))")
end

function calculate_costs(g::AbstractGraph{T}, s::Integer, node_cost::Float64) where {T}
    nbs = outneighbors(g, s)
    println("Source: $s nbs: $(collect(nbs))")
    for n in nbs
        node_cost += calculate_costs(g, n, 1.0) * g.weights[n, s]
        println("Source: $s Node: $n edge weight: $(g.weights[n, s]) total cost so far: $node_cost")
    end
    node_cost
end

function part1_sol()
    edges_tuple, nodes_to_ids = get_edges_and_nodes("day7_input.txt")
    g = build_graph(edges_tuple, nodes_to_ids)
    find_parent_bags(g, nodes_to_ids, "shiny gold")
    return g
end

function part2_sol()
    edges_tuple, nodes_to_ids = get_edges_and_nodes("day7_input.txt")
    g = build_graph(edges_tuple, nodes_to_ids)

    start_node_id = nodes_to_ids[String("shiny gold")]
    calculate_costs(g, start_node_id, 0.0)
end

part2_sol()


export simplify_expi

SymEngine.free_symbols(syms::Union{Real, Complex}) = Basic[]
SymEngine.free_symbols(syms::AbstractArray{T}) where {T<:Union{Real, Complex}} = Basic[]
SymEngine.free_symbols(syms::AbstractArray{T}) where {T<:Union{Basic, SymEngine.BasicType}} = SymEngine.free_symbols(Matrix(syms))
function rand_assign(syms...)
    fs = union(free_symbols.(syms)...)
    Dict(zip(fs, randn(length(fs))))
end

function _basic_approx(x, y; atol=1e-8)
    diff = x-y
    assign = rand_assign(x, y)
    length(assign) > 0 && (diff = subs.(diff, Ref.((assign...,))...))
    nres = ComplexF64.(diff)
    all(isapprox.(nres, 0; atol=atol))
end

Base.:≈(x::AbstractArray{<:Basic}, y::AbstractArray; atol=1e-8) = _basic_approx(x, y, atol=atol)
Base.:≈(x::AbstractArray, y::AbstractArray{<:Basic}; atol=1e-8) = _basic_approx(x, y, atol=atol)
Base.:≈(x::AbstractArray{<:Basic}, y::AbstractArray{<:Basic}; atol=1e-8) = _basic_approx(x, y, atol=atol)
function Base.Complex{T}(a::Basic) where T
    a = SymEngine.evalf(a)
    T(real(a)) + im*T(imag(a))
end

# simplification rules
function simplify_expi(expr::Basic)
    subs(expr, exp(im*Basic(π)/2)=>Basic(im), exp(im*Basic(π))=>-1, exp(im*Basic(π)*3/2)=>-Basic(im))
end


function prob(beta::Real, h::Real)
	return 1.0/(1.0+exp(-beta*h))    # beta = 1/T
end

function stochastic_sign{T<:Real}(p::T)
    return p > rand() ? one(T) : -one(T)
end

function activatefunc(beta::Float64, h::Float64)
    return stochastic_sign(prob(beta,h))
end

function activatefunc(beta::Float64, h::Complex{Float64})
    r_val = stochastic_sign(prob(beta,real(h)))
	i_val = stochastic_sign(prob(beta,imag(h)))
	return complex(r_val, i_val)
end

function activatefunc(beta::Float64, h::Quaternion{Float64})
    r_val = stochastic_sign(prob(beta,real(h)))
	i_val = stochastic_sign(prob(beta,imagi(h)))
	j_val = stochastic_sign(prob(beta,imagj(h)))
	k_val = stochastic_sign(prob(beta,imagk(h)))
	return quaternion(r_val, i_val, j_val, k_val)
end

function internal_potential(W::AbstractArray{Quaternion{Float64}}, x::AbstractVector{Quaternion{Float64}},cue_pattern::AbstractVector{Float64},pos::Integer,s::Float64,L::Integer)
	e1 = dot(W[:,pos],x)

	if L==2
		val = x[pos]
		idx = pos*2
		e2 = quaternion(cue_pattern[idx-1] * imagi(val),
						cue_pattern[idx-1] * real(val),
						cue_pattern[idx  ] * imagk(val),
						cue_pattern[idx  ] * imagj(val))
	elseif L==4
		val = x[pos]
		idx = pos*3
		e2 = quaternion(cue_pattern[idx-2] * imagi(val),
						cue_pattern[idx-2] * real(val),
						cue_pattern[idx-1] * real(val),
						cue_pattern[idx  ] * real(val))
	else
		@assert(false,"unimplemented")
	end

	return (e1 + s*e2)
end

function internal_potential(W::AbstractArray{Complex{Float64}}, x::AbstractVector{Complex{Float64}},cue_pattern::AbstractVector{Float64},pos::Integer,s::Float64,L::Integer)
	e1 = dot(W[:,pos],x)

	if L==2
		e2 = cue_pattern[pos] * complex(imag(x[pos]), real(x[pos]))
	else
		@assert(false,"unimplemented")
	end

	return (e1 + s*e2)
end

function internal_potential(W::AbstractArray{Float64}, x::AbstractVector{Float64},cue_pattern::AbstractVector{Float64},pos::Integer,s::Float64,L::Integer)
	e1 = dot(W[:,pos],x)

	if L==2
		n = length(cue_pattern)
		pos_cue = (pos -1)%n+1
		if pos > n
			e2 = cue_pattern[pos_cue] * x[pos_cue]
		else
			e2 = cue_pattern[pos_cue] * x[pos+n]
		end
	elseif L==4
		n = round(Int,length(x)/L)
		pos_cue = (pos -1)%n+1
		block = ifloor((pos-1)/n)
		@switch block begin
			0; e2 = cue_pattern[pos_cue] * x[pos_cue+1*n]
			1; e2 = cue_pattern[pos_cue] * x[pos_cue]
			2; e2 = cue_pattern[pos_cue+n] * x[pos_cue]
			3; e2 = cue_pattern[pos_cue+2*n] * x[pos_cue]
		end
	else
		@assert(false,"unimplemented")
	end

	return (e1 + s*e2)
end

function update_cue_async!(net::HopfieldNetwork, cue_pattern::AbstractVector,L::Integer, s::Real, beta::Real)
	ord = randperm(length(net.state))
	for i in ord
		h = internal_potential(net.weight,net.state,cue_pattern,i,s,L)
		net.state[i] = activatefunc(beta, h)
	end
end

function update_cue_sync!(net::HopfieldNetwork, cue_pattern::AbstractVector, L::Integer, s::Real, beta::Real)
	state = copy(net.state)
	for i = 1:length(net.state)
		h = internal_potential(net.weight,state,cue_pattern,i,s,L)
		net.state[i] = activatefunc(beta, h)
	end
end

function associate!{T1 <: Number, T2 <: Real}(net::HopfieldNetwork, init_pattern::AbstractVector{T1}, cue_pattern::AbstractVector{T2},L::Integer, s::Real, beta::Real, gamma::Real, iterations::Integer = 1_000, sync::Bool = false)
	copy!(net.state, init_pattern)
	associate!(net,cue_pattern,L,s,beta,gamma, iterations,sync)
end

function associate!{T <: Number}(net::HopfieldNetwork, cue_pattern::AbstractVector{T},L::Integer, s::Real, beta::Real, gamma::Real, iterations::Integer = 1_000, sync::Bool = false)
	if sync
		update! = update_cue_sync!
	else
		update! = update_cue_async!
	end
	beta_step = beta
	for i in 1:iterations
		update!(net, cue_pattern,L, s, beta_step)
		beta_step = beta_timestep(beta_step, gamma)
	end
	return copy(net.state)
end

function associate_step_async!{T <: Number}(net::HopfieldNetwork, cue_pattern::AbstractVector{T},L::Integer, s::Real, beta::Real, gamma::Real)
	update_cue_async!(net, cue_pattern, L, s, beta)
	return beta_timestep(beta, gamma)
end

function associate_step_sync!{T <: Number}(net::HopfieldNetwork, cue_pattern::AbstractVector{T},L::Integer, s::Real, beta::Real, gamma::Real)
	update_cue_sync!(net, cue_pattern, L, s, beta)
	return beta_timestep(beta, gamma)
end

function beta_timestep(beta::Real, gamma::Real)
	return gamma * beta;
end


#=

    This is a port of the SPPL model trueskill-poisson-binomial.pynb

    This model:
    * Prior()
    Summary Statistics
  parameters      mean       std   naive_se      mcse          ess      rhat   ess_per_sec
      Symbol   Float64   Float64    Float64   Float64      Float64   Float64       Float64

      perfB1   79.3452    0.9921     0.0099    0.0102    9771.5110    0.9999   105070.0105
      skillA   99.9423    9.4506     0.0945    0.0981    9873.7861    0.9999   106169.7428
      perfA1   89.9122    9.0726     0.0907    0.0948   10013.5210    0.9999   107672.2688
      result    0.8594    0.3476     0.0035    0.0032   10370.4888    1.0000   111510.6318

    * With observation result == 0
    Summary Statistics
    parameters      mean       std   naive_se      mcse         ess      rhat   ess_per_sec
        Symbol   Float64   Float64    Float64   Float64     Float64   Float64       Float64

        perfB1   79.4940    0.8266     0.0083    0.0249   1315.1090    1.0011      171.1045
        skillA   86.2807    4.2628     0.0426    0.1226   1134.3992    0.9999      147.5929
        perfA1   75.9473    3.2281     0.0323    0.0936   1280.7372    1.0003      166.6325
        result    0.0003    0.0173     0.0002    0.0003   3133.8451    1.0002      407.7342

=#

using Turing, StatsPlots, Distributions, StatsBase
using CSV
include("jl_utils.jl")

@model function trueskill_poisson_binomial()
    perfB1 ~ truncated(Binomial(100, 0.9), 0,80)
    skillA ~ truncated(Poisson(100), 77,125)
    perfA1 ~ Binomial(skillA, 0.9)

    if perfA1 > perfB1
       result ~ Dirac(1)
    else
        result ~ Dirac(0)
    end

    # Observation
    true ~ Dirac(result == 0)
end

model = trueskill_poisson_binomial()

# chns = sample(model, Prior(), 10_000)
# chns = sample(model, MH(), 10_000)
chns = sample(model, PG(5),  10_000)
# chns = sample(model, PG(5),  MCMCThreads(), 1_000, 4)
# chns = sample(model, SMC(), 10_000)
# chns = sample(model, IS(), 10_000)

# chns = sample(model, HMC(0.1,6), 1_000)
# chns = sample(model, NUTS(), 1_000)

display(chns)
# display(plot(chns))


@testset "laguerre" begin
    @test all(
        laguerrel(n, α, x) ≈ laguerre(n, α)(x)
        for n in 0:100, α in 0:100, x in -20:0.1:20
    )
end


using DistributionsLH, ModelObjectsLH
using Random, StatsBase, GLM, Statistics
using Test

@testset "DistributionsLH" begin
    include("uniform_test.jl");
    include("mv_normal_test.jl");
end

# ------------


using MomentClosure
using MomentClosure: define_μ, define_M
using Test
using Catalyst

@parameters t, c₁, c₂, c₃, c₄, Ω
@variables X(t), Y(t)

S_mat = [ 1 -1  1 -1;
         -1  1  0  0]
# NOTE: only holds if combinatoric_ratelaw=false
a = [c₁*X*Y*(X-1)/Ω^2, c₂*X, c₃*Ω, c₄*X]

# checking if ReactionSystem and ReactionSystemMod are consistent

rn1 = @reaction_network begin
    (c₁/Ω^2), 2X + Y → 3X
    (c₂), X → Y
    (Ω*c₃, c₄), 0 ↔ X
end c₁ c₂ c₃ c₄ Ω

rn2 = ReactionSystemMod(t, [X, Y], [c₁, c₂, c₃, c₄, Ω], a, S_mat)

@test isequal(species(rn1), species(rn2))
@test isequal(params(rn1), params(rn2))
@test isequal(speciesmap(rn1), speciesmap(rn2))
@test isequal(paramsmap(rn1), paramsmap(rn2))
@test numspecies(rn1) == numspecies(rn2) == 2
@test numreactions(rn1) == numreactions(rn2) == 4
@test numparams(rn1) == numparams(rn2) == 5
@test netstoichmat(rn1) == netstoichmat(rn2) == S_mat
@test isequal(propensities(rn1, combinatoric_ratelaw=false), propensities(rn2))


module QPointCellFieldsTests

using Test
using FillArrays
using Gridap.Helpers
using Gridap.Fields
using Gridap.CellData
using Gridap.TensorValues
using Gridap.ReferenceFEs
using Gridap.Integration

using Gridap.Fields: MockField, MockBasis, OtherMockBasis

ndofs = 4
c = fill(1.0,ndofs)
f = OtherMockBasis{2}(ndofs)

l = 10
cl = fill(c,l)
fl = Fill(f,l)
ϕl = lincomb(fl,cl)

degree = 3
quad = CellQuadrature(degree,[QUAD],Fill(1,l))
q = get_coordinates(quad)

s_q = [i*ones(size(qi)) for (i,qi) in enumerate(q)]
a = CellData.ArrayOfEvaluatedFields(s_q,q)
test_array_of_fields(a,q,s_q)

v = VectorValue(3.0,4.0)
cf = QPointCellField(v,ϕl,quad)
cf_q = evaluate(cf,q)
test_cell_field(cf,q,cf_q)

cf = CellField(v,ϕl,quad)
cf_q = evaluate(cf,q)
test_cell_field(cf,q,cf_q)

cf = CellField(1.0,ϕl,quad)
@test sum(integrate(cf,ϕl,quad)) ≈ 213.33333333333323

a = CellField(1.0,ϕl,quad)
b = CellField(1.0,ϕl,quad)
c = CellField(1.0,ϕl,quad)

function updater(a,b,c)
  b_new = 2*b
  c_new = 4*c
  true, b_new, c_new
end

@test evaluate(a,q)[1] == fill(1.0,4)
@test evaluate(b,q)[1] == fill(1.0,4)
@test evaluate(c,q)[1] == fill(1.0,4)

update_state_variables!(updater,quad,a,b,c)
@test evaluate(a,q)[1] == fill(1.0,4)
@test evaluate(b,q)[1] == fill(2.0,4)
@test evaluate(c,q)[1] == fill(4.0,4)

update_state_variables!(updater,quad,a,b,c)
@test evaluate(a,q)[1] == fill(1.0,4)
@test evaluate(b,q)[1] == fill(4.0,4)
@test evaluate(c,q)[1] == fill(16.0,4)

end # module


"""
    Quaternion{T<:Real} <: Number

Quaternionic number type with elements of type `T`.

`QuaternionF16`, `QuaternionF32` and `QuaternionF64` are aliases for `Quaternion{Float16}`,
`Quaternion{Float32}` and `Quaternion{Float64}` respectively.  See also [`Rotor`](@ref) and
[`QuatVec`](@ref).

The functions

    Quaternion(w, x, y, z)
    Quaternion(x, y, z)
    Quaternion(w)
    Quaternion{T}(w, x, y, z)

create a new quaternion with the given components.  The argument `w` is the
scalar component, and `x`, `y`, and `z` are the corresponding "vector"
components.  If any of these arguments is missing, it will be set to zero.  The
type of the returned quaternion will be inferred from the input arguments, or
can be specified, by passing the type parameter `T` as above.

Note that the constants [`imx`](@ref), [`imy`](@ref), and [`imz`](@ref) can
also be used like the complex `im` to create new `Quaternion` object.

# Examples

```jldoctest
julia> Quaternion(1, 2, 3, 4)
1 + 2𝐢 + 3𝐣 + 4𝐤
julia> Quaternion{Float64}(1, 2, 3, 4)
1.0 + 2.0𝐢 + 3.0𝐣 + 4.0𝐤
julia> Quaternion(1.0, 2.0, 3.0, 4.0)
1.0 + 2.0𝐢 + 3.0𝐣 + 4.0𝐤
julia> Quaternion(2, 3, 4)
0 + 2𝐢 + 3𝐣 + 4𝐤
julia> Quaternion(1)
1 + 0𝐢 + 0𝐣 + 0𝐤
```
"""
struct Quaternion{T<:Real} <: AbstractQuaternion{T}
    components::SVector{4, T}
    Quaternion{T}(a::SVector{4, T}) where {T<:Real} = new{T}(a)
    Quaternion{T}(a::A) where {T<:Real, A<:AbstractArray} = new{T}(SVector{4, T}(a))
end

@doc raw"""
    Rotor{T<:Real} <: Number

Quaternion of unit magnitude with elements of type `T`.  These objects can be
significantly faster *and* more accurate in certain operations representing
rotations.

A rotor is typically considered to be an element of the group
``\mathrm{Spin}(3) ≃ \mathrm{SU}(2)``, which can be thought of as the subgroup
of quaternions with norm 1.  They are particularly useful as representations of
rotations because a rotor ``R`` acts on a vector ``\vec{v}`` by "conjugation" as

```math
\vec{v}' = R\, \vec{v}\, R^{-1}.
```

This preserves the inner product between any two vectors conjugated in this
way, and so is a rotation.  Note that, because there are two factors of ``R``
here, the sign of ``R`` does not affect the result.  Therefore,
``\mathrm{Spin}(3)`` forms a *double* cover of the rotation group
``\mathrm{SO}(3)``.  For this reason, it will occasionally be useful to
disregard or arbitrarily change the sign of a `Rotor` (as in [`distance`](@ref)
functions) — though this is not generally the default, and may cause problems
if the input rotors change sign when the corresponding rotations are not so
different (cf. [`unflip`](@ref)).

`RotorF16`, `RotorF32` and `RotorF64` are aliases for `Rotor{Float16}`,
`Rotor{Float32}` and `Rotor{Float64}` respectively.  See also
[`Quaternion`](@ref) and [`QuatVec`](@ref).

The functions

    Rotor(w, x, y, z)
    Rotor(x, y, z)
    Rotor(w)
    Rotor{T}(w, x, y, z)

create a new rotor with the given components (where the components are as
described in [`Quaternion`](@ref)), automatically normalizing them on input.
If you would like to bypass this normalization step, you can call

    Rotor{T}(v)

where `v<:AbstractArray`, and can be converted to an `SVector{4, T}`.

However, once a `Rotor` is created, its norm will always be assumed to be
precisely 1.  So if its true norm is significantly different, you will like see
weird results — including vectors with very different lengths after "rotation"
by a non-unit `Rotor`.

Note that simply creating a `Quaternion` that happens to have norm 1 does not
make it a `Rotor`.  However, you can pass such a `Quaternion` to the `Rotor`
function and get the desired result.

# Examples

```jldoctest
julia> Rotor(1, 2, 3, 4)
0.18257418583505536 + 0.3651483716701107𝐢 + 0.5477225575051661𝐣 + 0.7302967433402214𝐤
julia> Rotor(Quaternion(1, 2, 3, 4))
0.18257418583505536 + 0.3651483716701107𝐢 + 0.5477225575051661𝐣 + 0.7302967433402214𝐤
julia> Rotor{Float16}(1, 2, 3, 4)
0.1826 + 0.3652𝐢 + 0.548𝐣 + 0.7305𝐤
julia> Rotor(2, 3, 4)
0.0 + 0.3713906763541037𝐢 + 0.5570860145311556𝐣 + 0.7427813527082074𝐤
julia> Rotor(1)
1 + 0𝐢 + 0𝐣 + 0𝐤
```
"""
struct Rotor{T<:Real} <: AbstractQuaternion{T}
    components::SVector{4, T}
    Rotor{T}(a::SVector{4, T}) where {T<:Real} = new{T}(a)
    Rotor{T}(a::A) where {T<:Real, A<:AbstractArray} = new{T}(SVector{4, T}(a))
end

"""
    QuatVec{T<:Real} <: Number

Pure-vector quaternion with elements of type `T`.  These objects can be
significantly faster *and* more accurate in certain operations than general
`Quaternion`s.

`QuatVecF16`, `QuatVecF32` and `QuatVecF64` are aliases for `QuatVec{Float16}`,
`QuatVec{Float32}` and `QuatVec{Float64}` respectively.  See also
[`Quaternion`](@ref) and [`Rotor`](@ref).

The functions

    QuatVec(w, x, y, z)
    QuatVec(x, y, z)
    QuatVec(w)
    QuatVec{T}(w, x, y, z)

create a new rotor with the given components (where the components are as
described in [`Quaternion`](@ref)), except that the scalar argument `w` is
always set to 0.

# Examples

```jldoctest
julia> QuatVec(1, 2, 3, 4)
0 + 2𝐢 + 3𝐣 + 4𝐤
julia> QuatVec(Quaternion(1, 2, 3, 4))
0 + 2𝐢 + 3𝐣 + 4𝐤
julia> QuatVec(2, 3, 4)
0 + 2𝐢 + 3𝐣 + 4𝐤
julia> QuatVec(1)
0 + 0𝐢 + 0𝐣 + 0𝐤
```
"""
struct QuatVec{T<:Real} <: AbstractQuaternion{T}
    components::SVector{4, T}
    QuatVec{T}(a::SVector{4, T}) where {T<:Real} = new{T}(a)
    QuatVec{T}(a::A) where {T<:Real, A<:AbstractArray} = new{T}(SVector{4, T}(a))
end

# Untyped constructor from SVector
# (::Type{QT})(v::SVector{4, T}) where {T<:Real, QT<:AbstractQuaternion} = QT{eltype(v)}(v)
Quaternion(v::SVector{4, T}) where {T<:Real} = Quaternion{eltype(v)}(v)
Rotor(v::SVector{4, T}) where {T<:Real} = Rotor{eltype(v)}(v)
QuatVec(v::SVector{4, T}) where {T<:Real} = QuatVec{eltype(v)}(v)

# Constructor from all 4 components
(::Type{QT})(w, x, y, z) where {QT<:AbstractQuaternion} = (v=SVector{4}(w, x, y, z); QT{eltype(v)}(v))
(::Type{QT})(w, x, y, z) where {T<:Real, QT<:AbstractQuaternion{T}} = QT(SVector{4, T}(w, x, y, z))
Rotor(w, x, y, z) = (n=√(w^2+x^2+y^2+z^2); Rotor(SVector{4}(w/n, x/n, y/n, z/n)))
Rotor{T}(w, x, y, z) where {T<:Real} = (n=√T(w^2+x^2+y^2+z^2); Rotor{T}(SVector{4, T}(w/n, x/n, y/n, z/n)))
QuatVec(w, x, y, z) = QuatVec(SVector{4}(oftype(w, false), x, y, z))
QuatVec{T}(w, x, y, z) where {T<:Real} = QuatVec{T}(SVector{4, T}(oftype(w, false), x, y, z))

# Constructor from vector components
(::Type{QT})(x, y, z) where {QT<:AbstractQuaternion} = (v=SVector{4}(false, x, y, z); QT{eltype(v)}(v))
(::Type{QT})(x, y, z) where {T<:Real, QT<:AbstractQuaternion{T}} = QT(SVector{4, T}(false, x, y, z))
Rotor(x, y, z) = (n=√(x^2+y^2+z^2); Rotor(SVector{4}(false, x/n, y/n, z/n)))
Rotor{T}(x, y, z) where {T<:Real} = (n=√T(x^2+y^2+z^2); Rotor{T}(SVector{4, T}(false, x/n, y/n, z/n)))
QuatVec(x, y, z) = QuatVec(SVector{4}(false, x, y, z))
QuatVec{T}(x, y, z) where {T<:Real} = QuatVec{T}(SVector{4, T}(false, x, y, z))

# Constructor from scalar component
# (::Type{QT})(w::Real) where {QT<:AbstractQuaternion} = (v=SVector{4}(w, false, false, false); QT{eltype(v)}(v))
# (::Type{QT})(w::Real) where {T<:Real, QT<:AbstractQuaternion{T}} = QT(SVector{4, T}(w, false, false, false))
Quaternion(w::Real) = Quaternion(SVector{4}(w, false, false, false))
Quaternion(w::Symbolics.Num) = Quaternion(SVector{4}(w, false, false, false))
Quaternion{T}(w::Real) where {T<:Real} = Quaternion{T}(SVector{4, T}(w, false, false, false))
Rotor(w::Real) = Rotor(SVector{4}(one(w), false, false, false))
Rotor(w::Symbolics.Num) = Rotor(SVector{4}(one(w), false, false, false))
Rotor{T}(w::Real) where {T<:Real} = Rotor{T}(SVector{4, T}(one(T), false, false, false))
QuatVec(w::Real) = QuatVec(SVector{4, typeof(w)}(false, false, false, false))
QuatVec(w::Symbolics.Num) = QuatVec(SVector{4, typeof(w)}(false, false, false, false))
QuatVec{T}(w::Real) where {T<:Real} = QuatVec{T}(SVector{4, T}(false, false, false, false))

# Copy constructor
# (::Type{QT})(q::AbstractQuaternion) where {QT<:AbstractQuaternion} = QT{eltype(q.components)}(q.components)
# (::Type{QT})(q::AbstractQuaternion{S}) where {T<:Real, S<:Real, QT<:AbstractQuaternion{T}} = QT(SVector{4, T}(q.components))
# (::Type{QT})(q::AbstractQuaternion{T}) where {T<:Real, QT<:AbstractQuaternion{T}} = QT(SVector{4, T}(q.components))
Quaternion(q::AbstractQuaternion{T}) where {T<:Real} = Quaternion(q.components...)
Quaternion{T}(q::AbstractQuaternion{S}) where {T<:Real, S<:Real} = Quaternion{T}(q.components...)
Rotor(q::QT) where {T<:Real, QT<:AbstractQuaternion{T}} = Rotor(q.components...)
Rotor{T}(q::AbstractQuaternion{S}) where {T<:Real, S<:Real} = Rotor{T}(q.components...)
QuatVec(q::AbstractQuaternion{T}) where {T<:Real} = QuatVec(q.components...)
QuatVec{T}(q::AbstractQuaternion{S}) where {T<:Real, S<:Real} = QuatVec{T}(q.components...)

# Type constructors
(::Type{QT})(::Type{T}) where {T<:Real, QT<:AbstractQuaternion} = QT{T}
(::Type{QT})(::Type{<:AbstractQuaternion{T}}) where {T<:Real, QT<:AbstractQuaternion} = QT{T}

# Handy aliases like `ComplexF64`, etc.
const QuaternionF64 = Quaternion{Float64}
const QuaternionF32 = Quaternion{Float32}
const QuaternionF16 = Quaternion{Float16}
const RotorF64 = Rotor{Float64}
const RotorF32 = Rotor{Float32}
const RotorF16 = Rotor{Float16}
const QuatVecF64 = QuatVec{Float64}
const QuatVecF32 = QuatVec{Float32}
const QuatVecF16 = QuatVec{Float16}

# Handy constants like `im`
"""
    imx

The quaternionic unit associated with rotation about the `x` axis.  Can also be entered as Unicode
bold: `𝐢`.

# Examples
```jldoctest
julia> imx * imx
-1 + 0𝐢 + 0𝐣 + 0𝐤
julia> 1.2imx
0.0 + 1.2𝐢 + 0.0𝐣 + 0.0𝐤
```
"""
const imx = QuatVec(false, true, false, false)
const 𝐢 = imx

"""
    imy

The quaternionic unit associated with rotation about the `y` axis.  Can also be entered as Unicode
bold: `𝐣`.

# Examples
```jldoctest
julia> imy * imy
-1 + 0𝐢 + 0𝐣 + 0𝐤
julia> 1.2imy
0.0 + 0.0𝐢 + 1.2𝐣 + 0.0𝐤
```
"""
const imy = QuatVec(false, false, true, false)
const 𝐣 = imy

"""
    imz

The quaternionic unit associated with rotation about the `z` axis.  Can also be entered as Unicode
bold: `𝐤`.

# Examples
```jldoctest
julia> imz * imz
-1 + 0𝐢 + 0𝐣 + 0𝐤
julia> 1.2imz
0.0 + 0.0𝐢 + 0.0𝐣 + 1.2𝐤
```
"""
const imz = QuatVec(false, false, false, true)
const 𝐤 = imz

# Essential constructors
Base.zero(::Type{QT}) where {T<:Real, QT<:AbstractQuaternion{T}} = QT(false, false, false, false)
Base.zero(q::QT) where {T<:Real, QT<:AbstractQuaternion{T}} = Base.zero(QT)
Base.zero(::Type{Rotor}) = throw(DomainError("Rotor", "Zero is not a possible rotor."))
Base.zero(::Type{Rotor{T}}) where T = throw(DomainError("Rotor", "Zero is not a possible rotor."))

Base.one(::Type{QT}) where {T<:Real, QT<:AbstractQuaternion{T}} = QT(true, false, false, false)
Base.one(q::QT) where {T<:Real, QT<:AbstractQuaternion{T}} = Base.one(QT)
Base.one(::Type{QuatVec}) = throw(DomainError("QuatVec", "One is not a possible 3-vector."))
Base.one(::Type{QuatVec{T}}) where T = throw(DomainError("QuatVec", "One is not a possible 3-vector."))

# Getting pieces of quaternions
@inline function Base.getproperty(q::AbstractQuaternion, sym::Symbol)
    @inbounds begin
        if sym === :w
            return q.components[1]
        elseif sym === :x
            return q.components[2]
        elseif sym === :y
            return q.components[3]
        elseif sym === :z
            return q.components[4]
        elseif sym === :re
            return q.components[1]
        elseif sym === :im
            return q.components[2:4]
        elseif sym === :vec
            return q.components[2:4]
        else # fallback to getfield
            return getfield(q, sym)
        end
    end
end
@inline Base.getindex(q::AbstractQuaternion, i::Int) = (@boundscheck checkbounds(q.components,i); q.components[i])
Base.@propagate_inbounds Base.getindex(q::AbstractQuaternion, I) = [q[i] for i in I]
Base.real(::Type{QT}) where {T<:Real, QT<:AbstractQuaternion{T}} = real(T)
Base.real(q::AbstractQuaternion{T}) where {T<:Real} = q.re
Base.imag(q::AbstractQuaternion{T}) where {T<:Real} = q.im

# Type games
wrapper(::T) where {T} = wrapper(T)
wrapper(T::UnionAll) = T
wrapper(T::Type{Q}) where {S<:Real, Q<:AbstractQuaternion{S}} = T.name.wrapper
wrapper(::Type{T}, ::Type{T}) where {T<:AbstractQuaternion} = wrapper(T)
wrapper(::Type{<:AbstractQuaternion}, ::Type{<:AbstractQuaternion}) = Quaternion

wrapper(::Type{<:AbstractQuaternion}, ::Val{OP}, ::Type{<:AbstractQuaternion}) where {OP} = Quaternion
wrapper(::Type{<:AbstractQuaternion}, ::Val{OP}, ::Type{<:Real}) where {OP} = Quaternion
wrapper(::Type{<:Real}, ::Val{OP}, ::Type{<:AbstractQuaternion}) where {OP} = Quaternion
wrapper(::Type{<:AbstractQuaternion}, ::Val{OP}, ::Type{<:Symbolics.Num}) where {OP} = Quaternion
wrapper(::Type{<:Symbolics.Num}, ::Val{OP}, ::Type{<:AbstractQuaternion}) where {OP} = Quaternion

wrapper(::Type{<:Rotor}, ::Val{*}, ::Type{<:Rotor}) = Rotor
wrapper(::Type{<:Rotor}, ::Val{/}, ::Type{<:Rotor}) = Rotor

wrapper(::Type{<:QuatVec}, ::Val{+}, ::Type{<:QuatVec}) = QuatVec
wrapper(::Type{<:QuatVec}, ::Val{-}, ::Type{<:QuatVec}) = QuatVec
wrapper(::Type{<:QuatVec}, ::Val{*}, ::Type{<:Real}) = QuatVec
wrapper(::Type{<:QuatVec}, ::Val{/}, ::Type{<:Real}) = QuatVec
wrapper(::Type{<:Real}, ::Val{*}, ::Type{<:QuatVec}) = QuatVec
wrapper(::Type{<:Real}, ::Val{/}, ::Type{<:QuatVec}) = QuatVec
wrapper(::Type{<:QuatVec}, ::Val{*}, ::Type{<:Symbolics.Num}) = QuatVec
wrapper(::Type{<:QuatVec}, ::Val{/}, ::Type{<:Symbolics.Num}) = QuatVec
wrapper(::Type{<:Symbolics.Num}, ::Val{*}, ::Type{<:QuatVec}) = QuatVec
wrapper(::Type{<:Symbolics.Num}, ::Val{/}, ::Type{<:QuatVec}) = QuatVec

Base.eltype(::Type{<:AbstractQuaternion{T}}) where {T} = T
Base.widen(::Type{Q}) where {Q<:AbstractQuaternion} = wrapper(Q){widen(eltype(Q))}
Base.float(::Type{Q}) where {Q<:AbstractQuaternion{<:AbstractFloat}} = Q
Base.float(::Type{Q}) where {Q<:AbstractQuaternion} = wrapper(Q){float(eltype(Q))}
Base.float(q::AbstractQuaternion{T}) where {T<:AbstractFloat} = q
Base.float(q::AbstractQuaternion{T}) where {T} = wrapper(q){float(T)}(float(q.components))

Base.big(::Type{Q}) where {Q<:AbstractQuaternion} = wrapper(Q){big(eltype(Q))}
Base.big(q::AbstractQuaternion{T}) where {T<:Real} = wrapper(q){big(T)}(q)

Base.promote_rule(::Type{Q}, ::Type{S}) where {Q<:AbstractQuaternion, S<:Real} =
    wrapper(Q){promote_type(eltype(Q),S)}
Base.promote_rule(::Type{QuatVec{T}}, ::Type{S}) where {T<:Real, S<:Real} =
    Quaternion{promote_type(T,S)}
Base.promote_rule(::Type{Q1}, ::Type{Q2}) where {Q1<:AbstractQuaternion, Q2<:AbstractQuaternion} =
    wrapper(wrapper(Q1), wrapper(Q2)){promote_type(eltype(Q1),eltype(Q2))}


"""
Differentiation
"""
function Δ(X,k,dt)
    return (X[:,k] - X[:, k-1])/dt
end

"""
Linspace
"""
function linspace(a,b,n)
    return [collect(range(a, stop=b, length=n))]
end

function linearpath(point_start, point_end, TotalPoint)

    initpath = point_start*ones(1,TotalPoint )

    for i in 1:size(initpath)[2]
        step = (point_end - point_start)*i/size(initpath)[2]
        initpath[:,i] = initpath[:,i] + step
    end
    return initpath
end