    using DifferentialEquations
    using OrdinaryDiffEq
    using DiffEqBase
    using Sundials
    using ODEInterfaceDiffEq

        # du -> Derivative
        # u -> State at time t
        # p-> paramter vector
        # t-> time as tuple (init, end)

        function Hes1ExampleODE!(du,u,p,t)
            m, p1, p2 = u;

            k1, P0, v, h, inps = p;

        du[1] = dm = -0.03*m + (1/(1+(p2/P0)^h))*inps;
        du[2] = dp1 = -0.03*p1 + v*m - k1*p1;
        du[3] = dp2 = -0.03*p2 + k1*p1;

        end

        # p -> parameter vector (plus inducer vector at the end)
        # I -> initial Y0 vector

        function Hes1ExampleSteadyState(p,I)

            m, p1, p2 = I;
               return(I)

        end

        # ts -> time vector going from t=0 to t=end every 1
        # p -> parameter vector
        # sp -> vector with switching times (t=0 and t=end needs to be included). Length of vector will be length(inducer steps)+1
        # inputs -> Vecor with the inducers (if more than 1 inducer in the system, the shape will be a1,b1,c1,..., a2,b2,c2,... aend,bend,cend,... where each leter represents a different inducer)
        # ivss -> Initial Y0 vector
        # pre -> Vector with the inducers in the ON. It can be empty

        function Hes1Example_solvecoupledODE(ts, p, sp, inputs, ivss, pre=[])

            maxtime = length(ts);
            Nsp = length(sp);
            Nevents = length(sp)-1;
            Neq = 3;
            if pre != []
                pSS = vcat(p, pre);
            else
                pSS = p;
            end

            final = zeros(maxtime, Neq);

            y_al = Hes1ExampleSteadyState(pSS,ivss) # Calculation of initial guesses for steady state

            y0 = y_al;
            initialV = y0;
            i = 1;

            for q in collect(1:Nevents)
                lts = length(ts[(sp[q]+1):sp[q+1]+1]);  # General way to define the number of elements in each event series
                Tevent = ts[(sp[q]+1):sp[q+1]+1];  # General way to extract the times of each event
                I = inputs[i:(i+(0))];
                pSte = vcat(p, I);

                if q == 1
                    prob = ODEProblem(Hes1ExampleODE!,initialV,(ts[(sp[q]+1)],ts[sp[q+1]+1]),pSte);
                    part1 = DifferentialEquations.solve(prob, Tsit5(),reltol=1.0e-9,abstol=1.0e-9,saveat=1);
                else
                    prob = ODEProblem(Hes1ExampleODE!,initialV,(ts[(sp[q]+1)],ts[sp[q+1]+1]),pSte);
                    part1 = DifferentialEquations.solve(prob, Tsit5(),reltol=1.0e-9,abstol=1.0e-9,saveat=1);
                end

                initialV = part1[lts];
                i+=1;

                for d in collect((sp[q]+1):(sp[q]+lts))
                    final[d,:] = part1[d-sp[q]];
                end
            end
            return(final)

        end

        # ts -> time vector going from t=0 to t=end every 1
        # pD -> Matrix containing all the draws for the parameters (can be a single vector if nly 1 draw of parameters is considered)
        # sp -> vector with switching times (t=0 and t=end needs to be included). Length of vector will be length(inducer steps)+1
        # inputs -> Vecor with the inducers (if more than 1 inducer in the system, the shape will be a1,b1,c1,..., a2,b2,c2,... aend,bend,cend,... where each leter represents a different inducer)
        # ivss -> Initial Y0 vector
        # samps -> Vector of Time points that want to be extracted from the simulations (it considers the initial time point as 0)
        # pre -> Vector with the inducers in the ON. It can be empty

        function Hes1Example_SolveAll(ts, pD, sp, inputs, ivss, samps, pre=[])

            if length(size(pD)) == 1
                pD = reshape(pD,size(pD)[1],1);
            end

            if size(pD)[2] != 4
                pD = pD';
            end

            if length(ivss)/3 > 1
                if size(ivss)[2] != 3
                    ivss = ivss';
                end
            end

            AllSolTest = zeros(length(samps), 3, length(pD[:,1]));

            for drawInd in collect(1:length(pD[:,1]))
                p = pD[drawInd,:];
                if length(ivss)/3 > 1
                    ivss2 = ivss[drawInd,:];
                else
                    ivss2 = ivss;
                end
                temp = Hes1Example_solvecoupledODE(ts, p, sp, inputs, ivss2, pre);
                AllSolTest[:,:,drawInd] = temp[convert.(Int,samps.+1),:];
            end

            return(AllSolTest);
        end


using MathProgComplex
using Base.Test

problems = [lasserre_ex1, lasserre_ex2, lasserre_ex5, test_gloptipoly_ex1, test_gloptipoly_ex2]

brokentestset = Set([(:MomentRelaxation, lasserre_ex2, 3), (:MomentRelaxation, lasserre_ex5, 2)])

@testset "SDP hierarchy base tests" begin

    @testset "$pbsolved tests" for pbsolved in [:SOSRelaxation, :MomentRelaxation]
        @testset "$problem_fct problem" for problem_fct in problems
            problem, order_to_obj, ε_abs = problem_fct()

            @testset "Order $d" for d in keys(order_to_obj)
                relax_ctx = set_relaxation(problem; hierarchykind=:Real,
                                            d = d,
                                            params = Dict(:opt_outlev=>0,
                                                        :opt_solver=>testsolver,
                                                        :opt_pbsolved=>pbsolved))

                primobj, dualobj = run_hierarchy(problem, relax_ctx)

                @show pbsolved, problem_fct, d
                @show d, primobj, dualobj, order_to_obj[d]

                if (pbsolved, problem_fct, d) in brokentestset
                    @test_broken primobj ≈ order_to_obj[d] atol=ε_abs
                else
                    @test primobj ≈ order_to_obj[d] atol=ε_abs
                    @test dualobj ≈ primobj atol=ε_abs*min(abs(primobj), abs(dualobj))
                end
            end
        end
    end
end


# ==================================
# Forward approximation from d/dt
# ==================================

"""
    ForwardDdt{EM, SO, SI, IT, BT} <: AbstractApproximationModel

Forward approximation model used in `d/dt`. It can be used for overapproximation
and underapproximation.

### Fields

- `oa`      -- (optional, default: `true`) flag to choose between
               overapproximation and underapproximation
- `exp`     -- (optional, default: `BaseExp`) exponentiation method
- `setops`  -- (optional, default: `:lazy`) set operations method
- `sih`     -- (optional, default: `:concrete`) way to compute the symmetric
               interval hull
- `inv`     -- (optional, default: `false`) if `true`, assume that the state
               matrix is invertible and use its inverse in the `Φ` functions
- `backend` -- (optional, default: `nothing`) used if the algorithm needs to
               apply concrete polyhedral computations

### Algorithm

The transformations are:

- ``Φ ← \\exp(Aδ)``,
- ``Ω_0 ← bloat(CH(\\mathcal{X}_0, Φ\\mathcal{X}_0), ε)``

where ``bloat(\\mathcal{X}, ε)`` bloats the set ``\\mathcal{X}`` with the value
``ε``. If `oa == false`, the bloating acts in an inverted way and shrinks the
set.

### Reference

E. Asarin, T. Dang, O. Maler, O. Bournez: *Approximate reachability analysis of
piecewise-linear dynamical systems*. HSCC 2000.
"""
struct ForwardDdt{EM, SO, SI, IT, BT} <: AbstractApproximationModel
    oa::Bool
    exp::EM
    setops::SO
    sih::SI
    inv::IT
    backend::BT
end

hasbackend(alg::ForwardDdt) = !isnothing(alg.backend)

# convenience constructor using symbols
function ForwardDdt(; oa::Bool=true, exp=BaseExp, setops=:lazy, sih=:concrete,
                      inv=false, backend=nothing)
    return ForwardDdt(oa, _alias(exp), _alias(setops), Val(sih), Val(inv), backend)
end

function Base.show(io::IO, alg::ForwardDdt)
    print(io, "`ForwardDdt` approximation model with: \n")
    print(io, "    - $(alg.oa ? "over" : "under")approximation \n")
    print(io, "    - exponentiation method: $(alg.exp) \n")
    print(io, "    - set operations method: $(alg.setops)\n")
    print(io, "    - symmetric interval hull method: $(alg.sih)\n")
    print(io, "    - invertibility assumption: $(alg.inv)")
    print(io, "    - polyhedral computations backend: $(alg.backend)")
end

Base.show(io::IO, m::MIME"text/plain", alg::ForwardDdt) = print(io, alg)

# ------------------------------------------------------------
# ForwardDdt Approximation: Homogeneous case
# ------------------------------------------------------------

# if A == |A|, then Φ can be reused in the computation of Φ₂(|A|, δ)
function discretize(ivp::IVP{<:CLCS, <:LazySet}, δ, alg::ForwardDdt)
    A = state_matrix(ivp)
    X0 = initial_state(ivp)

    Φ = _exp(A, δ, alg.exp)
    Y = ConvexHull(X0, Φ * X0)
    Y = _apply_setops(Y, alg)

    # bloating
    ε = _estimate_bloating_value_ddt(A, X0, δ)
    if !alg.oa
        ε = -ε
    end
    Ω0 = Bloating(Y, ε)

    X = stateset(ivp)
    Sdis = ConstrainedLinearDiscreteSystem(Φ, X)
    return InitialValueProblem(Sdis, Ω0)
end

function _estimate_bloating_value_ddt(A, X0, δ)
    norm_A = opnorm(A)
    u = 1/8 * norm_A^2 * δ^2
    v = exp(norm_A * δ) - 1 - (norm_A * δ) - (norm_A^2 * δ^2 / 2)
    return (u + v) * norm(X0)
end

# outer bloating a polytope into a polytope
function _bloat(P::LazySet, ε)
    A, b = tosimplehrep(P)
    bε = copy(b)
    for i in 1:length(b)
        bε[i] += ε * opnorm(transpose(A[i, :]))
    end
    return HPolytope(A, bε)
end


# Saving auxillary function which I need over multiple files

# Regression functions used in paper
function friedman(x::Matrix, errors::Matrix)

    Y = 10 .* sin.(π .* x[:,1] .* x[:,2]) .+ 20 .* (x[:,3] .- 0.5).^2 .+ 10 .* x[:,4] + 5 .* x[:,5] .+ errors

    return Y
end

function dp3(x::Matrix, errors::Matrix)

    Y = 4 .* (x[:, 1] .- 2 + 8 .* x[:, 2] .- 8 .* x[:,2].^2).^2 + (3 .- 4 .* x[:,2]).^2 + 16 .* (x[:, 3] .+ 1).^(0.5) .* (2 .* x[:,3] .- 1).^2 + errors

    return Y
end

function dp8(x::Matrix, errors::Matrix)

    Y = 4 .* (x[:, 1] .- 2 + 8 .* x[:, 2] .- 8 .* x[:,2].^2).^2 + (3 .- 4 .* x[:,2]).^2 + 16 .* (x[:, 3] .+ 1).^(0.5) .* (2 .* x[:,3] .- 1).^2 .+ sum(x[:, 4:8] .* (1 .+ log.(sum(x[:,1:3], dims=2))), dims=2) + errors

    return Y
end

function robot(x::Matrix, errors::Matrix)

    x[:, 1:4] = x[:, 1:4] .* 2 .* π     # Scaling inputs

    u = sum(x[:, 5:8] .* cos.(sum(x[:, 1:4], dims=2)) ,dims=2)
    v = sum(x[:, 5:8] .* sin.(sum(x[:, 1:4], dims=2)) ,dims=2)

    Y = (u.^2 .+ v.^2).^(0.5) + errors
    return Y
end

# Function for generating the data for given parameters
function make_data(n, d, func, σ)

    # Draw regressors
    x_train = rand(Uniform(0, 1), n, d)
    x_test = rand(Uniform(0, 1), n, d)

    # Drawing errors from N(0, σ)
    d = Normal(0, σ)
    td = truncated(d, -Inf, Inf)

    errors_train = rand(td, n, 1)
    errors_test = zeros(n, 1)

    if func=="friedman"
        y_train = friedman(x_train, errors_train)
        y_test = friedman(x_test, errors_test)
    elseif func=="dp3"
        y_train = dp3(x_train, errors_train)
        y_test = dp3(x_test, errors_test)
    elseif func=="dp8"
        y_train = dp8(x_train, errors_train)
        y_test = dp8(x_test, errors_test)
    elseif func=="robot"
        y_train = robot(x_train, errors_train)
        y_test = robot(x_test, errors_test)
    else
        error("Provide function to compute Y")
    end

    return x_train, x_test, y_train, y_test
end

# Get MSE for a vector of predictions and real values
function get_mse(pred, y)
    bias = abs(mean(pred .- y))
    variance = var(pred)
    mse = mean((pred .- mean(y)).^2)

    return bias, variance, mse
end

# Function used for me to get differences between models
function compare_models(rf1::RFR, rf2::RFR, x_test::Matrix, y_test::Matrix)

    pred1 = predict(rf1, x_test)
    pred2 = predict(rf2, x_test)

    println("Model 1: ", round.(get_mse(pred1, y_test), digits=5))
    println("Model 2: ", round.(get_mse(pred2, y_test), digits=5), "\n")
    println("Change 1 to 2: ", round.(get_mse(pred1, y_test)./get_mse(pred2, y_test), digits=5))
end


# export    UpdatableQR,
#           add_column!,
#           add_column_householder!,
#           remove_column!,
#           update_views,
#           qraddcol!
#
# # Credit to https://github.com/mpf/QRupdate.jl/blob/master/src/QRupdate.jl
#
# # function qraddcol(A::AbstractMatrix{T}, Rin::AbstractMatrix{T}, a::Vector{T}, β::T = 0.0) where {T <:Real}
#
# function qraddcol!(A::Array{Float64,2}, R::UpperTriangular{Float64,Array{Float64,2}}, a::Array{Float64,1}, β::Float64 = 0.0)
#
#   m, n = size(A)
#   anorm  = norm(a)
#   anorm2 = anorm^2
#   β2  = β^2
#   if β != 0
#       anorm2 = anorm2 + β2
#       anorm  = sqrt(anorm2)
#   end
#
#   if n == 0
#       return reshape([anorm], 1, 1)
#   end
#
#   # R = UpperTriangular(Rin)
#
#   c      = A'*a           # = [A' β*I 0]*[a; 0; β]
#   u      = R'\c
#   unorm2 = norm(u)^2
#   d2     = anorm2 - unorm2
#
#   if d2 > anorm2 #DISABLE 0.01*anorm2     # Cheap case: γ is not too small
#       γ = sqrt(d2)
#   else
#       z = R\u          # First approximate solution to min ||Az - a||
#       r = a - A*z
#       c = A'*r
#       if β != 0
#           c = c - β2*z
#       end
#       du = R'\c
#       dz = R\du
#       z  += dz          # Refine z
#     # u  = R*z          # Original:     Bjork's version.
#       u  += du          # Modification: Refine u
#       r  = a - A*z
#       γ = norm(r)       # Safe computation (we know gamma >= 0).
#       if β != 0
#           γ = sqrt(γ^2 + β2*norm(z)^2 + β2)
#       end
#   end
#
#   # This seems to be faster than concatenation, ie:
#   # [ Rin         u
#   #   zeros(1,n)  γ ]
#   # Rout = Array{T}(n+1, n+1)
#   # Rout[1:n,1:n] = R
#   # Rout[1:n,n+1] = u
#   # Rout[n+1,n+1] = γ
#   # Rout[n+1,1:n] = 0.0
#   R = UpperTriangular(hcat(vcat(R, zeros(1,n)),zeros(n+1,1)))
#   @show size(R)
#   # Rout = UpperTriangular(zeros(n+1, n+1))
#
#   # view(Rout,1:n,1:n) .= R.data
#   view(R,1:n,n+1) .= u
#   R[n+1,n+1] = γ
#   # fill!(view(Rout, n+1,1:n), 0.0)
#   # Rout[n+1,1:n] = 0.0
#
#   return R
#   # nothing
# end
#
#
# # Credit to https://github.com/oxfordcontrol/GeneralQP.jl/blob/master/src/linear_algebra.jl
# mutable struct UpdatableQR{T} <: Factorization{T}
#     """
#     Gives the qr factorization an (n, m) matrix as Q1*R1
#     Q2 is such that Q := [Q1 Q2] is orthogonal and R is an (n, n) matrix where R1 "views into".
#     """
#     Q::Matrix{T}
#     R::Matrix{T}
#     n::Int
#     m::Int
#
#     Q1::SubArray{T, 2, Matrix{T}, Tuple{Base.Slice{Base.OneTo{Int}}, UnitRange{Int}}, true}
#     Q2::SubArray{T, 2, Matrix{T}, Tuple{Base.Slice{Base.OneTo{Int}}, UnitRange{Int}}, true}
#     R1::UpperTriangular{T, SubArray{T, 2, Matrix{T}, Tuple{UnitRange{Int},UnitRange{Int}}, false}}
#
#     function UpdatableQR(A::AbstractMatrix{T}) where {T}
#         n, m = size(A)
#         @assert(m <= n, "Too many columns in the matrix.")
#
#         F = qr(A)
#         Q = F.Q*Matrix(I, n, n)
#         R = zeros(T, n, n)
#         R[1:m, 1:m] .= F.R
#
#         new{T}(Q, R, n, m,
#             view(Q, :, 1:m), view(Q, :, m+1:n),
#             UpperTriangular(view(R, 1:m, 1:m)))
#     end
# end
#
# function add_column!(F::UpdatableQR{T}, a::AbstractVector{T}) where {T}
#     a1 = F.Q1'*a;
#     a2 = F.Q2'*a;
#
#     x = copy(a2)
#     for i = length(x):-1:2
#         G, r = givens(x[i-1], x[i], i-1, i)
#         lmul!(G, x)
#         lmul!(G, F.Q2')
#     end
#
#     F.R[1:F.m, F.m+1] .= a1
#     F.R[F.m+1, F.m+1] = x[1]
#
#     F.m += 1; update_views!(F)
#
#     return a2
# end
#
# function add_column_householder!(F::UpdatableQR{T}, a::AbstractVector{T}) where {T}
#     a1 = F.Q1'*a;
#     a2 = F.Q2'*a;
#
#     Z = qr(a2)
#     LAPACK.gemqrt!('R','N', Z.factors, Z.T, F.Q2) # Q2 .= Q2*F.Q
#     F.R[1:F.m, F.m+1] .= a1
#     F.R[F.m+1, F.m+1] = Z.factors[1, 1]
#     F.m += 1; update_views!(F)
#
#     return Z
# end
#
# function remove_column!(F::UpdatableQR{T}, idx::Int) where {T}
#     Q12 = view(F.Q, :, idx:F.m)
#     R12 = view(F.R, idx:F.m, idx+1:F.m)
#
#     for i in 1:size(R12, 1)-1
#         G, r = givens(R12[i, i], R12[i + 1, i], i, i+1)
#         lmul!(G, R12)
#         rmul!(Q12, G')
#     end
#
#     for i in 1:F.m, j in idx:F.m-1
#         F.R[i, j] = F.R[i, j+1]
#     end
#     F.R[:, F.m] .= zero(T)
#
#     F.m -= 1; update_views!(F)
#
#     return nothing
# end
#
# function update_views!(F::UpdatableQR{T}) where {T}
#     F.R1 = UpperTriangular(view(F.R, 1:F.m, 1:F.m))
#     F.Q1 = view(F.Q, :, 1:F.m)
#     F.Q2 = view(F.Q, :, F.m+1:F.n)
# end


using LightGraphs: cycle_graph

n = 7
G = cycle_graph(n)
S = S0Graph(G)
@time λ = dsw(S, eye(n), eps=solver_eps)[1]
@test λ ≈ n*cos(pi/n) / (1 + cos(pi/n))  atol=tol


@testset "integration" begin
    prec = 64
    a = Acb(0, prec = prec)
    b = Acb(1, prec = prec)

    # Test with just a plain method
    f1 = sin
    f1! = Arblib.sin!
    res1 = "[0.459697694131860282 +/- 7.22e-19]"
    @test string(Arblib.integrate(f1, a, b, prec = prec)) == res1
    @test string(Arblib.integrate!(f1!, Acb(prec = prec), a, b)) == res1

    # Test with a method that accepts precision as a keyword argument
    f2 = (x; prec) -> Arblib.sin!(Acb(), x, prec = prec)
    f2! = Arblib.sin!
    res2 = "[0.459697694131860282 +/- 7.22e-19]"
    @test string(Arblib.integrate(f2, a, b, take_prec = true, prec = prec)) == res2
    @test string(Arblib.integrate!(f2!, Acb(prec = prec), a, b, take_prec = true)) == res2

    # Test with a method that accepts analytic as a keyword argument
    f3 = (x; analytic) -> Arblib.real_abs!(Acb(prec = prec), x, analytic)
    f3! = (res, x; analytic) -> Arblib.real_abs!(res, x, analytic)
    # FIXME: These are supposed to be identical but due to a bug in
    # Arb the produce slightly different results, see
    # https://github.com/kalmarek/Arblib.jl/issues/70. Once Arb is
    # updated so that they start to produce identical results this can
    # be updated and the issue closed.
    res3 = "[0.50000000000000000 +/- 2.73e-18]"
    @test string(Arblib.integrate(f3, a, b, check_analytic = true, prec = prec)) == res3
    @test string(Arblib.integrate!(f3!, Acb(prec = prec), a, b, check_analytic = true)) ==
          res3

    # Test with a method that accepts both precision and analytic as
    # a keyword arguments
    f4 = (x; analytic, prec) -> Arblib.real_abs!(Acb(), x, analytic, prec = prec)
    f4! = (res, x; analytic, prec) -> Arblib.real_abs!(res, x, analytic, prec = prec)
    # FIXME: See above
    res4 = "[0.50000000000000000 +/- 2.73e-18]"
    @test string(
        Arblib.integrate(f4, a, b, check_analytic = true, take_prec = true, prec = prec),
    ) == res4
    @test string(
        Arblib.integrate!(
            f4!,
            Acb(prec = prec),
            a,
            b,
            check_analytic = true,
            take_prec = true,
        ),
    ) == res4

    # Test with set tolerance
    f5 = x -> sin(exp(x))
    f5! = (res, x) -> Arblib.sin!(res, Arblib.exp!(res, x))
    res5 = "[0.8750 +/- 5.14e-5]"
    res5! = "[0.87495720 +/- 2.77e-9]"
    @test string(Arblib.integrate(f5, a, b, prec = prec, rtol = 1e-4)) == res5
    @test string(Arblib.integrate!(f5!, Acb(prec = prec), a, b, atol = 1e-8)) == res5!
end


# Property (*) means that the step size α*d is small -> β = 0
struct CGRestart{T} end
struct FletcherReeves{T} end
# using an "intial vectors" function we can initialize s if necessary or nothing if not to save on vectors
function optimize(obj::OnceDiffed, x0, cg::FletcherReeves, linesearch = Backtracking(); maxiter=100)
    T=eltype(x0)
    fx, ∇fx = obj(x0, copy(x0))
    Tf = typeof(fx)
    fz = fx
    d = -copy(∇fx)
    ∇fz = copy(∇fx)
    k = 0
    x = copy(x0)
    while norm(∇fx, Inf) ≥ 1e-8 && k ≤ maxiter
        α, f_α, ls_success = NLSolvers.find_steplength(linesearch, obj, d, x, Tf(1.0), fx, ∇fx)
        z=x+α*d
        fz, ∇fz = obj(z, ∇fz)

        βkp1 = calc_β(cg, ∇fz, ∇fx)
        fx = fz
        x .= z
        ∇fx .= ∇fz

        d .= -∇fx + βkp1*d
        k = k+1
    end
    x, fx
end
cg_restart(::FletcherReeves{<:CGRestart{true}}) = true

function calc_β(cg, ∇fz, ∇fx)
    T = eltype(∇fz)
    βkp1 = dot(∇fz, ∇fz)/dot(∇fx, ∇fx)
    if cg_restart(cg)
        βkp1 = max(βkp1, T(0))
    end
    βkp1
end
# function calc_β(method, ∇fx, ∇fz)
#
# if method == :fr
# return dot(∇fz, ∇fz)/dot(∇fx, ∇fx)
# elseif method == :pr
# β = dot(∇fz, s)/dot(∇fx, ∇fx) # dot(∇fx, ∇fx) can be saved between iteratoins only s is needed
# return    max(β, 0) # We need this according to Powell 1986 Gilbert Powell 1980 show it is globally  convergenct for exact and inexact
# # pr+ max(beta, 0)
# elseif method ==:hs
# dot(∇fz, s)/dot(s,p)
# end


using Analytical
using Test

# Running Analytical approximation with default parameters
@testset "Analytical.jl" begin

	adap = Analytical.parameters(N=500,n=661,alLow=0.2)
	cnv = Analytical.binomialDict()
	Analytical.binomOp!(adap,cnv.bn)
	Analytical.setThetaF!(adap)
	Analytical.setPpos!(adap)

	@test adap.pposH == 0.0002641360146660052
	@test adap.thetaF == 1.6433217979109007e-6

end


using Oceananigans.Architectures: architecture
using Oceananigans.Grids: interior_parent_indices
using Statistics: norm, dot
using LinearAlgebra

mutable struct PreconditionedConjugateGradientSolver{A, G, L, T, F, M, P}
               architecture :: A
                       grid :: G
          linear_operation! :: L
                  tolerance :: T
         maximum_iterations :: Int
                  iteration :: Int
                       ρⁱ⁻¹ :: T
    linear_operator_product :: F
           search_direction :: F
                   residual :: F
              precondition! :: M
     preconditioner_product :: P
end

no_precondition!(args...) = nothing

initialize_precondition_product(precondition, template_field) = similar(template_field)
initialize_precondition_product(::Nothing, template_field) = nothing

maybe_precondition(::Nothing, ::Nothing, r, args...) = r

function maybe_precondition(precondition!, z, r, args...)
    precondition!(z, r, args...)
    return z
end

"""
    PreconditionedConjugateGradientSolver(linear_operation;
                                          template_field,
                                          maximum_iterations = size(template_field.grid),
                                          tolerance = 1e-13,
                                          precondition = nothing)

Returns a PreconditionedConjugateGradientSolver that solves the linear equation
``A x = b`` using a iterative conjugate gradient method with optional preconditioning.
The solver is used by calling

```
solve!(x, solver::PreconditionedConjugateGradientOperator, b, args...)
```

for `solver`, right-hand side `b`, solution `x`, and optional arguments `args...`.

Arguments
=========

* `template_field`: Dummy field that is the same type and size as `x` and `b`, which
                    is used to infer the `architecture`, `grid`, and to create work arrays
                    that are used internally by the solver.

* `linear_operation`: Function with signature `linear_operation!(p, y, args...)` that calculates
                     `A*y` and stores the result in `p` for a "candidate solution `y`. `args...`
                     are optional positional arguments passed from `solve!(x, solver, b, args...)`.

* `maximum_iterations`: Maximum number of iterations the solver may perform before exiting.

* `tolerance`: Tolerance for convergence of the algorithm. The algorithm quits when
               `norm(A * x - b) < tolerance`.

* `precondition`: Function with signature `preconditioner!(z, y, args...)` that calculates
                  `P * y` and stores the result in `z` for linear operator `P`.
                  Note that some precondition algorithms describe the step
                  "solve `M * x = b`" for precondition `M`"; in this context,
                  `P = M⁻¹`.

See [`solve!`](@ref) for more information about the preconditioned conjugate-gradient algorithm.
"""
function PreconditionedConjugateGradientSolver(linear_operation;
                                               template_field::AbstractField,
                                               maximum_iterations = prod(size(template_field)),
                                               tolerance = 1e-13, #sqrt(eps(eltype(template_field.grid))),
                                               precondition = nothing)

    arch = architecture(template_field)
    grid = template_field.grid

    # Create work arrays for solver
    linear_operator_product = similar(template_field) # A*xᵢ = qᵢ
    search_direction = similar(template_field) # pᵢ
            residual = similar(template_field) # rᵢ

    # Either nothing (no precondition) or P*xᵢ = zᵢ
    precondition_product = initialize_precondition_product(precondition, template_field)

    return PreconditionedConjugateGradientSolver(arch,
                                                 grid,
                                                 linear_operation,
                                                 tolerance,
                                                 maximum_iterations,
                                                 0,
                                                 0.0,
                                                 linear_operator_product,
                                                 search_direction,
                                                 residual,
                                                 precondition,
                                                 precondition_product)
end

"""
    solve!(x, solver::PreconditionedConjugateGradientSolver, b, args...)

Solve `A * x = b` using an iterative conjugate-gradient method, where `A * x` is
determined by `solver.linear_operation`

See figure 2.5 in

> The Preconditioned Conjugate Gradient Method in "Templates for the Solution of Linear Systems: Building Blocks for Iterative Methods" Barrett et. al, 2nd Edition.

Given:
  * Linear Preconditioner operator `M!(solution, x, other_args...)` that computes `M * x = solution`
  * A matrix operator `A` as a function `A()`;
  * A dot product function `norm()`;
  * A right-hand side `b`;
  * An initial guess `x`; and
  * Local vectors: `z`, `r`, `p`, `q`

This function executes the algorithm

```
β  = 0
r = b - A(x)
iteration  = 0

Loop:
     if iteration > maximum_iterations
        break
     end

     ρ = r ⋅ z

     z = M(r)
     β = ρⁱ⁻¹ / ρ
     p = z + β * p
     q = A(p)

     α = ρ / (p ⋅ q)
     x = x + α * p
     r = r - α * q

     if |r| < tolerance
        break
     end

     iteration += 1
     ρⁱ⁻¹ = ρ
```
"""
function solve!(x, solver::PreconditionedConjugateGradientSolver, b, args...)

    # Initialize
    solver.iteration = 0

    # q = A*x
    q = solver.linear_operator_product
    solver.linear_operation!(q, x, args...)

    # r = b - A*x
    solver.residual .= b .- q

    @debug "PreconditionedConjugateGradientSolver, |b|: $(norm(b))"
    @debug "PreconditionedConjugateGradientSolver, |A(x)|: $(norm(q))"

    while iterating(solver)
        iterate!(x, solver, b, args...)
    end

    fill_halo_regions!(x, solver.architecture) # blocking

    return x
end

function iterate!(x, solver, b, args...)
    r = solver.residual
    p = solver.search_direction
    q = solver.linear_operator_product

    @debug "PreconditionedConjugateGradientSolver $(solver.iteration), |r|: $(norm(r))"

    # Preconditioned:   z = P * r
    # Unpreconditioned: z = r
    z = maybe_precondition(solver.precondition!, solver.preconditioner_product, r, args...)
    ρ = dot(z, r)

    @debug "PreconditionedConjugateGradientSolver $(solver.iteration), ρ: $ρ"
    @debug "PreconditionedConjugateGradientSolver $(solver.iteration), |z|: $(norm(z))"

    if solver.iteration == 0
        p .= z
    else
        β = ρ / solver.ρⁱ⁻¹
        p .= z .+ β .* p

        @debug "PreconditionedConjugateGradientSolver $(solver.iteration), β: $β"
    end

    # q = A * p
    solver.linear_operation!(q, p, args...)
    α = ρ / dot(p, q)

    @debug "PreconditionedConjugateGradientSolver $(solver.iteration), |q|: $(norm(q))"
    @debug "PreconditionedConjugateGradientSolver $(solver.iteration), α: $α"

    x .+= α .* p
    r .-= α .* q

    solver.iteration += 1
    solver.ρⁱ⁻¹ = ρ

    return nothing
end

function iterating(solver)
    # End conditions
    solver.iteration >= solver.maximum_iterations && return false
    norm(solver.residual) <= solver.tolerance && return false
    return true
end

function Base.show(io::IO, solver::PreconditionedConjugateGradientSolver)
    print(io, "Oceananigans-compatible preconditioned conjugate gradient solver.\n")
    print(io, " Problem size = "  , size(solver.q), '\n')
    print(io, " Grid = "  , solver.grid)
    return nothing
end


function digamma(x::Float64) # Use Int instead?
  ccall(("digamma", :libRmath), Float64, (Float64, ), x)
end

function trigamma(x::Float64) # Use Int instead?
  ccall(("trigamma", :libRmath), Float64, (Float64, ), x)
end

function fit(::Type{Bernoulli}, x::Array)
  for i in 1:length(x)
    if !insupport(Bernoulli(), x[i])
      error("Bernoulli observations must be 0/1 values")
    end
  end
  return Bernoulli(mean(x))
end

function fit(::Type{Beta}, x::Array)
  for i in 1:length(x)
    if !insupport(Beta(), x[i])
      error("Bernoulli observations must be in [0,1]")
    end
  end
  x_bar = mean(x)
  v_bar = var(x)
  a = x_bar * (((x_bar * (1. - x_bar)) / v_bar) - 1.)
  b = (1. - x_bar) * (((x_bar * (1. - x_bar)) / v_bar) - 1.)
  return Beta(a, b)
end

function fit(::Type{Binomial}, x::Real, n::Real)
  if x > n || x < 0
    error("For binomial observations, x must lie in [0, n]")
  end
  return Binomial(int(n), x / n)
end

# Categorical

# Cauchy

# Chisq

# Dirichlet

function fit(::Type{DiscreteUniform}, x::Array)
  DiscreteUniform(min(x), max(x))
end

function fit(::Type{Exponential}, x::Array)
  for i in 1:length(x)
    if !insupport(Exponential(), x[i])
      error("Exponential observations must be non-negative values")
    end
  end
  return Exponential(mean(x))
end

# FDist

function fit(::Type{Gamma}, x::Array)
  a_old = 0.5 / (log(mean(x)) - mean(log(x)))
  a_new = 0.5 / (log(mean(x)) - mean(log(x)))

  z = 1.0 / a_old + (mean(log(x)) - log(mean(x)) + log(a_old) - digamma(a_old)) / (a_old^2 * (1.0 / a_old - trigamma(a_old)))
  a_new, a_old = (1.0 / z, a_new)

  iterations = 1

  while abs(a_old - a_new) > 10e-16 && iterations <= 10_000
    z = 1.0 / a_old + (mean(log(x)) - log(mean(x)) + log(a_old) - digamma(a_old)) / (a_old^2 * (1.0 / a_old - trigamma(a_old)))
    a_new, a_old = (1.0 / z, a_new)
  end

  if iterations >= 10_000
    error("Parameter estimation failed to converge in 10,000 iterations")
  end

  Gamma(a_new, mean(x) / a_new)
end

function fit(::Type{Geometric}, x::Array)
  Geometric(1.0 / mean(convert(Array{Int}, x)))
end

# HyperGeometric

# Logistic

function fit(::Type{Laplace}, x::Array)
  a = median(x)
  deviations = 0.0
  for i in 1:length(x)
    deviations += abs(x[i] - a)
  end
  b = deviations / length(x)
  Laplace(a, b)
end

# logNormal

# function fit(Multinomial)
# end

# Assumes observations are rows
# Positive-definite errors
# function fit{T <: Real}(x::Matrix{T}, d::MultivariateNormal)
#   MultivariateNormal(reshape(mean(x, 1), size(x, 2)), cov(x))
# end

# fit([1.0 -0.1; -2.0 0.1], MultivariateNormal())

# NegativeBinomial

# NoncentralBeta

# NoncentralChisq

# NoncentralF

# NoncentralT

# Normal
# Not strict MLE
function fit(::Type{Normal}, x::Array)
  Normal(mean(x), std(x))
end

# Poisson
function fit(::Type{Poisson}, x::Array)
  for i in 1:length(x)
    if !insupport(Poisson(), x[i])
      error("Poisson observations must be non-negative integers")
    end
  end
  Poisson(mean(x))
end

# TDist

# Uniform
function fit(::Type{Uniform}, x::Array)
  Uniform(min(x), max(x))
end

# Weibull


using TensorDSL

function sign(per1, per2)
    changes = 0
    for i in 1:length(per1)
        if per1[i] != per2[i]
            changes += 1
            j = findfirst(x -> (x == per1[i]), per2[(i+1):end]) + i
            per2[j] = per2[i]
        end
    end
    return iseven(changes) ? 1 : -1
end

function getvalue(index)
    upperindices = []
    lowerindices = []
    for i in 1:length(index)
        if isodd(i)
            push!(upperindices, index[i])
        else
            push!(lowerindices, index[i])
        end
    end
    return (sort(upperindices) == sort(lowerindices) &&
            length(Set(upperindices)) == length(upperindices)) ? sign(lowerindices, upperindices) : 0
end

function shapetodelta(Vs::Vararg{AbstractVectorSpace})
    # Check the shape is legal
    DeltaTensor(Vs...)
    shape = [dim(V) for V in Vs]
    delta = Array{Int}(undef, shape...)
    indices = CartesianIndices(tuple(shape...))
    for index in indices
        delta[index] = getvalue(index)
    end
    return delta
end


"""
    - free cone
        x ∈ Rⁿ

    A Semismooth Newton Method for Fast, Generic Convex Programming
        https://arxiv.org/abs/1705.00772
"""

# free cone
function κ_free(z)
    z
end

# free cone Jacobian
function Jκ_free(z)
    p = ones(length(z))
    return Diagonal(p)
end


"""
    QuadratureVI(;ϵ::T=1e-5, nGaussHermite::Integer=20, clipping=Inf, natural::Bool=true, optimiser=Momentum(0.0001))

Variational Inference solver by approximating gradients via numerical integration via Quadrature

## Keyword arguments
- `ϵ::T` : convergence criteria
- `nGaussHermite::Int` : Number of points for the integral estimation
- `clipping::Real` : Limit the gradients values to avoid overshooting
- `natural::Bool` : Use natural gradients
- `optimiser` : Optimiser used for the variational updates. Should be an Optimiser object from the [Flux.jl](https://github.com/FluxML/Flux.jl) library, see list here [Optimisers](https://fluxml.ai/Flux.jl/stable/training/optimisers/) and on [this list](https://github.com/theogf/AugmentedGaussianProcesses.jl/tree/master/src/inference/optimisers.jl). Default is `Momentum(0.0001)`
"""
mutable struct QuadratureVI{T} <: NumericalVI{T}
    n_points::Int # Number of points for the quadrature
    nodes::Vector{T} # Nodes locations
    weights::Vector{T} # Weights for each node
    clipping::T # Clipping value of the gradient
    ϵ::T #Convergence criteria
    n_iter::Integer #Number of steps performed
    stoch::Bool #Use of mini-batches
    batchsize::Int #Size of mini-batches
    ρ::T #Stochastic Coefficient
    NaturalGradient::Bool
    HyperParametersUpdated::Bool # Flag for updating kernel matrices
    vi_opt::NVIOptimizer

    function QuadratureVI{T}(
        ϵ::T,
        n_points::Integer,
        optimiser,
        Stochastic::Bool,
        clipping::Real,
        batchsize::Int,
        natural::Bool,
    ) where {T}
        gh = gausshermite(n_points)
        return new{T}(
            n_points,
            gh[1] .* sqrt2,
            gh[2] ./ sqrtπ,
            clipping,
            ϵ,
            0,
            Stochastic,
            batchsize,
            one(T),
            natural,
            true,
            NVIOptimizer(optimiser),
        )
    end
end

function QuadratureVI(;
    ϵ::T=1e-5,
    nGaussHermite::Integer=100,
    optimiser=Momentum(1e-5),
    clipping::Real=0.0,
    natural::Bool=true,
) where {T<:Real}
    return QuadratureVI{T}(ϵ, nGaussHermite, optimiser, false, clipping, 1, natural)
end

"""
    QuadratureSVI(nMinibatch::Int; ϵ::T=1e-5, nGaussHermite::Int=20, clipping=Inf, natural=true, optimiser=Momentum(0.0001))

Stochastic Variational Inference solver by approximating gradients via numerical integration via Gaussian Quadrature.
See [`QuadratureVI`](@ref) for a more detailed reference.

## Arguments
-`nMinibatch::Integer` : Number of samples per mini-batches

## Keyword arguments
- `ϵ::T` : convergence criteria, which can be user defined
- `nGaussHermite::Int` : Number of points for the integral estimation (for the QuadratureVI)
- `natural::Bool` : Use natural gradients
- `optimiser` : Optimiser used for the variational updates. Should be an Optimiser object from the [Flux.jl](https://github.com/FluxML/Flux.jl) library, see list here [Optimisers](https://fluxml.ai/Flux.jl/stable/training/optimisers/) and on [this list](https://github.com/theogf/AugmentedGaussianProcesses.jl/tree/master/src/inference/optimisers.jl). Default is `Momentum(0.0001)`
"""
function QuadratureSVI(
    nMinibatch::Integer;
    ϵ::T=1e-5,
    nGaussHermite::Integer=100,
    optimiser=Momentum(1e-5),
    clipping::Real=0.0,
    natural=true,
) where {T<:Real}
    return QuadratureVI{T}(ϵ, nGaussHermite, optimiser, true, clipping, nMinibatch, natural)
end

function expec_loglikelihood(
    l::AbstractLikelihood, i::QuadratureVI, y, μ::AbstractVector, diagΣ::AbstractVector
)
    return mapreduce(+, y, μ, diagΣ) do y, μ, σ²
        apply_quad(i, y, μ, σ², l)
    end
end

function apply_quad(i::QuadratureVI, y::Real, μ::Real, σ²::Real, l::AbstractLikelihood)
    xs = i.nodes * sqrt(σ²) .+ μ
    return dot(i.weights, loglikelihood.(Ref(l), y, xs))
    # return mapreduce((w, x) -> w * Distributions.loglikelihood(l, y, x), +, i.weights, xs)# loglikelihood.(l, y, x))
end

function grad_expectations!(m::AbstractGPModel{T,L,<:QuadratureVI}, state, y) where {T,L}
    for (gp, opt_state, kernel_matrices) in zip(m.f, state.opt_state, state.kernel_matrices)
        μ = mean_f(gp, kernel_matrices)
        Σ = var_f(gp, kernel_matrices)
        for i in 1:batchsize(inference(m))
            opt_state.ν[i], opt_state.λ[i] = grad_quad(
                likelihood(m), y[i], μ[i], Σ[i], inference(m)
            )
        end
    end
end

# Compute the first and second derivative of the log-likelihood using the quadrature nodes
function grad_quad(l::AbstractLikelihood, y::Real, μ::Real, σ²::Real, i::AbstractInference)
    x = i.nodes * sqrt(max(σ², zero(σ²))) .+ μ
    Edloglike = dot(i.weights, ∇loglikehood.(l, y, x))
    Ed²loglike = dot(i.weights, hessloglikehood.(l, y, x))
    if i.clipping != 0
        return (
            abs(Edloglike) > i.clipping ? sign(Edloglike) * i.clipping : -Edloglike,
            abs(Ed²loglike) > i.clipping ? sign(Ed²loglike) * i.clipping : -Ed²loglike,
        )
    else
        return -Edloglike, Ed²loglike
    end
end


using StatsBase
using LinearAlgebra, Random, Test

@testset "StatsBase.Covariance" begin
weight_funcs = (weights, aweights, fweights, pweights)

@testset "$f" for f in weight_funcs
    X = randn(3, 8)

    Z1 = X .- mean(X, dims = 1)
    Z2 = X .- mean(X, dims = 2)

    w1 = rand(3)
    w2 = rand(8)

    # varcorrection is negative if sum of weights is smaller than 1
    if f === fweights
        w1[1] += 1
        w2[1] += 1
    end

    wv1 = f(w1)
    wv2 = f(w2)

    Z1w = X .- mean(X, wv1, 1)
    Z2w = X .- mean(X, wv2, 2)

    ## reference results

    S1 = Z1'Z1
    S2 = Z2 * Z2'

    Sz1 = X'X
    Sz2 = X * X'

    S1w = Z1w' * Matrix(Diagonal(w1)) * Z1w
    S2w = Z2w * Matrix(Diagonal(w2)) * Z2w'

    Sz1w = X' * Matrix(Diagonal(w1)) * X
    Sz2w = X * Matrix(Diagonal(w2)) * X'

    @testset "Scattermat" begin
        @test scattermat(X)    ≈ S1
        @test scattermat(X, 2) ≈ S2

        @test StatsBase.scattermatm(X, 0)    ≈ Sz1
        @test StatsBase.scattermatm(X, 0, 2) ≈ Sz2

        @test StatsBase.scattermatm(X, mean(X, dims = 1))    ≈ S1
        @test StatsBase.scattermatm(X, mean(X, dims = 2), 2) ≈ S2

        @test StatsBase.scattermatm(X, zeros(1,8))  ≈ Sz1
        @test StatsBase.scattermatm(X, zeros(3), 2) ≈ Sz2

        @testset "Weighted" begin
            @test scattermat(X, wv1)    ≈ S1w
            @test scattermat(X, wv2, 2) ≈ S2w

            @test StatsBase.scattermatm(X, 0, wv1)    ≈ Sz1w
            @test StatsBase.scattermatm(X, 0, wv2, 2) ≈ Sz2w

            @test StatsBase.scattermatm(X, mean(X, wv1, 1), wv1)    ≈ S1w
            @test StatsBase.scattermatm(X, mean(X, wv2, 2), wv2, 2) ≈ S2w

            @test StatsBase.scattermatm(X, zeros(1,8), wv1)  ≈ Sz1w
            @test StatsBase.scattermatm(X, zeros(3), wv2, 2) ≈ Sz2w
        end
    end

    @testset "Uncorrected" begin
        @testset "Weighted Covariance" begin
            @test cov(X, wv1; corrected=false)    ≈ S1w ./ sum(wv1)
            @test cov(X, wv2, 2; corrected=false) ≈ S2w ./ sum(wv2)

            @test StatsBase.covm(X, 0, wv1, 1; corrected=false) ≈ Sz1w ./ sum(wv1)
            @test StatsBase.covm(X, 0, wv2, 2; corrected=false) ≈ Sz2w ./ sum(wv2)

            @test StatsBase.covm(X, mean(X, wv1, 1), wv1, 1; corrected=false) ≈ S1w ./ sum(wv1)
            @test StatsBase.covm(X, mean(X, wv2, 2), wv2, 2; corrected=false) ≈ S2w ./ sum(wv2)

            @test StatsBase.covm(X, zeros(1,8), wv1, 1; corrected=false) ≈ Sz1w ./ sum(wv1)
            @test StatsBase.covm(X, zeros(3), wv2, 2; corrected=false)   ≈ Sz2w ./ sum(wv2)
        end

        @testset "Mean and covariance" begin
            (m, C) = mean_and_cov(X; corrected=false)
            @test m == mean(X, dims = 1)
            @test C == cov(X, dims = 1, corrected=false)

            (m, C) = mean_and_cov(X, 1; corrected=false)
            @test m == mean(X, dims = 1)
            @test C == cov(X, dims = 1, corrected = false)

            (m, C) = mean_and_cov(X, 2; corrected=false)
            @test m == mean(X, dims = 2)
            @test C == cov(X, dims = 2, corrected = false)

            (m, C) = mean_and_cov(X, wv1; corrected=false)
            @test m == mean(X, wv1, 1)
            @test C == cov(X, wv1, 1, corrected=false)

            (m, C) = mean_and_cov(X, wv1, 1; corrected=false)
            @test m == mean(X, wv1, 1)
            @test C == cov(X, wv1, 1, corrected=false)

            (m, C) = mean_and_cov(X, wv2, 2; corrected=false)
            @test m == mean(X, wv2, 2)
            @test C == cov(X, wv2, 2, corrected=false)
        end
        @testset "Conversions" begin
            std1 = std(X, wv1, 1; corrected=false)
            std2 = std(X, wv2, 2; corrected=false)

            cov1 = cov(X, wv1, 1; corrected=false)
            cov2 = cov(X, wv2, 2; corrected=false)

            cor1 = cor(X, wv1, 1)
            cor2 = cor(X, wv2, 2)

            @testset "cov2cor" begin
                @test cov2cor(cov(X, dims = 1), std(X, dims = 1)) ≈ cor(X, dims = 1)
                @test cov2cor(cov(X, dims = 2), std(X, dims = 2)) ≈ cor(X, dims = 2)
                @test cov2cor(cov1, std1) ≈ cor1
                @test cov2cor(cov2, std2) ≈ cor2
            end
            @testset "cor2cov" begin
                @test cor2cov(cor(X, dims = 1), std(X, dims = 1)) ≈ cov(X, dims = 1)
                @test cor2cov(cor(X, dims = 2), std(X, dims = 2)) ≈ cov(X, dims = 2)
                @test cor2cov(cor1, std1) ≈ cov1
                @test cor2cov(cor2, std2) ≈ cov2
            end
        end
    end

    @testset "Corrected" begin
        @testset "Weighted Covariance" begin
            if isa(wv1, Weights)
                @test_throws ArgumentError cov(X, wv1; corrected=true)
            else
                var_corr1 = StatsBase.varcorrection(wv1, true)
                var_corr2 = StatsBase.varcorrection(wv2, true)

                @test cov(X, wv1; corrected=true)    ≈ S1w .* var_corr1
                @test cov(X, wv2, 2; corrected=true) ≈ S2w .* var_corr2

                @test StatsBase.covm(X, 0, wv1, 1; corrected=true) ≈ Sz1w .* var_corr1
                @test StatsBase.covm(X, 0, wv2, 2; corrected=true) ≈ Sz2w .* var_corr2

                @test StatsBase.covm(X, mean(X, wv1, 1), wv1, 1; corrected=true) ≈ S1w .* var_corr1
                @test StatsBase.covm(X, mean(X, wv2, 2), wv2, 2; corrected=true) ≈ S2w .* var_corr2

                @test StatsBase.covm(X, zeros(1,8), wv1, 1; corrected=true) ≈ Sz1w .* var_corr1
                @test StatsBase.covm(X, zeros(3), wv2, 2; corrected=true)   ≈ Sz2w .* var_corr2
            end
        end
        @testset "Mean and covariance" begin
            (m, C) = mean_and_cov(X; corrected=true)
            @test m == mean(X, dims =1)
            @test C == cov(X, dims = 1, corrected = true)

            (m, C) = mean_and_cov(X, 1; corrected=true)
            @test m == mean(X, dims = 1)
            @test C == cov(X, dims = 1, corrected = true)

            (m, C) = mean_and_cov(X, 2; corrected=true)
            @test m == mean(X, dims = 2)
            @test C == cov(X, dims = 2, corrected = true)

            if isa(wv1, Weights)
                @test_throws ArgumentError mean_and_cov(X, wv1; corrected=true)
            else
                (m, C) = mean_and_cov(X, wv1; corrected=true)
                @test m == mean(X, wv1, 1)
                @test C == cov(X, wv1, 1; corrected=true)

                (m, C) = mean_and_cov(X, wv1, 1; corrected=true)
                @test m == mean(X, wv1, 1)
                @test C == cov(X, wv1, 1; corrected=true)

                (m, C) = mean_and_cov(X, wv2, 2; corrected=true)
                @test m == mean(X, wv2, 2)
                @test C == cov(X, wv2, 2; corrected=true)
            end
        end
        @testset "Conversions" begin
            if !isa(wv1, Weights)
                std1 = std(X, wv1, 1; corrected=true)
                std2 = std(X, wv2, 2; corrected=true)

                cov1 = cov(X, wv1, 1; corrected=true)
                cov2 = cov(X, wv2, 2; corrected=true)

                cor1 = cor(X, wv1, 1)
                cor2 = cor(X, wv2, 2)

                @testset "cov2cor" begin
                    @test cov2cor(cov(X, dims = 1), std(X, dims = 1)) ≈ cor(X, dims = 1)
                    @test cov2cor(cov(X, dims = 2), std(X, dims = 2)) ≈ cor(X, dims = 2)
                    @test cov2cor(cov1, std1) ≈ cor1
                    @test cov2cor(cov2, std2) ≈ cor2
                end

                @testset "cov2cor!" begin
                    tmp_cov1 = copy(cov1)
                    @test !(tmp_cov1 ≈ cor1)
                    StatsBase.cov2cor!(tmp_cov1, std1)
                    @test tmp_cov1 ≈ cor1

                    tmp_cov2 = copy(cov2)
                    @test !(tmp_cov2 ≈ cor2)
                    StatsBase.cov2cor!(tmp_cov2, std2)
                    @test tmp_cov2 ≈ cor2
                end

                @testset "cor2cov" begin
                    @test cor2cov(cor(X, dims = 1), std(X, dims = 1)) ≈ cov(X, dims = 1)
                    @test cor2cov(cor(X, dims = 2), std(X, dims = 2)) ≈ cov(X, dims = 2)
                    @test cor2cov(cor1, std1) ≈ cov1
                    @test cor2cov(cor2, std2) ≈ cov2
                end

                @testset "cor2cov!" begin
                    tmp_cor1 = copy(cor1)
                    @test !(tmp_cor1 ≈ cov1)
                    StatsBase.cor2cov!(tmp_cor1, std1)
                    @test tmp_cor1 ≈ cov1

                    tmp_cor2 = copy(cor2)
                    @test !(tmp_cor2 ≈ cov2)
                    StatsBase.cor2cov!(tmp_cor2, std2)
                    @test tmp_cor2 ≈ cov2
                end
            end
        end
    end

    @testset "Correlation" begin
        @test cor(X, f(ones(3)), 1) ≈ cor(X, dims = 1)
        @test cor(X, f(ones(8)), 2) ≈ cor(X, dims = 2)

        cov1 = cov(X, wv1, 1; corrected=false)
        std1 = std(X, wv1, 1; corrected=false)
        cov2 = cov(X, wv2, 2; corrected=false)
        std2 = std(X, wv2, 2; corrected=false)
        expected_cor1 = StatsBase.cov2cor!(cov1, std1)
        expected_cor2 = StatsBase.cov2cor!(cov2, std2)

        @test cor(X, wv1, 1) ≈ expected_cor1
        @test cor(X, wv2, 2) ≈ expected_cor2
    end
end
end # @testset "StatsBase.Covariance"


using BenchmarkTools
using NiBundleAdjustment
using NiBundleAdjustment: vec2cam, vec2scam, compute_reproj_err
using StaticArrays

function load(b, n, m, p)
    dir_in = joinpath(dirname(dirname(@__FILE__)), "data", "ba")
    fn = "ba$(b)_n$(n)_m$(m)_p$(p)"
    fn_in = joinpath(dir_in, fn)
    NiBundleAdjustment.read_ba_instance(string(fn_in,".txt"))
end

cams, X, w, obs, feats = load(4, 372, 47423, 204472)
CAMS = [vec2cam(cams[:,i]) for i = 1:size(cams,2)]
XX = [P3(X[:,i]...) for i=1:size(X, 2)]
FEATS = [P2(feats[:,i]...) for i=1:size(feats, 2)]
println("Normal Objective")
SCAMS = [NiBundleAdjustment.vec2scam(cams[:,i]) for i = 1:size(cams,2)]
display(@benchmark NiBundleAdjustment.ba_objective($SCAMS, $X, $w, $obs, $feats))

println()
println("Reversible Objective")
reproj_err! = zeros(P2{Float64}, size(feats, 2))
w_err! = zero(w)
reproj_err_cache! = zeros(P2{Float64}, size(feats, 2))
display(@benchmark NiBundleAdjustment.ba_objective!($reproj_err!, $w_err!, $reproj_err_cache!, $CAMS, $XX, $w, $obs, $FEATS))

println("ForwardDiff Gradient")
display(@benchmark compute_ba_J(Val(:ForwardDiff), $cams, $X, $w, $obs, $feats))

println()
println("NiLang Gradient")
display(@benchmark compute_ba_J(Val(:NiLang), $CAMS, $XX, $w, $obs, $FEATS))

using KernelAbstractions
using CuArrays
println()
println("NiLang-CUDA Gradient")
display(@benchmark compute_ba_J_multithread(Val(:NiLang), $CAMS, $XX, $w, $obs, $FEATS))
display(@benchmark compute_ba_J_cuda(Val(:NiLang), $(CuArray(CAMS)), $(CuArray(XX)), $(CuArray(w)), $(CuArray(obs)), $(CuArray(FEATS))))
println()


"""
   ograv_pk(pk, z, Ωm; kmin = 5e-4, kmax = 3d1)

Comoving density parameter of gravitational binding energy of large-scale structure of the Universe.

*Reference*: Equation (60) of Fukugita & Peebles, ApJ, 616, 643 (2004)
- With the redshift evolution factor of Equation (14) of Chiang, Makiya, Komatsu & Ménard, arXiv:2007.01679

# Arguments
- `pk`(k): a function which returns a matter power spectrum with the argument k being the comoving wavenumber.
    - **pk times k^3 must be dimensionless**. For example, if k is in units of h/Mpc, `pk` must be in units of Mpc^3/h^3.
- `z::Real`: redshift.
- `Ωm::Real`: present-day matter density parameter.

# Optional keyword arguments
- `kmin::Real=5e-4`: minimum wavenumber for integration, ``∫_{kmin}^{kmax} dk P(k)``.
- `kmax::Real=3e1`: maximum wavenumber for integration, ``∫_{kmin}^{kmax} dk P(k)``.
"""
function ograv_pk(pk, z::Real, Ωm::Real; kmin = 5e-4, kmax = 3e1)
   res, err = hquadrature(pk, kmin, kmax)
   halfW = -3 * Ωm * (1 + z) / 2998^2 / 16 / π^2 * res
   Ωgrav = halfW * Ωm
end

"""
   ograv_halo(pk, z, Ωm[, Ωcb=Ωm]; Mmin = 1e11, Mmax = 5e15, virial = false, t10MF = false)

Comoving density parameter of gravitational binding energy of gravitationally collapsed structures (halos).

*Reference*: Equation (17) of Chiang, Makiya, Komatsu & Ménard, arXiv:2007.01679

# Arguments
- `pk`(k): a function which returns a matter power spectrum with the argument k being the comoving wavenumber.
   - `pk` is in units of (Mpc/h)^3 and the argument k is in units of h/Mpc.
- `z::Real`: redshift.
- `Ωm::Real`: present-day total matter density parameter.

# Optional Arguments
- `Ωcb::Real=Ωm`: present-day baryon + cold dark matter density parameter.

# Optional keyword arguments
- `Mmin::Real=1e11`: minimum mass for integration, ``∫_{Mmin}^{Mmax} dM dn/dM Ag GM^2/R``.
- `Mmax::Real=5e15`: maximum mass for integration, ``∫_{Mmin}^{Mmax} dM dn/dM Ag GM^2/R``.
- `virial::Bool=false`: if `true`, use the virial overdensity `Δvir`. If `false` (the default), use `Δm=200`.
- `t10MF::Bool=false`: if `true`, use `tinker10MF` for the halo multiplicity function. If `false` (the default), use `tinker08MF`.
"""
function ograv_halo(
   pk,
   z::Real,
   Ωm::Real,
   Ωcb = Ωm;
   Mmin = 1e11,
   Mmax = 5e15,
   virial = false,
   t10MF = false,
)
   ρc = 2.775e11 # in units of h^2 M⊙/Mpc^3
   GN = 3 / 8π / 2998^2 / ρc # = 4.786e-20 Mpc/M⊙ is Newton's constant
   if virial # Determine Δm used by the mass function
      E2 = Ωm * (1 + z)^3 + 1 - Ωm # flat Universe
      Ωmz = Ωm * (1 + z)^3 / E2
      Δvir = 18 * π^2 + 82 * (Ωmz - 1) - 39 * (Ωmz - 1)^2
      Δm = Δvir / Ωmz
   else
      Δm = 200
   end
   nmass = 100
   func = zeros(nmass)
   lnMh = range(log(Mmin), length = nmass, log(Mmax))
   for i = 1:nmass
      func[i] = dndlnMh(pk, z, Δm, Ωcb, lnMh[i], t10MF = t10MF)
   end
   spl = Spline1D(lnMh, func)
   # %% Compute ρgrav = (1/2) * \int dlnM dn/dlnM Ag*GM^2/R, in units of h^2 M⊙/Mpc^3
   # c(M) relation is due to Equation (2) and "NFW, F, 0-2" of Table 1 in
   # Duffy et al., MNRAS, 390, L64 (2008)
   function dρdlnMh(lnMh)
      if virial # For Δ = Δvir(z)
         RΔh = cbrt(exp(lnMh) * 3 / 4π / (ρc * E2 * Δvir)) # in h^-1 Mpc
         A0, B0, C0 = 7.85, -0.081, -0.71 # Avir, Bvir, Cvir in Table 1 of Duffy et al.
      else # For Δ = Δm = 200
         RΔh = cbrt(exp(lnMh) * 3 / 4π / (ρc * Ωcb * Δm)) / (1 + z) # in h^-1 Mpc
         A0, B0, C0 = 10.14, -0.081, -1.01 # Amean, Bmean, Cmean in Table 1 of Duffy et al.
      end
      c = A0 * (exp(lnMh) / 2e12)^B0 * (1 + z)^C0
      dρdlnMh = 0.5 * spl(lnMh) * GN * exp(2 * lnMh) / RΔh * Ag(c)
   end
   ρgrav, err = hquadrature(dρdlnMh, log(Mmin), log(Mmax))
   redshift = z
   if t10MF
         fnu(x) = tinker10MF(x, redshift, Δm)
         factor, err = hquadrature(fnu, -50, 5)
         ρgrav /= factor
   end
   halfW = -ρgrav / ρc / Ωcb
   Ωgrav = halfW * Ωcb
end

"""
   Ag(c:Real; method = "poly4")

Integral of the squared Fourier transform of an NFW density profile, ``(1/π)∫_0^∞ dx |u(x,c)|^2``.

*Reference*: Equation (18) of Chiang, Makiya, Komatsu & Ménard, arXiv:2007.01679

# Arguments
- `c::Real`: concentration parameter.

# Optional keyword arguments
- `method::String="poly4"`: approximation method.
   - `poly4` (the default) or other values for 4th order polynomial fit.
   - `poly3` for 3rd order polynomial fit.
   - `exact` for numerical integration.
"""
function Ag(c::Real; method = "poly4")
   if method == "poly3" # 3rd order polynomial fit
      y = c - 5
      Ag = 1.0184 + 5.389e-2 * y - 1.514e-3 * y^2 + 3.845e-5 * y^3
   elseif method == "exact" # exact integration
      f = log(1 + c) - c / (1 + c)
      u2(x) =
         (
            sin(x / c) * (sinint((1 / c + 1) * x) - sinint(x / c)) +
            cos(x / c) * (cosint((1 / c + 1) * x) - cosint(x / c)) -
            sin(x) / (1 / c + 1) / x
         )^2
      res, err = quadde(u2, 0, 100)
      Ag = res / π / f^2
   else # 4th order polynomial fit
      y = c - 5
      Ag =
         1.0202 + 5.376e-2 * y - 1.842e-3 * y^2 + 9.489e-5 * y^3 -
         2.352e-6 * y^4
   end
end


using Ledger
using Test

@testset "Ledger.jl" begin

    # Extending the interface for my own types/valuation basis

    struct MyValuationBasis <: Ledger.ValuationBasis end

    struct MyBond <: Ledger.Asset
        par
    end

    Ledger.value(::MyValuationBasis,b::MyBond) = b.par

    gl = GeneralLedger()

    addAsset!(gl,MyBond(100))

    bal = Balance(MyValuationBasis(),gl)

    @test sum(bal.assets) == 100.0
    @test length(bal.liabilities) == 0
    @test sum(bal.equity) == 100.0

    struct myDebt <: Ledger.Liability
        balance
        int_rate
        years_due
    end

    Ledger.value(::MyValuationBasis,b::myDebt) = b.balance / (1+ b.int_rate) ^ b.years_due

    addLiability!(gl,myDebt(50,0.05,3))

    bal = Balance(MyValuationBasis(),gl)

    @test sum(bal.assets) == 100.0
    @test sum(bal.liabilities) == 50 / 1.05 ^ 3
    @test sum(bal.equity) == 100.0 - 50 / 1.05 ^ 3

end


# 225x225 complex dense
# Compare to S&P, Fig. 60.2

using Pseudospectra, Test

@testset "Complex dense" begin
    opts = Dict{Symbol,Any}(:ax => [-1.1,1.2,-1.1,1.1], :npts => 60,
                            :levels => -5:0.5:-1)
    A = Pseudospectra.landau_fox_li(225,16)

    ps_data = new_matrix(A,opts)
    driver!(ps_data,opts,gs)
    @test iscomputed(ps_data)
    modeplot(ps_data,0,0.0+1.0im)
    modeplot(ps_data,1,0.0+1.0im)

end


using AugmentedGaussianProcesses
using Distributions, LinearAlgebra
using Random: seed!;
seed!(123);
using Plots;
gr();

N = 100;
noise = 1e-1;
N_test = 1000
X = sort(rand(N))
X_test = range(0, 1; length=N_test)
k = SqExponentialKernel() ∘ ScaleTransform(10.0)
K = kernelmatrix(k, X) + noise * I
f = rand(MvNormal(K))
y = sign.(f)

logitmodel = VGP(X, y, k, LogisticLikelihood(), AnalyticVI());
train!(logitmodel, 300)
logitpred = (proba_y(logitmodel, X_test)[1] .- 0.5) .* 2

svmmodel = VGP(X, y, k, BayesianSVM(), AnalyticVI());
train!(svmmodel, 50)#,callback=intplot)
svmpred = (proba_y(svmmodel, X_test)[1] .- 0.5) .* 2

##
default(; legendfontsize=14.0, xtickfontsize=10.0, ytickfontsize=10.0)
p = plot(X, y; t=:scatter, lab="Training Points", ylims=(-1.5, 3.0), legend=:topright)
plot!(X_test, logitpred; lab="Logistic Prediction", lw=7.0)
plot!(X_test, svmpred; lab="BayesianSVM Prediction", lw=7.0)

display(p)
savefig(p, joinpath(@__DIR__, "Classification.png"))


# building fixed rate bonds
using QuantLib
using Dates
using Printf

function build_bonds(bond_mats::Vector{Date}, bond_rates::Vector{Float64}, tenor::QuantLib.Time.TenorPeriod, conv::QuantLib.Time.BusinessDayConvention,
                    rule::QuantLib.Time.DateGenerationRule, calendar::QuantLib.Time.BusinessCalendar, dc::QuantLib.Time.DayCount, freq::QuantLib.Time.Frequency, issue_date::Date)
  bonds = Vector{FixedRateBondHelper}(undef, length(bond_mats))
  pricing_engine = DiscountingBondEngine()

  for i =1:length(bond_mats)
    term_date = bond_mats[i]
    # rate = bond_rates[i] / 100
    rate = bond_rates[i]
    sched = QuantLib.Time.Schedule(issue_date, term_date, tenor, conv, conv, rule, true)
    bond_help = FixedRateBondHelper(Quote(100.0), FixedRateBond(0, 100.0, sched, rate, dc, conv, 100.0, issue_date, calendar, pricing_engine))
    bonds[i] = bond_help
  end

  return bonds
end

function build_depos(depo_quotes::Vector{Float64}, depo_tenors::Vector{P}, dc::QuantLib.Time.DayCount, conv::QuantLib.Time.BusinessDayConvention, calendar::QuantLib.Time.BusinessCalendar, fixing_days::Int) where {P <: Dates.Period}
  depos = Vector{DepositRateHelper}(undef, length(depo_quotes))
  for i = 1:length(depo_quotes)
    depo_quote = Quote(depo_quotes[i])
    depo_tenor = QuantLib.Time.TenorPeriod(depo_tenors[i])
    depo = DepositRateHelper(depo_quote, depo_tenor, fixing_days, calendar, conv, true, dc)
    depos[i] = depo
  end

  return depos
end

function build_swaps(swap_quotes::Vector{Float64}, swap_tenors::Vector{P}, fixed_dc::QuantLib.Time.DayCount, fixed_conv::QuantLib.Time.BusinessDayConvention, calendar::QuantLib.Time.BusinessCalendar, fixed_freq::QuantLib.Time.Frequency, float_index::IborIndex, forward_start::Int) where {P <: Dates.Period}
  forward_start_period = Dates.Day(forward_start)
  swaps = Vector{SwapRateHelper}(undef, length(swap_quotes))
  pricing_engine = DiscountingSwapEngine()

  for i = 1:length(swap_quotes)
    swaps[i] = SwapRateHelper(swap_quotes[i], swap_tenors[i], calendar, fixed_freq, fixed_conv, fixed_dc, float_index, 0.0, forward_start_period, pricing_engine)
  end

  return swaps
end

function get_npvs(bonds, issue_date, calendar, dc, freq)
  pricing_engine = DiscountingBondEngine()
  rate_quote = Quote(0.05)
  compounding = CompoundedCompounding()
  yts = FlatForwardTermStructure(0, issue_date, calendar, rate_quote, dc, compounding, freq)
  npvs = zeros(length(bonds))
  for i=1:length(bonds)
    npvs[i] = calculate(pricing_engine, yts, bonds[i])
  end

  return npvs
end

function par_rate(yts::YieldTermStructure, dates::Vector{Date}, dc::QuantLib.Time.DayCount)
  sum = 0.0
  for i = 2:length(dates)
    dt = QuantLib.Time.year_fraction(dc, dates[i - 1], dates[i])
    sum += discount(yts, dates[i]) * dt
  end

  result = discount(yts, dates[1]) - discount(yts, dates[end])
  return result/sum
end

function setup()
  today = now()
  issue_date = Date(Dates.Year(today), Dates.Month(today), Dates.Day(today))
  bond_mats = [issue_date + Dates.Year(i) for i in range(2, step=2, length=15)]
  # bond_rates = [5.75, 6.0, 6.25, 6.5, 6.75, 6.80, 7.00, 7.1, 7.15, 7.2]
  bond_rates = [0.0200, 0.0225, 0.0250, 0.0275, 0.0300, 0.0325, 0.0350, 0.0375, 0.0400, 0.0425, 0.0450, 0.0475, 0.0500, 0.0525, 0.0550]
  set_eval_date!(settings, issue_date)

  freq = QuantLib.Time.Annual()
  tenor = QuantLib.Time.TenorPeriod(freq)
  conv = QuantLib.Time.Unadjusted()
  rule = QuantLib.Time.DateGenerationBackwards()
  calendar = QuantLib.Time.USGovernmentBondCalendar()
  dc = QuantLib.Time.SimpleDayCount()
  bonds = build_bonds(bond_mats, bond_rates, tenor, conv, rule, calendar, dc, freq, issue_date)

  return issue_date, bonds, dc, calendar
end

function piecewise_yld_curve()
  issue_date, bonds, dc, calendar = setup()

  interp = QuantLib.Math.LogLinear()
  trait = Discount()
  bootstrap = IterativeBootstrap()

  yts = PiecewiseYieldCurve(issue_date, bonds, dc, interp, trait, 0.00000000001, bootstrap)

  solver = QuantLib.Math.BrentSolver()
  solver2 = QuantLib.Math.FiniteDifferenceNewtonSafe()
  calculate!(IterativeBootstrap(), yts, solver2, solver)

  # println(yts.data)

  for bond in bonds
    date_vec = Vector{Date}(undef, length(bond.cashflows.coupons) + 1)
    date_vec[1] = issue_date
    for (i, cf) in enumerate(bond.cashflows.coupons)
      date_vec[i+1] = date(cf)
    end
    par = par_rate(yts, date_vec, dc)
    println(100.0 * par)
  end
end

  # npvs = get_npvs(bonds, issue_date, calendar, dc, freq)
function fitted_bond_curve_exp()
  issue_date, bonds, dc, calendar = setup()

  esf = ExponentialSplinesFitting(true, length(bonds))
  fitted_curve = FittedBondDiscountCurve(0, issue_date, calendar, bonds, dc, esf, 1e-10, 5000, 1.0)
  initialize!(fitted_curve)
  calculate!(fitted_curve)

  println(fitted_curve.fittingMethod.commons.minimumCostValue)
  println(fitted_curve.fittingMethod.commons.numberOfIterations)
  println(fitted_curve.fittingMethod.commons.guessSolution)

  disc = discount(fitted_curve, 1.0)
  println(disc)
end

function fitted_bond_curve_simp()
  issue_date, bonds, dc, calendar = setup()

  spf = SimplePolynomialFitting(true, 3, length(bonds))
  fitted_curve = FittedBondDiscountCurve(0, issue_date, calendar, bonds, dc, spf, 1e-10, 5000, 1.0)
  initialize!(fitted_curve)
  calculate!(fitted_curve)

  # println(fitted_curve.fittingMethod.minimumCostValue)
  println(fitted_curve.fittingMethod.commons.numberOfIterations)
  # println(fitted_curve.fittingMethod.guessSolution)

  disc = discount(fitted_curve, 1.0)
  println(disc)
end

function fitted_bond_curve_ns()
  issue_date, bonds, dc, calendar = setup()

  nsf = NelsonSiegelFitting(length(bonds))
  fitted_curve = FittedBondDiscountCurve(0, issue_date, calendar, bonds, dc, nsf, 1e-10, 5000, 1.0)
  initialize!(fitted_curve)
  calculate!(fitted_curve)

  # println(fitted_curve.fittingMethod.minimumCostValue)
  println(fitted_curve.fittingMethod.commons.numberOfIterations)
  # println(fitted_curve.fittingMethod.guessSolution)

  disc = discount(fitted_curve, 1.0)
  println(disc)
end

function fitted_bond_curve_sven()
  issue_date, bonds, dc, calendar = setup()

  sf = SvenssonFitting(length(bonds))
  fitted_curve = FittedBondDiscountCurve(0, issue_date, calendar, bonds, dc, sf, 1e-10, 5000, 1.0)
  initialize!(fitted_curve)
  calculate!(fitted_curve)

  # println(fitted_curve.fittingMethod.minimumCostValue)
  println(fitted_curve.fittingMethod.commons.numberOfIterations)
  # println(fitted_curve.fittingMethod.guessSolution)

  disc = discount(fitted_curve, 1.0)
  println(disc)
end

function fitted_bond_curve_cbspline()
  issue_date, bonds, dc, calendar = setup()

  knots = [-30.0, -20.0, 0.0, 5.0, 10.0, 15.0, 20.0, 25.0, 30.0, 40.0, 50.0]
  cbsf = CubicBSplinesFitting(true, knots, length(bonds))
  fitted_curve = FittedBondDiscountCurve(0, issue_date, calendar, bonds, dc, cbsf, 1e-10, 5000, 1.0)
  initialize!(fitted_curve)
  calculate!(fitted_curve)

  println(fitted_curve.fittingMethod.commons.numberOfIterations)
  disc = discount(fitted_curve, 1.0)
  println(disc)
end

function fitted_bond_curve_all()
  tic()
  issue_date, bonds, dc, calendar = setup()

  esf = ExponentialSplinesFitting(true, length(bonds))
  esf_fitted_curve = FittedBondDiscountCurve(0, issue_date, calendar, bonds, dc, esf, 1e-10, 5000, 1.0)
  initialize!(esf_fitted_curve)
  calculate!(esf_fitted_curve)

  spf = SimplePolynomialFitting(true, 3, length(bonds))
  spf_fitted_curve = FittedBondDiscountCurve(0, issue_date, calendar, bonds, dc, spf, 1e-10, 5000, 1.0)
  initialize!(spf_fitted_curve)
  calculate!(spf_fitted_curve)

  nsf = NelsonSiegelFitting(length(bonds))
  nsf_fitted_curve = FittedBondDiscountCurve(0, issue_date, calendar, bonds, dc, nsf, 1e-10, 5000, 1.0)
  initialize!(nsf_fitted_curve)
  calculate!(nsf_fitted_curve)

  knots = [-30.0, -20.0, 0.0, 5.0, 10.0, 15.0, 20.0, 25.0, 30.0, 40.0, 50.0]
  cbsf = CubicBSplinesFitting(true, knots, length(bonds))
  cbsf_fitted_curve = FittedBondDiscountCurve(0, issue_date, calendar, bonds, dc, cbsf, 1e-10, 5000, 1.0)
  initialize!(cbsf_fitted_curve)
  calculate!(cbsf_fitted_curve)

  sf = SvenssonFitting(length(bonds))
  sf_fitted_curve = FittedBondDiscountCurve(0, issue_date, calendar, bonds, dc, sf, 1e-10, 5000, 1.0)
  initialize!(sf_fitted_curve)
  calculate!(sf_fitted_curve)
  toc()

  println("Exponential Splines: $(esf_fitted_curve.fittingMethod.commons.numberOfIterations)")
  println("Simple Polynomial: $(spf_fitted_curve.fittingMethod.commons.numberOfIterations)")
  println("Nelson Siegel: $(nsf_fitted_curve.fittingMethod.commons.numberOfIterations)")
  println("Cubic B-Splines: $(cbsf_fitted_curve.fittingMethod.commons.numberOfIterations)")
  println("Svensson Fitting: $(sf_fitted_curve.fittingMethod.commons.numberOfIterations)")
end

function generate_floatingrate_bond(yts::YieldTermStructure, disc_yts::YieldTermStructure, cap_vol::OptionletVolatilityStructure)
  # Floating Rate bond
  settlement_days = 3
  face_amount = 100.0
  fb_issue_date = Date(2005, 10, 21)
  bond_engine = DiscountingBondEngine(disc_yts)
  fb_dc = QuantLib.Time.Actual360()
  conv = QuantLib.Time.ModifiedFollowing()
  fb_schedule = QuantLib.Time.Schedule(Date(2005, 10, 21), Date(2010, 10, 21), QuantLib.Time.TenorPeriod(QuantLib.Time.Quarterly()),
                                        QuantLib.Time.Unadjusted(), QuantLib.Time.Unadjusted(), QuantLib.Time.DateGenerationBackwards(), false,
                                        QuantLib.Time.USNYSECalendar())
  fixing_days = 2
  in_arrears = true
  gearings = ones(length(fb_schedule.dates) - 1)
  spreads = fill(0.001, length(fb_schedule.dates) - 1)
  libor_3m = usd_libor_index(QuantLib.Time.TenorPeriod(Dates.Month(3)), yts)
  floating_bond = FloatingRateBond(settlement_days, face_amount, fb_schedule, libor_3m, fb_dc, conv, fixing_days, fb_issue_date, bond_engine, in_arrears, 100.0, gearings, spreads, cap_vol=cap_vol)

  return floating_bond
end

function generate_fixedrate_bond(yts::YieldTermStructure)
  settlement_days = 3
  face_amount = 100.0

  fx_schedule = QuantLib.Time.Schedule(Date(2007, 5, 15), Date(2017, 5, 15), QuantLib.Time.TenorPeriod(QuantLib.Time.Semiannual()),
                                        QuantLib.Time.Unadjusted(), QuantLib.Time.Unadjusted(), QuantLib.Time.DateGenerationBackwards(), false,
                                        QuantLib.Time.USGovernmentBondCalendar())

  pe = DiscountingBondEngine(yts)

  fixedrate_bond = FixedRateBond(settlement_days, face_amount, fx_schedule, 0.045, QuantLib.Time.ISMAActualActual(), QuantLib.Time.ModifiedFollowing(),
                                100.0, Date(2007, 5, 15), fx_schedule.cal, pe)

  return fixedrate_bond
end

function generate_discounting_ts(sett::Date)
  settlement_date = sett
  freq = QuantLib.Time.Semiannual()
  tenor = QuantLib.Time.TenorPeriod(freq)
  conv = QuantLib.Time.Unadjusted()
  conv_depo = QuantLib.Time.ModifiedFollowing()
  rule = QuantLib.Time.DateGenerationBackwards()
  calendar = QuantLib.Time.USGovernmentBondCalendar()
  dc_depo = QuantLib.Time.Actual365()
  dc = QuantLib.Time.ISDAActualActual()
  dc_bond = QuantLib.Time.ISMAActualActual()
  fixing_days = 3

  # build depos
  depo_rates = [0.0096, 0.0145, 0.0194]
  depo_tens = [Dates.Month(3), Dates.Month(6), Dates.Month(12)]

  # build bonds
  issue_dates = [Date(2005, 3, 15), Date(2005, 6, 15), Date(2006, 6, 30), Date(2002, 11, 15), Date(1987, 5, 15)]
  mat_dates = [Date(2010, 8, 31), Date(2011, 8, 31), Date(2013, 8, 31), Date(2018, 8, 15), Date(2038, 5, 15)]

  coupon_rates = [0.02375, 0.04625, 0.03125, 0.04000, 0.04500]
  market_quotes = [100.390625, 106.21875, 100.59375, 101.6875, 102.140625]

  insts = Vector{BootstrapHelper}(undef, length(depo_rates) + length(issue_dates))
  for i = 1:length(depo_rates)
    depo_quote = Quote(depo_rates[i])
    depo_tenor = QuantLib.Time.TenorPeriod(depo_tens[i])
    depo = DepositRateHelper(depo_quote, depo_tenor, fixing_days, calendar, conv_depo, true, dc_depo)
    insts[i] = depo
  end

  pricing_engine = DiscountingBondEngine()

  for i =1:length(coupon_rates)
    term_date = mat_dates[i]
    # rate = bond_rates[i] / 100
    rate = coupon_rates[i]
    issue_date = issue_dates[i]
    market_quote = market_quotes[i]
    sched = QuantLib.Time.Schedule(issue_date, term_date, tenor, conv, conv, rule, true)
    bond = FixedRateBondHelper(Quote(market_quote), FixedRateBond(3, 100.0, sched, rate, dc_bond, conv, 100.0, issue_date, calendar, pricing_engine))
    insts[i + length(depo_rates)] = bond
  end

  interp = QuantLib.Math.LogLinear()
  trait = Discount()
  bootstrap = IterativeBootstrap()

  yts = PiecewiseYieldCurve(settlement_date, insts, dc, interp, trait, 0.00000000001, bootstrap)

  # solver = QuantLib.Math.BrentSolver()
  # solver2 = QuantLib.Math.FiniteDifferenceNewtonSafe()
  calculate!(yts)

  return yts
end

function main()
  issue_date, bonds, dc, calendar = setup()

  println("Today's date: $issue_date")
  println("Calculating fit for 15 bonds....")

  interp = QuantLib.Math.LogLinear()
  trait = Discount()
  bootstrap = IterativeBootstrap()

  yts = PiecewiseYieldCurve(issue_date, bonds, dc, interp, trait, 0.00000000001, bootstrap)

  # solver = QuantLib.Math.BrentSolver()
  # solver2 = QuantLib.Math.FiniteDifferenceNewtonSafe()
  calculate!(yts)

  esf = ExponentialSplinesFitting(true, length(bonds))
  esf_fitted_curve = FittedBondDiscountCurve(0, issue_date, calendar, bonds, dc, esf, 1e-10, 5000, 1.0)
  # @code_warntype initialize!(esf_fitted_curve)
  initialize!(esf_fitted_curve)
  calculate!(esf_fitted_curve)

  println("(a) exponential splines")
  println("reference date : ", esf_fitted_curve.referenceDate)
  println("number of iterations : ", esf_fitted_curve.fittingMethod.commons.numberOfIterations)

  spf = SimplePolynomialFitting(true, 3, length(bonds))
  spf_fitted_curve = FittedBondDiscountCurve(0, issue_date, calendar, bonds, dc, spf, 1e-10, 5000, 1.0)
  initialize!(spf_fitted_curve)
  calculate!(spf_fitted_curve)

  println("(b) simple polynomial")
  println("reference date : ", spf_fitted_curve.referenceDate)
  println("number of iterations : ", spf_fitted_curve.fittingMethod.commons.numberOfIterations)

  nsf = NelsonSiegelFitting(length(bonds))
  nsf_fitted_curve = FittedBondDiscountCurve(0, issue_date, calendar, bonds, dc, nsf, 1e-10, 5000, 1.0)
  initialize!(nsf_fitted_curve)
  calculate!(nsf_fitted_curve)

  println("(c) Nelson-Siegel")
  println("reference date : ", nsf_fitted_curve.referenceDate)
  println("number of iterations : ", nsf_fitted_curve.fittingMethod.commons.numberOfIterations)

  knots = [-30.0, -20.0, 0.0, 5.0, 10.0, 15.0, 20.0, 25.0, 30.0, 40.0, 50.0]
  cbsf = CubicBSplinesFitting(true, knots, length(bonds))
  cbsf_fitted_curve = FittedBondDiscountCurve(0, issue_date, calendar, bonds, dc, cbsf, 1e-10, 5000, 1.0)
  initialize!(cbsf_fitted_curve)
  calculate!(cbsf_fitted_curve)

  println("(d) cubic B-splines")
  println("reference date : ", cbsf_fitted_curve.referenceDate)
  println("number of iterations : ", cbsf_fitted_curve.fittingMethod.commons.numberOfIterations)

  sf = SvenssonFitting(length(bonds))
  sf_fitted_curve = FittedBondDiscountCurve(0, issue_date, calendar, bonds, dc, sf, 1e-10, 5000, 1.0)
  initialize!(sf_fitted_curve)
  calculate!(sf_fitted_curve)

  println("(e) Svensson")
  println("reference date : ", sf_fitted_curve.referenceDate)
  println("number of iterations : ", sf_fitted_curve.fittingMethod.commons.numberOfIterations)

  println("Output par rates for each curve.  In this case, par rates should equal coupons for these par bonds")

  println(" tenor | coupon | bstrap |  (a)  |  (b)  |  (c)  |  (d)  |  (e)  ")

  for bh in bonds
    bond = bh.bond
    date_vec = Vector{Date}(undef, length(bond.cashflows.coupons) + 1)
    date_vec[1] = issue_date
    for (i, cf) in enumerate(bond.cashflows.coupons)
      date_vec[i+1] = date(cf)
    end

    if bh.bond.cashflows.redemption != nothing
      push!(date_vec, date(bh.bond.cashflows.redemption))
    end
    tenor = QuantLib.Time.year_fraction(dc, issue_date, date(bond.cashflows.coupons[end]))
    println(@sprintf(" %.2f  | %.3f | %.3f | %.3f | %.3f | %.3f | %.3f | %.3f ",
            tenor,
            100.0 * bond.cashflows.coupons[end].rate.rate,
            100.0 * par_rate(yts, date_vec, dc),
            100.0 * par_rate(esf_fitted_curve, date_vec, dc),
            100.0 * par_rate(spf_fitted_curve, date_vec, dc),
            100.0 * par_rate(nsf_fitted_curve, date_vec, dc),
            100.0 * par_rate(cbsf_fitted_curve, date_vec, dc),
            100.0 * par_rate(sf_fitted_curve, date_vec, dc)
            )
      )
  end
end

# function main2()
#   settlement_date = Date(2008, 9, 18)
#   calendar = QuantLib.Time.USGovernmentBondCalendar()
#   set_eval_date!(settings, settlement_date - Base.Dates.Day(3))
#
#   yts = generate_discounting_ts(settlement_date)
#   pe = DiscountingBondEngine(yts)
#
#   # build zero coupon bond
#   zcb = ZeroCouponBond(3, calendar, 100.0, Date(2013, 8, 15), QuantLib.Time.Following(), 116.92, Date(2003, 8, 15), pe)
#
#   # println(npv(zcb, pricing_engine, yts))
#   # println(clean_price(zcb))
#   # println(dirty_price(zcb))
#   return npv(zcb, zcb.pricingEngine), clean_price(zcb), dirty_price(zcb)
# end

function main3()
  settlement_date = Date(2008, 9, 18)
  set_eval_date!(settings, settlement_date - Dates.Day(3))
  cal = QuantLib.Time.TargetCalendar()
  dc = QuantLib.Time.ISDAActualActual()

  # Build deposits
  depo_quotes = [0.043375, 0.031875, 0.0320375, 0.03385, 0.0338125, 0.0335125]
  depo_tenors = [Dates.Week(1), Dates.Month(1), Dates.Month(3), Dates.Month(6), Dates.Month(9), Dates.Year(1)]
  deposit_dc = QuantLib.Time.Actual360()
  fixing_days = 3
  biz_conv = QuantLib.Time.ModifiedFollowing()

  depos = build_depos(depo_quotes, depo_tenors, deposit_dc, biz_conv, cal, fixing_days)

  # build swaps
  fixedLegFreq = QuantLib.Time.Annual()
  fixedLegConv = QuantLib.Time.Unadjusted()
  fixedLegDC = QuantLib.Time.EuroThirty360()
  floatingLegIndex = euribor_index(QuantLib.Time.TenorPeriod(Dates.Month(6)))
  forwardStart = 1

  swap_quotes = [0.0295, 0.0323, 0.0359, 0.0412, 0.0433]
  swap_tenors = [Dates.Year(2), Dates.Year(3), Dates.Year(5), Dates.Year(10), Dates.Year(15)]

  swaps = build_swaps(swap_quotes, swap_tenors, fixedLegDC, fixedLegConv, cal, fixedLegFreq, floatingLegIndex, forwardStart)

  insts = Vector{BootstrapHelper}(undef, length(swap_quotes) + length(depo_quotes))

  insts[1:length(depo_quotes)] = depos
  insts[length(depo_quotes) + 1: end] = swaps

  interp = QuantLib.Math.LogLinear()
  trait = Discount()
  bootstrap = IterativeBootstrap()

  yts = PiecewiseYieldCurve(settlement_date, insts, dc, interp, trait, 1e-15, bootstrap)

  # solver = QuantLib.Math.BrentSolver()
  # solver2 = QuantLib.Math.FiniteDifferenceNewtonSafe()
  calculate!(yts)

  disc_yts = generate_discounting_ts(settlement_date)

  cap_vol = ConstantOptionVolatility(3, cal, QuantLib.Time.ModifiedFollowing(), 0.0, QuantLib.Time.Actual365())
  # Floating Rate Bond
  fb = generate_floatingrate_bond(yts, disc_yts, cap_vol)

  # fb.iborIndex.ts = yts
  # fb.pricingEngine.yts = disc_yts

  # Zero Coupon Bond
  zcb_pe = DiscountingBondEngine(disc_yts)
  zcb_cal = QuantLib.Time.USGovernmentBondCalendar()

  # build zero coupon bond
  zcb = ZeroCouponBond(3, zcb_cal, 100.0, Date(2013, 8, 15), QuantLib.Time.Following(), 116.92, Date(2003, 8, 15), zcb_pe)

  # fixed rate bond
  fxb = generate_fixedrate_bond(disc_yts)
  # fxb.pricingEngine.yts = disc_yts

  println("Today's date: ", settings.evaluation_date)
  println("Settlement date: ", settlement_date)
  println("")
  println("              ZC      Fixed    Floating  ")
  println("-----------------------------------------")
  println(@sprintf("  NPV       %.2f   %.2f   %.2f", npv(zcb), npv(fxb), npv(fb)))
  println(@sprintf(" Clean      %.2f   %.2f   %.2f", clean_price(zcb), clean_price(fxb), clean_price(fb)))
  println(@sprintf(" Dirty      %.2f   %.2f   %.2f", dirty_price(zcb), dirty_price(fxb), dirty_price(fb)))
  println(@sprintf("accrued      %.2f     %.2f     %.2f", accrued_amount(zcb, settlement_date), accrued_amount(fxb, settlement_date), accrued_amount(fb, settlement_date)))
  println(@sprintf(" Next C      N/A      %.2f%%    %.2f%%", next_coupon_rate(fxb.cashflows, settlement_date) * 100.0, next_coupon_rate(fb.cashflows, settlement_date) * 100.0))
  println(@sprintf(" Yield      %.2f%%    %.2f%%    %.2f%%",
          QuantLib.yield(zcb, clean_price(zcb), QuantLib.Time.Actual360(), CompoundedCompounding(), QuantLib.Time.Annual(), settlement_date) * 100.0,
          QuantLib.yield(fxb, clean_price(fxb), QuantLib.Time.Actual360(), CompoundedCompounding(), QuantLib.Time.Annual(), settlement_date) * 100.0,
          QuantLib.yield(fb, clean_price(fb), QuantLib.Time.Actual360(), CompoundedCompounding(), QuantLib.Time.Annual(), settlement_date) * 100.0))

  # return npv(fb, fb.pricingEngine), clean_price(fb), dirty_price(fb), QuantLib.yield(fb, clean_price(fb), QuantLib.Time.Actual360(), CompoundedCompounding(), QuantLib.Time.Annual(), settlement_date), next_coupon_rate(fb.cashflows, settlement_date)
end


using Distributions, Parameters
import Base.rand
import Distributions: logpdf

struct LNR{T1,T2,T3} <: ContinuousUnivariateDistribution
    μ::T1
    σ::T2
    ϕ::T3
end

Broadcast.broadcastable(x::LNR) = Ref(x)

LNR(;μ, σ, ϕ) = LNR(μ, σ, ϕ)

function rand(dist::LNR)
    @unpack μ,σ,ϕ = dist
    x = @. rand(LogNormal(μ, σ)) + ϕ
    rt,resp = findmin(x)
    return resp,rt
end

rand(dist::LNR, N::Int) = [rand(dist) for i in 1:N]

function logpdf(d::T, r::Int, t::Float64) where {T<:LNR}
    @unpack μ,σ,ϕ = d
    LL = 0.0
    for (i,m) in enumerate(μ)
        if i == r
            LL += logpdf(LogNormal(m, σ), t-ϕ)
        else
            LL += log(1-cdf(LogNormal(m, σ), t-ϕ))
        end
    end
    return LL
end

logpdf(d::LNR, data::Tuple) = logpdf(d, data...)

@model lnrModel(data, idx, N, Nsub, ::Type{T}=Vector{Float64},
    ::Type{T1}=Array{Float64,2}) where {T,T1}  = begin
    Nr = 3
    mu_g = T(undef,Nr)
    sigma_g ~ Truncated(Cauchy(0, 1), 0, Inf)
    mu_g ~ [Normal(0,3)]
    mu = T1(undef,Nsub,Nr)
    for s in 1:Nsub
        for i in 1:Nr
            mu[s,i] ~ Normal(mu_g[i], sigma_g)
        end
    end
    sigma ~ Truncated(Cauchy(0, 1), 0, Inf)
    for i in 1:N
        data[i] ~ LNR(mu[idx[i],:], sigma, 0.0)
    end
end

function simulateLNR(; Nd=1, Nsub, mu_g = [-1.0,-2.0,-.5], sigma_g = 1,
     sigma = 1, kwargs...)
    N = Nd * Nsub
    choice = fill(0, N)
    rt = fill(0.0, N)
    idx = similar(choice)
    i = 0
    for s in 1:Nsub
        mu = @. rand(Normal(mu_g, sigma_g))
        dist = LNR(μ=mu, σ=sigma, ϕ=0.0)
        for _ in 1:Nd
            i += 1
            idx[i] = s
            choice[i],rt[i] = rand(dist)
        end
    end
    return (choice=choice, rt=rt, idx=idx, N=N, Nsub=Nsub)
 end

function AHMClnr(choice, rt, idx, N, Nsub)
    data=[(c,r) for (c,r) in zip(choice,rt)]

    return lnrModel(data, idx, N, Nsub)
end

AHMCconfig = Turing.NUTS(1000,.85)

CmdStanLNR = "
functions{

    real LNR_LL(row_vector mus,real sigma,real ter,real v,int c){
      real LL;
      LL = 0;
      for(i in 1:num_elements(mus)){
        if(i == c){
          LL += lognormal_lpdf(v-ter|mus[i],sigma);
      }else{
        LL += log(1-lognormal_cdf(v-ter,mus[i],sigma));
      }
    }
      return LL;
    }

    }

data{
     int N;
     int Nr;
     int Nsub;
     vector[N] rt;
     vector[N] choice;
}

parameters {
     vector[Nr] mu_g;
     matrix[Nsub,Nr] mu;
     real<lower=0> sigma_g;
     real<lower=0> sigma;
}

model {
     mu_g ~ normal(0, 3);
     sigma_g ~ cauchy(0,1);
     for(s in 1:Nsub){
         for(i in 1:Nr){
            mu[s,i] ~ normal(mu_g[i], sigma_g);
        }
     }
     sigma ~ cauchy(0,1);
     for(i in 1:N){
         target += LNR_LL(mu[idx[i],:], sigma, 0.0, rt[i], choice[i]);
      }
}
"

CmdStanConfig = Stanmodel(name="CmdStanLNR", model=CmdStanLNR, nchains=1,
   Sample(num_samples=1000, num_warmup=1000, adapt=CmdStan.Adapt(delta=0.8),
   save_warmup=true))

   struct LNRProb{T}
      data::T
      N::Int
      Nc::Int
  end

  function (problem::LNRProb)(θ)
      @unpack data=problem
      @unpack v,A,k,tau=θ
      d=LBA(ν=v, A=A, k=k, τ=tau)
      minRT = minimum(x->x[2], data)
      logpdf(d,data)+sum(logpdf.(TruncatedNormal(0, 3, 0, Inf), v)) +
      logpdf(TruncatedNormal(.8, .4, 0, Inf),A)+logpdf(TruncatedNormal(.2, .3 ,0 ,Inf), k)+
      logpdf(TruncatedNormal(.4, .1, 0, minRT), tau)
  end

function sampleDHMC(choice, rt, N, Nc, nsamples)
    data = [(c, r) for (c, r) in zip(choice, rt)]
    return sampleDHMC(data, N, Nc, nsamples)
end

# Define problem with data and inits.
function sampleDHMC(data, N, Nc, nsamples, autodiff)
    p = LBAProb(data, N, Nc)
    p((v=fill(.5, Nc), A=.8, k=.2, tau=.4))
    # Write a function to return properly dimensioned transformation.
    trans = as((v=as(Array, asℝ₊, Nc),A=asℝ₊, k=asℝ₊, tau=asℝ₊))
    # Use Flux for the gradient.
    P = TransformedLogDensity(trans, p)
    ∇P = ADgradient(autodiff, P)
    # FSample from the posterior.
    n = dimension(trans)
    results = mcmc_with_warmup(Random.GLOBAL_RNG, ∇P, nsamples;
        q = zeros(n), p = ones(n),reporter = NoProgressReport())
    # Undo the transformation to obtain the posterior from the chain.
    posterior = transform.(trans, results.chain)
    chns = nptochain(results, posterior)
    return chns
end

function simulateLBA(;Nd, v=[1.0,1.5,2.0], A=.8, k=.2, tau=.4, kwargs...)
    return (rand(LBA(ν=v, A=A, k=k, τ=tau), Nd)..., N=Nd, Nc=length(v))
end


using Distributions
using LinearAlgebra
using Random

# Define the DGP
function simulate_sample(N0, N1)
    dy1 = LogNormal(0.3, 0.7);
    dy2 = MixtureModel(LogNormal, [(0.0, 1.0), (0.3, 1.2)], [0.6, 0.4]);
    X0 = [ones(N0) rand([0, 1], N0, 4) rand(N0)]
    X1 = kron([ones(2, 4) [1, 0] [0, 0]], ones(N1))
    y0 = [X0[i, end] == 1 ? rand(dy2) : rand(dy1) for i in 1:N0]
    y1 = LinRange(0, 6, N1) |> collect |> x -> repeat(x, 2)
    c0 = rand(Exponential(12), N0)
    z0 = deepcopy(y0)
    return dy1, dy2, z0, c0, y0, X0, y1, X1
end

# Not run
Random.seed!(1);
N0, N1 = 500, 0;
dy1, dy2, z0, c0, y0, X0, y1, X1 = simulate_sample(N0, N1);
mean(z0 .> c0)


## Show the wikipedia categories graph
using SparseArrays
using LinearAlgebra
using DelimitedFiles
ei,ej = copy.(eachcol(Int.(readdlm("wiki-cats.edges"))))
xy = readdlm("wiki-cats.xy")
n = size(xy,1)
A = sparse(ei,ej,1,n,n) # we only store half the edges here
A = max.(A,A')
dropzeros!(A)
##
issymmetric(A)
##
ei,ej = findnz(triu(A,1))
using GraphRecipes
using Plots
graphplot(ei, ej, x =xy[:,1], y=xy[:,2],
  markercolor=:black, markerstrokecolor=:white,
  size=(1200,1200),dpi=400,
  markersize=2, linecolor=1, linealpha=0.2, linewidth=0.5,
  markeralpha=0.3,
  axis_buffer=0.02, background=nothing)
#savefig()
##
using MatrixNetworks
pr = pagerank(A,0.85)
##
p=sortperm(pr)
scatter(xy[p,1],xy[p,2],marker_z=-abs.(log10.(pr[p])),markerstrokewidth=0)
## Now show the entire graph
p=sortperm(pr)
ei, ej = findnz(triu(A,1))
graphplot(ei, ej, x =xy[:,1], y=xy[:,2],
  markercolor=:black, markerstrokecolor=:white,
  size=(1200,1200),dpi=300,
  markersize=0, linecolor=1, linealpha=0.2, linewidth=0.5,
  colorbar=false,
  axis_buffer=0.02, background=nothing)
scatter!(xy[p,1],xy[p,2],marker_z=-abs.(log10.(pr[p])),markerstrokewidth=0,
  markersize=-abs.(log10.(pr[p])).+6)
savefig("wiki-cats-pagerank.png")
##


using Distributions, SpecialFunctions
import Distributions: mean

""" Exact Bayes (Adams and MacKay, 2007) for the Gaussian task (section 2.3.1)
"""
struct GBayesFilter{D, A} <: BayesFilter
    posteriors::Array{D, 1}
    prior::D
    probabilities::Array{Float64, 1}
    sigma::Float64
    changeprobability::Float64
    cutoff::Float64
    approximation::A
    mu_n::Array{Float64, 1}
end
function GBayesFilter(;sigma = .1, mu_0 = 0., sigma_0 = 1.,
                    changeprobability = 0.01, cutoff = 1e-16)

    prior = GaussKnownVariance(mu_0, sigma_0, sigma)
    approximation = CutoffApprox(cutoff)
    GBayesFilter([deepcopy(prior)], prior, [1.], sigma, changeprobability,
                cutoff, approximation, [0.])
end
meanparameter(f::Union{GBayesFilter, GMPN, GSOR}) = sum(postmean.(f.posteriors) .* f.probabilities)
function update!(f::Union{GBayesFilter, GMPN, GSOR}, x)
    update_posterior!(f, x)
    approximate!(f.approximation, f)
    parameter_estimate = meanparameter(f)       # E[θ | x_{1:t}]
    update_prior!(f)
    f.mu_n[1] = copy(parameter_estimate)
end

mutable struct GaussKnownVariance
    mu_0::Float64
    sigma_0::Float64
    sigma::Float64
    ν::Float64
    χ::Float64
end
function GaussKnownVariance(mu_0, sigma_0, sigma)
    ν_0 = calcν(sigma, sigma_0)
    χ_0 = calcχ(mu_0, ν_0, sigma)

    GaussKnownVariance(mu_0, sigma_0, sigma, ν_0, χ_0)
end
dataprobability(d::GaussKnownVariance, x) = gauss_χ_ν(d, x)
function gauss_χ_ν(d::GaussKnownVariance, x)
    expterm = - x^2/(2. * d.sigma^2)
    expterm -= d.χ^2 / (2. * d.ν)
    expterm += (d.χ + x/d.sigma)^2 / (2. * (d.ν + 1.))
    if isnan(expterm)
        error("expterm is NaN")
    end
    sqrtterm = sqrt(d.ν/(d.ν + 1.)) * (1. / (sqrt(2π) * d.sigma) )
    dataprobability = exp(expterm) * sqrtterm
    if isinf(dataprobability)
         error("dataprobability is Inf")
    end
    dataprobability
end

function update!(d::Vector{<:GaussKnownVariance}, x)
    @inbounds for i in 1:length(d)
        d[i].ν += 1.
        d[i].χ += x/d[i].sigma
    end
end
vargauss(d::GaussKnownVariance) = 1/(1/(d.sigma_0^2) + d.ν/(d.sigma^2))
mean(d::GaussKnownVariance) = vargauss(d) * (d.mu_0/(d.sigma_0^2) + d.χ/(d.sigma^2))

function postmean(d::GaussKnownVariance)
    mu = d.sigma * d.χ / d.ν
end


"""
    Conv(kernel, in=>out, σ; stride = 1, pading = 0, set_w = "Xavier")
This is the traditional convolution layer. `kernel` is a tuple of integers that specifies the kernel size, it must have one or two elements. And, `in` and `out` specifies number of input and out channels.

The input data must have a dimensions WHCB(weight, width, channel, batch). If you want to use a data which has a dimentions WHC, you must be add a dimentioins of B.

`stride` and `padding` are single integers or tuple(stride is tuple of 2 elements, padding is tuple of 2 elements), and if you specifies `same` to padding, we adjust sizes of input and return a matrix which has the same sizes.
`set_w` is `Xavier` or `He`, it decide a method to create a first parameter. This parameter is the same as `Dense()`.

# Example
```
julia> C = Conv((2, 2), 2=>2, relu)
Convolution(k:(2, 2), IO:2 => 2, σ:relu)

julia> C(rand(10, 10, 2, 5)) |> size
(9, 9, 2, 5)
```

!!! warning
    When you specidies `same` to `padding`, in some cases, it will be returned one size smaller.
    Because of its expression.
```

julia> C = Conv((2, 2), 2=>2, relu, padding = "same")
Convolution(k:(2, 2), IO:2 => 2, σ:relu

julia> C(rand(10, 10, 2, 5)) |> size
(9, 9, 2, 5)
```
"""
struct Conv <: Param
    σ
    w
    W
    b
    stride::Tuple
    padding::Union{Tuple, String}
    shape::IOShape
end

function Conv(k::NTuple{N, Int}, io::Pair{<:Integer, <:Integer}, σ; stride = 1, padding = 0, set_w = "Xavier") where {N}
    weight = conv_w(k, io, set_w)
    M = io[2]
    w = reshape(weight, M, div(length(weight), M))
    n = ndims(weight)
    padding = convert(2N, padding)
    stride = convert(N, stride)
    b = create_bias(weight, size(weight, n))
    return Conv(σ, w, weight, b, stride, padding, IOShape())
end

function Base.show(io::IO, C::Conv)
    k = size(C.W)[1:end-2]
    i, o = size(C.W)[end-1:end]
    σ = C.σ
    print("Convolution(k:$k, IO:$i => $o, σ:$σ)")
end

function (C::Conv)(x::AbstractArray)
    IC = Im2Col(x, size(C.W)[1:2], C.stride, C.padding)
    if size(C.w)[2] != size(IC.x)[1]
        kc = size(C.w)[end-1]
        ic = IC.c
        throw(DimensionMismatch("the number of channels must match between the kernel and the input!(now kernel: $kc, input: $ic)"))
    end
    σ, b = C.σ, C.b
    M = size(C.W)[end]
    x = reshape(σ.(C.w * IC.x .+ b), IC.Oh, IC.Ow, M, IC.b)
    return x
end

function (C::Conv)(x::AbstractArray, record::Dict, i)
    Ih, Iw = size(x)[1:2]
    IC = Im2Col(x, size(C.W)[1:2], C.stride, C.padding)
    C.shape(Ih, Iw, IC.Oh, IC.Ow)
    if size(C.w)[2] != size(IC.x)[1]
        kc = size(C.w)[end-1]
        ic = IC.c
        throw(DimensionMismatch("the number of channels must match between the kernel and the input!(now kernel: $kc, input: $ic)"))
    end
    σ, b = C.σ, C.b
    M = size(C.W)[end]
    z = C.w * IC.x
    _z = reshape(z, IC.Oh, IC.Ow, M, IC.b)
    a = σ.(z .+ b)
    a = reshape(a, IC.Oh, IC.Ow, M, IC.b)
    record[i+1] = (_z, a)
    return a
end

function (C::Conv)(Δ, z, back::Bool)
    z = Im2Col(z, size(C.W)[1:2], C.stride, C.padding).x
    w, σ = C.w, C.σ
    Δ = w' * Δ .* σ.(z, true)
    Fh, Fw,  = size(C.W)
    Δ = Col2Im(Δ, size(C.W)[1:2], C.shape, C.stride, C.padding).x
    #I have to use Col2Im to Δ.
    return Δ
end


using Mimi
using MimiFUND
using Plots
using Statistics
using Distributions

# get the model
m = MimiFUND.get_model()
run(m)
explore(m)

sd = @defsim begin
    # Define random variables
    climatedynamics.climatesensitivity = truncated(Normal(2, 1), 0.1, 5) # truncated to be stricly positive
    save(climatedynamics.temp)
end

# Run 100 trials
si = run(sd, m, 100)

test = getdataframe(si, :climatedynamics, :temp)

test = filter("time" => ==(2100), test)

histogram(test[:,:temp], label = nothing, xlabel = "Temperature Increase", ylabel = "Nber of trials")


"""
    feynman_kac(𝕋 [; t, f, ψ, v])

With direction = :backward
Solve the following PDE:
u(x, t[end]) = ψ(x)
0 = u_t + 𝕋u - v(x, t)u + f(x, t)

With direction = :forward
Solve the following PDE:
u(x, t[1]) = ψ(x)
u_t = 𝕋u - v(x, t)u + f(x, t)
"""
function feynman_kac(𝕋;
    t::AbstractVector = range(0, 100, step = 1/12),
    f::Union{AbstractVector, AbstractMatrix} = zeros(size(𝕋, 1)),
    ψ::AbstractVector = ones(size(𝕋, 1)),
    v::Union{AbstractVector, AbstractMatrix} = zeros(size(𝕋, 1)),
    direction= :backward)
    if direction == :backward
        u = zeros(size(𝕋, 1), length(t))
        u[:, end] = ψ
        if isa(f, AbstractVector) && isa(v, AbstractVector)
            if isa(t, AbstractRange)
                dt = step(t)
                B = factorize(I + (Diagonal(v) - 𝕋) * dt)
                for i in (length(t)-1):(-1):1
                    ψ = ldiv!(B, u[:, i+1] .+ f .* dt)
                    u[:, i] = ψ
                end
            else
                for i in (length(t)-1):(-1):1
                    dt = t[i+1] - t[i]
                    B = I + (Diagonal(v) - 𝕋) * dt
                    u[:, i] = B \ (u[:, i+1] .+ f .* dt)
                end
            end
        elseif isa(f, AbstractMatrix) && isa(v, AbstractMatrix)
            for i in (length(t)-1):(-1):1
                dt = t[i+1] - t[i]
                B = I + (Diagonal(view(v, :, i)) - 𝕋) * dt
                u[:, i] = B \ (u[:, i+1] .+ f[:, i] .* dt)
            end
        else
            error("f and v must be both AbstractVectors or both AbstractMatrices")
        end
        return u
    elseif direction == :forward
        u = feynman_kac(𝕋; t = - reverse(t), ψ = ψ, f = f, v = v, direction = :backward)
        return u[:,end:-1:1]
    else
        error("Direction must be :backward or :forward")
    end
end


using ReachabilityAnalysis

const μ = 3.986e14 * 60^2
const r = 42164.0e3
const r² = r^2
const mc = 500.0
const n² = μ / r^3
const n = sqrt(n²)

const two_n = 2*n
const μ_r² = μ/r²;

const K₁ = [-28.8287 0.1005 -1449.9754 0.0046 0.0;
            -0.087 -33.2562 0.00462 -1451.5013 0.0]
const K₂ = [-288.0288 0.1312 -9614.9898 0.0 0.0;
            -0.1312 -288.0 0.0 -9614.9883   0.0]

const K₁mc = K₁/mc
const K₂mc = K₂/mc

function mymul!(v, A, x)
    @inbounds for ind = 1:length(v)
        v[ind] = zero(x[1])
        for jind = 1:length(x)
            v[ind] += A[ind, jind] * x[jind]
        end
    end
    return nothing
end

@taylorize function spacecraft_approaching!(du, u, p, t)
    x, y, vx, vy, t = u

    rx = r + x
    rx² = rx^2
    y² = y^2
    rc = sqrt(rx² + y²)
    rc³ = rc^3
    μ_rc³ = μ / rc³

    uxy = Vector{typeof(x)}(undef, 2)
    mymul!(uxy, K₁mc, u)

    # x' = vx
    du[1] = vx

    # y' = vy
    du[2] = vy

    # vx' = n²x + 2n*vy + μ/(r^2) - μ/(rc^3)*(r+x) + ux/mc
    du[3] = (n²*x + two_n*vy) + ((μ_r² - μ_rc³*rx) + uxy[1])

    # vy' = n²y - 2n*vx - μ/(rc^3)y + uy/mc
    du[4] = (n²*y - two_n*vx) - (μ_rc³*y - uxy[2])

    # t' = 1
    du[5] = one(x)

    return du
end

@taylorize function spacecraft_attempt!(du, u, p, t)
    x, y, vx, vy, t = u

    rx = r + x
    rx² = rx^2
    y² = y^2
    rc = sqrt(rx² + y²)
    rc³ = rc^3
    μ_rc³ = μ / rc³

    uxy = Vector{typeof(x)}(undef, 2)
    mymul!(uxy, K₂mc, u)

    # x' = vx
    du[1] = vx

    # y' = vy
    du[2] = vy

    # vx' = n²x + 2n*vy + μ/(r^2) - μ/(rc^3)*(r+x) + ux/mc
    du[3] = (n²*x + two_n*vy) + ((μ_r² - μ_rc³*rx) + uxy[1])

    # vy' = n²y - 2n*vx - μ/(rc^3)y + uy/mc
    du[4] = (n²*y - two_n*vx) - (μ_rc³*y - uxy[2])

    # t' = 1
    du[5] = one(x)

    return du
end

@taylorize function spacecraft_aborting!(du, u, p, t)
    x, y, vx, vy, t = u

    rx = r + x
    rx² = rx^2
    y² = y^2
    rc = sqrt(rx² + y²)
    rc³ = rc^3
    μ_rc³ = μ / rc³

    # x' = vx
    du[1] = vx

    # y' = vy
    du[2] = vy

    # vx' = n²x + 2n*vy + μ/(r^2) - μ/(rc^3)*(r+x)
    du[3] = (n²*x + two_n*vy) + (μ_r² - μ_rc³*rx)

    # vy' = n²y - 2n*vx - μ/(rc^3)y
    du[4] = (n²*y - two_n*vx) - μ_rc³*y

    # t' = 1
    du[5] = one(x)

    return du
end

using Symbolics

const var = @variables x y vx vy t

function spacecraft(; abort_time=(120.0, 150.0))

    n = 4 + 1  # number of variables
    t_abort_lower, t_abort_upper = abort_time

    automaton = LightAutomaton(3)

    # mode 1 "approaching"
    invariant = HalfSpace(x <= -100, var)
    approaching = @system(x' = spacecraft_approaching!(x), dim:5, x ∈ invariant)

    # mode 2 "rendezvous attempt"
    invariant = HalfSpace(x >= -100, var)
    attempt = @system(x' = spacecraft_attempt!(x), dim:5, x ∈ invariant)

    # mode 3 "aborting"
    invariant = Universe(n)
    aborting = @system(x' = spacecraft_aborting!(x), dim: 5, x ∈ invariant)

    # transition "approach" -> "attempt"
    add_transition!(automaton, 1, 2, 1)
    guard = HalfSpace(x >= -100, var)
    t1 = @map(x -> x, dim: n, x ∈ guard)

    # transition "approach" -> "abort"
    add_transition!(automaton, 1, 3, 2)
    guard_time = HPolyhedron([t >= t_abort_lower, t <= t_abort_upper], var)
    t2 = @map(x -> x, dim: n, x ∈ guard_time)

    # transition "attempt" -> "abort"
    add_transition!(automaton, 2, 3, 3)
    t3 = @map(x -> x, dim: n, x ∈ guard_time)

    H = HybridSystem(automaton=automaton,
                     modes=[approaching, attempt, aborting],
                     resetmaps=[t1, t2, t3])

     # initial condition in mode 1
     X0 = Hyperrectangle([-900., -400., 0., 0., 0.], [25., 25., 0., 0., 0.])
     init = [(1, X0)]

     return InitialValueProblem(H, init)
end

const tan30 = tand(30)

LineOfSightCone = HPolyhedron([x >= -100, y >= x*tan30, -y >= x*tan30], var)

Target = Hyperrectangle(zeros(2), [0.2, 0.2]);

function line_of_sight(sol)
    all_idx = findall(x -> x == 2, location.(sol))  # attempt
    for idx in all_idx
        verif = all(set(R) ⊆ LineOfSightCone for R in sol[idx][2:end])
        !verif && return false
    end
    return true
end

function absolute_velocity(R)
    vx, vy = 3, 4
    vx2 = set(overapproximate(project(R, vars=(vx,)), Interval))
    vy2 = set(overapproximate(project(R, vars=(vy,)), Interval))
    return sqrt(vx2.dat + vy2.dat)
end

function velocity_constraint(sol)
    all_idx = findall(x -> x == 2, location.(sol)) # attempt mode
    for idx in all_idx
        # maximum velocity measured in m/min
        verif = all(absolute_velocity(R) < 0.055 * 60. for R in sol[idx])
        !verif && return false
    end
    return true
end

function target_avoidance(sol)
    all_idx = findall(x -> x == 3, location.(sol)) # aborting mode
    for idx in all_idx
        verif = all(is_intersection_empty(set(Projection(R, [x, y])), Target) for R in sol[idx])
        !verif && return false
    end
    return true
end

using Plots

function solve_spacecraft(prob; k=25, s=missing)

    # template directions
    boxdirs = BoxDirections(5)

    # transition from mode 1 to mode 2
    sol12 = solve(prob,
                tspan=(0.0, 200.0),
                alg=TMJets20(abstol=1e-5, maxsteps=10_000, orderT=5, orderQ=1, disjointness=BoxEnclosure()),
                max_jumps=1,
                intersect_source_invariant=false,
                intersection_method=TemplateHullIntersection(boxdirs),
                clustering_method=LazyClustering(1),
                disjointness_method=BoxEnclosure())

    sol12jump = overapproximate(sol12[2](120 .. 150), Zonotope)
    t0 = tstart(sol12jump[1])
    sol12jump_c = cluster(sol12jump, 1:length(sol12jump), BoxClustering(k, s))

    # transition from mode 2 to mode 3
    H = system(prob)
    sol3 = solve(IVP(mode(H, 3), [set(X) for X in sol12jump_c]),
                 tspan=(t0, 200.0),
                 alg=TMJets20(abstol=1e-10, orderT=7, orderQ=1, disjointness=BoxEnclosure()))
    d = Dict{Symbol, Any}(:loc_id => 3)

    F12 = [fp for fp in sol12.F]
    F23 = [Flowpipe(fp.Xk, d) for fp in sol3.F]
    return HybridFlowpipe(vcat(F12, F23))
end

prob = spacecraft()
sol = solve_spacecraft(prob; k=25, s=missing)
solz = overapproximate(sol, Zonotope);

plot(solz, vars=(1, 2), xlab="x", ylab="y")

idx_approaching = findall(x -> x == 1, location.(solz))
idx_attempt = findall(x -> x == 2, location.(solz))
idx_aborting = findall(x -> x == 3, location.(solz))

using Plots, Plots.PlotMeasures, LaTeXStrings

fig = plot(legend=:bottomright, tickfont=font(10, "Times"), guidefontsize=15,
           xlab=L"x", ylab=L"y", lw=0.0,
           xtick=[-750, -500, -250, 0, 250.], ytick=[-400, -300, -200, -100, 0.],
           xlims=(-1000.0, 400.0), ylims=(-450.0, 0.0), size=(600, 600),
           bottom_margin=6mm, left_margin=2mm, right_margin=4mm, top_margin=3mm)

plot!(fig, solz[idx_approaching[1]], vars=(1, 2), lw=0.0, color=:lightgreen, alpha=1., lab="Approaching")
for k in 2:length(idx_approaching)
    plot!(fig, solz[idx_approaching[k]], vars=(1, 2), lw=0.0, color=:lightgreen, alpha=1.)
end

plot!(fig, solz[idx_attempt[1]], vars=(1, 2), lw=0.0, color=:red, alpha=1., lab="Attempt")
for k in 2:length(idx_attempt)
    plot!(fig, solz[idx_attempt[k]], vars=(1, 2), lw=0.0, color=:red, alpha=1.)
end

plot!(fig, solz[idx_aborting[1]], vars=(1, 2), lw=0.0, color=:cyan, alpha=1., lab="Aborting")
for k in 2:length(idx_aborting)
    plot!(fig, solz[idx_aborting[k]], vars=(1, 2), lw=0.0, color=:cyan, alpha=1.)
end

fig