text:
export learn_circuit,
       learn_circuit_miss

using LogicCircuits: split_step, struct_learn
using Statistics: mean
using Random

"""
Learn structure of a single structured decomposable circuit
"""
function learn_circuit(train_x;
        pick_edge="eFlow", pick_var="vMI", depth=1,
        pseudocount=1.0,
        entropy_reg=0.0,
        sanity_check=true,
        maxiter=100,
        seed=nothing,
        return_vtree=false,
        verbose=true,
        max_circuit_nodes=nothing)

    # Initial Structure
    pc, vtree = learn_chow_liu_tree_circuit(train_x)

    learn_circuit(train_x, pc, vtree; pick_edge, pick_var,
                           depth, pseudocount, sanity_check,
                           maxiter, seed, return_vtree, entropy_reg, verbose, max_circuit_nodes)
end

"""
Learn structure of a single structured decomposable circuit from missing data.

Missing feature are denoted by `missing`.
Median Imputation is used by default for initial structure (set `impute_method` for other options).
"""
function learn_circuit_miss(train_x;
        impute_method::Symbol=:median,
        pick_edge="eFlow", depth=1,
        pseudocount=1.0,
        entropy_reg=0.0,
        sanity_check=true,
        maxiter=100,
        seed=nothing,
        return_vtree=false,
        verbose=true,
        max_circuit_nodes=nothing,
        max_learning_time=nothing)

    # Initial Structure
    train_x_impute = impute(train_x; method=impute_method)
    pc, vtree = learn_chow_liu_tree_circuit(train_x_impute)

    # Only vRand supported for missing data
    pick_var="vRand"

    learn_circuit(train_x, pc, vtree; pick_edge, pick_var,
                  depth, pseudocount, sanity_check,
                  maxiter, seed, return_vtree, entropy_reg,
                  max_circuit_nodes, verbose,
                  max_learning_time,
                  has_missing=true)
end

function learn_circuit(train_x, pc, vtree;
        pick_edge="eFlow", pick_var="vMI", depth=1,
        pseudocount=1.0,
        sanity_check=true,
        maxiter=100,
        seed=nothing,
        return_vtree=false,
        batch_size=0,
        splitting_data=nothing,
        use_gpu=false,
        entropy_reg=0.0,
        verbose=true,
        max_circuit_nodes=nothing,
        max_learning_time=nothing, # in seconds
        has_missing::Bool=false
        )

    if seed !== nothing
        Random.seed!(seed)
    end

    if has_missing
        estimate_parameters_func = estimate_parameters_em
        likelihood_avg_func = marginal_log_likelihood_avg
        LogLikelihood_str = "Marginal LogLikelihood"
    else
        estimate_parameters_func = estimate_parameters
        likelihood_avg_func = log_likelihood_avg
        LogLikelihood_str = "LogLikelihood"
    end

    # structure_update
    loss(circuit) = heuristic_loss(circuit,
                                    splitting_data == nothing ? train_x : splitting_data;
                                    pick_edge=pick_edge,
                                    pick_var=pick_var,
                                    miss_data=has_missing)

    init_time_sec = time()

    pc_split_step(circuit) = begin
        r = split_step(circuit; loss=loss, depth=depth, sanity_check=sanity_check)
        if isnothing(r) return nothing end
        c, = r
        if batch_size > 0
            estimate_parameters_func(c, batch(train_x, batch_size); pseudocount, use_gpu, entropy_reg)
        else
            estimate_parameters_func(c, train_x; pseudocount, use_gpu, entropy_reg)
        end
        return c, missing
    end
    iter = 0
    log_per_iter(circuit) = begin
        if batch_size > 0
            ll = likelihood_avg_func(circuit, batch(train_x, batch_size); use_gpu)
        else
            ll = likelihood_avg_func(circuit, train_x; use_gpu)
        end
        verbose && println("Iteration $iter/$maxiter. $(LogLikelihood_str) = $(ll); nodes = $(num_nodes(circuit)); edges =  $(num_edges(circuit)); params = $(num_parameters(circuit))")
        iter += 1

        if !isnothing(max_circuit_nodes) && num_nodes(circuit) > max_circuit_nodes
            verbose && println("Stopping early, circuit node count ($(num_nodes(circuit))) is above max threshold $(max_circuit_nodes).");
            return true; # stop
        end

        total_time_taken = time() - init_time_sec
        if !isnothing(max_learning_time) && total_time_taken > max_learning_time
            verbose && println("Stopping early, time limit reached ($(max_learning_time)) is above max threshold $(max_learning_time).");
            return true; # stop
        end

        false
    end

    # Log Before Learning
    log_per_iter(pc)

    pc = struct_learn(pc;
        primitives=[pc_split_step], kwargs=Dict(pc_split_step=>()),
        maxiter=maxiter, stop=log_per_iter, verbose=verbose)

    if return_vtree
        pc, vtree
    else
        pc
    end
end


text:
@doc raw"""
    ChambollePock(
        M, N, cost, x0, ξ0, m, n, prox_F, prox_G_dual, adjoint_linear_operator;
        forward_operator=missing,
        linearized_forward_operator=missing,
        evaluation=AllocatingEvaluation()
    )

Perform the Riemannian Chambolle–Pock algorithm.

Given a `cost` function $\mathcal E:\mathcal M → ℝ$ of the form
```math
\mathcal E(x) = F(x) + G( Λ(x) ),
```
where $F:\mathcal M → ℝ$, $G:\mathcal N → ℝ$,
and $Λ:\mathcal M → \mathcal N$. The remaining input parameters are

* `x,ξ` primal and dual start points $x∈\mathcal M$ and $ξ∈T_n\mathcal N$
* `m,n` base points on $\mathcal M$ and $\mathcal N$, respectively.
* `adjoint_linearized_operator` the adjoint $DΛ^*$ of the linearized operator $DΛ(m): T_{m}\mathcal M → T_{Λ(m)}\mathcal N$
* `prox_F, prox_G_Dual` the proximal maps of $F$ and $G^\ast_n$

note that depenting on the [`AbstractEvaluationType`](@ref) `evaluation` the last three parameters
as well as the fordward_operator `Λ` and the `linearized_forward_operator` can be given as
allocating functions `(Manifolds, parameters) -> result`  or as mutating functions
`(Manifold, result, parameters)` -> result` to spare allocations.

By default, this performs the exact Riemannian Chambolle Pock algorithm, see the opional parameter
`DΛ` for ther linearized variant.

For more details on the algorithm, see[^BergmannHerzogSilvaLouzeiroTenbrinckVidalNunez2020].

# Optional Parameters

* `acceleration` – (`0.05`)
* `dual_stepsize` – (`1/sqrt(8)`) proximnal parameter of the primal prox
* `evaluation` ([`AllocatingEvaluation`](@ref)`()) specify whether the proximal maps and operators are
  allocating functions `(Manifolds, parameters) -> result`  or given as mutating functions
  `(Manifold, result, parameters)` -> result` to spare allocations.
* `Λ` (`missing`) the (forward) operator $Λ(⋅)$ (required for the `:exact` variant)
* `linearized_forward_operator` (`missing`) its linearization $DΛ(⋅)[⋅]$ (required for the `:linearized` variant)
* `primal_stepsize` – (`1/sqrt(8)`) proximnal parameter of the dual prox
* `relaxation` – (`1.`)
* `relax` – (`:primal`) whether to relax the primal or dual
* `variant` - (`:exact` if `Λ` is missing, otherwise `:linearized`) variant to use.
  Note that this changes the arguments the `forward_operator` will be called.
* `stopping_criterion` – (`stopAtIteration(100)`) a [`StoppingCriterion`](@ref)
* `update_primal_base` – (`missing`) function to update `m` (identity by default/missing)
* `update_dual_base` – (`missing`) function to update `n` (identity by default/missing)
* `retraction_method` – (`ExponentialRetraction()`) the rectraction to use
* `inverse_retraction_method` - (`LogarithmicInverseRetraction()`) an inverse retraction to use.
* `vector_transport_method` - (`ParallelTransport()`) a vector transport to use

[^BergmannHerzogSilvaLouzeiroTenbrinckVidalNunez2020]:
    > R. Bergmann, R. Herzog, M. Silva Louzeiro, D. Tenbrinck, J. Vidal-Núñez:
    > _Fenchel Duality Theory and a Primal-Dual Algorithm on Riemannian Manifolds_,
    > Foundations of Computational Mathematics, 2021.
    > doi: [10.1007/s10208-020-09486-5](http://dx.doi.org/10.1007/s10208-020-09486-5)
    > arXiv: [1908.02022](http://arxiv.org/abs/1908.02022)
"""
function ChambollePock(
    M::Manifold,
    N::Manifold,
    cost::Function,
    x::P,
    ξ::T,
    m::P,
    n::Q,
    prox_F::Function,
    prox_G_dual::Function,
    adjoint_linear_operator::Function;
    Λ::Union{Function,Missing}=missing,
    linearized_forward_operator::Union{Function,Missing}=missing,
    kwargs...,
) where {P,T,Q}
    x_res = allocate(x)
    copyto!(x_res, x)
    ξ_res = allocate(ξ)
    copyto!(ξ_res, ξ)
    m_res = allocate(m)
    copyto!(m_res, m)
    n_res = allocate(n)
    copyto!(n_res, n)
    return ChambollePock!(
        M,
        N,
        cost,
        x_res,
        ξ_res,
        m_res,
        n_res,
        prox_F,
        prox_G_dual,
        adjoint_linear_operator;
        Λ=Λ,
        linearized_forward_operator=linearized_forward_operator,
        kwargs...,
    )
end
@doc raw"""
    ChambollePock(M, N, cost, x0, ξ0, m, n, prox_F, prox_G_dual, adjoint_linear_operator)

Perform the Riemannian Chambolle–Pock algorithm in place of `x`, `ξ`, and potenitally `m`,
`n` if they are not fixed. See [`ChambollePock`](@ref) for details and optional parameters.
"""
function ChambollePock!(
    M::Manifold,
    N::Manifold,
    cost::Function,
    x::P,
    ξ::T,
    m::P,
    n::Q,
    prox_F::Function,
    prox_G_dual::Function,
    adjoint_linear_operator::Function;
    Λ::Union{Function,Missing}=missing,
    linearized_forward_operator::Union{Function,Missing}=missing,
    acceleration=0.05,
    dual_stepsize=1 / sqrt(8),
    primal_stepsize=1 / sqrt(8),
    relaxation=1.0,
    relax::Symbol=:primal,
    stopping_criterion::StoppingCriterion=StopAfterIteration(200),
    update_primal_base::Union{Function,Missing}=missing,
    update_dual_base::Union{Function,Missing}=missing,
    retraction_method::RM=ExponentialRetraction(),
    inverse_retraction_method::IRM=LogarithmicInverseRetraction(),
    vector_transport_method::VTM=ParallelTransport(),
    variant=ismissing(Λ) ? :exact : :linearized,
    return_options=false,
    kwargs...,
) where {
    P,
    Q,
    T,
    RM<:AbstractRetractionMethod,
    IRM<:AbstractInverseRetractionMethod,
    VTM<:AbstractVectorTransportMethod,
}
    p = PrimalDualProblem(
        M,
        N,
        cost,
        prox_F,
        prox_G_dual,
        adjoint_linear_operator;
        linearized_forward_operator=linearized_forward_operator,
        Λ=Λ,
    )
    o = ChambollePockOptions(
        m,
        n,
        x,
        ξ,
        primal_stepsize,
        dual_stepsize;
        acceleration=acceleration,
        relaxation=relaxation,
        stopping_criterion=stopping_criterion,
        relax=relax,
        update_primal_base=update_primal_base,
        update_dual_base=update_dual_base,
        variant=variant,
        retraction_method=retraction_method,
        inverse_retraction_method=inverse_retraction_method,
        vector_transport_method=vector_transport_method,
    )
    o = decorate_options(o; kwargs...)
    resultO = solve(p, o)
    if return_options
        return resultO
    else
        return get_solver_result(resultO)
    end
end

function initialize_solver!(::PrimalDualProblem, ::ChambollePockOptions) end

function step_solver!(p::PrimalDualProblem, o::ChambollePockOptions, iter)
    primal_dual_step!(p, o, Val(o.relax))
    o.m = ismissing(o.update_primal_base) ? o.m : o.update_primal_base(p, o, iter)
    if !ismissing(o.update_dual_base)
        n_old = deepcopy(o.n)
        o.n = o.update_dual_base(p, o, iter)
        vector_transport_to!(p.N, o.ξ, n_old, o.ξ, o.n, o.vector_transport_method)
        vector_transport_to!(p.N, o.ξbar, n_old, o.ξbar, o.n, o.vector_transport_method)
    end
    return o
end
#
# Variant 1: primal relax
#
function primal_dual_step!(p::PrimalDualProblem, o::ChambollePockOptions, ::Val{:primal})
    dual_update!(p, o, o.xbar, Val(o.variant))
    if ismissing(p.Λ!!)
        ptξn = o.ξ
    else
        ptξn = vector_transport_to(
            p.N, o.n, o.ξ, forward_operator(p, o.m), o.vector_transport_method
        )
    end
    xOld = o.x
    o.x = get_primal_prox!(
        p,
        o.x,
        o.primal_stepsize,
        retract(
            p.M,
            o.x,
            vector_transport_to(
                p.M,
                o.m,
                -o.primal_stepsize * (adjoint_linearized_operator(p, o.m, o.n, ptξn)),
                o.x,
                o.vector_transport_method,
            ),
            o.retraction_method,
        ),
    )
    update_prox_parameters!(o)
    retract!(
        p.M,
        o.xbar,
        o.x,
        -o.relaxation * inverse_retract(p.M, o.x, xOld, o.inverse_retraction_method),
        o.retraction_method,
    )
    return o
end
#
# Variant 2: dual relax
#
function primal_dual_step!(p::PrimalDualProblem, o::ChambollePockOptions, ::Val{:dual})
    if ismissing(p.Λ!!)
        ptξbar = o.ξbar
    else
        ptξbar = vector_transport_to(
            p.N, o.n, o.ξbar, forward_operator(p, o.m), o.vector_transport_method
        )
    end
    get_primal_prox!(
        p,
        o.x,
        o.primal_stepsize,
        retract(
            p.M,
            o.x,
            vector_transport_to(
                p.M,
                o.m,
                -o.primal_stepsize * (adjoint_linearized_operator(p, o.m, o.n, ptξbar)),
                o.x,
                o.vector_transport_method,
            ),
            o.retraction_method,
        ),
    )
    ξ_old = deepcopy(o.ξ)
    dual_update!(p, o, o.x, Val(o.variant))
    update_prox_parameters!(o)
    o.ξbar = o.ξ + o.relaxation * (o.ξ - ξ_old)
    return o
end
#
# Dual step: linearized
# depending on whether its primal relaxed or dual relaxed we start from start=o.x or start=o.xbar here
#
function dual_update!(
    p::PrimalDualProblem, o::ChambollePockOptions, start::P, ::Val{:linearized}
) where {P}
    # (1) compute update direction
    ξ_update = linearized_forward_operator(
        p, o.m, inverse_retract(p.M, o.m, start, o.inverse_retraction_method)
    )
    # (2) if p.Λ is missing, we assume that n = Λ(m) and do  not PT, otherwise we do
    (!ismissing(p.Λ!!)) && vector_transport_to!(
        p.N,
        ξ_update,
        forward_operator(p, o.m),
        ξ_update,
        o.n,
        o.vector_transport_method,
    )
    # (3) to the dual update
    get_dual_prox!(p, o.ξ, o.n, o.dual_stepsize, o.ξ + o.dual_stepsize * ξ_update)
    return o
end
#
# Dual step: exact
# depending on whether its primal relaxed or dual relaxed we start from start=o.x or start=o.xbar here
#
function dual_update!(
    p::PrimalDualProblem, o::ChambollePockOptions, start::P, ::Val{:exact}
) where {P}
    ξ_update = inverse_retract(
        p.N, o.n, forward_operator(p, start), o.inverse_retraction_method
    )
    get_dual_prox!(p, o.ξ, o.n, o.dual_stepsize, o.ξ + o.dual_stepsize * ξ_update)
    return o
end

@doc raw"""
    update_prox_parameters!(o)
update the prox parameters as described in Algorithm 2 of Chambolle, Pock, 2010, i.e.

1. ``θ_{n} = \frac{1}{\sqrt{1+2γτ_n}}``
2. ``τ_{n+1} = θ_nτ_n``
3. ``σ_{n+1} = \frac{σ_n}{θ_n}``
"""
function update_prox_parameters!(o::O) where {O<:PrimalDualOptions}
    if o.acceleration > 0
        o.relaxation = 1 / sqrt(1 + 2 * o.acceleration * o.primal_stepsize)
        o.primal_stepsize = o.primal_stepsize * o.relaxation
        o.dual_stepsize = o.dual_stepsize / o.relaxation
    end
    return o
end

ext: jl


text:
include("../../src/ToyModelsMD.jl")
using .ToyModelsMD
using Flux
using DelimitedFiles
using Plots

function Afunc(system::AbstractSystem,mm::MixedModel,ind::Int64,dt::Float64)
  lam::Float64 = 2.0e3
  pot::Float64 = potential(system,mm.potentials[1])
  ind==1500 && system.x[2]>0.7 && pot<-145.0 ? lam : 0.0
end

### define system ###
dt = 0.0001
steps = 1500
tf = steps*dt

# mueller-brown potential
mb = MullerBrown()
# neural network model
nn = NNModel(2,Chain(Dense(3,10),Dense(10,2)))

# mixed model
model = MixedModel(2,[mb,nn])

# thermostat and system
langevin = Langevin(1.0,1.0)
system = ThermostattedSystem(2,0.0,[0.63,0.03],zeros(Float64,2),langevin)

# define value baseline function
vbl = NNValueBaseline(Chain(Dense(3,10),Dense(10,1)))

# define callback function
mcvb = MCVBCallback(Afunc,nn,vbl)
# set up integrator and run!
integrator = StochasticEuler(system,model)
runtraj!(steps,dt,integrator,mcvb)
println(mcvb.dkl)
writedlm("gradf.txt",mcvb.gradf)
writedlm("gradv.txt",mcvb.gradv*dt)


text:
module  TestAlleleGenoFreq

using PopGenCore
using DataFrames
using Test

cats = @nancycats ;
g = cats.genodata.genotype[1:10]
sharks = @gulfsharks ;

@testset "Allele Frequncies" begin
    @testset "Basic" begin
        @test allelefreq(cats) isa NamedTuple
        @test allelefreq(g) isa Dict
        @test allelefreq(g[5]) isa Dict
        @test allelefreq_vec(g) isa AbstractVector
        a = allelefreq(g)
        b = allelefreq(cats.genodata.genotype[11:20])
        @test avg_allelefreq([a,b]) isa Dict
        @test avg_allelefreq([a,b], 2) isa Dict
        @test allelefreq(cats, "fca8") isa Dict
        @test allelefreq(cats, "fca8", population = true) isa DataFrame
        @test allelefreq(133, g) isa AbstractFloat
    end

    @testset "Nuanced" begin
        @test length(allelefreq(cats)) == 9
        @test eltype(allelefreq(cats)) == Dict{Int16,Float64}
        @test allelefreq(cats.genodata.genotype) isa Dict{Int16,Float64}
        df_cats = DataFrames.combine(
            groupby(cats.genodata, :population),
            :genotype => allelefreq => :genos
        )
        @test size(df_cats) == (17,2)
        @test eltype(df_cats.genos) == Dict{Int16,Float64}
        @test length(allelefreq(cats, "fca8")) == 16
        @test allelefreq(cats, "fca8") isa Dict{Int16,Float64}

        @test length(allelefreq(sharks)) == 2209
        @test eltype(allelefreq(sharks)) == Dict{Int8,Float64}
        @test allelefreq(sharks.genodata.genotype) isa Dict{Int8,Float64}
        df_sharks = DataFrames.combine(
            groupby(sharks.genodata, :population),
            :genotype => allelefreq => :genos
        )
        @test size(df_sharks) == (7,2)
        @test eltype(df_sharks.genos) == Dict{Int8,Float64}
        @test length(allelefreq(sharks, "contig_2784")) == 2
        @test allelefreq(sharks, "contig_2784") isa Dict{Int8,Float64}

        @test allelefreq_vec([(1,1), (2,2), (2,1)]) == [0.5, 0.5]
        @test allelefreq_vec(missing) === missing
    end
end

@testset "Genotype Frequncies" begin
    @testset "Basic" begin
        @test genofreq(g) isa Dict
        @test genofreq(cats, "fca8") isa Dict
        @test genofreq(cats, "fca8", population = true) isa DataFrame
        @test genofreq_expected(g) isa Dict
        @test genofreq_expected(cats, "fca8") isa Dict
        @test genofreq_expected(cats, "fca8", population = true) isa DataFrame
    end

    @testset "Nuanced" begin
        @test length(genofreq(cats.genodata.genotype)) == 295
        @test typeof(genofreq(cats.genodata.genotype)) <: Dict{<:Tuple,Float64}
        @test length(genofreq(cats, "fca8")) == 51
        @test typeof(genofreq(cats, "fca8")) <: Dict{<:Tuple,Float64}
        @test size(genofreq(cats, "fca8", population = true)) == (17,2)

        @test length(genofreq_expected(cats.genodata.genotype)) == 6241
        @test typeof(genofreq_expected(cats.genodata.genotype)) <: Dict{<:Tuple,Float64}
        @test length(genofreq_expected(cats, "fca8")) == 256
        @test typeof(genofreq_expected(cats, "fca8")) <: Dict{<:Tuple,Float64}
        @test size(genofreq_expected(cats, "fca8", population = true)) == (17,2)

        @test length(genofreq(sharks.genodata.genotype)) == 12
        @test length(genofreq(sharks, "contig_2784")) == 2
        @test typeof(genofreq(sharks, "contig_2784")) <: Dict{<:Tuple,Float64}
        @test size(genofreq(sharks, "contig_2784", population = true)) == (7,2)

        @test length(genofreq_expected(sharks.genodata.genotype)) == 25
        @test typeof(genofreq_expected(sharks.genodata.genotype)) <: Dict{<:Tuple,Float64}
        @test length(genofreq_expected(sharks, "contig_2784")) == 4
        @test typeof(genofreq_expected(sharks, "contig_2784")) <: Dict{<:Tuple,Float64}
        @test size(genofreq_expected(sharks, "contig_2784", population = true)) == (7,2)
    end
end

@testset "Genotype Counts" begin
    @testset "Basic" begin
        @test genocount_observed(g) isa Dict
        @test genocount_observed([missing, missing]) === missing
        @test genocount_expected(g) isa Dict
    end

    @testset "Nuanced" begin
        @test length(genocount_observed(cats.genodata.genotype)) == 295
        @test typeof(genocount_observed(cats.genodata.genotype)) <: Dict{<:Tuple,Int64}
        @test length(genocount_expected(cats.genodata.genotype)) == 6241
        @test typeof(genocount_expected(cats.genodata.genotype)) <: Dict{<:Tuple,Float64}

        @test length(genocount_observed(sharks.genodata.genotype)) == 12
        @test typeof(genocount_observed(sharks.genodata.genotype)) <: Dict{<:Tuple,Int64}
        @test length(genocount_expected(sharks.genodata.genotype)) == 25
        @test typeof(genocount_expected(sharks.genodata.genotype)) <: Dict{<:Tuple,Float64}
    end
end

end


text:
mutable struct Arellano_params{T<:Real, I<:Int64}
    #prefernces
    β :: T
    θ :: T
    #grid_point definitions
    grid :: Array{T}
    #Markov chian
    P :: Array{T}
    ϵ :: Array{T}
    #interest rate
    r :: T
    #defult output parametr
    y_hat :: T
    #maximal number of iterations
    n :: I

end

#outer constructor
function Arellano_params(ρ, η, step, ϵ_n, B_min, B_max)
    #constructor for given earning process and assetd grid with step and income grid with \epsilon_n  points
    grid = collect(B_min:step:B_max)
    MC = tauchen(ϵ_n, ρ, η)
    P_ϵ = MC.p
    ϵ = exp.(MC.state_values)
    Arellano_params(0.946, 0.280, grid,P_ϵ, ϵ, 0.017, 0.962, 1000)
end

utility(c) = c^(1.0-2.0)/(1.0-2.0) #utility with sigma = 2.0
"""
Compute the output if default, given the possible value of output in continuation (y) and the function parameter (y_hat)
"""
function y_d(y,y_hat)
    if(y<y_hat)
        return y
    else
        return y_hat
    end
end

#detrend data for staistics
function detrend(time, Y)
    β = inv(time'*time)*time'*Y
    return Y - β*time
end

"""
Make one iteration for the Value function iterations
    INPUTS:
    V_f:  current value function
    q: price matrix
    β: disount factor
    y_hat: paramter for the output in cse of default
    P: stochatsic matrix ofthe Markov Chain
    ϵ: values of the income schock for the given Markov Chain
    θ: paramter for the exclusion form the finacial arket after default
    grid: gird for assets
    OUTPUT:
    V_final: new value function
    Def_Mat: default region
"""
function Value_function_update_Arellano(V_f, q, β, y_hat,P, ϵ, θ, grid)
    #first read size of the value functions
    a_n = size(V_f)[2] #number of points in the asset grid
    ϵ_n = size(V_f,1) #number of the income process states

    #allocate memory
    E_V_f_prim = zeros(ϵ_n, a_n) #expected value next period for continuation
    E_V_left_prim  =  zeros(ϵ_n) #expected value next period for default

    V_f_prim = zeros(ϵ_n, a_n) #value for continuation
    V_def = zeros(ϵ_n)  #value for default
    C_vec = zeros(a_n) #possible consumption
    Default_mat = zeros(ϵ_n, a_n) #default decsion

    #find index with no assets
    zero_ind = searchsortedfirst(grid, 0.0)

    #compute expected values next-period for continuation
    for j in 1:ϵ_n
        for i in 1: a_n
            futur_val = 0
            for jj in 1: ϵ_n
                futur_val = copy(futur_val)+P[j,jj]*V_f[jj,i]
            end
            E_V_f_prim[j,i] = futur_val
        end
    end

    #compute optimal savings rule for continuation
    for j in 1: ϵ_n
        for i in 1:a_n
            C_vec = zeros(a_n)
            for ii in 1:a_n
                C_vec[ii] = grid[i] + ϵ[j] - q[j,ii]*grid[ii]
            end
            C_vec = max.(copy(C_vec), 1e-7)
            V_f_prim[j,i] = maximum(utility.(C_vec) + β*copy(E_V_f_prim[j,:]))
        end
    end

    #compute default value
    V_def = inv(I(ϵ_n) - β*(1.0-θ)*P)*(utility.(y_d.(ϵ,y_hat)) + β*θ*copy(E_V_f_prim[:,zero_ind]))

    #now choose the maximum between value of continuation and default
    V_final = V_f_prim
    for j in 1: ϵ_n
        for i in 1:a_n
           if(V_f_prim[j,i]<= V_def[j] && i< zero_ind)
              Default_mat[j,i] = 1.0
              V_final[j,i] = V_def[j]
            end
         end
    end
    return V_final, Default_mat
end

"""
Find optimal savings policy, this is to call arg_max function only once (after value and default function is found), beside this is the same function as Value_function_update_Arellano
"""

function Value_function_get_policy_Arellano(V_f, q, β, y_hat,P, ϵ, θ, grid)
    #first read size of the value functions
    a_n = size(V_f)[2]
    ϵ_n = size(V_f,1)

    #allocate memory
    E_V_f_prim = zeros(ϵ_n, a_n) #expected value next period for continuation
    E_V_left_prim  =  zeros(ϵ_n) #expected value next period for default

    V_f_prim = zeros(ϵ_n, a_n) #value for continuation
    V_def = zeros(ϵ_n)  #value for default
    C_vec = zeros(a_n) #possible consumption
    Default_mat = zeros(ϵ_n, a_n) #default decsion

    Policy = zeros(ϵ_n, a_n) #policy function
    C_vec = zeros(a_n) #possible consumption

    #find index with no assets
    zero_ind = searchsortedfirst(grid, 0.0)

     #compute expected values next-period for continuation
    for j in 1:ϵ_n
        for i in 1: a_n
            futur_val = 0
            for jj in 1: ϵ_n
                futur_val = copy(futur_val)+P[j,jj]*V_f[jj,i]
            end
            E_V_f_prim[j,i] = futur_val
        end
    end

    #compute optimal savings rule for continuation
    for j in 1: ϵ_n
        for i in 1:a_n
            C_vec = zeros(a_n)
            for ii in 1:a_n
                C_vec[ii] = grid[i] + ϵ[j] - q[j,ii]*grid[ii]
            end
            C_vec = max.(copy(C_vec), 1e-7)
            V_f_prim[j,i] = maximum(utility.(C_vec) + β*copy(E_V_f_prim[j,:]))
            Policy[j,i] = argmax(utility.(C_vec) + β*copy(E_V_f_prim[j,:]))

        end
    end
    #vaue of default
    V_def = inv(I(ϵ_n) - β*(1.0-θ)*P)*(utility.(y_d.(ϵ,y_hat)) + β*θ*copy(E_V_f_prim[:,zero_ind]))

    #now choose the maximum between value of continuation and default
    V_final = V_f_prim
    for j in 1: ϵ_n
        for i in 1:a_n
           if(V_f_prim[j,i]< V_def[j] && i< zero_ind)
              Default_mat[j,i] = 1.0
              V_final[j,i] = V_def[j]
              Policy[j,i] = zero_ind
            end
         end
    end

    #convert policy to the int64 matrix
    Policy_int = convert(Array{Int64}, Policy)
    return Policy_int
end

"""
Solve the Bellman equation given the price matrix q

INPUTS:
n: maximal number of iterations
q: price matrix
V_f: value function initial guess
β: disount factor
y_hat: paramter for the output in cse of default
P: stochatsic matrix ofthe Markov Chain
ϵ: values of the income schock for the given Markov Chain
θ: paramter for the exclusion form the finacial arket after default
grid: gird for assets

OUTPUTS:
V_final: value function
Def_Mat: default region
"""
function Solve_Bellman_Arellano(n, q, V_f,P, ϵ, θ, y_hat, β, grid)
    #check size of the asset and income grids
    a_n = size(V_f)[2]
    ϵ_n = size(V_f,1)

    #Allocate memory
    V_f_new  = ones(ϵ_n, a_n) #
    Default_mat = zeros(ϵ_n, a_n)

    #initialize iterations
    iter = 0

    while(maximum(abs.(V_f_new-V_f))>=1e-8 && iter <n)
        if(iter !=1)
         V_f = copy(V_f_new)
        end
         V_f_new, Default_mat = Value_function_update_Arellano( V_f, q, β, y_hat,P, ϵ, θ, grid)

         iter = iter+1
         #println(maximum(abs.(V_f_new-V_f)))
    end

    return(V_f, Default_mat)
end

"""
Solve the Arellano model:
INPUTS:
n: maximal number of iterations of outer loop (for q)
P: stochatsic matrix ofthe Markov Chain
ϵ: values of the income schock for the given Markov Chain
grid: gird for assets
r: interest rate
OUTPUTS:
q_new : price matrix
D_new: default region
Policy: savings policy
V_f_new: value function
"""
function Solve_eq_Arellano(Model::Arellano_params)

    #unpack params
    n = Model.n
    ϵ  = Model.ϵ
    P = Model.P
    grid = Model.grid
    r = Model.r
    y_hat = Model.y_hat
    θ = Model.θ
    β = Model.β

    #read size of the asset and income grids
    ϵ_n = length(ϵ)
    a_n = length(grid)

    #alloctate memory
    V_f  = zeros(ϵ_n, a_n) # value function, rather dull starting guess TODO
    D = zeros(ϵ_n, a_n) #default region
    D_new = zeros(ϵ_n, a_n) #new default region
    V_f_new  = zeros(ϵ_n, a_n) # new value function
    Policy = zeros(ϵ_n, a_n) # savings policy function
    δ = zeros(ϵ_n, a_n) #default probability
    q_new = zeros(ϵ_n, a_n) #new q price matrix
    #initilaize parameters
    q = 1.0/(1.0+r)*ones(ϵ_n, a_n) #initialize price matrix q

    #find index with no assets
    zero_ind = searchsortedfirst(grid, 0.0)
    #initialize number of iterations
    iter = 0

    #OUTER LOOP STARTS HERE
    while(maximum(abs.(q_new.-q))>=1e-7 && iter <n)
        #initilize the paramters
        if(iter != 0)
            q = copy(q_new)
        end

        #solve the Bellman equation for given q
        V_f_new, D_new = Solve_Bellman_Arellano(1000, q, copy(V_f_new), P, ϵ, θ, y_hat, β, grid)
        Policy =  Value_function_get_policy_Arellano(copy(V_f_new), q, β, y_hat,P, ϵ, θ, grid)

        #compute default probability
        for j in 1:ϵ_n
            for i in 1:a_n
                default_prob = 0.0
                for jj in 1:ϵ_n
                    default_prob = copy(default_prob) + P[j,jj]*D_new[jj,i]
                end
                δ[j,i] = copy(default_prob)

            end
        end

        #updatethe q
        q_new = 0.99*(1.0/(1.0+r)*(1.0.-copy(δ)))+0.01*copy(q)

        iter = iter+1
        #print iteration
        println("Iteration ", iter)
        println("Number of defaults ", sum(D_new))
        println("Maximal difference in q ", maximum(abs.(q_new.-q))," at ", argmax(abs.(q_new.-q)))
        println(" ")
    end
    return (q_new, D_new, Policy, V_f_new), grid
end

#simulate income process from markov chain
function mc_sample_path(P; init = 11, sample_size = 11000)
    @assert size(P)[1] == size(P)[2] # square required
    N = size(P)[1] # should be square

    # create vector of discrete RVs for each row
    dists = [Categorical(P[i, :]) for i in 1:N]

    # setup the simulation
    X = fill(0, sample_size) # allocate memory, or zeros(Int64, sample_size)
    init_prob = fill(1/N, N)
    X[1] = rand(Categorical(init_prob)) # set the initial state

    for t in 2:sample_size
        dist = dists[X[t-1]] # get discrete RV from last state's transition distribution
        X[t] = rand(dist) # draw new value
    end
    SIM = convert(Array{Int64},X)
    return SIM
end

#simulate if the country continue to be exclulded after default
function simulate_exclusion(θ)
    dists = Categorical([θ , 1.0- θ])
    sim = rand(dists, 1)
    return sim[1]
end
"""
Simulate the Arellano model, given the policy functions
INPUTS: (instructure of Model solution and model params)

P: stochatsic matrix ofthe Markov Chain
ϵ: values of the income schock for the given Markov Chain
grid: gird for assets
r: interest rate
θ: paramter for the exclusion form the finacial arket after default
grid: gird for assets
Def_mat: default region
q: price matrix
Policy: savings policy function

Params:
n_sim: number of independent simulatioins
t_sim: size of each simulation
burnout: number of first obervations to ignore
"""
function simulate_Arellano(Model::Arellano_params, Model_solution; burnout = 1000, t_sim =11000, n_sim = 1000)
    #unpack policy functions and paramters

    q = Model_solution[1]
    Def_mat = Model_solution[2]
    Policy = Model_solution[3]

    ϵ  = Model.ϵ
    P = Model.P
    grid = Model.grid
    r = Model.r
    y_hat = Model.y_hat
    θ = Model.θ

    #allocate memory
    A = ones(n_sim, t_sim) #assets history
    A_int = zeros(Int64, n_sim, t_sim) #position of assets on the grid
    Y = zeros(n_sim, t_sim) # output history
    Trade_B = zeros(n_sim, t_sim) # trade balance history
    C = zeros(n_sim, t_sim) # consumption history
    D = zeros(n_sim, t_sim) #Defaults histiry
    D_state = zeros(n_sim, t_sim) #exclusion from the financial market histiry
    R = zeros(n_sim, t_sim) #q hiatory
    SIM = zeros(Int64, n_sim, t_sim) #simulated income process

    #find index with no assets
    zero_ind = searchsortedfirst(grid, 0.0)

    #define starting values
    y_0 = 11 # start with 11th state
    Y_0 = ϵ[y_0]*ones(n_sim) #output start value
    A_0 = zeros(n_sim) #start with 0 assets

    Y[:,1] = Y_0
    A[:,1] = A_0
    A_int[:,1] .= zero_ind

    Default_flag = 0.0 # flag if the country was excluded from financial market in the previous period
    #simulate income process
    for i in 1:n_sim
        SIM[i,:] = mc_sample_path(P)
    end

    #now use simulated income process to find trade balance, output, assets and q history
    for t in 2:t_sim
        for i in 1:n_sim
            #case when country was excluded from financial market in the previous period
           if(Default_flag==1.0)
                Default_flag = simulate_exclusion(θ) -1.0
                if(Default_flag==1.0)
                    Y[i,t] = y_d(ϵ[SIM[i,t]], y_hat)
                    C[i,t] = Y[i,t]
                    Trade_B[i,t] = (Y[i,t] - C[i,t])/Y[i,t]
                    A_int[i,t] = zero_ind
                    A[i,t] = 0.0
                    R[i,t] = -1000 #just to
                    D[i,t] = 0.0
                    D_state[i,t] = 1.0
                else

                    Y[i,t] = ϵ[SIM[i,t]]

                    A_int[i,t] = Policy[SIM[i,t], zero_ind]
                    A[i,t] = grid[A_int[i,t]]
                    C[i,t] = Y[i,t] + A[i,t-1] - q[SIM[i,t],A_int[i,t]]*A[i,t]
                    R[i,t] = 1.0/q[SIM[i,t],A_int[i,t]] - 1.0
                    D[i,t] = 0.0
                    D_state[i,t] = 0.0
                    Trade_B[i,t] = (Y[i,t] - C[i,t])/Y[i,t]
                end
            else
                if(Def_mat[SIM[i,t], A_int[i,t-1]]==1.0) #case of default
                    Y[i,t] = y_d(ϵ[SIM[i,t]] , y_hat)
                    C[i,t] = Y[i,t]
                    A_int[i,t] = zero_ind
                    A[i,t] = 0.0
                    R[i,t] = -1000 #1.0/q[SIM[i,t],A_int[i,t]] - 1.0
                    D[i,t] = 1.0
                    D_state[i,t] = 1.0
                    Default_flag = 1.0
                    Trade_B[i,t] = (Y[i,t] - C[i,t])/Y[i,t]
                else
                    Y[i,t] = ϵ[SIM[i,t]]
                    A_int[i,t] = Policy[SIM[i,t], A_int[i,t-1]]
                    A[i,t] = grid[A_int[i,t]]
                    C[i,t] = Y[i,t] + A[i,t-1] - q[SIM[i,t],A_int[i,t]]*A[i,t]
                    R[i,t] = 1.0/q[SIM[i,t],A_int[i,t]] - 1.0
                    D[i,t] = 0.0
                    D_state[i,t] = 0.0
                    Trade_B[i,t] = (Y[i,t] - C[i,t])/Y[i,t]
                end
           end
        end
    end

    #compute stats for defaults
    n_defaults = sum(D[:, burnout:t_sim])

    non_defaults = D_state[:,burnout:t_sim].<1.0
    #default probability
    Def_prob  =  1.0 - (1.0-n_defaults/(sum(non_defaults)))^4.0

    #stats after default:
    Y_ab = Y[:,burnout:t_sim]
    R_ab = R[:,burnout:t_sim]
    A_ab = A[:,burnout:t_sim]
    D_ab = D[:,burnout:t_sim]
    Trade_B_ab = Trade_B[:,burnout:t_sim]
    D_state_ab = D_state[:,burnout:t_sim]
    C_ab = C[:,burnout:t_sim]

    #choose number of defaults to compute statistics
    n_chosen_def = 1000
    #allocate memory for the staoistics
    pre_def_stats_Y = zeros(n_chosen_def , 75)
    pre_def_stats_R = zeros(n_chosen_def , 75)
    pre_def_stats_TB = zeros(n_chosen_def , 75)
    pre_def_stats_C = zeros(n_chosen_def , 75)
    pre_def_stats_B = zeros(n_chosen_def , 75)
    pre_def_stats_NB = zeros(n_chosen_def , 75)

    #detrended staistics
    detrend_Y_stats = zeros(n_chosen_def , 74)
    detrend_R_stats = zeros(n_chosen_def , 74)
    detrend_C_stats = zeros(n_chosen_def , 74)
    detrend_TB_stats = zeros(n_chosen_def , 74)

    #standard deviations and correlations
    std_Y = zeros(n_chosen_def)
    std_R = zeros(n_chosen_def)
    std_C = zeros(n_chosen_def)
    std_TB = zeros(n_chosen_def)

    corr_R_Y = zeros(n_chosen_def)
    corr_R_C = zeros(n_chosen_def)
    corr_R_TB = zeros(n_chosen_def)

    corr_Y_R = zeros(n_chosen_def)
    corr_Y_C = zeros(n_chosen_def)
    corr_Y_TB = zeros(n_chosen_def)

    #compute statistics for 74 periods before the default, for n_def defaults (without any defults in the 74 periods before)
    iter = 0
    for i in 1:n_sim, t in 76:t_sim-burnout
            if(D_ab[i,t]==1 && sum(D_state_ab[i,t-75:t-1 ])== 0.0) #if default happen
                 iter =iter+1
                 pre_def_stats_Y[iter,:] = Y_ab[i, t-74:t]
                 pre_def_stats_R[iter,:] = (1.0.+R_ab[i, t-74:t]).^4.0.- (1.0+r).^4.0
                 pre_def_stats_B[iter,:] = A_ab[i, t-75:t-1]
                 pre_def_stats_NB[iter,:] = A_ab[i, t-74:t] #new debt
                 pre_def_stats_TB[iter,:] = Trade_B_ab[i, t-74:t]
                 pre_def_stats_C[iter,:] = C_ab[i, t-74:t]
                #detrend some of teh staistucs
                 detrend_Y_stats[iter,:] = detrend(ones(74),pre_def_stats_Y[iter,1:74])
                 detrend_R_stats[iter,:] = pre_def_stats_R[iter,1:74]
                 detrend_C_stats[iter,:] = detrend(ones(74),pre_def_stats_C[iter,1:74])
                 detrend_TB_stats[iter,:] = pre_def_stats_TB[iter,1:74]

                 std_Y[iter] = std(detrend_Y_stats[iter,:])
                 std_R[iter] = std(detrend_R_stats[iter,:])
                 std_C[iter] = std(detrend_C_stats[iter,:])
                 std_TB[iter] = std(detrend_TB_stats[iter,:])

                 corr_R_Y[iter] = cor(detrend_Y_stats[iter,:], detrend_R_stats[iter,:])
                 corr_R_C[iter] = cor(detrend_C_stats[iter,:], detrend_R_stats[iter,:])
                 corr_R_TB[iter] = cor(detrend_TB_stats[iter,:], detrend_R_stats[iter,:])

                 corr_Y_R[iter] = corr_R_Y[iter]
                 corr_Y_C[iter] = cor(detrend_C_stats[iter,:], detrend_Y_stats[iter,:])
                 corr_Y_TB[iter] = cor(detrend_TB_stats[iter,:], detrend_Y_stats[iter,:])

            end
            if(iter == n_chosen_def)

                break
            end
    end

    Calibration_target = [("Default probability", Def_prob), ("Trade balance voltality", std(Trade_B_ab)), ("Debt service to GDP", -mean(pre_def_stats_B[:,1:74] ./pre_def_stats_Y[:,1:74]))]

    Table_4 = hcat([mean(pre_def_stats_R[74]),mean(std_R), mean(corr_R_Y), nothing],
    [mean(pre_def_stats_Y[:,75]), mean(std_Y), nothing, mean(corr_R_Y)],
    [mean(pre_def_stats_C[:,75]), mean(std_C), mean(corr_Y_C) , mean(corr_R_C)],
    [mean(pre_def_stats_TB[:,74]),mean(std_TB), mean(corr_R_TB), mean(corr_Y_TB)])

    #decrease of the output due to defult
    OUT_D_Def = mean(log.(pre_def_stats_Y[:,75]).-mean( log.(pre_def_stats_Y[:,1:74])))
    Debt_GDP = mean(pre_def_stats_NB[:,1:74]./pre_def_stats_Y[:,1:74] )
    mean_spread = mean(detrend_R_stats)
    Table_4_b = [OUT_D_Def, Debt_GDP, mean_spread, Def_prob]

    println(Calibration_target)

    return (Calibration_target, Table_4, Table_4_b)
end

"""
SImulate the economy behaviour for a given shocks sequence
"""
function History_simulation(Model::Arellano_params, Model_solution, data)
    ϵ  = Model.ϵ
    P = Model.P
    grid = Model.grid
    r = Model.r
    y_hat = Model.y_hat
    θ = Model.θ

    q = Model_solution[1]
    Def_mat = Model_solution[2]
    Policy = Model_solution[3]
    t_sim = length(data)
    #find index with no assets
    zero_ind = searchsortedfirst(grid, 0.0)

    #allocate memory
    A = zeros(t_sim) #assets history
    A_int = zero_ind*ones(Int64, t_sim) #position of assets on the grid
    Y = zeros(t_sim) # output history
    Trade_B = zeros(t_sim) # trade balance history
    C = zeros(t_sim) # consumption history
    D = zeros(t_sim) #Defaults histiry
    D_state = zeros(t_sim) #exclusion from the financial market histiry
    R = zeros(t_sim) #q hiatory
    SIM = data[:,1] #simulated income process
    println(SIM)
    Y[1] = ϵ[SIM[1]]
    A_int[1] = Policy[SIM[1], zero_ind]
    A[1] = grid[A_int[1]]
    R[1] = 0.0

    Default_flag = 0.0 # flag if the country was excluded from financial market in the previous period
    #simulate income process

    #now use simulated income process to find trade balance, output, assets and q history
    for t in 2:t_sim

            #case when country was excluded from financial market in the previous period
           if(Default_flag==1.0)
                Default_flag = simulate_exclusion(θ) -1.0
                if(Default_flag==1.0)
                    Y[t] = y_d(ϵ[SIM[t]], y_hat)
                    C[t] = Y[t]
                    Trade_B[t] = (Y[t] - C[t])/Y[t]
                    A_int[t] = zero_ind
                    A[t] = 0.0
                    R[t] = 0.017 #just to
                    D[t] = 0.0
                    D_state[t] = 1.0
                else

                    Y[t] = ϵ[SIM[t]]

                    A_int[t] = Policy[SIM[t], zero_ind]
                    A[t] = grid[A_int[t]]
                    C[t] = Y[t] + A[t-1] - q[SIM[t],A_int[t]]*A[t]
                    R[t] = 1.0/q[SIM[t],A_int[t]] - 1.0
                    D[t] = 0.0
                    D_state[t] = 0.0
                    Trade_B[t] = (Y[t] - C[t])/Y[t]
                end
            else
                if(Def_mat[SIM[t], A_int[t-1]]==1.0) #case of default
                    Y[t] = y_d(ϵ[SIM[t]] , y_hat)
                    C[t] = Y[t]
                    A_int[t] = zero_ind
                    A[t] = 0.0
                    R[t] = 0.017 #1.0/q[SIM[i,t],A_int[i,t]] - 1.0
                    D[t] = 1.0
                    D_state[t] = 1.0
                    Default_flag = 1.0
                    Trade_B[t] = (Y[t] - C[t])/Y[t]
                else
                    Y[t] = ϵ[SIM[t]]
                    A_int[t] = Policy[SIM[t], A_int[t-1]]
                    A[t] = grid[A_int[t]]
                    C[t] = Y[t] + A[t-1] - q[SIM[t],A_int[t]]*A[t]
                    R[t] = 1.0/q[SIM[t],A_int[t]] - 1.0
                    D[t] = 0.0
                    D_state[t] = 0.0
                    Trade_B[t] = (Y[t] - C[t])/Y[t]
                end
           end

    end
    R = (1.0.+copy(R)).^4.0 .-(1.0+r)^4.0
    println(A)
    return Y, Trade_B, R, D
end
"""
Function for saving the equilbrium for the replication
"""
function Save_Arellano(Model::Arellano_params, Model_solution, Sim_results, Grid)
    #write tables
    println("Calibration")
    println(Sim_results[1])
    println("table 4")
    println(Sim_results[2])
    println("table 4b")
    println( Sim_results[3])

    q = Model_solution[1]
    Def_mat = Model_solution[2]
    Policy = Model_solution[3]
    V_f = Model_solution[4]

    #write plots

    p1 = plot(Grid[140:180], hcat(q[9,140:180], q[13,140:180]),label = ["y_low" "y_high"], ylabel = "q", xlabel = "B'")
    r_y_low = min.(1.1, 1.0./q[9,140:180]) #cut the r for very high values
    r_y_high = min.(1.1, 1.0./q[13,140:180])

    p2 = plot(Grid[140:180], hcat(r_y_low.-1.0, r_y_high.-1.0),label = ["y_low" "y_high"], ylabel = "1/q", xlabel = "B'")
    plot(p1,p2,layout = (1, 2))
    png("plot_q")

    p3 = plot(Grid[140:180], hcat([Grid[x] for x in Policy[9,140:180]], [Grid[x] for x in Policy[13,140:180]]),label = ["y_low" "y_high"], ylabel = "B'", xlabel = "B")

    p4 = plot(Grid[140:180], hcat(V_f[9,140:180], V_f[13,140:180]),label = ["y_low" "y_high"], ylabel = "V_f", xlabel = "B'")
    plot(p3,p4,layout = (1, 2))
    png("figure_4")
    #read the file for figure 5 replication
    input = CSV.File("shocks_y.csv",  header=false)
    data_input = zeros(35)
    for  i in 1:35
        data_input[i] = input[i,1][1]
    end

    input_int = convert(Array{Int64},data_input)
    Y, TB, R, D = History_simulation(Model, Model_solution, input_int)
    println("default series ", D)

   #write the data for figure 5
    outfile = "results.txt"
    open(outfile, "w") do f
        for i in 1:length(Y)
            println(f, Y[i], " ", R[i], " ", TB[i])
        end
    end

end

function Solve_Arellano(ρ, η, step, ϵ_n, B_min, B_max)
    Model = Arellano_params(ρ, η, step, ϵ_n, B_min, B_max)
    SolveRAEq(Model, Solve_eq_Arellano, simulate_Arellano, Save_Arellano )

end


text:
module DISSEQT

# temporary submodule
include("Levenshtein/Levenshtein.jl")

# submodules
include("AnnotatedArrays/AnnotatedArrays.jl")
include("BamReader/BamReader.jl")
include("Kruskal/Kruskal.jl")
include("SynapseTools/SynapseTools.jl")
include("AlignUtils/AlignUtils.jl")
include("Plots/Plots.jl")

using JLD
using StatsBase # for sample
using Statistics
using LinearAlgebra
using DataStructures
using DataFrames
using FASTX

using .AnnotatedArrays
using .BamReader

using .Kruskal

# TODO: get rid of these dependencies - only used by estimatesamplemix
using Convex
using SCS

include("fasta.jl")
include("consensus.jl")
include("bamcounts.jl")
include("codoncounts.jl")
include("codonfrequency.jl")
include("nuccounts.jl")
include("nucfrequency.jl")
include("swarm.jl")
include("timeseries.jl")
include("readcoverage.jl")
include("features.jl")
include("limitofdetection.jl")
include("estimatesamplemix.jl")
include("sparsepca.jl")
include("floydwarshall.jl")
include("mds.jl")
include("hierarchicalclustering.jl")
include("medianfilter.jl")
include("fitnessmodels.jl")
include("isomap.jl")

export
	Swarm,
	SwarmSegment,
	consensus,
	loadfasta,
	savefasta,
	seq2codons,
	codon2aa,
	codonqualitycount,
	nucqualitycount,
	qualityfilter!,
	removeambiguous!,
	mlcodonfreqs,
	mlnucfreqs,
	positions,
	coverage,
	positionfilter!,
	segments,
	segment,
	segmentinfo,
	segmentnames,
	segmentlengths,
	consensus,
	meta,
	processingmeta,
	loadswarm,
	timeseriesindices,
	timeseries,
	timeseriesat,
	timeseriesmean,
	distancecurve,
	distancecurveequals,
	bifurcationtime,
	evaluatedistancecurve,
	distancecurvemean,
	readcoverage,
	featurepositions,
	featuresegment,
	featureat,
	codonposition,
	annotatevariants!,
	variantdescriptions,
	limitofdetection,
	estimatemix,
	sparsepcal1,
	floydwarshall!,
	floydwarshall,
	mds,
	hierarchicalclustering,
	samplesinclusters,
	medianfilter,
	squareddistances,
	LandscapeModel,
	landscapekernelwidth,
	leaveoneoutlandscapekernelwidth,
	NearestNeighborModel,
	GroupModel,
	predictfitness,
	leaveoneoutpredict,
	varianceunexplained,
	kruskalmds,
	makeconnected!,
	isomapmatrix,
	isomap,
	kruskalisomap

end


text:
# plotting tools
include("generateVTK.jl")
include("vm2d_particle.jl")

# # set freestream
# function Uinf(x::Array{Float64,1})
#     Uinf = [0.0,0.0,0.0]
#     return Uinf
# end

# plot x-y view
function plot_xy(particles::Array{Particle,1}; velocity::Bool=false, xs=nothing, ys=nothing, zs=nothing, circulation::Bool=false, title::String="", save::Bool=false, name::String="default_sim", pause::Float64=-1.0)
    # plot
    plt.clf()
    plt.cool()
    for particle in particles
        plt.scatter(particle.x[1],particle.x[2],c=la.norm(particle.gamma))#la.norm(particle.gamma))
    end

    # velocity field
    if velocity
        gridxv, gridyv, gridzv, vx, vy, vz = vfield(particles, xs, ys, zs)
        plt.quiver(gridxv, gridyv, vx, vy)
    end

    # circulation vectors
    if circulation
        gridxc, gridyc, gridzc, gx, gy, gz = gfield(particles)
        plt.quiver(gridxc, gridyc, gx, gy)
    end

    plt.title(title)
    plt.xlabel("x")
    plt.ylabel("y")
    plt.axis("equal")

    # save frame
    if save==true
        plt.savefig("sim/"*name*".png")
    end

    # pause frame
    if pause > 0.0
        plt.pause(pause)
    end
end

# plot y-z view
function plot_yz(particles::Array{Particle,1}; velocity::Bool=false, xs=nothing, ys=nothing, zs=nothing, circulation::Bool=false, title::String="", save::Bool=false, name::String="default_sim", pause::Float64=-1.0)
    # plot
    plt.clf()
    plt.cool()
    for particle in particles
        plt.scatter(particle.x[2],particle.x[3],c=la.norm(particle.gamma))
    end

    # velocity field
    if velocity
        gridxv, gridyv, gridzv, vx, vy, vz = vfield(particles, xs, ys, zs)
        plt.quiver(gridyv, gridzv, vy, vz)
    end

    # circulation vectors
    if circulation
        gridxc, gridyc, gridzc, gx, gy, gz = gfield(particles)
        plt.quiver(gridyc, gridzc, gy, gz)
    end

    plt.title(title)
    plt.xlabel("y")
    plt.ylabel("z")
    plt.axis("equal")

    # save frame
    if save==true
        plt.savefig("sim/"*name*".png")
    end

    # pause frame
    if pause > 0.0
        plt.pause(pause)
    end
end

# plot y-z view
function plot_2d(particles::Array{Particle,1}, dims::String; velocity::Bool=false, xs=nothing, ys=nothing, zs=nothing, circulation::Bool=false, title::String="", save::Bool=false, name::String="default_sim", pause::Float64=-1.0)
    # set up axes
    xaxis = collect(dims)[1]
    yaxis = collect(dims)[2]
    if xaxis == 'x'
        xindex = 1
    elseif xaxis == 'y'
        xindex = 2
    elseif xaxis == 'z'
        xindex = 3
    end
    if yaxis == 'x'
        yindex = 1
    elseif yaxis == 'y'
        yindex = 2
    elseif yaxis == 'z'
        yindex = 3
    end

    # plot
    plt.clf()
    plt.cool()
    for particle in particles
        plt.scatter(particle.x[xindex],particle.x[yindex],c=la.norm(particle.gamma))
    end

    # velocity field
    if velocity
        gridxv, gridyv, gridzv, vx, vy, vz = vfield(particles, xs, ys, zs)
        gridv = [gridxv, gridyv, gridzv]
        v = [vx, vy, vz]
        plt.quiver(gridv[xindex], gridv[yindex], v[xindex], v[yindex])
    end

    # circulation vectors
    if circulation
        gridxc, gridyc, gridzc, gx, gy, gz = gfield(particles)
        gridc = [gridxc, gridyc, gridzc]
        g = [gx, gy, gz]
        plt.quiver(gridc[xindex], gridc[yindex], g[xindex], g[yindex])
    end

    plt.title(title)
    plt.xlabel(string(xaxis))
    plt.ylabel(string(yaxis))
    plt.axis("equal")

    # save frame
    if save==true
        plt.savefig("sim/"*name*".png")
    end

    # pause frame
    if pause > 0.0
        plt.pause(pause)
    end
end

# plot 3-D view
function plot_3d(particles::Array{Particle,1}; title::String="", save::Bool=false, name::String="default_sim", pause::Float64=-1.0)
    # plot
    plt.clf()
    plt.cool()
    for particle in particles
        plt.scatter(particle.x[2],particle.x[3],c=la.norm(particle.gamma))
    end
    plt.title(title)
    plt.xlabel("y")
    plt.ylabel("z")
    plt.axis("equal")

    # save frame
    if save==true
        plt.savefig("sim/"*name*".png")
    end

    # pause frame
    if pause > 0.0
        plt.pause(pause)
    end
end

Uinf(x) = [0.0,0.0,0.0] # default freestream velocity function
# animate
function animate(particles,timestart,timestep,numtimesteps; Uinf=Uinf)
    currenttime = timestart
    # set up plot
    plt.plot()
    plt.cool()
    plt.axis("equal")
    for particle in particles
        plt.scatter(particle.x[1],particle.x[2],c=particle.gamma[3])
    end
    plt.title("time: "*string(round(currenttime; digits=3)))
    # animate
    for i in range(1,length=numtimesteps)
        ## advance timestep
        advance(particles,timestep,Uinf)
        currenttime += timestep
        ## update plot
        plot_xy(particles; title="time: "*string(round(currenttime; digits=3)))
    end
end

"""
`packVTK(particles::Array{Particle,1},xs,ys,zs)`

Uses Eduardo Alvarez' `generateVTK()` function to generate VTK text files for visualizations with paraview.

Optional arguments `xs`, `ys`, and `zs` define the cartesian grid for visualizing the velocity field.
"""
function packVTK(filename, particles::Array{Particle,1}; xs=Nothing, ys=Nothing, zs=Nothing, num=nothing)
    # make array of points defining vortex particles
    points = Array{Float64,1}[]
    for particle in particles
        push!(points,particle.x)
    end

    # compile circulation and particle velocity vector data
    gammavecs = Array{Float64,1}[]
    particlevelocities = Array{Number,1}[]
    for particle in particles
        push!(gammavecs, particle.gamma)
        push!(particlevelocities, particle.v)
    end
    gammadata = Dict(
        "field_name" => "Circulation",
        "field_type" => "vector",
        "field_data" => gammavecs
    )
    particlevdata = Dict(
        "field_name" => "ParticleVelocity",
        "field_type" => "vector",
        "field_data" => particlevelocities
    )

    # compile all vtk data for particles
    point_data = [gammadata, particlevdata]

    # generate VTK file for particles
    generateVTK(filename, points; point_data = point_data, path = "vtk", num=num)

    if xs != Nothing
        # compile velocity field data
        gridxv, gridyv, gridzv, vx, vy, vz = vfield(particles, xs, ys, zs)
        velocityvecs = Array{Number,1}[]
        velocitycoordinates = Array{Number,1}[]
        for i in range(1, length=length(gridxv))
            push!(velocityvecs, [vx[i], vy[i], vz[i]])
            push!(velocitycoordinates, [gridxv[i], gridyv[i], gridzv[i]])
        end
        velocitydata = Dict(
            "field_name" => "Velocity",
            "field_type" => "vector",
            "field_data" => velocityvecs
        )

        # compile all vtk data for velocity field
        velocity_data = [velocitydata]

        # generate VTK file for velocity field
        generateVTK(filename*"_velocityfield", velocitycoordinates; point_data = velocity_data, path = "vtk", num=num)
    end

end


text:
using BioSeq
using Base.Test

@test nucleotide('A') == nucleotide(65) == nucleotide(0x41)
@test aminoacid('A') == aminoacid(65) == aminoacid(0x41)

@test nt"AC" == nucleotide([65; 67]) == [ nucleotide('A'); nucleotide('C') ]
@test aa"AC" == aminoacid([65; 67]) == [ aminoacid('A'); aminoacid('C') ]

@test nt"ACTG" == nucleotide("ACTG")
@test aa"ACTG" == aminoacid("ACTG")

@test aminoacid('A') + 32 == 32 + aminoacid('A') == aminoacid('a')
@test nucleotide('A') + 32 == 32 + nucleotide('A') == nucleotide('a')

@test aminoacid('a') - 32 == -32 + aminoacid('a') == aminoacid('A')
@test nucleotide('a') - 32 == -32 + nucleotide('a') == nucleotide('A')

# # but string interpolation is only supported for Julia 0.1
# seq = nt"TG"
# @test nt"AC$seq" == nt"ACTG"
#
# seq = aa"HM"
# @test aa"AC$seq" == aa"ACHM"

@test nt"ACTG" == nucleotide("ACTG")
@test aa"ACTG" == aminoacid("ACTG")

seq   = nt"ACTG"
seqII = nt"TGAC"

@test [ seq'; seqII' ] == [ seq seqII ]' == vcat( seq', seqII' ) == nucleotide([ "ACTG".data'; "TGAC".data' ])

@test nt"ACUG" == dna2rna(nt"ACTG")
@test nt"ACTG" == rna2dna(nt"ACUG")

@test nt"TGAC" == complement(nt"ACTG",DNA_COMPLEMENT)
@test nt"UGAC" == complement(nt"ACUG",RNA_COMPLEMENT)

@test nt"CAGT" == reversecomplement(nt"ACTG",DNA_COMPLEMENT)
@test nt"CAGU" == reversecomplement(nt"ACUG",RNA_COMPLEMENT)

@test aa"acmh" == lowercase(aa"ACmh")
@test nt"actg" == lowercase(nt"ACtg")

@test aa"ACMH" == uppercase(aa"ACmh")
@test nt"ACTG" == uppercase(nt"ACtg")

@test aminoacid('A') == uppercase(aminoacid('a'))
@test nucleotide('A') == uppercase(nucleotide('a'))

@test aminoacid('a') == lowercase(aminoacid('A'))
@test nucleotide('a') == lowercase(nucleotide('A'))

@test ismatch(aar"X",aa"ACMH") == true
@test ismatch(ntr"N",nt"ACTG") == true

@test match(aar"X",aa"ACMH").match == "A"
@test match(ntr"N",nt"ACTG").match == "A"

@test match(aar"X",aa"ACMH",2).match == "C"
@test match(ntr"N",nt"ACTG",2).match == "C"

@test search(aa"ACAC",'A') == 1
@test search(nt"ACAC",'A') == 1

@test search(aa"ACAC",'A',2) == 3
@test search(nt"ACAC",'A',2) == 3

@test search(aa"ACAC",aar"A") == 1:1
@test search(nt"ACAC",ntr"A") == 1:1

@test search(aa"ACAC",aar"A",2) == 3:3
@test search(nt"ACAC",ntr"A",2) == 3:3

@test [x.match for x in [eachmatch(aar"A",aa"ACAC")...]] == [ "A"; "A" ]
@test [x.match for x in [eachmatch(ntr"A",nt"ACAC")...]] == [ "A"; "A" ]

@test repeat(aa"A",3) == aa"AAA"
@test repeat(nt"A",3) == nt"AAA"

@test repeat(aa"AC",2) == aa"ACAC"
@test repeat(nt"AC",2) == nt"ACAC"

@test swap(aa"ACMH",'H','X') == aa"ACMX"
@test swap(nt"ACTG",'T','U') == nt"ACUG"

@test replace(aa"ACAC",r"A",aa"M") == aa"MCMC"
@test replace(aa"ACAC",r"AC",aa"M") == aa"MM"
@test replace(aa"ACAC",r"ACA",aa"M") == aa"MC"
@test replace(aa"ACAC",r"A",aa"TH") == aa"THCTHC"

@test sum(isin(nt"ACTGactg", NUCLEIC_IUPAC )) == 8
@test sum(isin(nt"ECTGactg", NUCLEIC_IUPAC )) == 7

@test swap(nt"ACTG", DNA_COMPLEMENT) == nt"TGAC"
@test swap(aa"MHAC", AMINO_1LETTER_TO_3) == [ "MET"; "HYS"; "ALA"; "CYS" ]

## Test for Nucleotide2bitSeq ##

##	A	C	T	G
## b1	0	0	1	1
## b2	0	1	0	1

@test nt2"ACTG".b1 == [false, false, true, true]
@test nt2"ACTG".b2 == [false, true, false, true]

@test nt2"ACTG" == nucleotide2bit("ACTG")
@test nt2"ACTG"[2] == 'C'

@test nt2"ACTG" == nt2"actg" == nt2"ACUG" == nt2"acug"

@test nucleotide(Nucleotide2bitSeq(10)) == nt"AAAAAAAAAA"

@test length(Nucleotide2bitSeq(10)[1:3])==3

@test complement(nt2"ACTG") == nt2"TGAC"
@test reversecomplement(nt2"ACTG") == nt2"CAGT"

@test percentGC(nt2"AACC") == 0.5

seq = nt2"ACTG"
ntseq = nucleotide(seq)
@test isadenine(seq) == (ntseq .== 'A')
@test iscytosine(seq) == (ntseq .== 'C')
@test isthymine(seq) == (ntseq .== 'T')
@test isguanine(seq) == (ntseq .== 'G')
@test isweak(seq) == ((ntseq .== 'A') | (ntseq .== 'T'))
@test isstrong(seq) == ((ntseq .== 'C') | (ntseq .== 'G'))
@test ispyrimidine(seq) == ((ntseq .== 'C') | (ntseq .== 'T'))
@test ispurine(seq) == ((ntseq .== 'A') | (ntseq .== 'G'))

## Test for Prosite Patterns ##

@test ismatch(prosite"M-[ALT]",aa"HML")
@test ismatch(prosite"M-{ALT}",aa"HMM")
@test ismatch(prosite"M-{ALT}",aa"HMM")
@test ismatch(prosite"M-x(2,10)-M",aa"MALHM")
@test ismatch(prosite"x(3)",aa"ALM")
@test ismatch(prosite"x(2,4)",aa"ALM")
@test ismatch(prosite"x(3)",aa"AAA")
@test ismatch(prosite"<A-L",aa"ALMH")
@test ismatch(prosite"<A-L",aa"LALMH") == false
@test ismatch(prosite"F-[GSTV]-P-R-L-[G>]",aa"FVPRLGH")
@test ismatch(prosite"F-[GSTV]-P-R-L-[G>]",aa"FVPRL")
@test ismatch(prosite"F-[GSTV]-P-R-L-[G>]",aa"FVPRLLL") == false
@test ismatch(prosite"M-A-S-K-E",aa"MASKE")
@test ismatch(prosite"MASKE",aa"MASKE")
@test ismatch(prosite"<{C}*>",aa"FVLRPHM")
@test ismatch(prosite"<{C}*>",aa"FVLRPHCM") == false
@test ismatch(prosite"[AC]-x-V-x(4)-{ED}",aa"AHVRLPLG")
@test ismatch(prosite"<A-x-[ST](2)-x(0,1)-V",aa"AHTSV")

## Test Translate ##

@test codon2aa(nt"GTG",1) == 'V'
@test isstart(nt"GTG",2)
@test isstop(nt"TGA",1)
@test isstart(nt"GTGGTAATGTGAAAGTAG",2) == reverse( isstop(nt"GTGGTAATGTGAAAGTAG",1) )
@test translateCDS(nt"GTGGTAATGTGAAAGTAG",2) == aa"MVMWK"
@test translatetostop(nt"GTGGTAATGTGAAAGTAG",1) == aa"VVM"

@test translate(nt"GTGGTAATGTGAAAGTAG",1) == aa"VVM*K*"
@test translate(nt2"GTGGTAATGTGAAAGTAG",1) == aa"VVM*K*"
@test translate(nt8"GTGGTAATGTGAAAGTAG",1) == aa"VVM*K*"

## Test for 8 bit Bit-Level Coding Scheme Nucleotides ##

@test nt8"ACTG" == nt8"actg" == nt8"ACUG" == nt8"acug"
@test nt8"ACTG" == nt"ACTG"

@test nt8"ACTG" == nucleotide8bit("ACTG")
@test nt8"ACTG"[2] == 'C'

@test nucleotide(nt8"ACTG") == nt"ACTG"

@test length(nt8"AAACCCTTT"[1:3])==3

@test percentGC(nt8"AACC") == 0.5

seq = nt8"ACTG"
ntseq = nucleotide(seq)
@test isadenine(seq) == (ntseq .== 'A')
@test iscytosine(seq) == (ntseq .== 'C')
@test isthymine(seq) == (ntseq .== 'T')
@test isguanine(seq) == (ntseq .== 'G')
@test ispyrimidine(seq) == ((ntseq .== 'C') | (ntseq .== 'T'))
@test ispurine(seq) == ((ntseq .== 'A') | (ntseq .== 'G'))

# # TO DO:
# @test nt8"ACTG" == nt"ACUG"
# @test nt8"ACTG" == nt2"ACTG"
# @test complement(nt8"ACTG") == nt8"TGAC"
# @test reversecomplement(nt8"ACTG") == nt8"CAGT"
# @test isweak(seq) == ((ntseq .== 'A') | (ntseq .== 'T'))
# @test isstrong(seq) == ((ntseq .== 'C') | (ntseq .== 'G'))


text:
export dropoutForward
import Base: unsafe_convert
using Knet.KnetArrays: DevArray
using AutoGrad: AutoGrad, @primitive1
using CUDA: CuArray

using CUDA.CUDNN: handle,
    cudnnDropoutForward,
    cudnnDropoutBackward,
    cudnnDropoutGetStatesSize,
    cudnnDropoutGetReserveSpaceSize,
    cudnnDropoutDescriptor_t,
        cudnnCreateDropoutDescriptor,
        unsafe_cudnnSetDropoutDescriptor,
        cudnnGetDropoutDescriptor,
        cudnnRestoreDropoutDescriptor,
        cudnnDestroyDropoutDescriptor

mutable struct DropoutDescriptor; ptr; end
unsafe_convert(::Type{<:Ptr}, dd::DropoutDescriptor)=dd.ptr

# cudnnDropoutForward() doc says:
# "This function should not be running concurrently with another cudnnDropoutForward() function using the same states."
# So I am going to assume using a single buffer is fine until we have to deal with concurrency.
# TODO: fix this based on default_rng in Random
GLOBAL_DROPOUT_STATE = Ref{CuArray{Int,1}}()

# To debug gradients set GLOBAL_DROPOUT_SEED[] >= 0 for all dropout operations
GLOBAL_DROPOUT_SEED = Ref{Int}(-1)

# cudnnSetDropoutDescriptor is expensive, so let's cache the descriptors based on dropout probability
GLOBAL_DROPOUT_DESCRIPTORS = Dict{Cfloat,DropoutDescriptor}()

@primitive1 dropoutForward(x; dropout=0.5, reserveSpace=Ref{Any}(), o...),dy,y  dropoutBackward(dy; dropout=dropout, reserveSpace=reserveSpace)
@primitive1 dropoutBackward(dy;o...)  throw(MethodError(back,dropoutBackward))

function dropoutForward(x::R; dropout=0.5, reserveSpace=Ref{Any}()) where {T,R<:DevArray{T}}
    y, td, dd = similar(x), TD(T, (1,1,length(x),1)), DD(dropout)
    if GLOBAL_DROPOUT_SEED[] >= 0
        # This is a very expensive call (40x dropout), so only use for debugging
        @warn "Knet.CUDNN.GLOBAL_DROPOUT_SEED >= 0: calling expensive cudnnSetDropoutDescriptor" maxlog=1
        @cudnn_retry unsafe_cudnnSetDropoutDescriptor(dd, handle(), dropout, GLOBAL_DROPOUT_STATE[], sizeof(GLOBAL_DROPOUT_STATE[]), GLOBAL_DROPOUT_SEED[])
    end
    if !isassigned(reserveSpace)
        # reserveSpace is 1/8 of tensor size and is specific to this one call
        rss = Csize_t[0]; cudnnDropoutGetReserveSpaceSize(td, rss)
        reserveSpace[] = CuArray{Int}(undef, (rss[1]-1)÷sizeof(Int)+1)
    end
    cudnnDropoutForward(handle(), dd, td, x, td, y, reserveSpace[], sizeof(reserveSpace[]))
    return y
end

function dropoutBackward(dy::R; dropout, reserveSpace) where {T,R<:DevArray{T}}
    #@show summary.((dy,dropoutDesc[],reserveSpace[]))
    dx, td, dd = similar(dy), TD(T, (1,1,length(dy),1)), DD(dropout)
    cudnnDropoutBackward(handle(), dd, td, dy, td, dx, reserveSpace[], sizeof(reserveSpace[]))
    return dx
end

function DD(dropout::Real=0.5)
    get!(GLOBAL_DROPOUT_DESCRIPTORS, Cfloat(dropout)) do
        if !isassigned(GLOBAL_DROPOUT_STATE)
            ssize = Csize_t[0]; cudnnDropoutGetStatesSize(handle(), ssize)
            GLOBAL_DROPOUT_STATE[] = CuArray{Int}(undef, (ssize[1]-1)÷sizeof(Int)+1)
        end
        seed = floor(Culonglong,time())
        ptr = cudnnDropoutDescriptor_t[C_NULL]
        cudnnCreateDropoutDescriptor(ptr)
        @cudnn_retry unsafe_cudnnSetDropoutDescriptor(ptr[1], handle(), dropout, GLOBAL_DROPOUT_STATE[], sizeof(GLOBAL_DROPOUT_STATE[]), seed)
        dd = DropoutDescriptor(ptr[1])
        finalizer(x->cudnnDestroyDropoutDescriptor(x.ptr), dd)
        return dd
    end
end


text:
# Concrete types

# Abstract types from abstract_types.jl
abstract type Asset end

abstract type Property <: Asset end
abstract type Investment <: Asset end
abstract type Cash <: Asset end

abstract type House <: Property end
abstract type Apartment <: Property end

abstract type FixedIncome <: Investment end
abstract type Equity <: Investment end

# Define immutable/composite types
struct Stock <: Equity
    symbol::String
    name::String
end

# Create an instance
stock = Stock("AAPL", "Apple, Inc.")

#= REPL
julia> stock = Stock("AAPL", "Apple, Inc.")
Stock("AAPL", "Apple, Inc.")
=#

# implementing the contract from abstract type
function describe(s::Stock)
    return s.symbol * "(" * s.name * ")"
end

# It's immutable!
#= REPL
julia> stock.name = "Apple LLC"
ERROR: setfield! immutable struct of type Stock cannot be changed
=#

# Create another type
# Note that types are specified here
struct BasketOfStocks
    stocks::Vector{Stock}
    reason::String
end

many_stocks = [
    Stock("AAPL", "Apple, Inc."),
    Stock("IBM", "IBM")
]

# Create a basket of stocks
basket = BasketOfStocks(many_stocks, "Anniversary gift for my wife")

#= REPL
julia> many_stocks = [
           Stock("AAPL", "Apple, Inc."),
           Stock("IBM", "IBM")
       ]
2-element Array{Stock,1}:
 Stock("AAPL", "Apple, Inc.")
 Stock("IBM", "IBM")

julia> basket = BasketOfStocks(many_stocks, "Anniversary gift for my wife")
BasketOfStocks(Stock[Stock("AAPL", "Apple, Inc."), Stock("IBM", "IBM")], "Anniversary gift for my wife")

julia> pop!(basket.stocks)
Stock("IBM", "IBM")

julia> basket
BasketOfStocks(Stock[Stock("AAPL", "Apple, Inc.")], "Anniversary gift for my wife")
=#

# ----------------------------------------
# mutable structs

# NOTE: restart REPL so we can define Stock again

mutable struct Stock <: Equity
    symbol::String
    name::String
end

#= REPL
julia> stock = Stock("AAPL", "Apple, Inc.")
StockHolding("AAPL", "Apple, Inc.")

julia> stock.name = "Apple LLC"
Apple LLC

julia> stock
Stock("AAPL", "Apple LLC")
=#


text:
using MITgcmTools, ClimateModels, MeshArrays, OceanStateEstimation
using Test

MITgcm_download()

@testset "MITgcmTools.jl" begin

    #format conversions
    (γ,Γ)=MeshArrays.GridOfOnes("CubeSphere",30,30)
    @test isa(convert2gcmfaces(Γ.XC),Array)
    @test isa(convert2array(Γ.XC),Array)

    #physical oceanography
    (ρP,ρI,ρR) = SeaWaterDensity(3.,35.5,3000.)
    @test isapprox(ρI,1041.83267, rtol=1e-6)

    D=collect(0.0:1.0:500.0); tmp=(1.0.-tanh.(5*(-1 .+ 2/D[end]*D)));
    T=2.0 .+ 8.0*tmp; S=34.0 .+ 0.5*tmp;
    (ρP,ρI,ρR) = SeaWaterDensity(T,S,D);
    mld=MixedLayerDepth(T,S,D,"BM");
    @test isapprox(mld,134.0)

    #running and reading in verification experiments
    exps=verification_experiments()
    @test isa(exps,Array)

    MC=MITgcm_config(configuration="advect_xy")
    setup(MC)

    fil=joinpath(MC.folder,string(MC.ID),"run","data")
    nml=read(fil,MITgcm_namelist())
    write(fil*"_new",nml)

    MITgcmTools.parse_param("1.0")
    MITgcmTools.parse_param(".TRUE.")
    MITgcmTools.parse_param(".false.")
    MITgcmTools.parse_param("10")

    @test isa(nml,MITgcm_namelist)
    @test nml.groups[1]==:PARM01
    @test nml.params[1][:implicitFreeSurface]

    #

    myexp="advect_cs"
    iexp=findall([exps[i].configuration==myexp for i in 1:length(exps)])[1]
    MC=exps[iexp]

    @test clean(MC)=="no task left in pipeline"
    @test build(MC)
    @test build(MC,"--allow-skip")
    @test compile(MC)
    @test setup(MC)

    push!(MC.status,("setup" => "ended"))

    launch(MC)
    pth=joinpath(MC.folder,string(MC.ID),"run")
    tmp=read_mdsio(pth,"XC.001.001")
    @test isa(tmp,Array)
    tmp=read_mdsio(pth,"XC")
    @test isa(tmp,Array)

    scan_rundir(pth)
    Γ=GridLoad_mdsio(MC)
    @test isa(Γ,NamedTuple)

    myexp="MLAdjust"
    iexp=findall([exps[i].configuration==myexp for i in 1:length(exps)])[1]
    MC=exps[iexp]
    setup(MC)
    build(MC,"--allow-skip")
    launch(MC)
    if isdir(joinpath(MC.folder,string(MC.ID),"run","mnc_test_0001"))
        Γ=GridLoad_mnc(MC)
    else
        Γ=GridLoad_mdsio(MC)
    end
    @test isa(Γ,NamedTuple)

    #read / write functions

    fil=joinpath(pth,"available_diagnostics.log")
    read_available_diagnostics("ETAN";filename=fil)

    readcube(xx::Array,x::MeshArray) = read(cube2compact(xx),x)
    function readcube(fil::String,x::MeshArray)
        p=dirname(fil)*"/"
        b=basename(fil)[1:end-5]
        xx=read_mdsio(p,b)
        read(cube2compact(xx),x)
    end
    writecube(x::MeshArray) = compact2cube(write(x))
    writecube(fil::String,x::MeshArray) = write(fil::String,x::MeshArray)

    γ=gcmgrid(pth,"CubeSphere",6,fill((32, 32),6), [192 32], Float64, readcube, writecube)
    Γ = GridLoad(γ)
    tmp1=writecube(Γ.XC)
    tmp2=readcube(tmp1,Γ.XC)

    @test isa(tmp2,MeshArray)

    ##

    γ=GridSpec("LatLonCap",MeshArrays.GRID_LLC90)
    findtiles(30,30,γ)
    findtiles(30,30,"LatLonCap",MeshArrays.GRID_LLC90)
    read_meta(MeshArrays.GRID_LLC90,"XC.meta")

    fil=joinpath(MeshArrays.GRID_LLC90,"XC.data")
    tmp1=read_bin(fil,γ)
    tmp2=convert2gcmfaces(tmp1)
    tmp3=convert2array(tmp1)
    tmp4=convert2array(tmp3,γ)
    read_bin(fil,γ.ioPrec,γ)
    read_bin(tmp2,tmp1)
    read_bin(tmp2,γ)

    function get_ecco_variable_if_needed(v::String)
        p=dirname(pathof(OceanStateEstimation))
        lst=joinpath(p,"../examples/nctiles_climatology.csv")
        pth=ECCOclim_path
        !isdir(pth*v) ? get_from_dataverse(lst,v,pth) : nothing
    end

    get_ecco_variable_if_needed("ETAN")
    tmp=read_nctiles(joinpath(ECCOclim_path,"ETAN/ETAN"),"ETAN",γ,I=(:,:,1))

    @test isa(tmp,MeshArray)

    ##

    MC=MITgcm_config(configuration="flt_example")
    tmp=testreport(MC)
    pth=MITgcm_path[1]*"verification/flt_example/run/"
    tmp=read_flt(pth,Float32)

    @test isa(tmp[1,1],Number)

end


text:
using DiffEqFlux, OrdinaryDiffEq, Flux, Optim, Plots
using DiffEqSensitivity
using Zygote
using ForwardDiff
using Interpolations
using LinearAlgebra
using Random
using Statistics
using ProgressBars, Printf
using Flux.Optimise: update!
using Flux.Losses: mae, mse
using BSON: @save, @load
using DelimitedFiles

is_restart = true;
n_epoch = 100000;
ntotal = 40;
batch_size = 32;
n_plot = 100;
grad_max = 1.e1;
maxiters = 10000;

nr = 10;
opt = ADAMW(0.005, (0.9, 0.999), 1.f-6);

atol = 1.e-8
rtol = 1.e-3
lb = atol;
ode_solver = AutoTsit5(Rosenbrock23(autodiff=false));
# ode_solver = Rosenbrock23(autodiff=false);

rawdata = readdlm("data/10atm_1300K_0.01.txt")';
tsteps = rawdata[1, :];
Tlist = rawdata[2, :];
Plist = rawdata[3, :];
ylabel = rawdata[4:end, :];
t_end = tsteps[end]
tspan = (0.0, tsteps[end]);
ns = size(ylabel)[1];

_tsteps = 10 .^ range(log10(t_end/100), log10(t_end/1.01), length=ntotal);
_tsteps[1] = 0;

itp = LinearInterpolation(tsteps, Tlist);
Tlist = itp.(_tsteps);
itp = LinearInterpolation(tsteps, Plist);
Plist = itp.(_tsteps);
_ylabel = zeros(ns, ntotal);
for i in 1:ns
    itp = LinearInterpolation(tsteps, ylabel[i, :])
    _ylabel[i, :] .= itp.(_tsteps)
end
tsteps = _tsteps;
ylabel = _ylabel
batch_size = minimum([batch_size, ntotal]);

varnames = ["C10H16", "H2", "CH4", "C2H2", "C2H4", "N2", "C4H81",  "H", "CH3"];
l_MW = [136.238, 2.016,  16.043,  26.038,  28.054,  28.014,  56.108, 1.008,  15.035][1:ns];  # kg/mol

E_C = [10, 0, 1, 2, 2, 0, 4, 0, 1];
E_H = [16, 2, 4, 2, 4, 0, 8, 1, 3];
E_N = [0, 0, 0, 0, 0, 2, 0, 0, 0];
E_ = hcat(E_C, E_H, E_N)[1:ns, :];
E_null = nullspace(E_')';
n_null = size(E_null)[1];

ymax = maximum(ylabel, dims=2);
ymin = minimum(ylabel, dims=2);

yscale = clamp.(ymax - ymin, lb, Inf);
normdata = ylabel
u0 = normdata[:, 1];

np = nr * (2 * ns + 3) + 1;
p = randn(Float64, np) .* 0.1;
p[end] = 0.1;  # slope

function p2vec(p)
    slope = p[end] .* 10.0
    w_b = p[1:nr] .* slope
    w_in_b = p[nr + 1:nr * 2]
    w_in_Ea = p[nr * 2 + 1:nr * 3] .* slope

    w_out = reshape(p[nr * 3 + 1:nr * (ns + 3)], ns, nr)
    w_in = reshape(p[nr * (ns + 3) + 1:nr * (ns * 2 + 3)], ns, nr);
    # w_out = (w_out * E_null)'
    w_out = @. -w_in * (10 ^ w_out)
    w_in = vcat(clamp.(w_in, 0.0, 2.5), w_in_Ea', w_in_b')
    return w_in, w_b, w_out
end

function display_p(p)
    w_in, w_b, w_out = p2vec(p)
    println("\nspecies (column) reaction (row)")
    println("w_in")
    show(stdout, "text/plain", round.(w_in', digits=2))
    println("\nw_out | log(A)")
    show(stdout, "text/plain", round.(hcat(w_out', w_b), digits=2))
    println("\n")
end
display_p(p);

itpT = LinearInterpolation(tsteps, Tlist);
itpP = LinearInterpolation(tsteps, Plist);

R = 1.98720425864083f-3; # kcal⋅K−1⋅mol−1
function Y2density(Y, P, T)
    return P / (8.31446261815324e3 * T * sum(Y ./ l_MW))
end
density = Y2density(u0, Plist[1], Tlist[1])

function Y2C(Y, density)
    return density * (Y ./ l_MW) * 1e3
end
Y2C(u0, density)

Clabel = Y2C(ylabel, density)
Cmax = maximum(Clabel, dims=2)[:, 1]

dydt_scale = (yscale[:, 1] / t_end)
function crnn!(du, u, p, t)
    P = itpP(t)
    T = itpT(t)
    Y = clamp.(u, lb, 10)
    density = Y2density(Y, P, T)
    C = Y2C(Y, density)
    logX = @. log(clamp(C, lb, 10))
    w_in_x = w_in' * vcat(logX, - 1 / R / T, log(T))
    wdot = w_out * (@. exp(w_in_x + w_b))
    @. du = wdot * l_MW / density * dydt_scale
end

prob = ODEProblem(crnn!, u0, tspan)
sense = BacksolveAdjoint(checkpointing=true; autojacvec=ZygoteVJP());
function predict_n_ode(p, sample=ntotal)
    global w_in, w_b, w_out = p2vec(p)
    _prob = remake(prob, p=p, tspan=[0, tsteps[sample]])
    pred = Array(solve(_prob, ode_solver, saveat=tsteps[1:sample],
                 atol=atol, rtol=rtol, sensalg=sense, maxiters=maxiters))
end
predict_n_ode(p)

function loss_n_ode(p, sample=ntotal)
    pred = predict_n_ode(p, sample)
    loss = mae(pred ./ yscale, normdata[:, 1:size(pred)[2]] ./ yscale)
    return loss
end
loss_n_ode(p)

list_loss = []
list_grad = []
iter = 1
cb = function (p, loss_mean, g_norm; doplot=true)
    global list_loss, list_grad, iter
    push!(list_loss, loss_mean)
    push!(list_grad, g_norm)

    if doplot & (iter % n_plot == 0)
        display_p(p)
        pred = predict_n_ode(p)

        list_plt = []
        for i in 1:ns
            plt = scatter(tsteps[2:end], normdata[i, 2:end], xscale=:log10, label="data")
            plot!(plt, tsteps[2:end], pred[i,2:end], label="pred")
            ylabel!(plt, "$(varnames[i])")
            xlabel!(plt, "Time [s]")
            push!(list_plt, plt)
        end
        plt_all = plot(list_plt..., legend=:false, size=(1000, 600))
        png(plt_all, "figs/pred.png")

        println("update plot")

        plt_loss = plot(list_loss, yscale=:log10, label="loss")
        plt_grad = plot(list_grad, yscale=:log10, label="grad_norm")
        xlabel!(plt_loss, "Epoch")
        xlabel!(plt_grad, "Epoch")
        ylabel!(plt_loss, "Loss")
        ylabel!(plt_grad, "Grad Norm")
        plt_all = plot([plt_loss, plt_grad]..., legend=:outertopright)
        png(plt_all, "figs/loss_grad")

        @save "./checkpoint/mymodel.bson" p opt list_loss list_grad iter
    end
    iter += 1
    return false
end

if is_restart
    @load "./checkpoint/mymodel.bson" p opt list_loss list_grad iter
    iter += 1
    # opt = ADAMW(0.0005, (0.9, 0.999), 1.f-6)
end

epochs = ProgressBar(iter:n_epoch);
for epoch in epochs
    global p
    sample = rand(batch_size:ntotal)
    loss = loss_n_ode(p, sample)
    grad = ForwardDiff.gradient(x -> loss_n_ode(x, sample), p);
    grad_norm = norm(grad, 2)

    if grad_norm > grad_max
        grad = grad ./ grad_norm .* grad_max
    end

    update!(opt, p, grad)

    set_description(epochs, string(@sprintf("Loss: %.4e grad: %.2e", loss, grad_norm)))
    cb(p, loss, grad_norm)
end


text:
################################################################################
# Utils
################################################################################\

joint_types = [
    :Fixed,
    :Prismatic,
    :Planar,
    :FixedOrientation,
    :Revolute,
    :Cylindrical,
    :PlanarAxis,
    :FreeRevolute,
    :Orbital,
    :PrismaticOrbital,
    :PlanarOrbital,
    :FreeOrbital,
    :Spherical,
    :CylindricalFree,
    :PlanarFree
    ]

################################################################################
# Test get and set position and velocities
################################################################################
@testset "Get and set position and velocities" begin
	@testset "Maximal coordinates" begin
		mech = Dojo.get_mechanism(:raiberthopper)
		timestep= mech.timestep
		joint1 = mech.joints[1]
		joint2 = mech.joints[2]
		pbody = mech.bodies[1]
		cbody = mech.bodies[2]
		tra2 = joint2.translational
		rot2 = joint2.rotational

		x = srand(1)
		Δx = Dojo.zerodimstaticadjoint(Dojo.nullspace_mask(tra2)) * x
		Δq = rand(QuatRotation).q
		Dojo.set_maximal_configurations!(pbody, cbody;
			parent_vertex=tra2.vertices[1],
			child_vertex=tra2.vertices[2],
			Δx=Δx,
			Δq=Δq)
		@test norm(Dojo.minimal_coordinates(tra2, pbody, cbody) - x[1], Inf) < 1.0e-8

		v = srand(1)
		Δv = Dojo.zerodimstaticadjoint(Dojo.nullspace_mask(tra2)) * v
		Δω = rand(3)
		Dojo.set_maximal_velocities!(pbody, cbody;
			parent_vertex=tra2.vertices[1],
			child_vertex=tra2.vertices[2],
			Δv=Δv,
			Δω=Δω)
		@test norm(Dojo.minimal_velocities(tra2, pbody, cbody, timestep) - v[1], Inf) < 1.0e-8
	end

	@testset "Minimal coordinates" begin
		for joint_type in joint_types
			mech = Dojo.get_snake(
					num_bodies=10,
					joint_type=joint_type)

			timestep = mech.timestep
			for joint in mech.joints
				joint.rotational.axis_offset = rand(QuatRotation).q
			end
			joint0 = mech.joints[1]
			tra0 = joint0.translational
			rot0 = joint0.rotational
			pnodes0 = [mech.origin; mech.bodies[1:end-1]]
			cnodes0 = mech.bodies

			Random.seed!(100)
			Δθ = rand(input_dimension(rot0))
			Δx = rand(input_dimension(tra0))
			Δϕ = rand(input_dimension(rot0))
			Δv = rand(input_dimension(tra0))
			for i = 1:10
				Dojo.set_minimal_coordinates!(rot0, pnodes0[i], cnodes0[i],  timestep,
					Δθ=Δθ)
				Δθ0 = Dojo.minimal_coordinates(rot0, pnodes0[i], cnodes0[i])
				@test norm(Δθ0 - Δθ, Inf) < 1.0e-7

				Dojo.set_minimal_coordinates!(tra0, pnodes0[i], cnodes0[i], timestep,
					Δx=Δx)
				Δx0 = Dojo.minimal_coordinates(tra0, pnodes0[i], cnodes0[i])
				@test norm(Δx0 - Δx, Inf) < 1.0e-7

				Dojo.set_minimal_velocities!(joint0, pnodes0[i], cnodes0[i], timestep,
					Δv=Δv,
					Δϕ=Δϕ)
				Δϕ0 = Dojo.minimal_velocities(rot0, pnodes0[i], cnodes0[i], timestep)
				Δv0 = Dojo.minimal_velocities(tra0, pnodes0[i], cnodes0[i], timestep)
				@test norm(Δϕ0 - Δϕ, Inf) < 1.0e-7
				@test norm(Δv0 - Δv, Inf) < 1.0e-7
			end
		end
	end
end

################################################################################
# Test min -> max -> min
################################################################################
@testset "Minimal to maximal to minimal" begin
	# raiberthopper
	@testset "Raibert hopper" begin
		mech = Dojo.get_mechanism(:raiberthopper);
		Random.seed!(100)
		nx = Dojo.minimal_dimension(mech)
		x0 = rand(nx)
		z0 = Dojo.minimal_to_maximal(mech, x0)
		x1 = Dojo.maximal_to_minimal(mech, z0)

		@test norm(x0[1:3] - x1[1:3], Inf) < 1.0e-8
		@test norm(x0[4:6] - x1[4:6]) < 1.0e-8
		@test norm(x0[7:9] - x1[7:9], Inf) < 1.0e-8
		@test norm(x0[11:12] - x1[11:12], Inf) < 1.0e-8
		@test norm(x0[13] - x1[13], Inf) < 1.0e-8
		@test norm(x0[14] - x1[14], Inf) < 1.0e-8
	end

	# box
	@testset "Box" begin
		mech = Dojo.get_mechanism(:box)
		Random.seed!(100)
		nx = Dojo.minimal_dimension(mech)
		x0 = rand(nx)
		z0 = Dojo.minimal_to_maximal(mech, x0)
		x1 = Dojo.maximal_to_minimal(mech, z0)
		@test norm(x0 - x1, Inf) < 1.0e-8
	end

	# pendulum
	@testset "Pendulum" begin
		mech = Dojo.get_mechanism(:pendulum)
		Random.seed!(100)
		nx = Dojo.minimal_dimension(mech)
		x0 = rand(nx)
		z0 = Dojo.minimal_to_maximal(mech, x0)
		x1 = Dojo.maximal_to_minimal(mech, z0)
		@test norm(x0 - x1, Inf) < 1.0e-8
	end

	# halfcheetah
	@testset "Halfcheetah" begin
		mech = Dojo.get_mechanism(:halfcheetah)
		Random.seed!(100)
		nx = Dojo.minimal_dimension(mech)
		x0 = rand(nx)
		z0 = Dojo.minimal_to_maximal(mech, x0)
		x1 = Dojo.maximal_to_minimal(mech, z0)
		@test norm(x0 - x1, Inf) < 1.0e-8
	end

	# nslider
	@testset "Nslider" begin
		Nb0 = 5
		mech = Dojo.get_mechanism(:nslider,
			num_bodies=Nb0)
		Random.seed!(100)
		nx = Dojo.minimal_dimension(mech)
		x0 = rand(nx)
		z0 = Dojo.minimal_to_maximal(mech, x0)
		x1 = Dojo.maximal_to_minimal(mech, z0)
		@test norm(x0 - x1, Inf) < 1.0e-8
	end

	# npendulum
	@testset "Npendulum" begin
		for joint_type in joint_types
			# @show joint_type
			Nb0 = 5
			mech = Dojo.get_mechanism(:npendulum,
				num_bodies=Nb0,
				joint_type=joint_type)
			Random.seed!(100)
			nx = Dojo.minimal_dimension(mech)
			x0 = rand(nx)
			z0 = Dojo.minimal_to_maximal(mech, x0)
			x1 = Dojo.maximal_to_minimal(mech, z0)
			@test norm(x0 - x1, Inf) < 1.0e-8
		end
	end

	# snake
	@testset "Snake" begin
		for joint_type in joint_types
			# @show joint_type
			Nb0 = 5
			mech = Dojo.get_mechanism(:snake,
				num_bodies=Nb0,
				joint_type=joint_type)
			mech = Dojo.get_mechanism(:snake,
				num_bodies=Nb0,
				joint_type=:Fixed)
			Random.seed!(100)
			nx = Dojo.minimal_dimension(mech)
			x0 = rand(nx)
			z0 = Dojo.minimal_to_maximal(mech, x0)
			x1 = Dojo.maximal_to_minimal(mech, z0)
			@test norm(x0 - x1, Inf) < 1.0e-8
		end
	end

	# twister
	@testset "Twister" begin
		for joint_type in joint_types
			# @show joint_type
			Nb0 = 5
			mech = Dojo.get_mechanism(:twister,
				num_bodies=Nb0,
				joint_type=joint_type)
			Random.seed!(100)
			nx = Dojo.minimal_dimension(mech)
			x0 = rand(nx)
			z0 = Dojo.minimal_to_maximal(mech, x0)
			x1 = Dojo.maximal_to_minimal(mech, z0)
			@test norm(x0 - x1, Inf) < 1.0e-8
		end
	end

	# humanoid
	@testset "Humanoid" begin
		mech = Dojo.get_mechanism(:humanoid)
		Random.seed!(100)
		nx = Dojo.minimal_dimension(mech)
		x0 = rand(nx)
		z0 = Dojo.minimal_to_maximal(mech, x0)
		x1 = Dojo.maximal_to_minimal(mech, z0)
		@test norm(x0 - x1, Inf) < 1.0e-8
	end

	# quadruped
	@testset "Quadruped" begin
		mech = Dojo.get_mechanism(:quadruped)
		Random.seed!(100)
		nx = Dojo.minimal_dimension(mech)
		x0 = rand(nx)
		z0 = Dojo.minimal_to_maximal(mech, x0)
		x1 = Dojo.maximal_to_minimal(mech, z0)
		@test norm(x0 - x1, Inf) < 1.0e-8
	end

	# atlas
	@testset "Atlas" begin
		mech = Dojo.get_mechanism(:atlas,
			model_type=:simple,
			contact_feet=true,
			damper=10.0)

		Random.seed!(100)
		nx = Dojo.minimal_dimension(mech)
		x0 = rand(nx)
		z0 = Dojo.minimal_to_maximal(mech, x0)
		x1 = Dojo.maximal_to_minimal(mech, z0)
		@test norm(x0 - x1, Inf) < 1.0e-8
	end
end

################################################################################
# Test minimal coordinates and velocities Jacobians
################################################################################
@testset "Jacobians" begin
	@testset "Minimal velocity Jacobian" begin
		mech = Dojo.get_humanoid()
		timestep= mech.timestep
		for jointcon in mech.joints
			for joint in [jointcon.translational, jointcon.rotational]
				# generate random configuration in minimal space
				x = rand(minimal_dimension(mech))

				# convert to maximal
				z = Dojo.minimal_to_maximal(mech, x)

				# extract body states
				Ne = Dojo.length(mech.joints)
				if Dojo.get_body(mech, jointcon.parent_id).name == :origin
					zp = [mech.origin.state.x2; mech.origin.state.v15; Dojo.vector(mech.origin.state.q2); mech.origin.state.ϕ15]
				else
					zp = z[(jointcon.parent_id - Ne - 1) * 13 .+ (1:13)]
				end
				zc = z[(jointcon.child_id - Ne - 1) * 13 .+ (1:13)]

				xa = SVector{3}(zp[1:3])
				va = SVector{3}(zp[3 .+ (1:3)])
				qa = Quaternion(zp[6 .+ (1:4)]...)
				ωa = SVector{3}(zp[10 .+ (1:3)])

				xb = SVector{3}(zc[1:3])
				vb = SVector{3}(zc[3 .+ (1:3)])
				qb = Quaternion(zc[6 .+ (1:4)]...)
				ωb = SVector{3}(zc[10 .+ (1:3)])

				Dojo.minimal_velocities(joint, xa, va, qa, ωa, xb, vb, qb, ωb, timestep)

				# Jacobians
				∇0 = Dojo.minimal_velocities_jacobian_configuration(:parent, joint, xa, va, qa, ωa, xb, vb, qb, ωb, timestep)
				∇1 = FiniteDiff.finite_difference_jacobian(
					xq -> Dojo.minimal_velocities(joint, xq[Dojo.SUnitRange(1,3)], va, Quaternion(xq[4:7]...), ωa, xb, vb, qb, ωb, timestep),
					[xa; Dojo.vector(qa)]) * cat(I(3), Dojo.LVᵀmat(qa), dims=(1,2))
				@test norm(∇0 - ∇1, Inf) < 1.0e-5

				∇0 = Dojo.minimal_velocities_jacobian_configuration(:child, joint, xa, va, qa, ωa, xb, vb, qb, ωb, timestep)
				∇1 = FiniteDiff.finite_difference_jacobian(
					xq -> Dojo.minimal_velocities(joint, xa, va, qa, ωa, xq[Dojo.SUnitRange(1,3)], vb, Quaternion(xq[4:7]...), ωb, timestep),
					[xb; Dojo.vector(qb)]) * cat(I(3), Dojo.LVᵀmat(qb), dims=(1,2))
				@test norm(∇0 - ∇1, Inf) < 1.0e-5

				∇0 = Dojo.minimal_velocities_jacobian_velocity(:parent, joint, xa, va, qa, ωa, xb, vb, qb, ωb, timestep)
				∇1 = FiniteDiff.finite_difference_jacobian(
					vϕ -> Dojo.minimal_velocities(joint, xa, vϕ[Dojo.SUnitRange(1,3)], qa, vϕ[Dojo.SUnitRange(4,6)], xb, vb, qb, ωb, timestep),
					[va; ωa])
				@test norm(∇0 - ∇1, Inf) < 1.0e-5

				∇0 = Dojo.minimal_velocities_jacobian_velocity(:child, joint, xa, va, qa, ωa, xb, vb, qb, ωb, timestep)
				∇1 = FiniteDiff.finite_difference_jacobian(
					vϕ -> Dojo.minimal_velocities(joint, xa, va, qa, ωa, xb, vϕ[Dojo.SUnitRange(1,3)], qb, vϕ[Dojo.SUnitRange(4,6)], timestep),
					[vb; ωb])
				@test norm(∇0 - ∇1, Inf) < 1.0e-5
			end
		end
	end

	@testset "Minimal coordinate Jacobian" begin
		mech = Dojo.get_humanoid()
		for jointcon in mech.joints
			for joint in [jointcon.translational, jointcon.rotational]
				# generate random configuration in minimal space
				x = rand(Dojo.minimal_dimension(mech))

				# convert to maximal
				z = Dojo.minimal_to_maximal(mech, x)

				# extract body states
				Ne = Dojo.length(mech.joints)
				if Dojo.get_body(mech, jointcon.parent_id).name == :origin
					zp = [mech.origin.state.x2; mech.origin.state.v15; Dojo.vector(mech.origin.state.q2); mech.origin.state.ϕ15]
				else
					zp = z[(jointcon.parent_id - Ne - 1) * 13 .+ (1:13)]
				end
				zc = z[(jointcon.child_id - Ne - 1) * 13 .+ (1:13)]

				xa = SVector{3}(zp[1:3])
				# va = SVector{3}(zp[3 .+ (1:3)])
				qa = Quaternion(zp[6 .+ (1:4)]...)
				# ωa = SVector{3}(zp[10 .+ (1:3)])

				xb = SVector{3}(zc[1:3])
				# vb = SVector{3}(zc[3 .+ (1:3)])
				qb = Quaternion(zc[6 .+ (1:4)]...)
				# ωb = SVector{3}(zc[10 .+ (1:3)])

				Dojo.minimal_coordinates(joint, xa, qa, xb, qb)

				∇0 = Dojo.minimal_coordinates_jacobian_configuration(:parent, joint, xa, qa, xb, qb)
				∇1 = FiniteDiff.finite_difference_jacobian(
					xq -> Dojo.minimal_coordinates(joint, xq[1:3], Quaternion(xq[4:7]...), xb, qb),
					[xa; Dojo.vector(qa)]) * cat(I(3), Dojo.LVᵀmat(qa), dims=(1,2))
				@test norm(∇0 - ∇1, Inf) < 1.0e-6

				∇0 = Dojo.minimal_coordinates_jacobian_configuration(:child, joint, xa, qa, xb, qb)
				∇1 = FiniteDiff.finite_difference_jacobian(
					xq -> Dojo.minimal_coordinates(joint, xa, qa, xq[1:3], Quaternion(xq[4:7]...)),
					[xb; Dojo.vector(qb)]) * cat(I(3), Dojo.LVᵀmat(qb), dims=(1,2))
				@test norm(∇0 - ∇1, Inf) < 1.0e-6
			end
		end
	end

	@testset "Minimal to maximal Jacobian" begin
		function maximal_to_minimal_jacobian_fd(mechanism::Mechanism, z)
			J = FiniteDiff.finite_difference_jacobian(y -> Dojo.maximal_to_minimal(mechanism, y), z)
			G = attitude_jacobian(z, length(mechanism.bodies))
			return J * G
		end

		function minimal_to_maximal_jacobian_fd(mechanism::Mechanism, x)
			J = FiniteDiff.finite_difference_jacobian(y -> Dojo.minimal_to_maximal(mechanism, y), x)
			z = minimal_to_maximal(mechanism, x)
			G = attitude_jacobian(z, length(mechanism.bodies))
			return G' * J
		end

		function ctrl!(mechanism, k)
			Dojo.set_input!(mechanism, 0.1 * srand(Dojo.input_dimension(mechanism)))
		end

		# n-pendulum
		mechanism = Dojo.get_mechanism(:npendulum,
			timestep=0.1,
			gravity=-9.81,
			num_bodies=1)

		base_angle = 0.3 * π
		Dojo.initialize!(mechanism, :npendulum,
			base_angle=base_angle)
		storage = Dojo.simulate!(mechanism, 1.0,
			record=true,
			verbose=false)

		x = Dojo.get_minimal_state(mechanism)
		z = Dojo.get_maximal_state(mechanism)
		u = zeros(Dojo.input_dimension(mechanism))

		@test norm(minimal_to_maximal(mechanism, x) - z) < 1.0e-5
		@test norm(Dojo.maximal_to_minimal(mechanism, z) - x) < 1.0e-5

		M_fd = maximal_to_minimal_jacobian_fd(mechanism, z)
		M_a = Dojo.maximal_to_minimal_jacobian(mechanism, z)
		@test size(M_fd) == size(M_a)
		@test norm(M_fd - M_a, Inf) < 1.0e-5

		N_fd = minimal_to_maximal_jacobian_fd(mechanism, Dojo.maximal_to_minimal(mechanism, z))
		N_a = Dojo.minimal_to_maximal_jacobian(mechanism, Dojo.maximal_to_minimal(mechanism, z))
		@test size(N_fd) == size(N_a)
		@test norm(N_fd - N_a, Inf) < 1.0e-5
		@test abs(sum(diag(M_fd * N_fd)) - Dojo.minimal_dimension(mechanism)) < 1.0e-5
		@test abs(sum(diag(M_a * N_a)) - Dojo.minimal_dimension(mechanism)) < 1.0e-5

		# n-pendulum
		mechanism = Dojo.get_mechanism(:npendulum,
			timestep=0.1,
			gravity=-9.81,
			num_bodies=2)
		base_angle = 0.3 * π
		Dojo.initialize!(mechanism, :npendulum,
			base_angle=base_angle)
		storage = Dojo.simulate!(mechanism, 1.0,
			record=true,
			verbose=false)

		x = Dojo.get_minimal_state(mechanism)
		z = Dojo.get_maximal_state(mechanism)
		u = zeros(Dojo.input_dimension(mechanism))

		@test norm(minimal_to_maximal(mechanism, x) - z) < 1.0e-5
		@test norm(Dojo.maximal_to_minimal(mechanism, z) - x) < 1.0e-5

		M_fd = maximal_to_minimal_jacobian_fd(mechanism, z)
		M_a = Dojo.maximal_to_minimal_jacobian(mechanism, z)
		@test size(M_fd) == size(M_a)
		@test norm(M_fd - M_a, Inf) < 1.0e-5

		N_fd = minimal_to_maximal_jacobian_fd(mechanism, Dojo.maximal_to_minimal(mechanism, z))
		N_a = Dojo.minimal_to_maximal_jacobian(mechanism, Dojo.maximal_to_minimal(mechanism, z))
		@test size(N_fd) == size(N_a)
		@test norm(N_fd - N_a, Inf) < 1.0e-5
		@test abs(sum(diag(M_fd * N_fd)) - Dojo.minimal_dimension(mechanism)) < 1.0e-5
		@test abs(sum(diag(M_a * N_a)) - Dojo.minimal_dimension(mechanism)) < 1.0e-5

		# sphere
		mechanism = Dojo.get_mechanism(:sphere,
			timestep=0.01,
			gravity=-9.81)
		Dojo.initialize!(mechanism, :sphere)
		storage = Dojo.simulate!(mechanism, 1.0,
			record=true,
			verbose=false)

		z = Dojo.get_maximal_state(mechanism)
		x = Dojo.get_minimal_state(mechanism)
		u = zeros(Dojo.input_dimension(mechanism))

		@test norm(minimal_to_maximal(mechanism, x) - z) < 1.0e-5
		@test norm(Dojo.maximal_to_minimal(mechanism, z) - x) < 1.0e-5

		M_fd = maximal_to_minimal_jacobian_fd(mechanism, z)
		M_a = Dojo.maximal_to_minimal_jacobian(mechanism, z)
		@test size(M_fd) == size(M_a)
		@test norm(M_fd - M_a, Inf) < 1.0e-5

		N_fd = minimal_to_maximal_jacobian_fd(mechanism, Dojo.maximal_to_minimal(mechanism, z))
		N_a = Dojo.minimal_to_maximal_jacobian(mechanism, Dojo.maximal_to_minimal(mechanism, z))
		@test size(N_fd) == size(N_a)
		@test norm(N_fd - N_a, Inf) < 1.0e-6
		@test abs(sum(diag(M_fd * N_fd)) - Dojo.minimal_dimension(mechanism)) < 1.0e-5
		@test abs(sum(diag(M_a * N_a)) - Dojo.minimal_dimension(mechanism)) < 1.0e-5

		# half cheetah
		mechanism = Dojo.get_mechanism(:halfcheetah,
			timestep=0.01,
			gravity=-9.81)
		Dojo.initialize!(mechanism, :halfcheetah)
		storage = Dojo.simulate!(mechanism, 1.0, ctrl!,
			record=true,
			verbose=false)

		z = Dojo.get_maximal_state(mechanism)
		x = Dojo.get_minimal_state(mechanism)
		u = zeros(Dojo.input_dimension(mechanism))

		@test norm(minimal_to_maximal(mechanism, x) - z) < 1.0e-5
		@test norm(Dojo.maximal_to_minimal(mechanism, z) - x) < 1.0e-5

		M_fd = maximal_to_minimal_jacobian_fd(mechanism, z)
		M_a = Dojo.maximal_to_minimal_jacobian(mechanism, z)
		@test size(M_fd) == size(M_a)
		@test norm(M_fd - M_a, Inf) < 1.0e-5

		N_fd = minimal_to_maximal_jacobian_fd(mechanism, Dojo.maximal_to_minimal(mechanism, z))
		N_a = Dojo.minimal_to_maximal_jacobian(mechanism, Dojo.maximal_to_minimal(mechanism, z))
		@test size(N_fd) == size(N_a)
		@test norm(N_fd - N_a, Inf) < 1.0e-5

		@test abs(sum(diag(M_fd * N_fd)) - Dojo.minimal_dimension(mechanism)) < 1.0e-5
		@test abs(sum(diag(M_a * N_a)) - Dojo.minimal_dimension(mechanism)) < 1.0e-5
		@test abs(sum(diag(M_a * N_fd)) - Dojo.minimal_dimension(mechanism)) < 1.0e-5

		# atlas
		mechanism = Dojo.get_mechanism(:atlas,
			timestep=0.01,
			gravity=-9.81,
			friction_coefficient=0.5,
			damper=100.0,
			spring=1.0,
			contact_feet=true)
		Dojo.initialize!(mechanism, :atlas_stance,
			body_position=[0,0,0.5],
			body_orientation=[0.0,0.0,0.0])
		storage = Dojo.simulate!(mechanism, 1.0, ctrl!,
			record=true,
			verbose=false)

		z = Dojo.get_maximal_state(mechanism)
		x = Dojo.get_minimal_state(mechanism)
		u = zeros(Dojo.input_dimension(mechanism))

		@test norm(minimal_to_maximal(mechanism, x) - z) < 1.0e-5
		@test norm(Dojo.maximal_to_minimal(mechanism, z) - x) < 1.0e-5

		M_fd = maximal_to_minimal_jacobian_fd(mechanism, z)
		M_a = Dojo.maximal_to_minimal_jacobian(mechanism, z)
		@test size(M_fd) == size(M_a)
		@test norm(M_fd - M_a, Inf) < 1.0e-5

		N_fd = minimal_to_maximal_jacobian_fd(mechanism, Dojo.maximal_to_minimal(mechanism, z))
		N_a = Dojo.minimal_to_maximal_jacobian(mechanism, Dojo.maximal_to_minimal(mechanism, z))
		@test size(N_fd) == size(N_a)
		@test norm(N_fd - N_a, Inf) < 5.0e-5

		@test abs(sum(diag(M_fd * N_fd)) - Dojo.minimal_dimension(mechanism)) < 1.0e-5
		@test abs(sum(diag(M_a * N_a)) - Dojo.minimal_dimension(mechanism)) < 1.0e-5
	end

	@testset "Maximal to minimal Jacobian" begin
		function maximal_to_minimal_jacobian_fd(mechanism::Mechanism, z)
			J = FiniteDiff.finite_difference_jacobian(y -> maximal_to_minimal(mechanism, y), z)
			G = attitude_jacobian(z, length(mechanism.bodies))
			return J * G
		end

		# 5-link pendulum
		mech = Dojo.get_mechanism(:npendulum,
			timestep=0.01,
			gravity=-9.81,
			num_bodies=5)

		Random.seed!(100)
		base_angle = 0.3π
		Dojo.initialize!(mech, :npendulum,
			base_angle=base_angle)
		storage = Dojo.simulate!(mech, 1.0,
			record=true,
			verbose=false)

		Dojo.maximal_dimension(mech) == 13
		Dojo.minimal_dimension(mech) == 12
		z = Dojo.get_maximal_state(mech)

		attjac = Dojo.attitude_jacobian(z, length(mech.bodies))
		M_fd = maximal_to_minimal_jacobian_fd(mech, z)
		M_a = Dojo.maximal_to_minimal_jacobian(mech, z)
		@test size(M_fd) == size(M_a)
		@test norm(M_fd - M_a, Inf) < 1.0e-6

		# # sphere
		mech = Dojo.get_mechanism(:sphere,
			timestep=0.01,
			gravity=-9.81)
		Dojo.initialize!(mech, :sphere)
		storage = Dojo.simulate!(mech, 1.0,
			record=true,
			verbose=false)

		Dojo.maximal_dimension(mech)
		Dojo.minimal_dimension(mech)
		z = Dojo.get_maximal_state(mech)

		attjac = Dojo.attitude_jacobian(z, length(mech.bodies))
		M_fd = maximal_to_minimal_jacobian_fd(mech, z)
		M_a = Dojo.maximal_to_minimal_jacobian(mech, z)

		@test size(M_fd) == size(M_a)
		@test norm(M_fd - M_a, Inf) < 1.0e-6
	end
end


text:
using ..TopOpt.TopOptProblems: AbstractGrid
const Vec = Ferrite.Vec

# @params struct TrussGrid{xdim,N,M,C<:Ferrite.Cell{xdim,N,M},T} <: AbstractGrid{xdim, T}
struct TrussGrid{xdim,T,N,M,TG<:Ferrite.Grid{xdim,<:Ferrite.Cell{xdim,N,M},T}} <: AbstractGrid{xdim, T}
    grid::TG
    white_cells::BitVector
    black_cells::BitVector
    constant_cells::BitVector
    crosssecs::Vector{TrussFEACrossSec{T}}
end
# nels::NTuple{dim, Int} # num of elements in x,y,z direction in the ground mesh, not applicable to truss
# sizes::NTuple{dim, T}  # length of the ground mesh in x,y,z direction, not applicaiton to truss
# corners::NTuple{2, Vec{dim, T}} # corner of the ground mesh

nnodespercell(::TrussGrid{xdim,T,N,M}) where {xdim,T,N,M} = N
nfacespercell(::TrussGrid{xdim,T,N,M}) where {xdim,T,N,M} = M
nnodes(cell::Type{Ferrite.Cell{xdim,N,M}}) where {xdim, N, M} = N
nnodes(cell::Ferrite.Cell) = nnodes(typeof(cell))
Ferrite.getncells(tg::TrussGrid) = Ferrite.getncells(tg.grid)

function TrussGrid(node_points::Dict{iT, SVector{xdim, T}}, elements::Dict{iT, Tuple{iT, iT}},
        boundary::Dict{iT, SVector{xdim, fT}}; crosssecs=TrussFEACrossSec{T}(1.0)) where {xdim, T, iT, fT}
    grid = _LinearTrussGrid(node_points, elements, boundary)
    ncells = getncells(grid)
    if crosssecs isa Vector
        @assert length(crosssecs) == ncells
        crosssecs = convert(Vector{TrussFEACrossSec{T}}, crosssecs)
    elseif crosssecs isa TrussFEACrossSec
        crosssecs = [TrussFEACrossSec{T}(crosssecs) for i=1:ncells]
    else
        error("Invalid crossecs: $(crossecs)")
    end
    return TrussGrid(grid, falses(ncells), falses(ncells), falses(ncells), crosssecs)
end

function _LinearTrussGrid(node_points::Dict{iT, SVector{xdim, T}}, elements::Dict{iT, Tuple{iT, iT}},
        boundary::Dict{iT, SVector{xdim, fT}}) where {xdim, T, iT, fT}
    n_nodes = length(node_points)

    # * Generate cells, Line2d or Line3d
    @assert xdim ∈ [2,3]
    CellType = xdim == 2 ? Line2D : Line3D
    cells = Vector{CellType}(undef, length(elements))
    for (e_ind, element) in elements
        cells[e_ind] = CellType((element...,))
    end

    # * Generate nodes
    nodes = Vector{Node{xdim,T}}(undef, length(node_points))
    for (n_id, node_point) in node_points
        nodes[n_id] = Node((node_point...,))
    end

    return Grid(cells, nodes)
end

function Base.show(io::Base.IO, mime::MIME"text/plain", tg::TrussGrid)
    println(io, "TrussGrid:")
    print(io, "\t-")
    Base.show(io, mime, tg.grid)
    println(io,"")
    print(io, "\t-")
    println(io, "white cells:T|$(sum(tg.white_cells))|, black cells:T|$(sum(tg.black_cells))|, const cells:T|$(sum(tg.constant_cells))|")
end

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

# function Base.show(io::Base.IO, ::MIME"text/plain", grid::Ferrite.Grid{xdim, Union{Line2D, Line3D}}) where {xdim}
#     print(io, "$(typeof(grid)) with $(getncells(grid)) $(extra_celltypes[eltype(grid.cells)]) cells and $(getnnodes(grid)) nodes")
# end

# const extra_celltypes = Dict{DataType, String}(Line2D  => "Line2D",
#                                                Line3D  => "Line3D")