import MixedSubdivisions
const MS = MixedSubdivisions

@testset "Binomial System Test" begin

    if Sys.iswindows()
        @testset "GMP set Int64 on 32-bit systems" begin
            x = BigInt()
            a = round(UInt64, rand() * typemax(Int32))
            for y in [a, a + typemax(Int32), a + UInt64(4) * typemax(UInt32)]
                HC.set_ui64!(x, y)
                @test x == y
                HC.set_si64!(x, Int64(y))
                @test x == Int64(y)
                HC.set_si64!(x, -Int64(y))
                @test x == -Int64(y)
            end
            HC.set_si64!(x, Int64(typemin(Int32)))
            @test x == typemin(Int32)
            HC.set_si64!(x, typemin(Int64))
            @test x == typemin(Int64)
        end
    end

    @testset "Hermite Normal Form" begin
        A = [0 3 1 0; -2 2 -1 2; 1 -1 2 3; -3 3 3 2]
        H, U = HC.hnf(A)
        @test A * U == H

        D = zeros(Int, 1, 1)
        D[1, 1] = -2
        H, U = HC.hnf(D)
        @test H[1, 1] == 2
        @test U[1, 1] == -1

        # overflow A
        A = [
            2 4 4 2 -4 -3 -4 3 0 3
            -4 -2 -5 -5 1 0 -3 0 -1 -2
            -3 5 -3 1 4 3 -1 -4 -1 0
            1 3 3 4 -2 -3 -2 -5 -4 -3
            2 2 0 2 -4 4 3 -4 -2 -4
            -3 -3 1 1 -4 4 0 4 4 -4
            -3 -1 -1 -1 2 -4 -3 -4 4 -4
            0 1 -2 4 5 4 3 1 -5 2
            -5 -4 -5 3 1 5 0 -3 -3 -1
            2 -5 -3 -1 -1 5 -4 5 -3 1
        ]
        @test_throws OverflowError HC.hnf(A)
        H, U = HC.hnf(A, BigInt)
        @test A * U == H
    end

    @testset "Solving (unit b)" begin
        A = [0 1 1 0 -1; 0 0 0 1 -1; -1 0 -1 0 -1; 0 -1 0 -1 -1; 1 0 0 0 -1]
        b = [
            0.9053223983046926 + 0.4247250347316951im,
            0.7000429487508004 - 0.7141007421255663im,
            0.018811552539476483 - 0.9998230470893611im,
            -0.9983533473373446 + 0.057363698105327716im,
            0.9999409355072137 - 0.010868555421874721im,
        ]
        X = HC.solve!(HC.BinomialSystemSolver(A, b))
        @test maximum(eachcol(X)) do x
            maximum(abs.([prod(x .^ a) for a in eachcol(A)] - b) ./ abs.(b))
        end < 1e-12

        A = [
            3 2 3 4 2 5
            6 -3 8 -3 8 7
            -2 -5 7 3 6 5
            1 2 3 4 5 6
            1 0 2 0 5 0
            1 2 0 -3 0 5
        ]
        b = randn(ComplexF64, 6)
        b ./= abs.(b)
        X = HC.solve!(HC.BinomialSystemSolver(A, b))
        @test maximum(eachcol(X)) do x
            maximum(abs.([prod(x .^ a) for a in eachcol(A)] - b) ./ abs.(b))
        end < 1e-12
    end

    @testset "General solving" begin
        A = [0 1 1 0 -1; 0 0 0 1 -1; -1 0 -1 0 -1; 0 -1 0 -1 -1; 1 0 0 0 -1]
        b = randn(ComplexF64, size(A, 1))
        X = HC.solve!(HC.BinomialSystemSolver(A, b))
        @test maximum(eachcol(X)) do x
            maximum(abs.([prod(x .^ a) for a in eachcol(A)] - b) ./ abs.(b))
        end < 1e-12

        A = [
            3 2 3 4 2 5
            6 -3 8 -3 8 7
            -2 -5 7 3 6 5
            1 2 3 4 5 6
            1 0 2 0 5 0
            1 2 0 -3 0 5
        ]
        b = randn(ComplexF64, 6)
        X = HC.solve!(HC.BinomialSystemSolver(A, b))
        @test maximum(eachcol(X)) do x
            maximum(abs.([prod(x .^ a) for a in eachcol(A)] - b) ./ abs.(b))
        end < 1e-12
    end

    @testset "MixedCell" begin
        f = cyclic(5)
        supp = first.(ModelKit.exponents_coefficients.(f, Ref(variables(f))))

        cells, lift = MS.fine_mixed_cells(supp)
        coeffs = map(l -> randn(ComplexF64, length(l)), lift)

        cell = first(cells)
        BSS = HC.BinomialSystemSolver(5)
        X = HC.solve(BSS, supp, coeffs, cell)
        @test maximum(eachcol(X)) do x
            maximum(abs.([prod(x .^ a) for a in eachcol(BSS.A)] - BSS.b) ./ abs.(BSS.b))
        end < 1e-12
    end
end


@testset "Simplicial Complex" begin
    cplx = SimplicialComplex()
    @test eltype(cplx) == Simplex
    @test length(cells(cplx)) == 0
    @test length(cells(cplx,2)) == 0
    @test size(cplx) == ()

    cplx = SimplicialComplex(Simplex(1,2,3), Simplex(2,4), Simplex(4))

    @test eltype(cplx) == Simplex
    @test length(cplx) == 9
    @test length(cells(cplx)) == 3
    @test length(cells(cplx,2)) == 1
    @test length(cells(cplx,3)) == 0
    @test size(cplx) == (4,4,1)
    @test size(cplx,-1) == 0
    @test size(cplx,0) == 4
    @test size(cplx,1) == 4
    @test size(cplx,2) == 1
    @test size(cplx,3) == 0
    @test Simplex(4) ∈ cplx
    @test ComputationalHomology.position(cplx, Simplex(4)) == 4
    @test Simplex(10) ∉ cplx

    @test first(Iterators.drop(cplx, 5)) == Simplex(1,3)

    cplx = SimplicialComplex(Simplex(1,2,3),
                       Simplex(2,4),
                       Simplex(3,4),
                       Simplex(5,4),
                       Simplex(6))

    # (co)faces
    fs = [Simplex(4), Simplex(3)]
    @testset "Faces" for fi in faces(cplx, hash(Simplex(3,4)))
        @test cplx[fi,0] ∈ fs
        @test fi ∈ map(hash, fs)
    end
    @test length(faces(cplx, hash(Simplex(1)))) == 0
    fs = [Simplex(2,4), Simplex(3,4), Simplex(5,4)]
    @testset "Cofaces" for fi in cofaces(cplx, hash(Simplex(4)))
        @test cplx[fi,1] ∈ fs
        @test fi ∈ map(hash, fs)
    end
    @test length(cofaces(cplx, hash(Simplex(1,2,3)))) == 0

    # boundary matrix
    @test convert(Matrix, boundary(cplx, 2, Int)) == [1 -1 1 0 0 0]'

    # simplex values
    cplx = SimplicialComplex()

    splx = push!(cplx, Simplex(1))
    @test ComputationalHomology.position(cplx, splx[1]) == 1
    @test dim(splx[1]) == 0
    @test sum(size(cplx)) == 1
    @test indices(cells(cplx, 0)[1]) == [1]

    splx = push!(cplx, Simplex(1,2))
    @test splx[1] == Simplex(1,2)
    @test dim(splx[1]) == 1
    @test sum(size(cplx)) == 2
    @test sort!(indices(cells(cplx, 1)[1])) == [1,2]

    cplx = SimplicialComplex()
    splxs = push!(cplx, Simplex(1, 2), recursive=true)
    @test splxs[1] == Simplex(1,2)
    @test sort!(indices(splxs[1])) == [1, 2]
    @test splxs[3] == Simplex(2)
    @test indices(splxs[3]) == [2]
    @test splxs[2] == Simplex(1)
    @test indices(splxs[2]) == [1]
    @test sum(size(cplx)) == 3
    @test size(cplx, 0) == 2
    @test size(cplx, 1) == 1

    splxs = push!(cplx, Simplex(1,3), recursive=true)
    @test splxs[1] == Simplex(1,3)
    @test sort!(indices(splxs[1])) == [1,3]
    @test splxs[2] == Simplex(3)
    @test indices(splxs[2]) == [3]
    @test sum(size(cplx)) == 5
    @test size(cplx, 0) == 3
    @test size(cplx, 1) == 2

    @test Simplex(1,3) in cplx
    @test first(cells(cplx,1)) in cplx
    @test !(Simplex(2,3) in cplx)
    @test !(Simplex(2,3,3) in cplx)

    io = IOBuffer()
    write(io, cplx)
    seekstart(io)
    cplx2 = read(io, SimplicialComplex(), Simplex{0,Int})
    @test size(cplx2, 0) == 3
    @test size(cplx2, 1) == 2
    @test all(cells(cplx2,1) .== cells(cplx,1))

    A = ComputationalHomology.adjacency_matrix(ComputationalHomology.sphere(1), UInt8)
    @test collect(A) == UInt8[0x00 0x01; 0x01 0x00]
    @test eltype(A) == UInt8
end


using DifferentialEquations
using PyCall
pygui(:qt)
using PyPlot
pygui(true)

include("Dimer_Eq_of_M.jl")
Ham2ndSHO(p, q, t) = 8 * (p[1]^2 + p[2]^2) + (q[1]^2 + q[2]^2);
ω_0 = [1.,1e4]
tspan = (0., 100);
p0 = [1.0,0]
q0 = [0,1.0]
params_1 = [1,2]
params_2 = [1,sqrt(2)]
param = sqrt(2)
# prob = HamiltonianProblem{true}(Ham2ndSHO, p0, q0, tspan);
# sol = solve(prob, RK4(), maxiters = 1e20, reltol = 1e-8, abstol = 1e-10)

# prob_1 = HamiltonianProblem{true}(Ham2ndSHO, p0, q0, params_1, tspan);
# sol_1 = solve(prob_1, RK4(), maxiters = 1e20, reltol = 1e-8, abstol = 1e-10)

# prob_2 = HamiltonianProblem{true}(Ham2ndSHO, p0, q0, tspan);
# sol_2 = solve(prob_2, RK4(), maxiters = 1e20, reltol = 1e-8, abstol = 1e-10)

function condition_hits_PSS_q1(u, t, integrator) # Event when event_f(u,t) == 0
    u[1]
end

function affect_update_iterator_q1!(integrator)
    integrator.u[5] = integrator.u[5] + 1
end

u0=[0.,0,1,0]
callback_hits_PSS_q1_forward = ContinuousCallback(condition_hits_PSS_q1, nothing,  affect_update_iterator_q1!, rootfind = true)
cb_forward = CallbackSet(callback_hits_PSS_q1_forward)
prob=ODEProblem(Dimer_Eq_of_M,u0,tspan )
sol = solve(prob, RK4(), maxiters = 1e20, reltol = 1e-8, abstol = 1e-10, callback = cb_forward, save_start = false, save_end = true, save_everystep = false)
plot(sol[2,:],sol[4,:],".")


using ProgressMeter, KitBase
import KitML

begin
    # space
    x0 = -1.5
    x1 = 1.5
    y0 = -1.5
    y1 = 1.5
    nx = 50#100
    ny = 50#100
    dx = (x1 - x0) / nx
    dy = (y1 - y0) / ny

    pspace = KitBase.PSpace2D(x0, x1, nx, y0, y1, ny)

    # time
    tEnd = 1.0
    cfl = 0.5
    dt = cfl / 2 * (dx * dy) / (dx + dy)

    # quadrature
    quadratureorder = 5
    points, triangulation = KitBase.octa_quadrature(quadratureorder)
    weights = KitBase.quadrature_weights(points, triangulation)
    nq = size(points, 1)

    # particle
    SigmaS = 1 * ones(ny + 4, nx + 4)
    SigmaA = 0 * ones(ny + 4, nx + 4)
    SigmaT = SigmaS + SigmaA

    # moments
    L = 1
    ne = (L + 1)^2
    phi = zeros(ne, nx, ny)
    α = zeros(ne, nx, ny)
    m = KitBase.eval_spherharmonic(points, L)
end

begin
    s2 = 0.03^2
    flr = 1e-4
    init_field(x, y) = max(flr, 1.0 / (4.0 * pi * s2) * exp(-(x^2 + y^2) / 4.0 / s2))
    for j = 1:nx
        for i = 1:ny
            y = y0 + (i - 0.5) * dy
            x = x0 + (j - 0.5) * dx
            # only zeroth order moment is non-zero
            phi[1, i, j] = init_field(x, y)
        end
    end
end

global t = 0.0
flux1 = zeros(ne, nx + 1, ny)
flux2 = zeros(ne, nx, ny + 1)

@showprogress for iter = 1:20
    # regularization
    Threads.@threads for j = 1:ny
        for i = 1:nx
            res = KitBase.optimize_closure(
                α[:, i, j],
                m,
                weights,
                phi[:, i, j],
                KitBase.maxwell_boltzmann_dual,
            )
            α[:, i, j] .= res.minimizer

            phi[:, i, j] .= KitBase.realizable_reconstruct(
                res.minimizer,
                m,
                weights,
                KitBase.maxwell_boltzmann_dual_prime,
            )
        end
    end

    # flux
    fη1 = zeros(nq)
    for j = 1:ny
        for i = 2:nx
            KitBase.flux_kfvs!(
                fη1,
                KitBase.maxwell_boltzmann_dual.(α[:, i-1, j]' * m)[:],
                KitBase.maxwell_boltzmann_dual.(α[:, i, j]' * m)[:],
                points[:, 1],
                dt,
            )

            for k in axes(flux1, 1)
                flux1[k, i, j] = sum(m[k, :] .* weights .* fη1)
            end
        end
    end

    fη2 = zeros(nq)
    for i = 1:nx
        for j = 2:ny
            KitBase.flux_kfvs!(
                fη2,
                KitBase.maxwell_boltzmann_dual.(α[:, i, j-1]' * m)[:],
                KitBase.maxwell_boltzmann_dual.(α[:, i, j]' * m)[:],
                points[:, 2],
                dt,
            )

            for k in axes(flux2, 1)
                flux2[k, i, j] = sum(m[k, :] .* (weights .* fη2))
            end
        end
    end

    # update
    for j = 2:ny-1
        for i = 2:nx-1
            for q = 1:ne
                phi[q, i, j] =
                    phi[q, i, j] +
                    (flux1[q, i, j] - flux1[q, i+1, j]) / dx +
                    (flux2[q, i, j] - flux2[q, i, j+1]) / dy #+
                #(integral - phi[q, i, j]) * dt
            end
        end
    end

    global t += dt
end

X = zeros(Float32, 4, nx * ny)
Y = zeros(Float32, 4, nx * ny)
for j = 1:ny
    for i = 1:nx
        res = KitBase.optimize_closure(
            α[:, i, j],
            m,
            weights,
            phi[:, i, j],
            KitBase.maxwell_boltzmann_dual,
        )
        α[:, i, j] .= res.minimizer
        phi[:, i, j] .= KitBase.realizable_reconstruct(
            res.minimizer,
            m,
            weights,
            KitBase.maxwell_boltzmann_dual_prime,
        )

        X[:, (j-1)*nx+i] .= phi[:, i, j]
        Y[:, (j-1)*nx+i] .= α[:, i, j]
    end
end

using Flux
#data = [(X, Y)]
data = Flux.Data.DataLoader(X, Y, batchsize = 1, shuffle = true)

using Flux: @epochs
```nn = Chain(Dense(4, 16, tanh), Dense(16, 16, tanh), Dense(16, 16, tanh), Dense(16, 16, tanh), Dense(16, 4))```
nn = Chain(
    Dense(4, 16, relu),
    SkipConnection(Chain(Dense(16, 16, relu), Dense(16, 16, relu)), +),
    Dense(16, 4, relu),
)
ps = params(nn)

function loss(x, y)
    sum(abs2, nn(x) - y)
end
cb() = @show(loss(X, Y))
@epochs 10 Flux.train!(loss, ps, data, ADAM(), cb = Flux.throttle(cb, 1))

struct Shortcut{T}
    f::T
    σ::Function
end
Flux.@functor Shortcut
(m::Shortcut)(x) = m.σ.(m.f(x) .+ x)

nn = Chain(
    Dense(4, 16, tanh),
    Shortcut(Chain(Dense(16, 16, relu), Dense(16, 16, relu)), tanh),
    Dense(16, 4, tanh),
)
ps = params(nn)

using DiffEqFlux
ann = FastChain(
    FastDense(4, 16, tanh),
    FastDense(16, 16, tanh),
    FastDense(16, 16, tanh),
    FastDense(16, 4),
)
p = initial_params(ann)
function loss2(θ)
    sum(abs2, ann(X, θ) - Y)
end
callback = function (p, l)
    println("loss: $l")
    return false
end

res = DiffEqFlux.sciml_train(loss2, p, ADAM(), cb = callback, maxiters = 200)
res = DiffEqFlux.sciml_train(
    loss2,
    res.minimizer,
    ADAM(),
    cb = Flux.throttle(callback, 1),
    maxiters = 5000,
)

p, re = Flux.destructure(nn)
function loss3(θ)
    sum(abs2, re(θ)(X) - Y)
end
res = DiffEqFlux.sciml_train(loss3, p, ADAM(), cb = callback, maxiters = 200)


"""
    Weibull(α,θ)

The *Weibull distribution* with shape `α` and scale `θ` has probability density function

```math
f(x; \\alpha, \\theta) = \\frac{\\alpha}{\\theta} \\left( \\frac{x}{\\theta} \\right)^{\\alpha-1} e^{-(x/\\theta)^\\alpha},
    \\quad x \\ge 0
```

```julia
Weibull()        # Weibull distribution with unit shape and unit scale, i.e. Weibull(1, 1)
Weibull(α)       # Weibull distribution with shape α and unit scale, i.e. Weibull(α, 1)
Weibull(α, θ)    # Weibull distribution with shape α and scale θ

params(d)        # Get the parameters, i.e. (α, θ)
shape(d)         # Get the shape parameter, i.e. α
scale(d)         # Get the scale parameter, i.e. θ
```

External links

* [Weibull distribution on Wikipedia](http://en.wikipedia.org/wiki/Weibull_distribution)

"""
struct Weibull{T<:Real} <: ContinuousUnivariateDistribution
    α::T   # shape
    θ::T   # scale

    function Weibull{T}(α::T, θ::T) where {T <: Real}
        new{T}(α, θ)
    end
end

function Weibull(α::T, θ::T; check_args::Bool=true) where {T <: Real}
    check_args && @check_args(Weibull, α > zero(α) && θ > zero(θ))
    return Weibull{T}(α, θ)
end

Weibull(α::Real, θ::Real; check_args::Bool=true) = Weibull(promote(α, θ)...; check_args=check_args)
Weibull(α::Integer, θ::Integer; check_args::Bool=true) = Weibull(float(α), float(θ); check_args=check_args)
Weibull(α::Real=1.0) = Weibull(α, one(α); check_args=false)

@distr_support Weibull 0.0 Inf

#### Conversions

convert(::Type{Weibull{T}}, α::Real, θ::Real) where {T<:Real} = Weibull(T(α), T(θ))
convert(::Type{Weibull{T}}, d::Weibull{S}) where {T <: Real, S <: Real} = Weibull(T(d.α), T(d.θ), check_args=false)

#### Parameters

shape(d::Weibull) = d.α
scale(d::Weibull) = d.θ

params(d::Weibull) = (d.α, d.θ)
partype(::Weibull{T}) where {T<:Real} = T

#### Statistics

mean(d::Weibull) = d.θ * gamma(1 + 1/d.α)
median(d::Weibull) = d.θ * logtwo ^ (1/d.α)
mode(d::Weibull{T}) where {T<:Real} = d.α > 1 ? (iα = 1 / d.α; d.θ * (1 - iα)^iα) : zero(T)

var(d::Weibull) = d.θ^2 * gamma(1 + 2/d.α) - mean(d)^2

function skewness(d::Weibull)
    μ = mean(d)
    σ2 = var(d)
    σ = sqrt(σ2)
    r = μ / σ
    gamma(1 + 3/d.α) * (d.θ/σ)^3 - 3r - r^3
end

function kurtosis(d::Weibull)
    α, θ = params(d)
    μ = mean(d)
    σ = std(d)
    γ = skewness(d)
    r = μ / σ
    r2 = r^2
    r4 = r2^2
    (θ/σ)^4 * gamma(1 + 4/α) - 4γ*r - 6r2 - r4 - 3
end

function entropy(d::Weibull)
    α, θ = params(d)
    0.5772156649015328606 * (1 - 1/α) + log(θ/α) + 1
end

#### Evaluation

function pdf(d::Weibull{T}, x::Real) where T<:Real
    if x >= 0
        α, θ = params(d)
        z = x / θ
        (α / θ) * z^(α - 1) * exp(-z^α)
    else
        zero(T)
    end
end

function logpdf(d::Weibull{T}, x::Real) where T<:Real
    if x >= 0
        α, θ = params(d)
        z = x / θ
        log(α / θ) + (α - 1) * log(z) - z^α
    else
        -T(Inf)
    end
end

zval(d::Weibull, x::Real) = (max(x, 0) / d.θ) ^ d.α
xval(d::Weibull, z::Real) = d.θ * z ^ (1 / d.α)

cdf(d::Weibull, x::Real) = -expm1(- zval(d, x))
ccdf(d::Weibull, x::Real) = exp(- zval(d, x))
logcdf(d::Weibull, x::Real) = log1mexp(- zval(d, x))
logccdf(d::Weibull, x::Real) = - zval(d, x)

quantile(d::Weibull, p::Real) = xval(d, -log1p(-p))
cquantile(d::Weibull, p::Real) = xval(d, -log(p))
invlogcdf(d::Weibull, lp::Real) = xval(d, -log1mexp(lp))
invlogccdf(d::Weibull, lp::Real) = xval(d, -lp)

function gradlogpdf(d::Weibull{T}, x::Real) where T<:Real
    if insupport(Weibull, x)
        α, θ = params(d)
        (α - 1) / x - α * x^(α - 1) / (θ^α)
    else
        zero(T)
    end
end

#### Sampling

rand(rng::AbstractRNG, d::Weibull) = xval(d, randexp(rng))

#### Fit model

"""
    fit_mle(::Type{<:Weibull}, x::AbstractArray{<:Real};
    alpha0::Real = 1, maxiter::Int = 1000, tol::Real = 1e-16)

Compute the maximum likelihood estimate of the [`Weibull`](@ref) distribution with Newton's method.
"""
function fit_mle(::Type{<:Weibull}, x::AbstractArray{<:Real};
    alpha0::Real = 1, maxiter::Int = 1000, tol::Real = 1e-16)

    N = 0

    lnx = map(log, x)
    lnxsq = lnx.^2
    mean_lnx = mean(lnx)

    # first iteration outside loop, prevents type instabililty in α, ϵ

    xpow0 = x.^alpha0
    sum_xpow0 = sum(xpow0)
    dot_xpowlnx0 = dot(xpow0, lnx)

    fx = dot_xpowlnx0 / sum_xpow0 - mean_lnx - 1 / alpha0
    ∂fx = (-dot_xpowlnx0^2 + sum_xpow0 * dot(lnxsq, xpow0)) / (sum_xpow0^2) + 1 / alpha0^2

    Δα = fx / ∂fx
    α = alpha0 - Δα

    ϵ = abs(Δα)
    N += 1

    while ϵ > tol && N < maxiter

        xpow = x.^α
        sum_xpow = sum(xpow)
        dot_xpowlnx = dot(xpow, lnx)

        fx = dot_xpowlnx / sum_xpow - mean_lnx - 1 / α
        ∂fx = (-dot_xpowlnx^2 + sum_xpow * dot(lnxsq, xpow)) / (sum_xpow^2) + 1 / α^2

        Δα = fx / ∂fx
        α -= Δα

        ϵ = abs(Δα)
        N += 1
    end

    θ = mean(x.^α)^(1 / α)
    return Weibull(α, θ)
end


using ProximalOperators
using ProximalOperators: Zero

ProximalOperators.Conjugate(_::Zero) = IndZero()
ProximalOperators.Conjugate(_::IndZero) = Zero()
ProximalOperators.Conjugate(f::SqrNormL2) = SqrNormL2(1.0 / f.lambda)

# TODO: Add other useful functions and calculus rules such as translation


export
    FockState,
    NumberState,
    VacuumState,
    SinglePhotonState

function FockState(T::Type{<:Number}, n::Integer, dim::Integer, ::Type{StateVector})
    v = zeros(T, dim)
    v[n+1] = 1

    return StateVector{T}(v, dim)
end

function FockState(T::Type{<:Number}, n::Integer, dim::Integer, ::Type{StateMatrix})
    𝛒 = zeros(T, dim, dim)
    𝛒[n+1, n+1] = 1

    return StateMatrix{T}(𝛒, dim)
end

"""
    FockState(n; T=ComplexF64, dim=DIM, rep=StateVector)

Construct a Fock state in `rep` representation.

* `n`: Photon number of Fock state.
* `T`: Numeric data types, default is `ComplexF64`
* `dim`: Maximum photon number for truncate, default is $DIM
* `rep`: In which representation, default is `StateVector`.

# Examples
```jldoctest
julia> state = FockState(1);

julia> state = FockState(1, dim=100);

julia> state = FockState(1, rep=StateMatrix);
```
"""
FockState(T::Type{<:Number}, n; dim=DIM, rep=StateVector) = FockState(T, n, dim, rep)
FockState(n; dim=DIM, rep=StateVector) = FockState(ComplexF64, n, dim, rep)

"""
    NumberState(n; T=ComplexF64, dim=DIM, rep=StateVector)

Exactly the same with `FockState`.
"""
NumberState(T::Type{<:Number}, n; dim=DIM, rep=StateVector) = FockState(T, n, dim, rep)
NumberState(n; dim=DIM, rep=StateVector) = FockState(ComplexF64, n, dim, rep)

"""
    VacuumState(; T=ComplexF64, dim=DIM, rep=StateVector)

Construct a vacuum state in `rep` representation.

* `T`: Numeric data types, default is `ComplexF64`
* `dim`: Maximum photon number for truncate, default is $DIM
* `rep`: In which representation, default is `StateVector`.

# Examples
```jldoctest
julia> state = VacuumState();

julia> state = VacuumState(dim=100);

julia> state = VacuumState(rep=StateMatrix);
```
"""
VacuumState(T::Type{<:Number}; dim=DIM, rep=StateVector) = FockState(T, 0, dim, rep)
VacuumState(; dim=DIM, rep=StateVector) = FockState(ComplexF64, 0, dim, rep)

"""
    SinglePhotonState(; T=ComplexF64, dim=DIM, rep=StateVector)

Construct a single photon state in `rep` representation.

* `T`: Numeric data types, default is `ComplexF64`
* `dim`: Maximum photon number for truncate, default is $DIM
* `rep`: In which representation, default is `StateVector`.

# Examples
```jldoctest
julia> state = SinglePhotonState();

julia> state = SinglePhotonState(dim=100);

julia> state = SinglePhotonState(rep=StateMatrix);
```
"""
SinglePhotonState(T::Type{<:Number}; dim=DIM, rep=StateVector) = FockState(T, 1, dim, rep)
SinglePhotonState(; dim=DIM, rep=StateVector) = FockState(ComplexF64, 1, dim, rep)


#rump's formula

using Unums
include("../methods/unum_optimizer.jl")

vtypes = [Float16, Float32, Float64, BigFloat]

function type_rump(T, x, y)
  A = T(333.75)
  B = T(11)
  C = T(121)
  D = T(2)
  E = T(5.5)

  A * y^6 + (x^2) * (B * x^2 * y^2 - y^6 - C * y^4 - D) + E * y^8 + x / (D * y)
end

for T in vtypes
  x = convert(T, 77617)
  y = convert(T, 33096)

  z = Float64(type_rump(T, x, y))

  println("rump's formula result in $(T) : $(z)")
end

function u_rump(T)
  x = T(77617)
  y = T(33096)

  type_rump(T, x, y)
end

r = optimize(u_rump, 0.1, 1, 1, verbose = true)


mutable struct OrbitSystem{T, WT<:AbstractArray, ST<:AbstractArray,
                           SiT<:AbstractArray, GT<:AbstractArray,
                           TT<:AbstractArray}
    # sum(((W*f .- d)./err).^2) + (f .- mu)'*inv(alpha*S)*(f .- mu)
    W::WT
    d::Vector{T}
    err::Vector{T}
    S::ST
    S_inv::SiT
    G::GT
    T::TT
    mu::Vector{T}
    alpha::T
end

function Base.show(io::IO, OS::OrbitSystem)
    println(io, "OrbitSystem")
    println(io, " # of measurements: $(length(OS.d))")
    print(io,   " # of orbits: $(size(OS.W,2))")
end

function OrbitSystem(W, d, err, S; S_inv = inv(S), G = cholesky(S_inv).U,
                     T=Diagonal(ones(size(W,2))), mu = zeros(size(W,2)), alpha=1.0)

    OrbitSystem(W,d,err,S,S_inv,G,T,mu,alpha)
end

function lcurve_point(OS::OrbitSystem, f::Vector)
    W = OS.W
    d = OS.d
    err = OS.err
    mu = OS.mu
    S_inv = OS.S_inv

    chisq = sum(((W*f .- d)./err).^2)
    R = (f .- mu)'*S_inv*(f .- mu)
    return chisq, R
end

function lcurve_point(OS::OrbitSystem; kwargs...)
    f = solve(OS; kwargs...)
    return lcurve_point(OS,f)
end

function lcurve(OS; n = 50, log_bounds = (-6,6), kwargs...)
    alpha = 10.0.^range(log_bounds...,length=n)
    chisq = zeros(n)
    R = zeros(n)
    for i=1:n
        OS.alpha = alpha[i]
        chisq[i], R[i] = lcurve_point(OS; kwargs...)
    end
    return alpha, chisq, R
end

function marginal_loglike(OS::OrbitSystem; norm=1e18/size(OS.W,2),nonneg=false,max_iter=30*size(OS.W,1))

    d = OS.d
    err = OS.err
    Σ_d = Diagonal(err.^2)
    Σ_d_inv = inv(Σ_d)

    K = norm*OS.W

    mu_X = OS.mu/norm

    # Scale covariance matrices
    Σ_X = OS.alpha*OS.S
    Σ_X_inv = inv(OS.alpha)*OS.S_inv

    Σ_inv = Σ_X_inv .+ K'*Σ_d_inv*K

    if nonneg
        Γ = sqrt(inv(OS.alpha))*OS.G
        X = vec(nonneg_lsq(vcat(K./err, Γ), vcat(d./err, Γ*mu_X),
                           alg=:fnnls,use_parallel=false,max_iter=max_iter))
    else
        X = Σ_inv \ (K'*(Σ_d_inv*d) + Σ_X_inv*mu_X)
    end

    l = 0.0
    try
        dhat = K*X
        l = -logabsdet(Array(Σ_X))[1] - logabsdet(Array(Σ_inv))[1] -
            ((d .- dhat)'*Σ_d_inv*(d .- dhat))[1] - ((X .- mu_X)'*Σ_X_inv*(X .- mu_X))[1]
    catch er
        if isa(er,InterruptException)
            rethrow(er)
        end
        l = -Inf
    end

    return l
end

function optimize_alpha!(OS; log_bounds = (-6,6), kwargs...)
    f = x -> begin
        OS.alpha = 10.0^x
        return -marginal_loglike(OS; kwargs...)
    end

    op = optimize(f, log_bounds[1], log_bounds[2], Brent(),rel_tol=1e-3)

    OS.alpha = 10.0^Optim.minimizer(op)

    return Optim.minimum(op)
end

function estimate_rtol(S)
    s = filter(x -> x > 0, eigvals(S))
    sn = log10.(reverse(s)/s[end])
    n = length(sn)
    x = collect(range(0,1,length=n))
    f = PolyharmonicSpline(3,x,sn,s=0.1)
    df2(t) = ForwardDiff.derivative(tt->ForwardDiff.derivative(f,tt),t)
    roots = find_zeros(df2,0.2,0.8)
    return 10.0^(f(roots[1]))
end

function optimize_parameters(make_orbit_system::Function, lbounds, ubounds;
                             niter=10, nseed = 15, objective::Function = optimize_alpha!,
                             warmstart=false, checkpoint=true,
                             file="optim_progress.jld2",verbose=true, kwargs...)

    s = Sobol.SobolSeq(lbounds,ubounds)
    points = Vector{Float64}[]
    values = Float64[]

    if warmstart && isfile(file) && (filesize(file) != 0)
        s, points, values = load(file, "s", "points", "values")
    end

    if checkpoint
        touch(file)
    end

    while s.s.n < nseed
        p = Sobol.next!(s)
        OS = make_orbit_system(p)
        v = objective(OS; kwargs...)
        push!(points, p)
        push!(values, v)
        if checkpoint
            @save file s points values
        end
        if verbose
            println("Seed Number: $(s.s.n)")
            println("Seed Point:  ", repr(p))
            println("Seed Value:  ", v)
            println("Best Seed: ", repr(points[argmin(values)]))
            println("")
        end
    end

    for i=1:niter
        spl = PolyharmonicSpline(2, Array(hcat(points...)'), values)
        xstart = s.lb .+ s.ub.*rand(length(s.ub))
        op = optimize(x -> spl(x...), s.lb, s.ub, xstart, Fminbox(NelderMead()))
        p = Optim.minimizer(op)
        pv = Optim.minimum(op)
        if p in points
            p = Sobol.next!(s)
            pv = spl(p...)
        end
        if verbose
            println("Iteration: $i")
            println("Guess Point:  ",repr(p))
            println("Guess Value:  ", pv)
        end
        OS = make_orbit_system(p)
        v = objective(OS; kwargs...)
        if isfinite(v)
            push!(points, p)
            push!(values, v)
        end
        if verbose
            println("Actual Value: ", v)
            println("Best Point: ", repr(points[argmin(values)]))
            println("")
        end
        if checkpoint
            @save file s points values
        end
    end

    return points[argmin(values)]
end

function inv_chol(Σ; rtol=0.0)
    if rtol == 0.0
        Σ_inv = inv(Σ)
        Γ = cholesky(Hermitian(Σ_inv)).U
    else
        S = svd(Σ)
        smax = maximum(S.S)
        svals_inv = Diagonal([s >= smax*rtol ? 1/s : 0.0 for s in S.S])
        Σ_inv = (svals_inv*S.Vt)' * S.U'
        Γ = sqrt.(svals_inv)*S.Vt
    end
    return Σ_inv, Γ
end

function inv_chol(S::LinearAlgebra.SVD; rtol=0.0)
    if rtol == 0.0
        Σ_inv = inv(S)
        Γ = cholesky(Hermitian(Σ_inv)).U
    else
        smax = maximum(S.S)
        svals_inv = Diagonal([s >= smax*rtol ? 1/s : 0.0 for s in S.S])
        Σ_inv = (svals_inv*S.Vt)' * S.U'
        Γ = sqrt.(svals_inv)*S.Vt
    end
    return Σ_inv, Γ
end

function solve(OS::OrbitSystem; norm=1e18/size(OS.W,2), nonneg=true,max_iter=30*size(OS.W,1))

    d = OS.d
    err = OS.err
    Σ_d = Diagonal(err.^2)
    Σ_d_inv = inv(Σ_d)

    K = norm*OS.W

    mu_X = OS.mu/norm

    # Scale covariance matrices
    Σ_X = OS.alpha*OS.S
    Σ_X_inv = inv(OS.alpha)*OS.S_inv
    Γ = sqrt(inv(OS.alpha))*OS.G

    if nonneg
        try
            X = vec(nonneg_lsq(vcat(K./err, Γ), vcat(d./err, Γ*mu_X),
                               alg=:fnnls,use_parallel=false,max_iter=max_iter))
        catch er
            if isa(er,InterruptException)
                rethrow(er)
            end
            @warn "Non-negative Least Squares failed. Using MAP estimate"
            println(err)
            Σ_inv = Σ_X_inv .+ K'*Σ_d_inv*K
            X = Σ_inv \ (K'*(Σ_d_inv*d) + Σ_X_inv*mu_X)
        end
    else
        Σ_inv = Σ_X_inv .+ K'*Σ_d_inv*K
        X = Σ_inv \ (K'*(Σ_d_inv*d) + Σ_X_inv*mu_X)
    end

    return norm*X
end


using NeutralLandscapes
using Test

pl33 = rand(PlanarGradient(0.0), (3,3))
@test pl33 == [1 1 1; 1/2 1/2 1/2; 0 0 0]

testmask = [true false true; false true false; true true true]
pl33m = rand(PlanarGradient(0.0), (3,3); mask=testmask)
@test all(pl33m .=== [1 NaN 1; NaN 1/2 NaN; 0 0 0])

@test rand(PlanarGradient(0.0), (5, 5)) == rand(EdgeGradient(0.0), (5, 5))

mat = rand(Float64, (10, 10))
rand!(mat, PlanarGradient(90.0))


import Dierckx
import Calculus

using Formatting
using Crayons.Box

struct SplineFunction <: Function
    spline::Dierckx.Spline1D
    function_object::Function
    firstderivative_object::Function
    a::Float64
    b::Float64
    y_min::Float64
    y_max::Float64
    Δx_min::Float64
    Δx_max::Float64
    Δy_min::Float64
    Δy_max::Float64

    function SplineFunction(X::Vector{Float64}, Y::Vector{Float64})
        @assert(!isempty(X), "X must non be empty.")
        @assert(!isempty(Y), "Y must non be empty.")
        @assert(length(X) > 2, "The length of X must greater than 2.")
        @assert(length(Y) > 2, "The length of Y must greater than 2.")
        @assert(issorted(X), "The values in X must be sorted in the ascending order.")
        @assert(length(X) == length(Y), "X and Y must have equal sizes.")

        a = X[1]; b = X[end]
        y_min = minimum(Y); y_max = maximum(Y)
        Δx_min = minimum(X[2:end] .- X[1:end-1]); Δx_max = maximum(X[2:end] .- X[1:end-1])
        Δy_min = minimum(abs.(Y[2:end] .- Y[1:end-1])); Δy_max = maximum(abs.(Y[2:end] .- Y[1:end-1]))
        spline = Dierckx.Spline1D(X, Y; w=ones(length(X)), k=2, bc="extrapolate", s=0.0)

        this = new(spline,
            x::Number -> Dierckx.evaluate(spline, x),
            x::Number -> Dierckx.derivative(spline, x),
            a, b,
            y_min, y_max,
            Δx_min, Δx_max, Δy_min, Δy_max)

        return this
    end

    function (sfunction::SplineFunction)(x::Number)
        return sfunction.function_object(x)
    end
end
function Base.show(io::IO, sfunction::SplineFunction)
    compact = get(io, :compact, false)

    a = sfunction.a; b = sfunction.b
    y_min = sfunction.y_min; y_max = sfunction.y_max
    Δx_min = sfunction.Δx_min; Δx_max = sfunction.Δx_max
    Δy_min = sfunction.Δy_min; Δy_max = sfunction.Δy_max
    spline = sfunction.spline
    order = spline.k

    sb = IOBuffer()
    if compact
        printfmt(sb, "𝐒𝐩𝐥𝐢𝐧𝐞(𝐗 ∈ [{1:.6e}, {2:.6e}], 𝐘 ∈ [{3:.6e}, {4:.6e}];\n$(' '^15)𝚶 = {5}; Δxᵐⁱⁿ = {6:.6e}, Δxᵐᵃˣ = {7:.6e})",
            a, b, y_min, y_max, order, Δx_min, Δx_max)
    else
        printfmt(sb, "𝐒𝐩𝐥𝐢𝐧𝐞(𝐗 ∈ [{1:.6e}, {2:.6e}], 𝐘 ∈ [{3:.6e}, {4:.6e}];\n$(' '^15)𝚶 = {5}; Δxᵐⁱⁿ = {6:.6e}, Δxᵐᵃˣ = {7:.6e}; |Δyᵐⁱⁿ| = {8:.6e}, |Δyᵐᵃˣ| = {9:.6e})",
            a, b, y_min, y_max, order, Δx_min, Δx_max, Δy_min, Δy_max)
    end
    print(io, LIGHT_GRAY_FG(String(take!(sb))))
end
function Base.show(io::IO, ::MIME"text/plain", sfunction::SplineFunction)
    show(io, sfunction)
end

function ⋮(sfunction::SplineFunction, x::Number)
    return sfunction.firstderivative_object(x)
end
function d2(sfunction::SplineFunction, x::Number)
    return Calculus.derivative(sfunction.firstderivative_object, x)
end
function ∫(sfunction::SplineFunction, (a, b)::Tuple{Number, Number})
    return Dierckx.integrate(sfunction.spline, a, b)
end
function ⩪(sfunction::SplineFunction)
    return Dierckx.roots(sfunction.spline)
end


function sqarea_triangle(x::Array{T,1}) where {T <: Real}
# How to compute (twice) area squared of triangle with
# high precision (Goldberg 1991):
#a,b,c=sort([a,b,c],rev=true)
a=maximum(x); b=median(x); c=minimum(x)
area = (a+(b+c))*(c-(a-b))*(c+(a-b))*(a+(b-c))
return area
end

# Computes a limb-darkened transit light curve with the dependence:
# I(\mu) = 1-\sum_{n=1}^N u_n (1-\mu)^n
# where \mu = \cos{\theta} = z is the cosine of the angle
# angle from the sub-stellar point (or equivalently the
# height on the star relative to the sky plane if the radius
# of the star is unity.

#include("sn.jl")
include("s2.jl")
include("IJv_derivative.jl")
#include("area_triangle.jl")

function transit_poly(r::T,b::T,u_n::Array{T,1}) where {T <: Real}
# Transform the u_n coefficients to c_n, which are coefficients
# of the basis in which the P(G_n) functions are computed.
n = length(u_n)
#c_n = zeros(typeof(r),n+3)
#a_n = zeros(typeof(r),n+1)
#c_n = Array{typeof(r)}(n+3)
#c_n[n+2] = zero(T)
#c_n[n+3] = zero(T)
fill!(a_n,zero(T))
a_n[1] = one(r)  # Add in the first constant coefficient term
for i=1:n
  # Compute the contribution to a_n*\mu^n
  bcoeff= one(T)
  for j=0:i
#    a_n[j+1] -= u_n[i]*binomial(i,j)*(-1)^j
    a_n[j+1] -= u_n[i]*bcoeff*(-1)^j
    bcoeff *= (i-j)/(j+1)
#    println("i: ",i," j: ",j," a_i: ",a_n[j+1])
  end
end
# Now, compute the c_n coefficients:
for j=n:-1:2
  c_n[j+1] = a_n[j+1]/(j+2)+c_n[j+3]
end
c_n[2] = a_n[2]+3*c_n[4]
c_n[1] = a_n[1]+2*c_n[3]
#println("u_n: ",u_n)
#println("a_n: ",a_n)
#if typeof(r) == Float64
#  println("c_n: ",c_n)
#end
return transit_poly_c(r,b,c_n[1:n+1])
end

function transit_poly_c(r::T,b::T,c_n::Array{T,1}) where {T <: Real}
# Number of limb-darkening components to include (beyond 0 and 1):
N_c = length(c_n)-1
# We are parameterizing these with the function:
# g_n = c_n [(n+2) z^n - n z^{n-2}] for n >= 2
# while g_{0} = c_0 z^0 (uniform source) and g_{1} = c_1 z^1 (linear limb-darkening)
# which gives a Green's integral of:
# P(G_n) = \int_{\pi-\phi}^{2\pi+\phi} (1-r^2-b^2-2br s_\varphi)^{n/2} (r+b s_\varphi) d\varphi
# for which we have a solution in terms of I_v (for even n) and J_v (for odd n).

# Set up a vector for storing results of P(G_n)-Q(G_n); note that
# this is a different vector than the Starry case:
#sn = zeros(typeof(r),N_c+1)
#sn = Array{T}(N_c+1)
#fill!(sn,zero(T))

# Check for different cases:
if b >= 1+r
  # unobscured - return one:
  return one(r)
end
if r >= 1+b
  # full obscuration - return zero:
  return zero(r)
end
if b == 0.0
  # Annular eclipse - integrate around the full boundary of both bodies:
  flux = zero(r); sqrt1mr2 = sqrt(1-r^2)
  flux = (c_n[1]*(1-r^2)+2/3*c_n[2]*sqrt1mr2^3)
  fac= 2r^2*(1-r^2)
  for i=2:N_c
    flux += -c_n[i+1]*fac
    fac *= sqrt1mr2
  end
  return flux/(c_n[1]+2*c_n[2]/3)
else
# Next, compute k^2 = m:
  onembmr2=(r+1-b)*(1-r+b); fourbr = 4b*r
#  onembpr2 = (1-r-b)*(1+b+r); onembmr2=(r-b+1)*(1-r+b); fourbr = 4b*r
#  k2 = onembpr2/fourbr+1
  k2 = onembmr2/fourbr
  if k2 > 1
    if k2 > 2.0
      kc = sqrt(1.-inv(k2))
    else
#      kc2 = (1-(b+r)^2)/(1-(b-r)^2)
      kc2 = (1-r-b)*(1+b+r)/(1-b+r)/(1-r+b)
      kc = sqrt(kc2)
    end
  else
    if k2 > 0.5
#      kc2 = ((b+r)^2-1)/(4*b*r)
      kc2 = (r-1+b)*(b+r+1)/(4*b*r)
      kc = sqrt(kc2)
    else
      kc = sqrt(1.-k2)
    end
  end
end

# Compute the highest value of v in J_v or I_v that we need:
if iseven(N_c)
  v_max = round(Int64,N_c/2)+2
else
  v_max = round(Int64,(N_c-1)/2)+2
end
# Compute sn[1] and sn[2]:
# Uniform disk case:
if b <= 1-r
  lam = pi*r^2
  sn[1] = pi-lam
  kap0 = convert(typeof(r),pi); kck = zero(r)
else
  # Twice area of kite-shaped region connecting centers of circles & intersection points:
  kite_area2 = sqrt(sqarea_triangle([one(r),b,r]))
  # Angle of section for occultor:
  kap0  = atan2(kite_area2,(r-1)*(r+1)+b^2)
  # Angle of section for source:
  pimkap1 = atan2(kite_area2,(r-1)*(r+1)-b^2)
  # Flux of visible uniform disk:
  sn[1] = pimkap1 - r^2*kap0 + .5*kite_area2
  kck = kite_area2/(4*b*r)
end
sn[2] = s2(r,b)
#if typeof(r) == Float64
# Compute the J_v and I_v functions:
#Iv = zeros(typeof(k2),v_max+1); Jv = zeros(typeof(k2),v_max+1)
#Iv = Array{T}(v_max+1); Jv = Array{T}(v_max+1)
#Iv = Array{T}(v_max+1); Jv = Array{T}(v_max+1)
if k2 > 0
  if (k2 < 0.5 || k2 > 2.0) # && v_max > 3
# This computes I_v,J_v for the largest v, and then works down to smaller values:
    IJv_lower!(v_max,k2,kck,kc,kap0,Iv,Jv)
  else
# This computes I_0,J_0,J_1, and then works upward to larger v:
    IJv_raise!(v_max,k2,kck,kc,kap0,Iv,Jv)
  end
end

#nphi = 1000; dphi=2pi/nphi; phigrid = linspace(.5*dphi,1-.5*dphi,nphi)
# Next, loop over the Green's function components:
for n=2:N_c
  pofgn = zero(r)
  if iseven(n)
# For even values of n, sum over I_v:
    n0 = convert(Int64,n/2)
    coeff = (-fourbr)^n0
    # Compute i=0 term
    pofgn = coeff*((r-b)*Iv[n0+1]+2b*Iv[n0+2])
# For even n, compute coefficients for the sum over I_v:
#    println("n0: ",n0," i: ",0," coeff: ",coeff)
    for i=1:n0
      coeff *= -(n0-i+1)/i*k2
#      println("n0: ",n0," i: ",i," coeff: ",coeff)
      pofgn += coeff*((r-b)*Iv[n0-i+1]+2b*Iv[n0-i+2])
    end
    pofgn *= 2r
  else
# Now do the same for odd N_c in sum over J_v:
    n0 = convert(Int64,(n-3)/2)
    coeff = (-fourbr)^n0
    # Compute i=0 term
    pofgn = coeff*((r-b)*Jv[n0+1]+2b*Jv[n0+2])
#    println("n0: ",n0," i: ",0," coeff: ",coeff)
# For even n, compute coefficients for the sum over I_v:
    for i=1:n0
      coeff *= -(n0-i+1)/i*k2
#      println("n0: ",n0," i: ",i," coeff: ",coeff)
      pofgn += coeff*((r-b)*Jv[n0-i+1]+2b*Jv[n0-i+2])
    end
    pofgn *= 2r*onembmr2^1.5
  end
#  pofgn_num = sum(sqrt.(1-r^2-b^2-2*b*r*sin.(phigrid)).^n.*(r+b.*sin.(phigrid))*r*dphi)
#  println("n: ",n," P(G_n): ",pofgn," P(G_n),num: ",pofgn_num)
# Q(G_n) is zero in this case since on limb of star z^n = 0 at the stellar
# boundary for n > 0.
# Compute sn[n]:
  #println("n: ",n," P(G_n): ",pofgn)
  sn[n+1] = -pofgn
end
#  println("r: ",r," b: ",b," s2 error: ",convert(Float64,s2(big(r),big(b)))-sn[2])
#end
# That's it!
#println("s_n: ",sn)
#println("c_n*s_n: ",c_n.*sn)
flux = sum(c_n.*sn)/(pi*(c_n[1]+2*c_n[2]/3))  # for c_2 and above, the flux is zero.
return flux
end
# That's it!

# Now, for the versions of these functions which include derivatives:
function transit_poly!(r::T,b::T,u_n::Array{T,1},dfdrbu::Array{T,1}) where {T <: Real}
# Transform the u_n coefficients to c_n, which are coefficients
# of the basis in which the P(G_n) functions are computed.
# Compute the derivatives of the flux with respect to the u coefficients.
n = length(u_n)
# We define c_n with two extra elements which are zero:
#c_n = zeros(typeof(r),n+3)
#dfdrbc = zeros(typeof(r),n+3)
#a_n = zeros(typeof(r),n+1)
#dadu = zeros(typeof(r),n+1,n)
#dcdu = zeros(typeof(r),n+3,n)
#c_n = Array{T}(n+3)
c_n[n+2]=zero(T)
c_n[n+3]=zero(T)
#dfdrbc = Array{T}(n+3)
#a_n = Array{T}(n+1)
#dadu = Array{T}(n+1,n)
#dcdu = Array{T}(n+3,n)
#fill!(c_n,zero(r))
#fill!(dfdrbc,zero(r))
fill!(a_n,zero(T))
fill!(dadu,zero(T))
#fill!(dcdu,zero(r))
a_n[1] = one(r)  # Add in the first constant coefficient term
for i=1:n
  # Compute the contribution to a_n*\mu^n
  bcoeff = one(T)
  for j=0:i
#    a_n[j+1] -= u_n[i]*binomial(i,j)*(-1)^j
    a_n[j+1] -= u_n[i]*bcoeff*(-1)^j
#    dadu[j+1,i] -= binomial(i,j)*(-1)^j
    dadu[j+1,i] -= bcoeff*(-1)^j
    bcoeff *= (i-j)/(j+1)
#    println("i: ",i," j: ",j," a_i: ",a_n[j+1])
  end
end
# Now, compute the c_n coefficients and propagate derivatives:
for j=n:-1:2
  c_n[j+1] = a_n[j+1]/(j+2)+c_n[j+3]
  for i=1:n
    dcdu[j+1,i] = dadu[j+1,i]/(j+2) + dcdu[j+3,i]
  end
end
c_n[2] = a_n[2]+3*c_n[4]
for i=1:n
  dcdu[2,i] = dadu[2,i] + 3*dcdu[4,i]
end
c_n[1] = a_n[1]+2*c_n[3]
for i=1:n
  dcdu[1,i] = dadu[1,i] + 2*dcdu[3,i]
end
#println("a_n: ",a_n)
#if typeof(r) == Float64
#  println("u_n: ",u_n)
#  println("c_n: ",c_n)
#  println("dcdu: ",dcdu)
#end
# Pass c_n (without last two dummy values):
flux = transit_poly_c!(r,b,c_n[1:n+1],dfdrbc)
# Now, transform derivaties from c to u:
fill!(dfdrbu,zero(r))
dfdrbu[1] = dfdrbc[1]  # r derivative
dfdrbu[2] = dfdrbc[2]  # b derivative
# u_n derivatives:
for i=1:n, j=0:n
  dfdrbu[i+2] += dfdrbc[j+3]*dcdu[j+1,i]
end
#println("dcdu: ",dcdu," dfdrbc: ",dfdrbc," dfdrbu: ",dfdrbu)
return flux
end

function transit_poly_c!(r::T,b::T,c_n::Array{T,1},dfdrbc::Array{T,1}) where {T <: Real}
@assert((length(c_n)+2) == length(dfdrbc))
@assert(r > 0)
# Number of limb-darkening components to include (beyond 0 and 1):
N_c = length(c_n)-1
# We are parameterizing these with the function:
# g_n = c_n [(n+2) z^n - n z^{n-2}] for n >= 2
# while g_{0} = c_0 z^0 (uniform source) and g_{1} = c_1 z^1 (linear limb-darkening)
# which gives a Green's integral of:
# P(G_n) = \int_{\pi-\phi}^{2\pi+\phi} (1-r^2-b^2-2br s_\varphi)^{n/2} (r+b s_\varphi) d\varphi
# for which we have a solution in terms of I_v (for even n) and J_v (for odd n).
# Compute the derivative of the flux with respect to the different coefficients.

# Set up a vector for storing results of P(G_n)-Q(G_n); note that
# this is a different vector than the Starry case:
#sn = Array{T}(N_c+1)
#dsndr = Array{T}(N_c+1)
#dsndb = Array{T}(N_c+1)
#sn = zeros(typeof(r),N_c+1)
#dsndr = zeros(typeof(r),N_c+1)
#dsndb = zeros(typeof(r),N_c+1)
#fill!(sn,zero(r))
#fill!(dsndr,zero(r))
#fill!(dsndb,zero(r))
fill!(dfdrbc,zero(r))
# Check for different cases:
if b >= 1+r || r ==  0.0
  # unobscured - return one:
  return one(r)
end
if r >= 1+b
  # full obscuration - return zero:
  return zero(r)
end
if b == 0.0
  # Annular eclipse - integrate around the full boundary of both bodies:
  flux = zero(r); onemr2 = 1-r^2; sqrt1mr2 = sqrt(onemr2)
  den = inv(c_n[1]+2*c_n[2]/3)
  flux = (c_n[1]*onemr2+2/3*c_n[2]*sqrt1mr2^3)*den
  fac  = 2r^2*onemr2*den
  facd = -2r*den
  dfdrbc[1] = c_n[1]*facd + c_n[2]*facd*sqrt1mr2
  for i=2:N_c
    flux -= c_n[i+1]*fac
    dfdrbc[1] += c_n[i+1]*facd*(2*onemr2-i*r^2)
    dfdrbc[i+3] -= fac
    fac *= sqrt1mr2
    facd *= sqrt1mr2
  end
  #  dfdrbc[2]=0 since the derivative with respect to b is zero.
  dfdrbc[3] = (onemr2-flux)*den
  dfdrbc[4] = 2/3*(sqrt1mr2^3-flux)*den
  # Also need to compute derivatives [ ]
  return flux
else
# Next, compute k^2 = m:
  onembmr2=(r-b+1)*(1-r+b); fourbr = 4b*r
  k2 = onembmr2/fourbr;
  if k2 > 0
    k = sqrt(k2)
  else
    println("negative k2: ",k2," r: ",r," b: ",b)
  end
#  onembpr2 = (1-b-r)*(1+b+r); onembmr2=(r-b+1)*(1-r+b); fourbr = 4b*r
#  k2 = onembpr2/fourbr+1; k = sqrt(k2)
  dkdr = (b^2-r^2-1)/(8*k*b*r^2)
#  dkdr = ((b-r)*(b+r)-1)/(8*k*b*r^2)
  dkdb = (r^2-b^2-1)/(8*k*b^2*r)
#  dkdb = ((r-b)*(r+b)-1)/(8*k*b^2*r)
  if k2 > 1
    if k2 > 2.0
      kc = sqrt(1.-inv(k2))
    else
#      kc2 = (1-(b+r)^2)/(1-(b-r)^2)
      kc2 = (1-r-b)*(1+b+r)/(1+r-b)/(1-r+b)
      kc = sqrt(kc2)
    end
#    if typeof(kc) == Float64
#      kc_error = kc-convert(Float64,sqrt((big(1.0)-big(b)-big(r))*(big(1.0)+big(b)+big(r))/(big(1.0)-big(b)+big(r))/(big(1.0)-big(r)+big(b))))
#      println("kc: ",kc," error on kc: ",kc_error)
#    end
  else
    if k2 > 0.5
#      kc2 = ((b+r)^2-1)/(4*b*r)
      kc2 = (r-1+b)*(b+r+1)/(4*b*r)
      kc = sqrt(kc2)
    else
      kc2 = (r-1+b)*(b+r+1)/(4*b*r)
      kc = sqrt(kc2)
#      kc = sqrt(1.-k2)
    end
#    if typeof(kc) == Float64
#      kc_error = kc-convert(Float64,sqrt((big(b)+big(r)-big(1.0))*(big(1.0)+big(b)+big(r))/(big(4.0)*big(b)*big(r))))
#      println("kc: ",kc," error on kc: ",kc_error)
#    end
  end
end

# Compute the highest value of v in J_v or I_v that we need:
if iseven(N_c)
  v_max = round(Int64,N_c/2)+2
else
  v_max = round(Int64,(N_c-1)/2)+2
end
#println("v_max: ",v_max," N_c: ",N_c)
# Compute sn[1] and its derivatives:
if b <= 1-r  # k^2 > 1
  lam = pi*r^2
  sn[1] = pi-lam
  dsndr[1] = -2*pi*r
  dsndb[1] = 0.
  kap0 = convert(typeof(r),pi); kck = zero(r)
else
  # Twice area of kite-shaped region connecting centers of circles & intersection points:
  kite_area2 = sqrt(sqarea_triangle([one(r),b,r]))
  # Angle of section for occultor:
  kap0  = atan2(kite_area2,(r-1)*(r+1)+b^2)
  # Angle of section for source:
  pimkap1 = atan2(kite_area2,(r-1)*(r+1)-b^2)
  # Flux of visible uniform disk:
  sn[1] = pimkap1 - r^2*kap0 + .5*kite_area2
  dsndr[1]= -2*r*kap0
  dsndb[1]= kite_area2/b
  kck = kite_area2/(4*b*r)
end
# Compute sn[2] and its derivatives:
#s2_grad = zeros(typeof(r),2)
#s2_grad = Array{T}(2)
sn[2] = s2!(r,b,s2_grad)
dsndr[2] = s2_grad[1]
dsndb[2] = s2_grad[2]

# Compute the J_v and I_v functions:
#Iv = zeros(typeof(k2),v_max+1); Jv = zeros(typeof(k2),v_max+1)
#Iv = Array{T}(v_max+1); Jv = Array{T}(v_max+1)
#fill!(Iv,zero(k2)); fill!(Jv,zero(k2))
# And their derivatives with respect to k:
#dIvdk = zeros(typeof(k2),v_max+1); dJvdk = zeros(typeof(k2),v_max+1)
#dIvdk = Array{T}(v_max+1); dJvdk = Array{T}(v_max+1)
#fill!(dIvdk,zero(k2)); fill!(dJvdk,zero(k2))
if k2 > 0
  if (k2 < 0.5 || k2 > 2.0) # && v_max > 3
# This computes I_v,J_v for the largest v, and then works down to smaller values:
    dIJv_lower_dk!(v_max,k2,kck,kc,kap0,Iv,Jv,dIvdk,dJvdk)
  else
# This computes I_0,J_0,J_1, and then works upward to larger v:
    dIJv_raise_dk!(v_max,k2,kck,kc,kap0,Iv,Jv,dIvdk,dJvdk)
  end
end

#nphi = 1000; dphi=2pi/nphi; phigrid = linspace(.5*dphi,1-.5*dphi,nphi)
# Next, loop over the Green's function components:
for n=2:N_c
  pofgn = zero(r)
  dpdr = zero(r)
  dpdb = zero(r)
  dpdk = zero(r)
  if iseven(n)
# For even values of n, sum over I_v:
    n0 = convert(Int64,n/2)
    coeff = (-fourbr)^n0
#    dIv_fac = (1+(r-b)*(r+b))/(2r)
    # Compute i=0 term
    pofgn = coeff*((r-b)*Iv[n0+1]+2b*Iv[n0+2])
    dpdr = coeff*Iv[n0+1]
    dpdb = coeff*(-Iv[n0+1]+2*Iv[n0+2])
    dpdr += (n0+1)*pofgn/r
#    dpdr += coeff*((n0+2-(n0+1)*b/r)*Iv[n0+1]+2*(n0+1)*b/r*Iv[n0+2])
    dpdb += n0*pofgn/b
#    println("v: ",n0,"-Iv[v]: ",-Iv[n0+1]," 2Iv[v+1]: ",2Iv[n0+2]," diff: ",-Iv[n0+1]+2*Iv[n0+2])
#    dpdk = coeff*((r-b)*dIvdk[n0+1]+2b*dIvdk[n0+2])
#    dpdk = coeff*dIvdk[n0+1]*dIv_fac
# For even n, compute coefficients for the sum over I_v:
#    println("n0: ",n0," i: ",0," coeff: ",coeff)
    for i=1:n0
      coeff *= -(n0-i+1)/i*k2
#      println("n0: ",n0," i: ",i," coeff: ",coeff)
      term =  coeff*((r-b)*Iv[n0-i+1]+2b*Iv[n0-i+2])
      pofgn += term
      dpdr += coeff*Iv[n0-i+1]
      dpdb += coeff*(-Iv[n0-i+1]+2*Iv[n0-i+2])
      dpdr += term*(i*2*(b-r)/onembmr2+(n0+1-i)/r)
#      dpdr += coeff*((n0+1)+i*((b-r)*(r+b)-1)/onembmr2)*((1-b/r)*Iv[n0-i+1]+2b/r*Iv[n0-i+2])
      dpdb += term*(i*2*(r-b)/onembmr2+(n0-i)/b)
#      println("v: ",n0-i,"-Iv[v]: ",-Iv[n0-i+1]," 2Iv[v+1]: ",2Iv[n0-i+2]," diff: ",-Iv[n0-i+1]+2*Iv[n0-i+2])
#      dpdk += coeff*((r-b)*dIvdk[n0-i+1]+2b*dIvdk[n0-i+2])
#      dpdk += coeff*dIvdk[n0-i+1]*dIv_fac
#      dpdk += coeff*2*i/k*((r-b)*Iv[n0-i+1]+2b*Iv[n0-i+2])
    end
    pofgn *= 2r
    dpdr *= 2r
#    dpdr += (n0+1)*pofgn/r
#    dpdr += pofgn/r
    dpdb *= 2r
#    dpdb += n0*pofgn/b
#    dpdk *= 2r
  else
# Now do the same for odd N_c in sum over J_v:
    n0 = convert(Int64,(n-3)/2)
    coeff = (-fourbr)^n0
    # Compute i=0 term
    pofgn = coeff*((r-b)*Jv[n0+1]+2b*Jv[n0+2])
    dpdr = coeff*Jv[n0+1]
    dpdb = coeff*(-Jv[n0+1]+2*Jv[n0+2])
#    dpdr += n0*pofgn/r
#    dpdb += n0*pofgn/b
#    dpdr  += pofgn*(3*(b-r)/onembmr2+(n0+1)/r)
    dpdr  += pofgn*(3*(b-r)*r+(n0+1)*onembmr2)/(r*onembmr2)
#    dpdb  += pofgn*(3*(r-b)/onembmr2+n0/b)
    dpdb  += pofgn*(3*(r-b)*b+n0*onembmr2)/(b*onembmr2)
    dpdk = coeff*((r-b)*dJvdk[n0+1]+2b*dJvdk[n0+2])
#    println("n0: ",n0," i: ",0," coeff: ",coeff)
# For even n, compute coefficients for the sum over I_v:
    for i=1:n0
      coeff *= -(n0-i+1)/i*k2
#      println("n0: ",n0," i: ",i," coeff: ",coeff)
      term = coeff*((r-b)*Jv[n0-i+1]+2b*Jv[n0-i+2])
      pofgn += term
      dpdr  +=  coeff*Jv[n0-i+1]
      dpdb  +=  coeff*(-Jv[n0-i+1]+2*Jv[n0-i+2])
#      dpdr  += term*((i*2+3)*(b-r)/onembmr2+(n0+1-i)/r)
      dpdr  += term*((i*2+3)*(b-r)*r+(n0+1-i)*onembmr2)/(r*onembmr2)
#      dpdb  += term*((i*2+3)*(r-b)/onembmr2+(n0-i)/b)
      dpdb  += term*((i*2+3)*(r-b)*b+(n0-i)*onembmr2)/(onembmr2*b)
      dpdk  += coeff*((r-b)*dJvdk[n0-i+1]+2b*dJvdk[n0-i+2])
#      dpdk  += coeff*2*i/k*((r-b)*Jv[n0-i+1]+2b*Jv[n0-i+2])
    end
    pofgn *= 2r*onembmr2^1.5
    dpdr  *= 2r*onembmr2^1.5
    dpdb  *= 2r*onembmr2^1.5
#    dpdr += ((n0+1)/r+3*(b-r)/onembmr2)*pofgn
#    dpdr  += (1/r+3*(b-r)/onembmr2)*pofgn
#    dpdb += (n0/b-3*(b-r)/onembmr2)*pofgn
#    dpdb  += 3*(r-b)/onembmr2*pofgn
#    dpdb += (n*onembmr2-3*(b-r)*(r+b)-3)/(2*b*onembmr2)*pofgn
    dpdk  *= 2r*onembmr2^1.5
  end
#  pofgn_num = sum(sqrt.(1-r^2-b^2-2*b*r*sin.(phigrid)).^n.*(r+b.*sin.(phigrid))*r*dphi)
#  println("n: ",n," P(G_n): ",pofgn," P(G_n),num: ",pofgn_num)
# Q(G_n) is zero in this case since on limb of star z^n = 0 at the stellar
# boundary for n > 0.
# Compute sn[n]:
  sn[n+1] = -pofgn
#  println("n: ",n," sn: ",sn[n+1]," dpdr: ",dpdr," dpdk*dkdr: ",dpdk*dkdr," dpdb: ",dpdb," dpdk*dkdb: ",dpdk*dkdb)
  dsndr[n+1] = -(dpdr+dpdk*dkdr)
  dsndb[n+1] = -(dpdb+dpdk*dkdb)
end
#if typeof(r) == Float64
#  println("r: ",r," b: ",b," s2 error: ",convert(Float64,s2(big(r),big(b)))-sn[2])
#end
# That's it!
#println("s_n: ",sn)
#println("c_n*s_n: ",c_n.*sn)
# Compute derivatives with respect to the coefficients:
den = inv(pi*(c_n[1]+2*c_n[2]/3))
flux = zero(r)
dfdrbc[1]=zero(r)  # Derivative with respect to r
dfdrbc[2]=zero(r)  # Derivative with respect to b
for n=0:N_c
  # derivatives with respect to the coefficients:
  dfdrbc[n+3]= sn[n+1]*den
  # total flux:
  flux += c_n[n+1]*dfdrbc[n+3]
  # derivatives with respect to r and b:
  dfdrbc[1] += c_n[n+1]*dsndr[n+1]*den
  dfdrbc[2] += c_n[n+1]*dsndb[n+1]*den
end
# Include derivatives with respect to first two c_n parameters:
dfdrbc[3] -= flux*den*pi
dfdrbc[4] -= flux*den*2pi/3
#flux = sum(c_n.*sn)*den   # for c_2 and above, the flux integrated over the star is zero.
return flux
end


##########################
# FunctionCell Interface #
##########################
"""
	kmedoids(k, d=Distances.Euclidean() [;kwargs])

Constructs an untrained cell that when piped data inside, clusters the data using the
k-medoids algorithm where `k` is the desired number of clusters and `d` is the distance
that will be used to calculate distances from new samples to the clustering generated
from the untrained function cell.

# Keyword arguments (same as in `Clustering`)
  * `init` initialization algorithm; can be `:rand` for random, `:kmpp` for K-medoids++ or `:kmcen` for choosing samples with highest centrality (default `:kmpp`)
or a vector of size 'k' providing the indexes of the initial seeds.
  * `maxiter` maximum number of iterations (default `100`)
  * `tol` tolerable change in objective at convergence (default `1.0e-6`)
  * `display` Level of information; can be `:none`, `:final`, `:iter` (default `:none`)

Read the `Clustering.jl` documentation for more information.
"""
kmedoids(k::Int, d::Distances.PreMetric=Distances.Euclidean(); kwargs...) = FunctionCell(kmedoids, (k,d), ModelProperties(), kwtitle("K-medoids ($k clusters)", kwargs); kwargs...)

############################
# DataCell/Array Interface #
############################
"""
	kmedoids(x, k, d=Distances.Euclidean() [;kwargs])

Performs k-medoids clustering on the dataset `x` (square cost matrix) using `k` number of clusters and returns the clustering; when piping
data to the resulting clustering (e.g. trained cell), the distances from input data to the medoids are
returned, using `d` as distance.
"""
# Training
kmedoids(x::T where T<:CellData, k::Int, d::Distances.PreMetric=Distances.Euclidean(); kwargs...) = kmedoids(getx!(x), k, d; kwargs...)
kmedoids(x::T where T<:AbstractVector, k::Int, d::Distances.PreMetric=Distances.Euclidean(); kwargs...) =
	error("[kmedoids] Vectors are not supported as input data. Try using `v |> dist(v) |> kmedoids(n)`")
kmedoids(x::T where T<:AbstractMatrix, k::Int, d::Distances.PreMetric=Distances.Euclidean(); kwargs...) = begin

	# Perform clustering
	kmeddata = Clustering.kmedoids(x, k; kwargs...)

	@assert length(kmeddata.medoids) == k "[kmedoids] The number of medoids (cluster centers) differes from the desired one k=$k."

	# Build model properties
	modelprops = ModelProperties(length(kmeddata.acosts), k)

	FunctionCell(kmedoids, Model((d, x[:,sort(kmeddata.medoids)], kmeddata), modelprops), kwtitle("K-medoids ($k clusters)", kwargs) );
end

# Execution
kmedoids(x::T where T<:CellData, model::Model{<:Tuple{<:Distances.PreMetric, <:AbstractArray, <:Clustering.KmedoidsResult}}) =
	datacell(kmedoids(getx!(x), model), gety(x))
kmedoids(x::T where T<:AbstractVector, model::Model{<:Tuple{<:Distances.PreMetric, <:AbstractArray, <:Clustering.KmedoidsResult}}) =
	kmedoids(mat(x, LearnBase.ObsDim.Constant{2}()), model)
kmedoids(x::T where T<:AbstractMatrix, model::Model{<:Tuple{<:Distances.PreMetric, <:AbstractArray, <:Clustering.KmedoidsResult}}) =
	# Return distances from each input sample to the medoids
	Distances.pairwise(model.data[1], model.data[2], x)

##########################
# FunctionCell Interface #
##########################
"""
	kmedoids!(medoids, d=Distances.Euclidean() [;kwargs])

Constructs an untrained cell that when piped data inside, clusters the data using the
k-medoids algorithm where `medoids` contains the initial seed observation indices
(will be updated inplace) and `d` is the distance that will be used to calculate distances
from new samples to the clustering generated from the untrained function cell.

# Keyword arguments (same as in `Clustering`)
  * `maxiter` maximum number of iterations (default `100`)
  * `tol` tolerable change in objective at convergence (default `1.0e-6`)
  * `display` Level of information; can be `:none`, `:final`, `:iter` (default `:none`)
"""
kmedoids!(medoids::Vector{Int}, d::Distances.PreMetric=Distances.Euclidean(); kwargs...) =
	FunctionCell(kmedoids!, (medoids,d), ModelProperties(), kwtitle("K-medoids! ($(length(medoids)) clusters)", kwargs); kwargs...)

############################
# DataCell/Array Interface #
############################
"""
	kmedoids!(x, medoids, d=Distances.Euclidean() [;kwargs])

Performs k-medoids clustering on the dataset `x` using `medoids` as initial seed observation indices
returning a clustering; when piping data to the resulting clustering (e.g. trained cell),
the distances from input data to the new medoids are returned, using `d` as distance.
"""
# Training
kmedoids!(x::T where T<:CellData, medoids::Vector{Int}, d::Distances.PreMetric=Distances.Euclidean(); kwargs...) = kmedoids!(getx!(x), medoids, d; kwargs...)
kmedoids!(x::T where T<:AbstractVector, medoids::Vector{Int}, d::Distances.PreMetric=Distances.Euclidean() ; kwargs...) =
	error("[kmedoids!] Vectors are not supported as input data. Try using `v |> dist(v) |> kmedoids!(medoids)`")
kmedoids!(x::T where T<:AbstractMatrix, medoids::Vector{Int}, d::Distances.PreMetric=Distances.Euclidean() ; kwargs...) = begin

	# Perform clustering
	kmeddata = Clustering.kmedoids!(x, medoids; kwargs...)

	# Build model properties
	modelprops = ModelProperties(length(kmeddata.acosts), length(kmeddata.medoids))

	# Return a trained cell that uses `kmedoids` and not `kmedoids!` as execution function
	# since the only difference is in training (e.g. clustering) and not in
	# assigning new data to clusters.
	FunctionCell(kmedoids, Model((d, x[:,sort(kmeddata.medoids)], kmeddata), modelprops), kwtitle("K-medoids ($(length(medoids)) clusters)", kwargs) );
end


using PSCOPF

using Test
using JuMP
using Dates
using DataStructures

@testset verbose=true "test_energy_market_dp" begin

    #=
    TS: [11h]
    S: [S1]
                        bus 1
                        |
    (limitable) wind_1_1|load
    Pmin=0, Pmax=100    | S1: 55
    Csta=0, Cprop=1     |
    DP => 10h40         |
      S1: 20            |
                        |
    (pilotable) prod_1_1|
    Pmin=10, Pmax=100   |
    Csta=0, Cprop=10    |
    DP => 10h30         |
                        |
    (pilotable) prod_1_2|
     Pmin=10, Pmax=100  |
     Csta=0, Cprop=15   |
     DP => 10h45        |
                        |

    =#
    function create_instance(ech, next_ech, ts,
                            wind_1_level,
                            prod_1_level,
                            prod_2_level;
                            fo_length=Minute(5))
        TS = [ts]
        ECH = [DateTime("2015-01-01T10:00:00"), DateTime("2015-01-01T10:40:00")]
        network = PSCOPF.Networks.Network()
        # Buses
        PSCOPF.Networks.add_new_bus!(network, "bus_1")
        # Limitables
        PSCOPF.Networks.add_new_generator_to_bus!(network, "bus_1", "wind_1_1", PSCOPF.Networks.LIMITABLE,
                                                0., 100.,
                                                0., 1.,
                                                Dates.Second(20*60), Dates.Second(20*60))
        # Pilotables
        PSCOPF.Networks.add_new_generator_to_bus!(network, "bus_1", "prod_1_1", PSCOPF.Networks.PILOTABLE,
                                                10., 100.,
                                                0., 10.,
                                                Dates.Second(3*60*60), Dates.Second(30*60))
        PSCOPF.Networks.add_new_generator_to_bus!(network, "bus_1", "prod_1_2", PSCOPF.Networks.PILOTABLE,
                                                10., 100.,
                                                0., 15.,
                                                Dates.Second(3*60*60), Dates.Second(15*60))
        # Uncertainties
        uncertainties = PSCOPF.Uncertainties()
        PSCOPF.add_uncertainty!(uncertainties, ECH[1], "wind_1_1", TS[1], "S1", 20.)
        PSCOPF.add_uncertainty!(uncertainties, ECH[1], "bus_1", TS[1], "S1", 55.)
        PSCOPF.add_uncertainty!(uncertainties, ECH[2], "wind_1_1", TS[1], "S1", 20.)
        PSCOPF.add_uncertainty!(uncertainties, ECH[2], "bus_1", TS[1], "S1", 55.)
        # initial generators state : need to pay starting cost at TS[1]
        generators_init_state = SortedDict(
            "prod_1_1" => PSCOPF.OFF,
            "prod_1_2" => PSCOPF.OFF
        )
        #ManagementMode
        mode = PSCOPF.ManagementMode("test_mode", fo_length)

        context = PSCOPF.PSCOPFContext(network, TS, mode,
                                        generators_init_state,
                                        uncertainties, nothing)
        market = PSCOPF.EnergyMarket()
        firmness = PSCOPF.compute_firmness(market,
                                            ech, next_ech,
                                            TS, context)

        prod_1_definitive_state = (prod_1_level < 1e-09) ? PSCOPF.OFF : PSCOPF.ON
        prod_1_definitive_level = (PSCOPF.get_power_level_firmness(firmness, "prod_1_1", ts) == PSCOPF.DECIDED) ? prod_1_level : missing
        prod_2_definitive_state = (prod_2_level < 1e-09) ? PSCOPF.OFF : PSCOPF.ON
        prod_2_definitive_level = (PSCOPF.get_power_level_firmness(firmness, "prod_1_2", ts) == PSCOPF.DECIDED) ? prod_2_level : missing
        context.tso_schedule = PSCOPF.Schedule(PSCOPF.Market(), Dates.DateTime("2015-01-01T06:00:00"), SortedDict(
                "wind_1_1" => PSCOPF.GeneratorSchedule("wind_1_1",
                    SortedDict{Dates.DateTime, PSCOPF.UncertainValue{PSCOPF.GeneratorState}}(),
                    # SortedDict(Dates.DateTime("2015-01-01T11:00:00") => PSCOPF.UncertainValue{PSCOPF.GeneratorState}(missing,
                    #                                                                         SortedDict("S1"=>PSCOPF.ON))),
                    SortedDict(ts => PSCOPF.UncertainValue{Float64}(missing,
                                                                    SortedDict("S1"=>wind_1_level)))
                    ),
                "prod_1_1" => PSCOPF.GeneratorSchedule("prod_1_1",
                    SortedDict(ts => PSCOPF.UncertainValue{PSCOPF.GeneratorState}(prod_1_definitive_state,
                                                                    SortedDict("S1"=>prod_1_definitive_state))),
                    SortedDict(ts => PSCOPF.UncertainValue{Float64}(prod_1_definitive_level,
                                                                    SortedDict("S1"=>prod_1_level)))
                    ),
                "prod_1_2" => PSCOPF.GeneratorSchedule("prod_1_2",
                    SortedDict(ts => PSCOPF.UncertainValue{PSCOPF.GeneratorState}(prod_2_definitive_state,
                                                                    SortedDict("S1"=>prod_2_definitive_state))),
                    SortedDict(ts => PSCOPF.UncertainValue{Float64}(prod_2_definitive_level,
                                                                    SortedDict("S1"=>prod_2_level)))
                    )
                )
            )

        return context, market, firmness
    end

    #=
    TS: [11h]
    S: [S1]
                        bus 1
                        |
    (limitable) wind_1_1|load
    Pmin=0, Pmax=100    | S1: 55
    Csta=0, Cprop=1     |
    DP => 10h40         |
      S1: 20            |
                        |
    (pilotable) prod_1_1|
    Pmin=10, Pmax=100   |
    Csta=0, Cprop=10    |
    DP => 10h30         |
                        |
    (pilotable) prod_1_2|
     Pmin=10, Pmax=100  |
     Csta=0, Cprop=15   |
     DP => 10h45        |
                        |

    ECH1*         ECH2            ECH3                                 <---FO--->TS
    |             |               |
    10h           10h30           10h40          10h45                          11h
                  <------------------------------------------DP(prod1)----------->
                                  <--------------------------DP(wind1)----------->
                                                   <---------DP(prod2)----------->

    We can change the production level of a unit before reaching the DP
    ech = 10h
    TS = 11h
    DP(prod_1_1) = 30 mins => 10h30
    DP(prod_1_2) = 15 mins => 10h45

    We suppose that in the preceding step we decided :
    prod_1_1 : 0. (OFF)
    prod_1_2 : 37. (ON)

    In terms of cost, it is cheaper to use prod_1_1. But, since we are past DMO, we can no longer start prod_1_1.
    wind_1_1 : 20 MW
    remaining Demand : 35MW
    Since we haven't reached the DP yet, we can change the production level of prod_1_2 from 37 MW to 35 MW
    => the new decision :
    wind_1_1 : 20 MW
    prod_1_1 : OFF (unchanged)
    prod_1_2 : 35. (decreased from 37MW)
    =#
    @testset "energy_market_can_decrease_the_production_level_before_dp_if_ON" begin
        ech = DateTime("2015-01-01T10:00:00")
        next_ech = DateTime("2015-01-01T10:30:00")
        TS = [DateTime("2015-01-01T11:00:00")]
        context, market, firmness = create_instance(ech, next_ech, TS[1],
                                                    18., 0., 37.,
                                                    )
        expected_firmness = PSCOPF.Firmness(
                SortedDict("prod_1_1" => SortedDict(Dates.DateTime("2015-01-01T11:00:00") => PSCOPF.DECIDED),
                            "prod_1_2" => SortedDict(Dates.DateTime("2015-01-01T11:00:00") => PSCOPF.DECIDED), ),
                SortedDict("wind_1_1" => SortedDict(Dates.DateTime("2015-01-01T11:00:00") => PSCOPF.FREE),
                            "prod_1_1" => SortedDict(Dates.DateTime("2015-01-01T11:00:00") => PSCOPF.FREE),
                            "prod_1_2" => SortedDict(Dates.DateTime("2015-01-01T11:00:00") => PSCOPF.FREE), )
                )
        @test firmness == expected_firmness

        result = PSCOPF.run(market, ech, firmness,
                    PSCOPF.get_target_timepoints(context),
                    context)
        PSCOPF.update_market_schedule!(context, ech, result, firmness, market)

        # Solution is optimal
        @test PSCOPF.get_status(result) == PSCOPF.pscopf_OPTIMAL
        @test 20. ≈ PSCOPF.get_prod_value(context.market_schedule, "wind_1_1", TS[1], "S1")
        # prod_1_1 is OFF cause it already was OFF:
        @test PSCOPF.OFF == PSCOPF.get_commitment_value(context.market_schedule, "prod_1_1", TS[1], "S1")
        @test PSCOPF.get_prod_value(context.market_schedule, "prod_1_1", TS[1], "S1") < 1e-09
        # prod_1_2 was ON, it can change it's production : 37 => 35.
        @test PSCOPF.ON == PSCOPF.get_commitment_value(context.market_schedule, "prod_1_2", TS[1], "S1")
        @test 35. ≈ PSCOPF.get_prod_value(context.market_schedule, "prod_1_2", TS[1], "S1")
    end

    #=
    TS: [11h]
    S: [S1]
                        bus 1
                        |
    (limitable) wind_1_1|load
    Pmin=0, Pmax=100    | S1: 55
    Csta=0, Cprop=1     |
    DP => 10h40         |
      S1: 20            |
                        |
    (pilotable) prod_1_1|
    Pmin=10, Pmax=100   |
    Csta=0, Cprop=10    |
    DP => 10h30         |
                        |
    (pilotable) prod_1_2|
     Pmin=10, Pmax=100  |
     Csta=0, Cprop=15   |
     DP => 10h45        |
                        |

    ECH1*         ECH2            ECH3                                 <---FO--->TS
    |             |               |
    10h           10h30           10h40          10h45                          11h
                  <------------------------------------------DP(prod1)----------->
                                  <--------------------------DP(wind1)----------->
                                                   <---------DP(prod2)----------->

    We can change the production level of a unit before reaching the DP
    ech = 10h
    TS = 11h
    DP(prod_1_1) = 30 mins => 10h30
    DP(prod_1_2) = 15 mins => 10h45

    We suppose that in the preceding step we decided :
    prod_1_1 : 0. (OFF)
    prod_1_2 : 33. (ON)

    In terms of cost, it is cheaper to use prod_1_1. But, since we are past DMO, we can no longer start prod_1_1.
    wind_1_1 : 20 MW
    remaining Demand : 35MW
    Since we haven't reached the DP yet, we can change the production level of prod_1_2 from 33 MW to 35 MW
    => the new decision :
    wind_1_1 : 20 MW
    prod_1_1 : OFF (unchanged)
    prod_1_2 : 35. (increased from 33MW)
    =#
    @testset "energy_market_can_increase_the_production_level_before_dp_if_ON" begin
        ech = DateTime("2015-01-01T10:00:00")
        next_ech = DateTime("2015-01-01T10:30:00")
        TS = [DateTime("2015-01-01T11:00:00")]
        context, market, firmness = create_instance(ech, next_ech, TS[1],
                                                    22., 0., 33.,
                                                    )
        expected_firmness = PSCOPF.Firmness(
                SortedDict("prod_1_1" => SortedDict(Dates.DateTime("2015-01-01T11:00:00") => PSCOPF.DECIDED),
                            "prod_1_2" => SortedDict(Dates.DateTime("2015-01-01T11:00:00") => PSCOPF.DECIDED), ),
                SortedDict("wind_1_1" => SortedDict(Dates.DateTime("2015-01-01T11:00:00") => PSCOPF.FREE),
                            "prod_1_1" => SortedDict(Dates.DateTime("2015-01-01T11:00:00") => PSCOPF.FREE),
                            "prod_1_2" => SortedDict(Dates.DateTime("2015-01-01T11:00:00") => PSCOPF.FREE), )
                )
        @test firmness == expected_firmness

        result = PSCOPF.run(market, ech, firmness,
                    PSCOPF.get_target_timepoints(context),
                    context)
        PSCOPF.update_market_schedule!(context, ech, result, firmness, market)

        # Solution is optimal
        @test PSCOPF.get_status(result) == PSCOPF.pscopf_OPTIMAL
        @test 20. ≈ PSCOPF.get_prod_value(context.market_schedule, "wind_1_1", TS[1], "S1")
        # prod_1_1 is OFF cause it already was OFF:
        @test PSCOPF.OFF == PSCOPF.get_commitment_value(context.market_schedule, "prod_1_1", TS[1], "S1")
        @test PSCOPF.get_prod_value(context.market_schedule, "prod_1_1", TS[1], "S1") < 1e-09
        # prod_1_2 was ON, it can change it's production : 33 => 35.
        @test PSCOPF.ON == PSCOPF.get_commitment_value(context.market_schedule, "prod_1_2", TS[1], "S1")
        @test 35. ≈ PSCOPF.get_prod_value(context.market_schedule, "prod_1_2", TS[1], "S1")
    end

    #=
    ECH1          ECH2            ECH3*          ECH4                  <---FO--->TS
    |             |               |
    10h           10h30           10h40          10h45                          11h
                  <------------------------------------------DP(prod1)----------->
                                  <--------------------------DP(wind1)----------->
                                                   <---------DP(prod2)----------->

    We can no longer change the production level of a unit once DP is reached
    ech = 10h40
    TS = 11h
    DP(prod_1_1) = 30 mins => 10h30 => we are past DP
    DP(prod_1_2) = 15 mins => 10h45

    We suppose that in the preceding step we decided :
    prod_1_1 : 17. (ON)
    prod_1_2 : 20. (ON)

    In terms of cost, it is cheaper to use prod_1_1.
    But, since we are past DP, we can no longer increase the level of prod_1_1.

    wind_1_1 : 20 MW
        remaining Demand : 35MW
    prod_1_1 : 17. (unchanged)
    prod_1_2 : 18. (decreased from 20MW)
    =#
    @testset "energy_market_cannot_change_the_production_level_after_dp" begin
        # We are past DP[prod_1_1]
        ech = DateTime("2015-01-01T10:40:00")
        next_ech = DateTime("2015-01-01T10:45:00")
        TS = [DateTime("2015-01-01T11:00:00")]
        context, market, firmness = create_instance(ech, next_ech, TS[1],
                                                    18., 17., 20.,
                                                    )

        expected_firmness = PSCOPF.Firmness(
                                    SortedDict("prod_1_1" => SortedDict(TS[1] => PSCOPF.DECIDED),
                                                "prod_1_2" => SortedDict(TS[1] => PSCOPF.DECIDED),
                                                ),
                                    SortedDict("wind_1_1" => SortedDict(TS[1] => PSCOPF.TO_DECIDE),
                                                "prod_1_1" => SortedDict(TS[1] => PSCOPF.DECIDED),
                                                "prod_1_2" => SortedDict(TS[1] => PSCOPF.FREE),
                                                )
            )
        @test firmness == expected_firmness

        # prod_1_1 has a definitive value since we are past the DP

        result = PSCOPF.run(market, ech, firmness,
                    PSCOPF.get_target_timepoints(context),
                    context)
        PSCOPF.update_market_schedule!(context, ech, result, firmness, market)

        # Solution is optimal
        @test PSCOPF.get_status(result) == PSCOPF.pscopf_OPTIMAL
        @test 20. ≈ PSCOPF.get_prod_value(context.market_schedule, "wind_1_1", TS[1], "S1")
        # we are past DP[prod_1_] => we can't change the level of prod_1_1
        # prod_1_1 : 17.
        @test PSCOPF.ON == PSCOPF.get_commitment_value(context.market_schedule, "prod_1_1", TS[1], "S1")
        @test 17. ≈ PSCOPF.get_prod_value(context.market_schedule, "prod_1_1", TS[1], "S1")
        # prod_1_2 : 20 => 18.
        @test PSCOPF.ON == PSCOPF.get_commitment_value(context.market_schedule, "prod_1_2", TS[1], "S1")
        @test 18. ≈ PSCOPF.get_prod_value(context.market_schedule, "prod_1_2", TS[1], "S1")
    end

    #=
    at ECH2, we are at wind_1_1's DP,
    but the decision made will still be by scenario cause it depends on the uncertainties
    limitables's DP only applies to limitation decisions (=> applied in TSO)
    ECH1          ECH2            ECH3*          ECH4                  <---FO--->TS
    |             |               |
    10h           10h30           10h40          10h45                          11h
                  <------------------------------------------DP(prod1)----------->
                                  <--------------------------DP(wind1)----------->
                                                   <---------DP(prod2)----------->
    =#
    @testset "energy_market_wind_is_always_uncertain" begin
        # We are past DP[prod_1_1]
        ech = DateTime("2015-01-01T10:40:00")
        next_ech = DateTime("2015-01-01T10:45:00")
        TS = [DateTime("2015-01-01T11:00:00")]
        context, market, firmness = create_instance(ech, next_ech, TS[1],
                                                    18., 17., 20.,
                                                    )

        expected_firmness = PSCOPF.Firmness(
                                    SortedDict("prod_1_1" => SortedDict(TS[1] => PSCOPF.DECIDED),
                                                "prod_1_2" => SortedDict(TS[1] => PSCOPF.DECIDED),
                                                ),
                                    SortedDict("wind_1_1" => SortedDict(TS[1] => PSCOPF.TO_DECIDE),
                                                "prod_1_1" => SortedDict(TS[1] => PSCOPF.DECIDED),
                                                "prod_1_2" => SortedDict(TS[1] => PSCOPF.FREE),
                                                )
            )
        @test firmness == expected_firmness

        result = PSCOPF.run(market, ech, firmness,
                    PSCOPF.get_target_timepoints(context),
                    context)
        PSCOPF.update_market_schedule!(context, ech, result, firmness, market)

        # Solution is optimal
        @test PSCOPF.get_status(result) == PSCOPF.pscopf_OPTIMAL
        @test 20. ≈ PSCOPF.get_prod_value(context.market_schedule, "wind_1_1", TS[1], "S1")
        @test !PSCOPF.is_definitive(PSCOPF.get_prod_uncertain_value(context.market_schedule, "wind_1_1", TS[1]))
    end

    #=
    ECH1*         ECH2            ECH3           ECH4                  <---FO--->TS
    |             |               |
    10h           10h30           10h40          10h45                          11h
                  <------------------------------------------DP(prod1)----------->
                                  <--------------------------DP(wind1)----------->
                                                   <---------DP(prod2)----------->
    =#
    @testset "energy_market_cannot_change_the_production_level_before_dp_if_unit_is_off_and_past_DMO" begin
        ech = DateTime("2015-01-01T10:00:00")
        next_ech = DateTime("2015-01-01T10:30:00")
        TS = [DateTime("2015-01-01T11:00:00")]
        context, market, firmness = create_instance(ech, next_ech, TS[1],
                                                    35., 0., 25.,
                                                    )

        expected_firmness = PSCOPF.Firmness(
                                    SortedDict("prod_1_1" => SortedDict(TS[1] => PSCOPF.DECIDED),
                                                "prod_1_2" => SortedDict(TS[1] => PSCOPF.DECIDED),
                                                ),
                                    SortedDict("wind_1_1" => SortedDict(TS[1] => PSCOPF.FREE),
                                                "prod_1_1" => SortedDict(TS[1] => PSCOPF.FREE),
                                                "prod_1_2" => SortedDict(TS[1] => PSCOPF.FREE),
                                                )
            )
        @test firmness == expected_firmness

        result = PSCOPF.run(market, ech, firmness,
                    PSCOPF.get_target_timepoints(context),
                    context)
        PSCOPF.update_market_schedule!(context, ech, result, firmness, market)

        # Solution is optimal
        @test PSCOPF.get_status(result) == PSCOPF.pscopf_OPTIMAL
        @test 20. ≈ PSCOPF.get_prod_value(context.market_schedule, "wind_1_1", TS[1], "S1")
        # We are past DMO => unit state was already fixed
        # prod_1_1 : OFF
        @test PSCOPF.OFF == PSCOPF.get_commitment_value(context.market_schedule, "prod_1_1", TS[1], "S1")
        @test PSCOPF.get_prod_value(context.market_schedule, "prod_1_1", TS[1], "S1") < 1e-09
        # prod_1_2 : 25 => 35.
        @test PSCOPF.ON == PSCOPF.get_commitment_value(context.market_schedule, "prod_1_2", TS[1], "S1")
        @test 35. ≈ PSCOPF.get_prod_value(context.market_schedule, "prod_1_2", TS[1], "S1")
    end

end


using AlgebraOfGraphics, CategoricalArrays, CairoMakie, DataFrames
using LaTeXStrings, HypergeometricFunctions, SpecialFunctions
set_aog_theme!()
update_theme!(font = "Arial", fontsize = 28)
theme_minimal
function weight(j, s, xi)
    exp(
        logabsbinomial(j + s - 2, j - 1)[1] -
        log(j) +
        s * log(xi) +
        (j - 1) * log(1 - xi) +
        log(_₂F₁(j + s - 1, 1, j + 1, 1 - xi))
    )
end

df_j = DataFrame(j = 1:20);
df_ϕ = DataFrame(ϕ = [0.1, 0.3, 0.5]);
df_s = DataFrame(s = [1, 2, 3, 4]);
df = crossjoin(df_j, df_ϕ, df_s);
df[!, :w] = weight.(df[!, :j], df[!, :s], df[!, :ϕ]);
df[!, :s] = string.(df[!, :s]);
df[!, :ϕ] = string.(df[!, :ϕ]);
df[!, :ϕ] = LaTeXString.("φ(x_i ⋅ β) = " .* df[!, :ϕ]);
plt = data(df) *
    visual(Lines) *
    mapping(:j, :w, col = :ϕ, color = :s)
axis = (xlabel = L"j", ylabel = L"w_j(x)", color = L"s");
fg = draw(plt; axis, figure = (resolution = (900, 450), px_per_unit = 8))
fg
save("figures/final/lineplot-dgsb-weights.pdf", fg, px_per_unit = 3)


# generate the table found in
# https://arxiv.org/abs/2201.05967

using DyadicKDE

# specify parameters
n_data = 500
n_evals = 50
n_repeats = 2000
degeneracies = ["total", "partial", "none"]
kernel_names = ["epanechnikov_order_2", "epanechnikov_order_4"]
evals = collect(range(-2.0, stop=2.0, length=n_evals))
sdp_solver = "cosmo"
n_resample = 10000
significance_level = 0.05
results = []
ps = Dict(
    "total" => [0.5, 0.0, 0.5],
    "partial" => [0.25, 0.0, 0.75],
    "none" => [0.2, 0.2, 0.6],
)

println("Running experiments")
for degen in degeneracies
    for kernel_name in kernel_names

        global counter = 0
        Threads.@threads for rep in 1:n_repeats

            global counter += 1
            println("Degeneracy: $degen, Kernel: $kernel_name, Repeat $counter/$n_repeats")
            data = make_data(n_data, ps[degen])
            h_ROT = estimate_ROT_bandwidth(data, "epanechnikov_order_2")

            est = DyadicKernelDensityEstimator(
                kernel_name, h_ROT, significance_level,
                n_resample, sdp_solver, evals, data,
                Dict("degen" => degen, "p" => ps[degen]))

            fit(est)
            push!(results, est)

        end
    end
end

println("\nResults:")
for degen in degeneracies
    for kernel_name in kernel_names

        results_reduced = [est for est in results if
                               est.meta["degen"] == degen
                               && est.kernel_name == kernel_name]

        p = ps[degen]
        h_ROT = mean([est.bandwidth for est in results_reduced])
        RIMSE = mean([get_RIMSE(est) for est in results_reduced])
        ucb_coverage = mean([get_ucb_coverage(est) for est in results_reduced])
        pci_coverage = mean([get_pci_coverage(est) for est in results_reduced])
        ucb_width = mean([get_ucb_average_width(est) for est in results_reduced])
        pci_width = mean([get_pci_average_width(est) for est in results_reduced])

        println("π = $p")
        println("Degeneracy: $degen")
        println("h_ROT = $h_ROT")
        println("Kernel: $kernel_name")
        println("RIMSE: $RIMSE")
        println("UCB coverage: $ucb_coverage")
        println("UCB width: $ucb_width")
        println("PCI coverage: $pci_coverage")
        println("PCI width: $pci_width")
        println()

    end
end


module MultiFieldSparseMatrixAssemblersTests

using Gridap.Arrays
using Gridap.Algebra
using Gridap.Geometry
using Gridap.FESpaces
using Gridap.Fields
using Gridap.Integration
using SparseArrays
using Gridap.CellData
using Gridap.MultiField
using Gridap.ReferenceFEs
using Gridap.TensorValues

using Test

order = 2

domain = (0,1,0,1)
partition = (3,3)
model = CartesianDiscreteModel(domain,partition)

degree = order
trian = get_triangulation(model)
quad = CellQuadrature(trian,degree)

trian_Γ = SkeletonTriangulation(model)
quad_Γ = CellQuadrature(trian_Γ,degree)

V = TestFESpace(model,ReferenceFE(lagrangian,Float64,order);conformity=:H1)
Q = TestFESpace(model,ReferenceFE(lagrangian,Float64,order-1),conformity=:L2)

U = TrialFESpace(V)
P = TrialFESpace(Q)

Y = MultiFieldFESpace([V,Q])
X = MultiFieldFESpace([U,P])

free_values = rand(num_free_dofs(X))
xh = FEFunction(X,free_values)
uh, ph = xh

dy = get_cell_shapefuns(Y)
dx = get_cell_shapefuns_trial(X)
dv, dq = dy
du, dp = dx

cellmat = integrate(dv*du,quad)
cellvec = integrate(dv*2,quad)
cellids = get_cell_to_bgcell(trian)
cellmatvec = pair_arrays(cellmat,cellvec)

#cellmat_Γ = integrate(  jump(dv)*dp.⁺ + mean(dq)*jump(dp), quad_Γ)
cellmat_Γ = integrate(  jump(dv)*mean(du) + jump(∇(dq))⋅jump(∇(dp)), quad_Γ)

cellvec_Γ = integrate(  jump(dv) + mean(dq),quad_Γ)
cellmatvec_Γ = pair_arrays(cellmat_Γ,cellvec_Γ)
cellids_Γ = get_cell_to_bgcell(trian_Γ)

assem = SparseMatrixAssembler(SparseMatrixCSR{0,Float64,Int},X,Y)

matvecdata = ([cellmatvec,cellmatvec_Γ],[cellids,cellids_Γ],[cellids,cellids_Γ])
matdata = ([cellmat,cellmat_Γ],[cellids,cellids_Γ],[cellids,cellids_Γ])
vecdata = ([cellvec,cellvec_Γ],[cellids,cellids_Γ])
data = (matvecdata,matdata,vecdata)

test_assembler(assem,matdata,vecdata,data)

A = assemble_matrix(assem,matdata)

A = allocate_matrix(assem,matdata)

assem = SparseMatrixAssembler(X,Y)
test_assembler(assem,matdata,vecdata,data)

end # module


if !(VERSION.major==0&&VERSION.minor<=6)
	using Test
else
	using Base.Test
end
using FinancialToolbox

print_colored("Starting Hyper Dual Numbers Test\n",:green)

#Test Parameters
testToll=1e-14;
spot=10;K=10;r=0.02;T=2.0;sigma=0.2;d=0.01;
assert_(value,toll)=@test abs(value)<toll
#EuropeanCall Option
PriceCall=blsprice(spot,K,r,T,sigma,d);
DeltaCall=blsdelta(spot,K,r,T,sigma,d);
ThetaCall=blstheta(spot,K,r,T,sigma,d);
RhoCall=blsrho(spot,K,r,T,sigma,d);
LambdaCall=blslambda(spot,K,r,T,sigma,d);
VannaCall=blsvanna(spot,K,r,T,sigma,d);

#EuropeanPut Option
PricePut=blsprice(spot,K,r,T,sigma,d,false);
DeltaPut=blsdelta(spot,K,r,T,sigma,d,false);
ThetaPut=blstheta(spot,K,r,T,sigma,d,false);
RhoPut=blsrho(spot,K,r,T,sigma,d,false);
LambdaPut=blslambda(spot,K,r,T,sigma,d,false);
VannaPut=blsvanna(spot,K,r,T,sigma,d,false);

#Equals for both Options
Gamma=blsgamma(spot,K,r,T,sigma,d);
Vega=blsvega(spot,K,r,T,sigma,d);

########HYPER DUAL NUMBERS
using HyperDualNumbers;
DerToll=1e-13;
di=1e-15;
#Function definition
SpotHyper=hyper(spot,1.0,1.0,0.0);
rHyper=hyper(r,1.0,1.0,0.0);
THyper=hyper(T,1.0,1.0,0.0);
SigmaHyper=hyper(sigma,1.0,1.0,0.0);
KHyper=hyper(K,1.0,1.0,0.0);

#Function definition
#Call
F(spot)=blsprice(spot,K,r,T,sigma,d);
FF(spot)=blsdelta(spot,K,r,T,sigma,d);
G(r)=blsprice(spot,K,r,T,sigma,d);
H(T)=blsprice(spot,K,r,T,sigma,d);
L(sigma)=blsprice(spot,K,r,T,sigma,d);
P(spot)=blsprice(spot,K,r,T,sigma,d)*spot.value/blsprice(spot.value,K,r,T,sigma,d);
V(sigma)=blsdelta(spot,K,r,T,sigma,d);
VV(spot,sigma)=blsprice(spot,K,r,T,sigma,d);
#Put
F1(spot)=blsprice(spot,K,r,T,sigma,d,false);
FF1(spot)=blsdelta(spot,K,r,T,sigma,d,false);
G1(r)=blsprice(spot,K,r,T,sigma,d,false);
H1(T)=blsprice(spot,K,r,T,sigma,d,false);
P1(spot)=blsprice(spot,K,r,T,sigma,d,false)*spot.value/blsprice(spot.value,K,r,T,sigma,d,false);
V1(sigma)=blsdelta(spot,K,r,T,sigma,d,false);
V11(spot,sigma)=blsprice(spot,K,r,T,sigma,d,false);

#TEST
print_colored("--- European Call Sensitivities: HyperDualNumbers\n",:yellow)
print_colored("-----Testing Delta\n",:blue);
assert_(F(SpotHyper).epsilon1-DeltaCall,DerToll)
print_colored("-----Testing Gamma\n",:blue);
assert_(FF1(SpotHyper).epsilon1-Gamma,DerToll)
print_colored("-----Testing Gamma  2\n",:blue);
assert_(F(SpotHyper).epsilon12-Gamma,DerToll)
print_colored("-----Testing Rho\n",:blue);
assert_(G(rHyper).epsilon1-RhoCall,DerToll)
print_colored("-----Testing Theta\n",:blue);
assert_(-H(THyper).epsilon1-ThetaCall,DerToll)
print_colored("-----Testing Lambda\n",:blue);
assert_(P(SpotHyper).epsilon1-LambdaCall,DerToll)
print_colored("-----Testing Vanna\n",:blue);
assert_(V(SigmaHyper).epsilon1-VannaCall,DerToll)
print_colored("-----Testing Vanna  2\n",:blue);
assert_(VV(hyper(spot,1,0,0),hyper(sigma,0,1,0)).epsilon12-VannaCall,DerToll)

## Complex Test with Complex Step Approximation for European Put
#TEST
print_colored("--- European Put Sensitivities: HyperDualNumbers\n",:yellow)
print_colored("-----Testing Delta\n",:blue);
assert_(F1(SpotHyper).epsilon1-DeltaPut,DerToll)
print_colored("-----Testing Rho\n",:blue);
assert_(G1(rHyper).epsilon1-RhoPut,DerToll)
print_colored("-----Testing Theta\n",:blue);
assert_(-H1(THyper).epsilon1-ThetaPut,DerToll)
print_colored("-----Testing Lambda\n",:blue);
assert_(P1(SpotHyper).epsilon1-LambdaPut,DerToll)
print_colored("-----Testing Vanna\n",:blue);
assert_(V1(SigmaHyper).epsilon1-VannaPut,DerToll)
print_colored("-----Testing Vanna  2\n",:blue);
assert_(V11(hyper(spot,1,0,0),hyper(sigma,0,1,0)).epsilon12-VannaPut,DerToll)
print_colored("-----Testing Vega\n",:blue);
assert_(L(SigmaHyper).epsilon1-Vega,DerToll)
print_colored("Hyper Dual Numbers Test Passed\n\n",:green)

#TEST OF INPUT VALIDATION
print_colored("Starting Input Validation Test Hyper Dual Numbers\n",:magenta)
print_colored("----Testing Negative  Spot Price \n",:cyan)
@test_throws(ErrorException, blsprice(-SpotHyper,K,r,T,sigma,d))
@test_throws(ErrorException, blkprice(-SpotHyper,K,r,T,sigma))
@test_throws(ErrorException, blsdelta(-SpotHyper,K,r,T,sigma,d))
@test_throws(ErrorException, blsgamma(-SpotHyper,K,r,T,sigma,d))
@test_throws(ErrorException, blstheta(-SpotHyper,K,r,T,sigma,d))
@test_throws(ErrorException, blsrho(-SpotHyper,K,r,T,sigma,d))
@test_throws(ErrorException, blspsi(-SpotHyper,K,r,T,sigma,d))
@test_throws(ErrorException, blslambda(-SpotHyper,K,r,T,sigma,d))
@test_throws(ErrorException, blsvanna(-SpotHyper,K,r,T,sigma,d))
@test_throws(ErrorException, blsvega(-SpotHyper,K,r,T,sigma,d))

print_colored("----Testing Negative  Strike Price \n",:cyan)
@test_throws(ErrorException, blsprice(spot,-KHyper,r,T,sigma,d))
@test_throws(ErrorException, blkprice(spot,-KHyper,r,T,sigma))
@test_throws(ErrorException, blsdelta(spot,-KHyper,r,T,sigma,d))
@test_throws(ErrorException, blsgamma(spot,-KHyper,r,T,sigma,d))
@test_throws(ErrorException, blstheta(spot,-KHyper,r,T,sigma,d))
@test_throws(ErrorException, blsrho(spot,-KHyper,r,T,sigma,d))
@test_throws(ErrorException, blspsi(spot,-KHyper,r,T,sigma,d))
@test_throws(ErrorException, blslambda(spot,-KHyper,r,T,sigma,d))
@test_throws(ErrorException, blsvanna(spot,-KHyper,r,T,sigma,d))
@test_throws(ErrorException, blsvega(spot,-KHyper,r,T,sigma,d))

print_colored("----Testing Negative  Time to Maturity \n",:cyan)
@test_throws(ErrorException, blsprice(spot,K,r,-THyper,sigma,d))
@test_throws(ErrorException, blkprice(spot,K,r,-THyper,sigma))
@test_throws(ErrorException, blsdelta(spot,K,r,-THyper,sigma,d))
@test_throws(ErrorException, blsgamma(spot,K,r,-THyper,sigma,d))
@test_throws(ErrorException, blstheta(spot,K,r,-THyper,sigma,d))
@test_throws(ErrorException, blsrho(spot,K,r,-THyper,sigma,d))
@test_throws(ErrorException, blspsi(spot,K,r,-THyper,sigma,d))
@test_throws(ErrorException, blslambda(spot,K,r,-THyper,sigma,d))
@test_throws(ErrorException, blsvanna(spot,K,r,-THyper,sigma,d))
@test_throws(ErrorException, blsvega(spot,K,r,-THyper,sigma,d))

print_colored("----Testing Negative  Volatility \n",:cyan)
@test_throws(ErrorException, blsprice(spot,K,r,T,-SigmaHyper,d))
@test_throws(ErrorException, blkprice(spot,K,r,T,-SigmaHyper))
@test_throws(ErrorException, blsdelta(spot,K,r,T,-SigmaHyper,d))
@test_throws(ErrorException, blsgamma(spot,K,r,T,-SigmaHyper,d))
@test_throws(ErrorException, blstheta(spot,K,r,T,-SigmaHyper,d))
@test_throws(ErrorException, blsrho(spot,K,r,T,-SigmaHyper,d))
@test_throws(ErrorException, blspsi(spot,K,r,T,-SigmaHyper,d))
@test_throws(ErrorException, blslambda(spot,K,r,T,-SigmaHyper,d))
@test_throws(ErrorException, blsvanna(spot,K,r,T,-SigmaHyper,d))
@test_throws(ErrorException, blsvega(spot,K,r,T,-SigmaHyper,d))

print_colored("Hyper Dual Input Validation Test Passed\n",:magenta)


# using Optim
# using Elliptic
# using Struve
# using SpecialFunctions

function CalcFlux(Ae,Le,μr,N,V,ω)
    μ0 = 4*pi*1e-7
    MagnetizingInductance = N^2*μr*μ0*Ae /Le
    Impedance = ω*MagnetizingInductance
    I = V/Impedance
    Flux = N*I*μr*μ0*Ae /Le

    B = Flux/Ae
    return B
end

function CalcPermeability(AL,Ae,Le)
    #This function takes in the nominal inductance (inductance at one turn),effective length and area
    #And calculates the permeabily
    # L = N^2 μ_r*μ_o A_e /L_e ⟶ N=1, A_L = μA_e/L_e
    μ0 = 4*pi*1e-7
    return AL*Le/Ae/μ0
end

function BMax(V,f,N,A;k=(2*π/√2))
    return V ./ (k.*N.*A.*f)
end

function MaxTurns(CoreID,WireD)
    InnerWindingCircum = pi*(CoreID-WireD)
    return InnerWindingCircum/WireD
end

function RunMagnetizingFlux(WireD,CoreID,Ae,Le,AL,NRatio,VPk,ω)
    μ0 = 4*pi*1e-7
    μr = CalcPermeability(AL,Ae,Le)
    NTot  = MaxTurns(CoreID,WireD)
    NPrimary = NTot* (1/(1+NRatio))
    BCore = CalcFlux(Ae,Le,μr,NPrimary,VPk,ω)
    println("
    μr              = $(round(μr;sigdigits=3))
    Primary Turns   = $(round(NPrimary;sigdigits=3))
    Total Turns     = $(round(NTot;sigdigits=3))
    BCore(mT)       = $(round(BCore*1e3;sigdigits=3))
    ")
end

"""
Calculate Magnetizing Inductance, units in meters
"""
function MagInductance(Aₑ,Lₑ,N,μᵣ)
    μ₀ = 4*pi*1e-7
    L = N^2* μᵣ*μ₀ * Aₑ /Lₑ
    return L
end

"""
Calculate Magnetizing Inductance, units in meters
"""
function MagInductance_Gap(Aₑ,Lₑ,N,μᵣ,LGap)
    μ₀ = 4*pi*1e-7
    ℛgap = LGap/(μ₀*Aₑ)#\scrR
    ℛcore = Lₑ/(μ₀*μᵣ*Aₑ)#\scrR

    L = N^2 / (ℛcore+ℛgap)

    return L
end

function MagnetizingCurrent(L,f)
    return 2*π*f*L
end
""""
determine the resistance of a winding on a rectangular core with N layers of copper film winding

N is the number of layers
t the the thickness of copper
L is the length of the core
w is the width of the core (Inner dimension)
H is the height of the core
ρ is the resistivity of copper
"""
function FilmCoreResist_Rect(N,t,L,w,H,ρ)
    Length = N* 2*( (w+N/2*t) + (L+N/2*t))
    println(Length)
    R = ρ*Length / (H*t)
    return R
end

""""
determine the resistance of a winding on a rectangular core with Nₗ layers and Nₜ turns of copper circular wire

Nₗ is the number of layers
Nₜ is the number of turns
D the the diamter of copper wire
L is the length of the core
w is the width of the core (Inner dimension)
H is the height of the core
ρ is the resistivity of copper
"""
function WireWoundCoreResist_Rect(Nₗ,Nₜ,D,L,w,H,ρ)
    Length = Nₜ* 2*( (w+Nₗ/2*D) + (L+Nₗ/2*D))
    println(Length)
    R = ρ*Length / (π/4*D^2)
    return R
end

"""
The minimum allowable window area given some current density J₀, Maximum current, and turns (N)
"""
function MinWindowArea(Imax, J₀,N)
    return Imax*N/J₀
end

function MinWindowArea(Vmax,Imax, BMax, Acore,f, J₀)
    return Imax*MinTurns(Vmax, BMax, Acore,f) /J₀
end

function MinTurns(Vmax, BMax, Acore,f)
    return Vmax*√2 /(2*π*f* BMax * Acore)
end

function MeasuredInductance2Params(ZₚSecₒ,ZₚSecₛ,ZₛPriₒ,ZₛPriₛ)

    ZₚEqiv_SecLoad(Zₗ1,Zₘ1,Zₗ2,N,RTerm) = Zₗ1+ Par(Zₘ1,(1 / N^2 * (Zₗ2+RTerm))) #Primary referred impedance with some secondary termination (RTerm)
    ZₛEqiv_PriLoad(Zₗ1,Zₘ1,Zₗ2,N,RTerm) = Zₗ2+ N^2 * Par(Zₘ1,(Zₗ1+RTerm)) #Primary referred impedance with some secondary termination (RTerm)

    Cost(Zₗ1,Zₘ1,Zₗ2,N) = abs((ZₚSecₒ - ZₚEqiv_SecLoad(Zₗ1,Zₘ1,Zₗ2,N,1e9)) + (ZₚSecₛ - ZₚEqiv_SecLoad(Zₗ1,Zₘ1,Zₗ2,N,0)) +
                         (ZₛPriₒ - ZₛEqiv_PriLoad(Zₗ1,Zₘ1,Zₗ2,N,1e9)) + (ZₛPriₛ - ZₛEqiv_PriLoad(Zₗ1,Zₘ1,Zₗ2,N,0)))

    CostVecIn(Vec) = Cost(Vec[1],Vec[2],Vec[3],Vec[4])
    optimize(CostVecIn, [10.0 , 40.0, 10.0, 2.0])

end

# """
# Function to determine the leakage inductnace of a toroidal transformer
# Ref. `Calculation of self and mutual impedances between sections of tranformer windings`, Wilcox et al. IEE Proceedings, 1989

# """
# function CalcLeakageInductance_Toroid()

# end
# """
# Function to determine the leakage inductnace of a toroidal transformer
# Ref. `Calculation of self and mutual impedances between sections of tranformer windings`, Wilcox et al. IEE Proceedings, 1989

# """
# function CalcLₖₘ(Nₖ,Nₘ,r,a,z)
#     μ₀ = 4*π*1e-7
#     κ =κFunk(a,r,z)
#     Lₖₘ = μ₀*Nₖ*Nₘ*√(r*a)*2/κ*(  (1 - κ^2/2) * Elliptic.K(κ) - Elliptic.E(κ) )

# end

# function Z₁₍ₖₘ₎(Nₖ,Nₘ,b,μᵣ,ω,λ,ρ)

#     μ₁ = 4*π*1e-7 #for consistency with eqn. using 1 subscript
#     μₓ =4*π*1e-7*μᵣ
#     m = √(im*ω*μₓ/ρ)
#     Z = 1*im*ω*Nₖ*Nₘ *π*b^2/λ * (2*μₓ*besseli(1,m*b) / (m*b*besseli(0,m*b)) - μ₁)

# end

# """
# Equation 12,Ref. `Calculation of self and mutual impedances between sections of tranformer windings`, Wilcox et al. IEE Proceedings, 1989
# """
# function BigP₁(x,y,β)
#     (p₁(x) - p₁(y)) ./ (β.^2)

# end

# function p₁(α)
#     π*α/2 * (bessely(0,α)*Struve.L(0, α)+bessely(1,α)*Struve.L(1, α))
# end

# """
# Equation 10,Ref. `Calculation of self and mutual impedances between sections of tranformer windings`, Wilcox et al. IEE Proceedings, 1989
# """
# function κFunk(a,r,z)
#     √(4*a*r ./ (z^2 + (a+r)^2))
# end


include("testjacobi.jl")
include("testgauss.jl")


function _update_dual_constitutive!(H, B, C::Array{T,5}, nx, ny, nz) where T
    @turbo for k = 2:nz-1
        for j = 2:ny-1
            for i = 2:nx-1
                by = 0.25 * (B[i, j, k, 2] + B[i-1, j, k, 2] + B[i-1, j+1, k, 2] + B[i, j+1, k, 2])
                bz = 0.25 * (B[i, j, k, 3] + B[i, j, k+1, 3] + B[i-1, j, k+1, 3] + B[i-1, j, k, 3])
                H[i, j, k, 1] = C[i, j, k, 1, 1] * B[i, j, k, 1] + C[i, j, k, 1, 2] * by + C[i, j, k, 1, 3] * bz

                bx = 0.25 * (B[i, j, k, 1] + B[i, j-1, k, 1] + B[i+1, j-1, k, 1] + B[i+1, j, k, 1])
                bz = 0.25 * (B[i, j, k, 3] + B[i, j-1, k, 3] + B[i, j-1, k+1, 3] + B[i, j, k+1, 3])

                H[i, j, k, 2] = C[i, j, k, 2, 1] * bx + C[i, j, k, 2, 2] * B[i, j, k, 2] + C[i, j, k, 2, 3] * bz

                bx = 0.25 * (B[i, j, k, 1] + B[i+1, j, k, 1] + B[i+1, j, k-1, 1] + B[i, j, k-1, 1])
                by = 0.25 * (B[i, j, k, 2] + B[i, j+1, k, 2] + B[i, j+1, k-1, 2] + B[i, j, k-1, 2])

                H[i, j, k, 3] = C[i, j, k, 3, 1] * bx + C[i, j, k, 3, 2] * by + C[i, j, k, 3, 3] * B[i, j, k, 3]

            end
        end
    end

    @turbo for k = 1:1
        for j = 1:1
            for i = 2:nx-1
                by = 0.25 * (B[i, j, k, 2] + B[i-1, j, k, 2] + B[i-1, j+1, k, 2] + B[i, j+1, k, 2])
                bz = 0.25 * (B[i, j, k, 3] + B[i, j, k+1, 3] + B[i-1, j, k+1, 3] + B[i-1, j, k, 3])
                H[i, j, k, 1] = C[i, j, k, 1, 1] * B[i, j, k, 1] + C[i, j, k, 1, 2] * by + C[i, j, k, 1, 3] * bz

            end
        end
    end
    @turbo for k = 1:1
        for j = 2:ny-1
            for i = 1:1
                bx = 0.25 * (B[i, j, k, 1] + B[i, j-1, k, 1] + B[i+1, j-1, k, 1] + B[i+1, j, k, 1])
                bz = 0.25 * (B[i, j, k, 3] + B[i, j-1, k, 3] + B[i, j-1, k+1, 3] + B[i, j, k+1, 3])
                H[i, j, k, 2] = C[i, j, k, 2, 1] * bx + C[i, j, k, 2, 2] * B[i, j, k, 2] + C[i, j, k, 2, 3] * bz
            end
        end
    end
    @turbo for k = 2:nz-1
        for j = 1:1
            for i = 1:1
                bx = 0.25 * (B[i, j, k, 1] + B[i+1, j, k, 1] + B[i+1, j, k-1, 1] + B[i, j, k-1, 1])
                by = 0.25 * (B[i, j, k, 2] + B[i, j+1, k, 2] + B[i, j+1, k-1, 2] + B[i, j, k-1, 2])
                H[i, j, k, 3] = C[i, j, k, 3, 1] * bx + C[i, j, k, 3, 2] * by + C[i, j, k, 3, 3] * B[i, j, k, 3]
            end
        end
    end

end

function _update_dual_constitutive!(H, B, C::Array{T,4}, nx, ny, nz) where T
    @turbo for p in 1:3
        for k = 2:nz-1
            for j = 2:ny-1
                for i = 2:nx-1
                    H[i, j, k, p] = C[i, j, k, p] * B[i, j, k, p]
                end
            end
        end
    end

    @turbo for k = 1:1
        for j = 1:1
            for i = 2:nx-1
                H[i, j, k, 1] = C[i, j, k, 1] * B[i, j, k, 1]
            end
        end
    end
    @turbo for k = 1:1
        for j = 2:ny-1
            for i = 1:1
                H[i, j, k, 2] = C[i, j, k, 2] * B[i, j, k, 2]
            end
        end
    end
    @turbo for k = 2:nz-1
        for j = 1:1
            for i = 1:1
                H[i, j, k, 3] = C[i, j, k, 3] * B[i, j, k, 3]
            end
        end
    end
end


@safetestset "blocking" begin
    using DifferentialEvolutionMCMC, Test, Random, Distributions
    using MCMCChains
    Random.seed!(58122)

    data = rand(Normal(0.0, 1.0), 100)

    function prior_loglike(μ, σ)
        LL = 0.0
        LL += logpdf(Normal(0, 10), μ)
        LL += logpdf(truncated(Cauchy(0, 1), 0, Inf), σ)
        return LL
    end

    function sample_prior()
        μ = rand(Normal(0, 10))
        σ = rand(truncated(Cauchy(0, 1), 0, Inf))
        return [μ,σ]
    end

    bounds = ((-Inf,Inf),(0.0,Inf))

    data = rand(Normal(0.0, 1.0), 1000)

    function loglike(data, μ, σ)
        return sum(logpdf.(Normal(μ, σ), data))
    end

    bounds = ((-Inf,Inf),(0.0,Inf))

    names = (:μ,:σ)

    blocks = [[true,false],[false,true]]

    blocks = as_union(blocks)

    blocking_on = x -> true

    model = DEModel(;
        sample_prior,
        prior_loglike,
        loglike,
        data,
        names,
    )

    de = DE(;
        sample_prior,
        bounds,
        burnin = 1000,
        Np = 6,
        blocking_on,
        blocks
    )

    n_iter = 2000
    chains = sample(model, de, n_iter, progress=true)

    means = describe(chains)[1][:,:mean]
    rhat = describe(chains)[1][:,:rhat]

    @test means[1] ≈ 0.0 atol = .1
    @test means[2] ≈ 1.0 atol = .1
    @test rhat[1] ≈ 1.0 atol = .01
    @test rhat[2] ≈ 1.0 atol = .01

    chains = sample(model, de, MCMCThreads(), n_iter, progress=true)

    means = describe(chains)[1][:,:mean]
    rhat = describe(chains)[1][:,:rhat]

    @test means[1] ≈ 0.0 atol = .1
    @test means[2] ≈ 1.0 atol = .1
    @test rhat[1] ≈ 1.0 atol = .01
    @test rhat[2] ≈ 1.0 atol = .01
end