import Base.deepcopy_internal
deepcopy_internal(x::Module,stackdict::IdDict) = x

# Model Types.  The template args are required for inference for cache/perturbation solutions
struct PerturbationModel{MaxOrder,N_y,N_x,N_ϵ,N_z,N_p,HasΩ,T1,T2}
    mod::Module

    # Could extract from type, but here for simplicity
    max_order::Int64
    n_y::Int64
    n_x::Int64
    n_p::Int64
    n_ϵ::Int64
    n_z::Int64

    # Variations
    has_Ω::Bool
    η::T1
    Q::T2
end

# Construct from a module.  Inherently type unstable, so use function barrier from return type
function PerturbationModel(mod)
    return PerturbationModel{mod.max_order,mod.n_y,mod.n_x,mod.n_ϵ,mod.n_z,mod.n_p,
                             mod.has_Ω,typeof(mod.η),typeof(mod.Q)}(mod, mod.max_order,
                                                                    mod.n_y, mod.n_x,
                                                                    mod.n_p, mod.n_ϵ,
                                                                    mod.n_z, mod.has_Ω,
                                                                    mod.η, mod.Q)
end

# TODO: Add in latex stuff for the mod.H_latex,
function Base.show(io::IO, ::MIME"text/plain", m::PerturbationModel) where {T}
    return print(io,
                 "Perturbation Model: n_y = $(m.n_y), n_x = $(m.n_x), n_p = $(m.n_p), n_ϵ = $(m.n_ϵ), n_z = $(m.n_z)\n y = $(m.mod.y_symbols) \n x = $(m.mod.x_symbols) \n p = $(m.mod.p_symbols)")
end

# Buffers for the solvers to reduce allocations
# General rule for cache vs. buffers
# 1. If something should be used in multiple parts of the algorithm, put it in the cache
# 2. Otherwise, use the buffers, which you can "trash" with inplace operations as required
# 3. The cache should never be modified after it has been filled in a given sequence of events, buffers can be

# Note: When fully generic, this pattern has an issue when the exact number fields meets the number of arguments to the convenience constructor.  Forcing the first type bypasses the issue
Base.@kwdef struct FirstOrderSolverBuffers{ComplexMatrixType<:AbstractMatrix,RealMatrixType,
                                           UpperTriangularMatrixType}
    A::ComplexMatrixType
    B::ComplexMatrixType
    Z::RealMatrixType  # real version, transposed relative to the schur
    Z_ll::RealMatrixType
    S_bb::UpperTriangularMatrixType
    T_bb::UpperTriangularMatrixType
end
function FirstOrderSolverBuffers(n_y, n_x, n_p_d, n_ϵ, n_z)
    n = n_x + n_y
    return FirstOrderSolverBuffers(; A = zeros(Complex{Float64}, n, n),
                                   B = zeros(Complex{Float64}, n, n), Z = zeros(n, n),
                                   Z_ll = zeros(n_y, n_y),
                                   S_bb = UpperTriangular(zeros(n_x, n_x)),
                                   T_bb = UpperTriangular(zeros(n_x, n_x)))
end
Base.@kwdef struct FirstOrderDerivativeSolverBuffers{RealMatrixType}
    R::RealMatrixType
    A::RealMatrixType
    C::RealMatrixType
    D::RealMatrixType
    E::RealMatrixType
    dH::RealMatrixType
    bar::RealMatrixType
end
function FirstOrderDerivativeSolverBuffers(n_y, n_x, n_p_d, n_ϵ, n_z)
    n = n_x + n_y
    return FirstOrderDerivativeSolverBuffers(; R = zeros(2 * n, n_x), A = zeros(n, n),
                                             C = zeros(n, n), D = zeros(n_x, n_x),
                                             dH = zeros(n, 2n), bar = zeros(2n, 1),
                                             E = zeros(n, n_x))
end
Base.@kwdef struct SecondOrderSolverBuffers{RealMatrixType}
    A::RealMatrixType
    C::RealMatrixType
    B::RealMatrixType
    D::RealMatrixType
    E::RealMatrixType
    R::RealMatrixType
    A_σ::RealMatrixType
    R_σ::RealMatrixType
end
function SecondOrderSolverBuffers(n_y, n_x, n_p_d, n_ϵ, n_z)
    n = n_x + n_y
    return SecondOrderSolverBuffers(; A = zeros(n, n), B = zeros(n_x^2, n_x^2),
                                    C = zeros(n, n), D = zeros(n_x^2, n_x^2),
                                    E = zeros(n, n_x^2), R = zeros(2 * n, n_x),
                                    A_σ = zeros(n, n), R_σ = zeros(2 * n, n_x))
end
Base.@kwdef struct SecondOrderDerivativeSolverBuffers{RealMatrixType,VectorOfMatrixType,VectorOfVectorOfMatrixType}
    A::RealMatrixType
    B::RealMatrixType
    C::RealMatrixType
    D::RealMatrixType
    E::RealMatrixType
    R::RealMatrixType
    dH::RealMatrixType
    dΨ::VectorOfMatrixType
    gh_stack::RealMatrixType
    g_xx_flat::RealMatrixType
    Ψ_x_sum::VectorOfVectorOfMatrixType
    Ψ_y_sum::VectorOfVectorOfMatrixType
    bar::RealMatrixType
    kron_h_x::RealMatrixType
    R_p::RealMatrixType
    A_σ::RealMatrixType
    R_σ::RealMatrixType
end

function SecondOrderDerivativeSolverBuffers(n_y, n_x, n_p_d, n_ϵ, n_z)
    n = n_x + n_y
    return SecondOrderDerivativeSolverBuffers(; A = zeros(n, n), B = zeros(n_x^2, n_x^2),
                                              C = zeros(n, n), D = zeros(n_x^2, n_x^2),
                                              E = zeros(n, n_x^2), R = zeros(2n, n_x),
                                              dH = zeros(n, 2n), dΨ = [zeros(2n, 2n) for _ in 1:n],
                                              gh_stack = zeros(n, n_x^2), g_xx_flat = zeros(n_y, n_x^2),
                                              Ψ_x_sum =  [[zeros(2n, 2n) for _ in 1:n] for _ in 1:n_x],
                                              Ψ_y_sum =  [[zeros(2n, 2n) for _ in 1:n] for _ in 1:n_y],
                                              bar = zeros(2n, 1), kron_h_x = zeros(n_x^2, n_x^2),
                                              R_p = zeros(2n, n_x),
                                              A_σ = zeros(n, n), R_σ = zeros(2 * n, n_x),
                                              )
end

# The cache if for both 1st and 2nd order
# Constructors set values to nothing as appropriate
abstract type AbstractSolverCache{Order} end
Base.@kwdef struct SolverCache{Order,MatrixType,MatrixType2,MatrixType3,MatrixType4,
                               MatrixType5,VectorType,VectorType2,VectorOfVectorType,
                               VectorOfMatrixType,VectorOfMatrixType2,VectorOfMatrixType3,
                               VectorOrNothingType,VectorOfVectorOrNothingType,
                               MatrixScalingOrNothingType,SymmetricMatrixType,
                               SymmetricVectorOfMatrixType,VectorOfVectorOfMatrixType,
                               ThreeTensorType,CholeskyType,ChangeVarianceType,
                               VectorOfThreeTensorType,FirstOrderSolverBuffersType,
                               FirstOrderSolverDerivativesBuffersType,
                               SecondOrderSolverBuffersOrNothingType,
                               SecondOrderSolverDerivativesBuffersOrNothingType} <:
                   AbstractSolverCache{Order}
    order::Val{Order}  # allows inference in construction
    p_d_symbols::Vector{Symbol}
    H::VectorType
    H_yp::MatrixType
    H_y::MatrixType
    H_xp::MatrixType
    H_x::MatrixType
    H_yp_p::VectorOfMatrixType
    H_y_p::VectorOfMatrixType
    H_xp_p::VectorOfMatrixType
    H_x_p::VectorOfMatrixType
    H_p::VectorOfVectorType
    Γ::MatrixType2
    Γ_p::VectorOfMatrixType2
    Σ::SymmetricMatrixType
    Σ_p::SymmetricVectorOfMatrixType
    Ω::VectorOrNothingType
    Ω_p::VectorOfVectorOrNothingType
    Ψ::VectorOfMatrixType

    # Used in solution
    x::VectorType
    y::VectorType
    y_p::VectorOfVectorType
    x_p::VectorOfVectorType
    g_x::MatrixType4
    h_x::MatrixType4
    g_x_p::VectorOfMatrixType3
    h_x_p::VectorOfMatrixType3
    B::MatrixType2
    B_p::VectorOfMatrixType2
    Q::MatrixScalingOrNothingType
    η::MatrixType3
    A_1_p::VectorOfMatrixType3
    C_1::MatrixType4
    C_1_p::VectorOfMatrixType3
    V::CholeskyType
    V_p::ChangeVarianceType
    η_Σ_sq::SymmetricMatrixType

    # Additional for 2nd order
    Ψ_p::VectorOfVectorOfMatrixType
    Ψ_yp::VectorOfVectorOfMatrixType
    Ψ_y::VectorOfVectorOfMatrixType
    Ψ_xp::VectorOfVectorOfMatrixType
    Ψ_x::VectorOfVectorOfMatrixType
    g_xx::ThreeTensorType
    h_xx::ThreeTensorType
    g_σσ::VectorType2
    h_σσ::VectorType2
    g_xx_p::VectorOfThreeTensorType
    h_xx_p::VectorOfThreeTensorType
    g_σσ_p::MatrixType5
    h_σσ_p::MatrixType5

    # Additional for solution type 2nd order
    A_0_p::MatrixType5
    A_2_p::VectorOfThreeTensorType
    C_0::VectorType2
    C_2::ThreeTensorType
    C_0_p::MatrixType5
    C_2_p::VectorOfThreeTensorType

    # Buffers for additional calculations
    first_order_solver_buffer::FirstOrderSolverBuffersType
    first_order_solver_p_buffer::FirstOrderSolverDerivativesBuffersType
    second_order_solver_buffer::SecondOrderSolverBuffersOrNothingType
    second_order_solver_p_buffer::SecondOrderSolverDerivativesBuffersOrNothingType
    I_x::MatrixType  #dense identity matrices
    I_x_2::MatrixType
    zeros_x_x::MatrixType
    zeros_y_x::MatrixType
end

# The Val(2), etc. for the order required for inference to function
# Note that the n_p_d is the number of differentiated parameters to allocate for
# Extracts the field names from the order of the p_d argument (which could be dict, named tuple, etc.)
function SolverCache(m::PerturbationModel{MaxOrder,N_y,N_x,N_ϵ,N_z,N_p,HasΩ,T1,T2},
                     ::Val{Order},
                     p_d) where {Order,MaxOrder,N_y,N_x,N_ϵ,N_z,N_p,HasΩ,T1,T2}
    n_p_d = length(p_d)
    p_d_symbols = collect(Symbol.(keys(p_d)))
    return SolverCache(; order = Val(Order), p_d_symbols, H = zeros(N_x + N_y),
                       H_yp = zeros(N_x + N_y, N_y), H_y = zeros(N_x + N_y, N_y),
                       H_xp = zeros(N_x + N_y, N_x), H_x = zeros(N_x + N_y, N_x),
                       Γ = zeros(N_ϵ, N_ϵ), Ω = !HasΩ ? nothing : zeros(N_z),
                       Ψ = [zeros(2(N_x + N_y), 2(N_x + N_y)) for i in 1:(N_x + N_y)],
                       H_p = [zeros(N_x + N_y) for i in 1:n_p_d],
                       H_yp_p = [zeros(N_x + N_y, N_y) for i in 1:n_p_d],
                       H_y_p = [zeros(N_x + N_y, N_y) for i in 1:n_p_d],
                       H_xp_p = [zeros(N_x + N_y, N_x) for i in 1:n_p_d],
                       H_x_p = [zeros(N_x + N_y, N_x) for i in 1:n_p_d],
                       Γ_p = [zeros(N_ϵ, N_ϵ) for i in 1:n_p_d],
                       Ω_p = !HasΩ ? nothing : [zeros(N_z) for i in 1:n_p_d],
                       x = zeros(N_x), y = zeros(N_y), y_p = [zeros(N_y) for i in 1:n_p_d],
                       x_p = [zeros(N_x) for i in 1:n_p_d], g_x = zeros(N_y, N_x),
                       h_x = zeros(N_x, N_x), g_x_p = [zeros(N_y, N_x) for _ in 1:n_p_d],
                       h_x_p = [zeros(N_x, N_x) for _ in 1:n_p_d],
                       Σ = Symmetric(zeros(N_ϵ, N_ϵ)), η_Σ_sq = Symmetric(zeros(N_x, N_x)),
                       Σ_p = [Symmetric(zeros(N_ϵ, N_ϵ)) for _ in 1:n_p_d], m.Q, m.η,
                       B = zeros(N_x, N_ϵ), B_p = [zeros(N_x, N_ϵ) for _ in 1:n_p_d],
                       C_1 = zeros(N_z, N_x), C_1_p = [zeros(N_z, N_x) for _ in 1:n_p_d],
                       A_1_p = [zeros(N_x, N_x) for _ in 1:n_p_d],
                       V = cholesky(Array(I(N_x))),
                       V_p = [zeros(N_x, N_x) for _ in 1:n_p_d],

                       # Stuff for 2nd order
                       Ψ_p = (Order == 1) ? nothing :
                             [[zeros(2 * (N_x + N_y), 2 * (N_x + N_y))
                               for _ in 1:(N_x + N_y)] for _ in 1:n_p_d],
                       Ψ_yp = (Order == 1) ? nothing :
                              [[zeros(2 * (N_x + N_y), 2 * (N_x + N_y))
                                for _ in 1:(N_x + N_y)] for _ in 1:N_y],
                       Ψ_y = (Order == 1) ? nothing :
                             [[zeros(2 * (N_x + N_y), 2 * (N_x + N_y))
                               for _ in 1:(N_x + N_y)] for _ in 1:N_y],
                       Ψ_xp = (Order == 1) ? nothing :
                              [[zeros(2 * (N_x + N_y), 2 * (N_x + N_y))
                                for _ in 1:(N_x + N_y)] for _ in 1:N_x],
                       Ψ_x = (Order == 1) ? nothing :
                             [[zeros(2 * (N_x + N_y), 2 * (N_x + N_y))
                               for _ in 1:(N_x + N_y)] for _ in 1:N_x],
                       g_xx = (Order == 1) ? nothing : zeros(N_y, N_x, N_x),
                       h_xx = (Order == 1) ? nothing : zeros(N_x, N_x, N_x),
                       g_σσ = (Order == 1) ? nothing : zeros(N_y),
                       h_σσ = (Order == 1) ? nothing : zeros(N_x),
                       g_xx_p = (Order == 1) ? nothing :
                                [zeros(N_y, N_x, N_x) for _ in 1:n_p_d],
                       h_xx_p = (Order == 1) ? nothing :
                                [zeros(N_x, N_x, N_x) for _ in 1:n_p_d],
                       g_σσ_p = (Order == 1) ? nothing : zeros(N_y, n_p_d),
                       h_σσ_p = (Order == 1) ? nothing : zeros(N_x, n_p_d),
                       A_0_p = (Order == 1) ? nothing : zeros(N_x, n_p_d),
                       A_2_p = (Order == 1) ? nothing :
                               [zeros(N_x, N_x, N_x) for _ in 1:n_p_d],
                       C_0 = (Order == 1) ? nothing : zeros(N_z),
                       C_0_p = (Order == 1) ? nothing : zeros(N_z, n_p_d),
                       C_2 = (Order == 1) ? nothing : zeros(N_z, N_x, N_x),
                       C_2_p = (Order == 1) ? nothing :
                               [zeros(N_z, N_x, N_x) for _ in 1:n_p_d],

                       # buffers for algorithms
                       first_order_solver_buffer = FirstOrderSolverBuffers(N_y, N_x, n_p_d,
                                                                           N_ϵ, N_z),
                       first_order_solver_p_buffer = FirstOrderDerivativeSolverBuffers(N_y,
                                                                                N_x,
                                                                                n_p_d,
                                                                                N_ϵ,
                                                                                N_z),
                       second_order_solver_buffer = (Order == 1) ? nothing :
                                                    SecondOrderSolverBuffers(N_y, N_x,
                                                                             n_p_d, N_ϵ,
                                                                             N_z),
                       second_order_solver_p_buffer = (Order == 1) ? nothing :
                                                      SecondOrderDerivativeSolverBuffers(N_y,
                                                                                N_x,
                                                                                n_p_d,
                                                                                N_ϵ,
                                                                                N_z),
                       I_x = Matrix{Float64}(I(N_x)), I_x_2 = Matrix{Float64}(I(N_x^2)),
                       zeros_x_x = zeros(N_x, N_x), zeros_y_x = zeros(N_y, N_x))
end
Base.@kwdef struct PerturbationSolverSettings{T1,T2,T3,T4,T5,T6}
    print_level::Int64 = 1  # 0 is no output at all
    ϵ_BK::Float64 = 1e-6 # For checking Blanchard-Kahn condition
    nlsolve_method::Symbol = :trust_region
    nlsolve_iterations::Int64 = 1000
    nlsolve_show_trace::Bool = false
    nlsolve_ftol::Float64 = 1e-8
    use_solution_cache::Bool = true
    evaluate_functions_callback::T1 = nothing
    calculate_steady_state_callback::T2 = nothing
    solve_first_order_callback::T3 = nothing
    solve_first_order_p_callback::T4 = nothing
    solve_second_order_callback::T5 = nothing
    solve_second_order_p_callback::T6 = nothing
end

function nlsolve_options(s::PerturbationSolverSettings)
    return (method = s.nlsolve_method, iterations = s.nlsolve_iterations,
            show_trace = s.nlsolve_show_trace, ftol = s.nlsolve_ftol)
end

# State Space types
abstract type AbstractPerturbationSolution end
abstract type AbstractFirstOrderPerturbationSolution <: AbstractPerturbationSolution end

# For this, all are dense due to schur decomposition
# All are dense due to schur decomposition
Base.@kwdef struct FirstOrderPerturbationSolution{T1<:AbstractVector,T2<:AbstractVector,
                                                  T3<:AbstractMatrix,T4<:AbstractMatrix,
                                                  T5<:AbstractMatrix,
                                                  T6<:Union{Nothing,Distribution},
                                                  T7<:Union{Nothing,AbstractMatrix,
                                                            UniformScaling},
                                                  T8<:AbstractMatrix,T9<:AbstractMatrix,
                                                  T10<:Distribution,T11<:AbstractMatrix} <:
                   AbstractFirstOrderPerturbationSolution
    retcode::Symbol
    x_symbols::Vector{Symbol}
    y_symbols::Vector{Symbol}
    p_symbols::Vector{Symbol}
    p_d_symbols::Vector{Symbol}
    u_symbols::Vector{Symbol}
    # TODO: differentiated parameter ordering?
    n_y::Int64
    n_x::Int64
    n_p::Int64
    n_ϵ::Int64
    n_z::Int64
    y::T1
    x::T2
    g_x::T3
    A::T4
    B::T5
    D::T6  # current a matrix or nothing, later could make more general
    Q::T7  # can be nothing
    η::T8
    C::T9  # i.e. Q * g_x
    x_ergodic::T10
    Γ::T11
end

maybe_diagonal(x::AbstractVector) = TuringDiagMvNormal(zero(x), x)
maybe_diagonal(x) = x # otherwise, just return raw.  e.g. nothing

function FirstOrderPerturbationSolution(retcode, m::PerturbationModel, c::SolverCache)
    return FirstOrderPerturbationSolution(; retcode, m.mod.x_symbols, m.mod.y_symbols,
                                          m.mod.u_symbols, m.mod.p_symbols, c.p_d_symbols,
                                          m.n_x, m.n_y, m.n_p, m.n_ϵ, m.n_z, c.Q, c.η, c.y,
                                          c.x, c.B, D = maybe_diagonal(c.Ω), c.g_x,
                                          A = c.h_x, C = c.C_1,
                                          x_ergodic = TuringDenseMvNormal(zeros(m.n_x),
                                                                          c.V), c.Γ)
end

Base.@kwdef struct SecondOrderPerturbationSolution{T1<:AbstractVector,T2<:AbstractVector,
                                                   T3<:AbstractMatrix,T4<:AbstractMatrix,
                                                   T5<:AbstractMatrix,
                                                   T6<:Union{Nothing,Distribution},
                                                   T7<:Union{Nothing,AbstractMatrix,
                                                             UniformScaling},
                                                   T8<:AbstractMatrix,T9<:AbstractMatrix,
                                                   T10,T11<:AbstractArray,T12,
                                                   T13<:AbstractVector,T14<:AbstractMatrix,
                                                   T15<:AbstractVector,
                                                   T16<:AbstractArray} <:
                   AbstractPerturbationSolution
    retcode::Symbol
    x_symbols::Vector{Symbol}
    y_symbols::Vector{Symbol}
    p_symbols::Vector{Symbol}
    p_d_symbols::Vector{Symbol}
    u_symbols::Vector{Symbol}
    n_y::Int64
    n_x::Int64
    n_p::Int64
    n_ϵ::Int64
    n_z::Int64
    y::T1
    x::T2
    g_x::T3
    B::T5
    D::T6
    Q::T7  # can be nothing
    η::T8
    Γ::T9

    g_xx::T10
    g_σσ::T12
    A_0::T13
    A_1::T4
    A_2::T11

    C_1::T14
    C_0::T15
    C_2::T16
end

function SecondOrderPerturbationSolution(retcode, m::PerturbationModel, c::SolverCache)
    return SecondOrderPerturbationSolution(; retcode, m.mod.x_symbols, m.mod.y_symbols,
                                           m.mod.u_symbols, m.mod.p_symbols, c.p_d_symbols,
                                           m.n_x, m.n_y, m.n_p, m.n_ϵ, m.n_z, c.Q, c.η, c.y,
                                           c.x, c.B, D = maybe_diagonal(c.Ω), c.Γ, c.g_x,
                                           A_1 = c.h_x, c.g_xx, A_2 = 0.5 * c.h_xx, c.g_σσ,
                                           A_0 = 0.5 * c.h_σσ, c.C_1, c.C_0, c.C_2)
end


@testset "api" begin
    # drop database "quantum_states" if exists
    drop_db(string(QuantumStatesData))

    @test QuantumStateDB.QSDB == "quantum_states"

    create_all()

    @test string(QuantumStatesData) in from_sql("SELECT datname FROM pg_database;")[!, 1]

    # modify dbconfig
    dbconfig = copy(current_dbconfig())
    dbconfig[:dbname] = string(QuantumStatesData)

    column_names = [
        "id",
        "r", "theta",
        "dim", "rho",
        "n_points", "bhd",
        "w_range", "w",
    ]
    column_names_from_sql = from_sql(
        """
        SELECT column_name
            FROM information_schema.columns
            WHERE table_schema='public' AND table_name='$(string(SqueezedStatesData))';
        """,
        dbconfig=dbconfig
    )[!, 1]
    @test all(lowercase(col) in column_names_from_sql for col in column_names)

    column_names = [
        "id",
        "r", "theta", "nbar",
        "dim", "rho",
        "n_points", "bhd",
        "w_range", "w",
    ]
    column_names_from_sql = from_sql(
        """
        SELECT column_name
            FROM information_schema.columns
            WHERE table_schema='public' AND table_name='$(string(SqueezedThermalStatesData))';
        """,
        dbconfig=dbconfig
    )[!, 1]
    @test all(lowercase(col) in column_names_from_sql for col in column_names)
end

@testset "insert into SqueezedStatesData" begin
    r = 0.8; θ = π/2
    dim = 100; ρ = SqueezedState(r, θ, Matrix, dim=dim)
    np = 40960; ps = rand(GaussianStateBHD(ρ), np)
    w_range = 3; w = wigner(ρ, LinRange(-w_range, w_range, 101), LinRange(-w_range, w_range, 101))

    df = DataFrame([
        :r=>r, :theta=>θ,
        :dim=>dim, :rho=>hexbytes_str(ρ),
        :n_points=>np, :bhd=>hexbytes_str(ps),
        :w_range=> w_range, :w=>hexbytes_str(w.𝐰_surface),
    ])

    # modify dbconfig
    dbconfig = copy(current_dbconfig())
    dbconfig[:dbname] = string(QuantumStatesData)

    to_sql(df, SqueezedStatesData, dbconfig=dbconfig)

    df_from_sql = from_sql(SqueezedStatesData, dbconfig=dbconfig)

    postprocessor = [
        :r => Float64,
        :theta => Float64,
        :dim => Int,
        :rho => hexbytes2array(ComplexF64),
        :n_points => Int,
        :bhd => hexbytes2array(Float64),
        :w_range => Int,
        :w => hexbytes2array(Float64),
    ]
    apply!(df_from_sql, postprocessor...)

    @test df_from_sql[1, :r] == r
    @test df_from_sql[1, :theta] == θ
    @test df_from_sql[1, :dim] == dim
    @test all(reshape(df_from_sql[1, :rho], dim, dim) .== ρ)
    @test df_from_sql[1, :n_points] == np
    @test all(reshape(df_from_sql[1, :bhd], 2, np) .== ps)
    @test df_from_sql[1, :w_range] == w_range
    @test all(reshape(df_from_sql[1, :w], 101, 101) .== w.𝐰_surface)
end

@testset "insert into SqueezedThermalStatesData" begin
    r = 0.8; θ = π/2; n̄ = 0.3
    dim = 100; ρ = SqueezedThermalState(r, θ, n̄, dim=dim)
    np = 40960; ps = rand(GaussianStateBHD(ρ), np)
    w_range = 3; w = wigner(ρ, LinRange(-w_range, w_range, 101), LinRange(-w_range, w_range, 101))

    df = DataFrame([
        :r=>r, :theta=>θ, :nbar=>n̄,
        :dim=>dim, :rho=>hexbytes_str(ρ),
        :n_points=>np, :bhd=>hexbytes_str(ps),
        :w_range=> w_range, :w=>hexbytes_str(w.𝐰_surface),
    ])

    # modify dbconfig
    dbconfig = copy(current_dbconfig())
    dbconfig[:dbname] = string(QuantumStatesData)

    to_sql(df, SqueezedThermalStatesData, dbconfig=dbconfig)

    df_from_sql = from_sql(SqueezedThermalStatesData, dbconfig=dbconfig)

    postprocessor = [
        :r => Float64,
        :theta => Float64,
        :nbar => Float64,
        :dim => Int,
        :rho => hexbytes2array(ComplexF64),
        :n_points => Int,
        :bhd => hexbytes2array(Float64),
        :w_range => Int,
        :w => hexbytes2array(Float64),
    ]
    apply!(df_from_sql, postprocessor...)

    @test df_from_sql[1, :r] == r
    @test df_from_sql[1, :theta] == θ
    @test df_from_sql[1, :nbar] == n̄
    @test df_from_sql[1, :dim] == dim
    @test all(reshape(df_from_sql[1, :rho], dim, dim) .== ρ)
    @test df_from_sql[1, :n_points] == np
    @test all(reshape(df_from_sql[1, :bhd], 2, np) .== ps)
    @test df_from_sql[1, :w_range] == w_range
    @test all(reshape(df_from_sql[1, :w], 101, 101) .== w.𝐰_surface)
end

@testset "insert table into SqueezedStatesData" begin
    # modify dbconfig
    dbconfig = copy(current_dbconfig())
    dbconfig[:dbname] = string(QuantumStatesData)

    n = 10

    r = LinRange(0, 1, n); θ = LinRange(0, 2π, n)
    dim = 100; sq(r, θ) = SqueezedState(r, θ, Matrix, dim=dim); ρ = sq.(r, θ)
    np = 4096; p(ρ) = rand(GaussianStateBHD(ρ), np); ps = p.(ρ)
    w_range = 3; wf = ρ->wigner(ρ, LinRange(-w_range, w_range, 101), LinRange(-w_range, w_range, 101)).𝐰_surface ; w = wf.(ρ)

    df = DataFrame([
        :r=>r, :theta=>θ,
        :dim=>repeat([dim], n), :rho=>hexbytes_str.(ρ),
        :n_points=>repeat([np], n), :bhd=>hexbytes_str.(ps),
        :w_range=> repeat([w_range], n), :w=>hexbytes_str.(w),
    ])

    to_sql(df, SqueezedStatesData, dbconfig=dbconfig)

    df_from_sql = from_sql(SqueezedStatesData, dbconfig=dbconfig)

    postprocessor = [
        :r => Float64,
        :theta => Float64,
        :dim => Int,
        :rho => hexbytes2array(ComplexF64),
        :n_points => Int,
        :bhd => hexbytes2array(Float64),
        :w_range => Int,
        :w => hexbytes2array(Float64),
    ]
    apply!(df_from_sql, postprocessor...)

    for nᵢ in 1:n
        # index from db should +1 because we already insert 1 row in previous testset
        @test df_from_sql[nᵢ+1, :r] == r[nᵢ]
        @test df_from_sql[nᵢ+1, :theta] == θ[nᵢ]
        @test df_from_sql[nᵢ+1, :dim] == dim
        @test all(reshape(df_from_sql[nᵢ+1, :rho], dim, dim) .== ρ[nᵢ])
        @test df_from_sql[nᵢ+1, :n_points] == np
        @test all(reshape(df_from_sql[nᵢ+1, :bhd], 2, np) .== ps[nᵢ])
        @test df_from_sql[nᵢ+1, :w_range] == w_range
        @test all(reshape(df_from_sql[nᵢ+1, :w], 101, 101) .== w[nᵢ])
    end
end

@testset "insert table into SqueezedThermalStatesData" begin
    # modify dbconfig
    dbconfig = copy(current_dbconfig())
    dbconfig[:dbname] = string(QuantumStatesData)

    n = 10

    r = LinRange(0, 1, n); θ = LinRange(0, 2π, n); n̄ = LinRange(0, 0.3, n)
    dim = 100; sqth(r, θ, n̄) = SqueezedThermalState(r, θ, n̄, dim=dim); ρ = sqth.(r, θ, n̄)
    np = 4096; p(ρ) = rand(GaussianStateBHD(ρ), np); ps = p.(ρ)
    w_range = 3; wf = ρ->wigner(ρ, LinRange(-w_range, w_range, 101), LinRange(-w_range, w_range, 101)).𝐰_surface ; w = wf.(ρ)

    df = DataFrame([
        :r=>r, :theta=>θ, :nbar=>n̄,
        :dim=>repeat([dim], n), :rho=>hexbytes_str.(ρ),
        :n_points=>repeat([np], n), :bhd=>hexbytes_str.(ps),
        :w_range=> repeat([w_range], n), :w=>hexbytes_str.(w),
    ])

    to_sql(df, SqueezedThermalStatesData, dbconfig=dbconfig)

    df_from_sql = from_sql(SqueezedThermalStatesData, dbconfig=dbconfig)

    postprocessor = [
        :r => Float64,
        :theta => Float64,
        :nbar => Float64,
        :dim => Int,
        :rho => hexbytes2array(ComplexF64),
        :n_points => Int,
        :bhd => hexbytes2array(Float64),
        :w_range => Int,
        :w => hexbytes2array(Float64),
    ]
    apply!(df_from_sql, postprocessor...)

    for nᵢ in 1:n
        # index from db should +1 because we already insert 1 row in previous testset
        @test df_from_sql[nᵢ+1, :r] == r[nᵢ]
        @test df_from_sql[nᵢ+1, :theta] == θ[nᵢ]
        @test df_from_sql[nᵢ+1, :nbar] == n̄[nᵢ]
        @test df_from_sql[nᵢ+1, :dim] == dim
        @test all(reshape(df_from_sql[nᵢ+1, :rho], dim, dim) .== ρ[nᵢ])
        @test df_from_sql[nᵢ+1, :n_points] == np
        @test all(reshape(df_from_sql[nᵢ+1, :bhd], 2, np) .== ps[nᵢ])
        @test df_from_sql[nᵢ+1, :w_range] == w_range
        @test all(reshape(df_from_sql[nᵢ+1, :w], 101, 101) .== w[nᵢ])
    end
end

@testset "nrow" begin
    @test QuantumStateDB.nrow(SqueezedStatesData) == 11
    @test QuantumStateDB.nrow(SqueezedThermalStatesData) == 11
end

@testset "from_sql" begin
    @show nrow(from_sql(SqueezedStatesData, 5, order=:r, offset=6)) == 5
end


"""
    fit!(m::WSVarLmmModel,
    solver=IpoptSolver(print_level=0, mehrotra_algorithm="yes", max_iter=100);
    init=init_ls!(m), runs = 2)

Fit a `WSVarLmmModel` object using a weighted NLS method.

# Positional arguments
- `m::WSVarLmmModel`: Model to fit.
- `solver`: Nonlinear programming solver to use. Common choices include:
    - `Ipopt.IpoptSolver(print_level=0, mehrotra_algorithm="yes", warm_start_init_point="yes", max_iter=100)`.
    - `Ipopt.IpoptSolver(print_level=0, watchdog_shortened_iter_trigger=3, max_iter=100)`.
    - `Ipopt.IpoptSolver(print_level=0, max_iter=100)`.
    - `KNITRO.KnitroSolver(outlev=3)`. (Knitro is commercial software)
    - `NLopt.NLoptSolver(algorithm=:LD_MMA, maxeval=4000)`.
    - `NLopt.NLoptSolver(algorithm=:LD_LBFGS, maxeval=4000)`.

# Keyword arguments
- `init`: Initialization strategy. `fit!` will use `m.τ` and `m.Lγ` to set the
    weight matrices `Vi` and solve the weighted NLS to obtain an
    estimate for `m.β`, `m.τ`, and `m.Lγ`.  Choices for `init` include
    - `init_ls!(m)` (default): initialize by the least squares analytical solution.
    - `init_mom!(m)`: initialize by the unweighted NLS (MoM).
    - `m`: initilize from user supplied values in `m.τ` and `m.Lγ`.
- `runs::Integer`: Number of weighted NLS runs; default is 2. Each run will use the
    newest `m.τ` and `m.Lγ` to update the weight matrices `Vi` and solve the
    new weighted NLS.
- `parallel::Bool`: Multi-threading or not. Default is `false`.
- `verbose::Bool`: Verbose display or not, Default is `true`.
"""
function fit!(
    m        :: WSVarLmmModel{T},
    solver   :: MOI.AbstractOptimizer = Ipopt.Optimizer();
    solver_config::Dict =
        Dict("print_level"           => 0,
             "mehrotra_algorithm"    => "yes",
             "warm_start_init_point" => "yes",
             "max_iter"              => 100),
    init     :: WSVarLmmModel{T} = init_ls!(m),
    runs     :: Integer = 2,
    parallel :: Bool = false,
    verbose  :: Bool = true
    ) where {T<:BlasReal}
    solvertype = typeof(solver)
    solvertype <: Ipopt.Optimizer ||
        @warn("Optimizer object is $solvertype, `solver_config` may need to be defined.")
    # Pass options to solver
    config_solver(solver, solver_config)
    # set up NLP optimization problem
    npar = m.l + ◺(m.q)
    MOI.empty!(solver)
    lb   = T[]
    ub   = T[]

    NLPBlock = MOI.NLPBlockData(
        MOI.NLPBoundsPair.(lb, ub), m, true
    )

    par0 = Vector{T}(undef, npar)
    modelpar_to_optimpar!(par0, m)

    solver_pars = MOI.add_variables(solver, npar)
    for i in 1:npar
        MOI.set(solver, MOI.VariablePrimalStart(), solver_pars[i], par0[i])
    end

    MOI.set(solver, MOI.NLPBlock(), NLPBlock)
    MOI.set(solver, MOI.ObjectiveSense(), MOI.MIN_SENSE)

    # optimize weighted NLS
    m.iswtnls[1] = true
    m.ismthrd[1] = parallel
    xsol = similar(par0)
    βprev, τprev, Lγprev = similar(m.β), similar(m.τ), similar(m.Lγ)
    for run in 1:runs
        copyto!(βprev, m.β); copyto!(τprev, m.τ), copyto!(Lγprev, m.Lγ)
        tic = time() # start timing
        # update Vi, then β and residuals with WLS
        update_wtmat!(m)
        # update τ and Lγ by WNLS
        modelpar_to_optimpar!(par0, m)
        for i in 1:npar
            MOI.set(solver, MOI.VariablePrimalStart(), solver_pars[i], par0[i])
        end
        MOI.optimize!(solver)
        toc = time()
        optstat = MOI.get(solver, MOI.TerminationStatus())
        optstat in (MOI.LOCALLY_SOLVED, MOI.ALMOST_LOCALLY_SOLVED) ||
            @warn("Optimization unsuccessful; got $optstat")
        # Get solver solution values
        fill!(xsol, zero(T))
        for i in eachindex(xsol)
            xsol[i] = MOI.get(solver, MOI.VariablePrimal(), MOI.VariableIndex(i))
        end
        optimpar_to_modelpar!(m, xsol)
        verbose && @printf(
            "run = %d, ‖Δβ‖ = %f, ‖Δτ‖ = %f, ‖ΔL‖ = %f, status = %s, time(s) = %f\n",
            run,
            norm(m.β  - βprev ),
            norm(m.τ  - τprev ),
            norm(m.Lγ - Lγprev),
            optstat,
            toc - tic)
    end
    m.isfitted[1] = true
    # refresh objective, gradient, and Hessian
    mul!(m.Σγ, m.Lγ, transpose(m.Lγ))
    nlsv_obj!(m, true, true, false)
    # sandwich estimator
    sandwich!(m)
    m
end

"""
    modelpar_to_optimpar!(par, m)

Translate model parameters in `m` to optimization variables in `par`.
"""
function modelpar_to_optimpar!(
    par :: AbstractVector{T},
    m   :: WSVarLmmModel{T}
    ) where {T<:BlasReal}
    q, l = m.q, m.l
    # τ
    copyto!(par, m.τ)
    # Lγ
    offset = l + 1
    @inbounds for j in 1:q, i in j:q
        par[offset] = i == j ? log(max(m.Lγ[i, j], floatmin(T))) : m.Lγ[i, j]
        offset += 1
    end
    par
end

"""
    optimpar_to_modelpar!(m, par)

Translate optimization variables in `par` to the model parameters in `m`.
"""
function optimpar_to_modelpar!(
    m   :: WSVarLmmModel{T},
    par :: Vector{T}
    ) where {T<:BlasReal}
    q, l = m.q, m.l
    # τ
    copyto!(m.τ, 1, par, 1, l)
    # Lγ
    fill!(m.Lγ, 0)
    offset = l + 1
    @inbounds for j in 1:q, i in j:q
        m.Lγ[i, j] = i == j ? exp(par[offset]) : par[offset]
        offset += 1
    end
    m
end

function MOI.initialize(
    m                  :: WSVarLmmModel,
    requested_features :: Vector{Symbol}
    )
    for feat in requested_features
        if !(feat in MOI.features_available(m))
            error("Unsupported feature $feat")
        end
    end
end

MOI.features_available(m::WSVarLmmModel) = [:Grad, :Hess]

function MOI.eval_objective(
    m   :: WSVarLmmModel{T},
    par :: AbstractVector{T}
    ) where {T<:BlasReal}
    optimpar_to_modelpar!(m, par)
    nlsv_obj!(m, false, false, false)
end

function MOI.eval_objective_gradient(
    m    :: WSVarLmmModel{T},
    grad :: Vector{T},
    par  :: AbstractVector{T}
    ) where {T<:BlasReal}
    q, l = m.q, m.l
    optimpar_to_modelpar!(m, par)
    obj = nlsv_obj!(m, true, false, false)
    # gradient wrt τ
    copyto!(grad, m.∇τ)
    # gradient wrt Lγ
    offset = l + 1
    @inbounds for j in 1:q, i in j:q # traverse lower triangular of Lγ[i, j]
        grad[offset] = m.∇Lγ[i, j]
        i == j && (grad[offset] *= m.Lγ[i, j])
        offset += 1
    end
end

function MOI.eval_constraint(
    m   :: WSVarLmmModel,
    g   :: Vector{T},
    par :: AbstractVector{T}
    ) where {T<:BlasReal}
    return nothing
end

function MOI.hessian_lagrangian_structure(m::WSVarLmmModel)
    # our Hessian is a dense matrix, work on the upper triangular part
    npar = m.l + ◺(m.q)
    arr1 = Vector{Int}(undef, ◺(npar))
    arr2 = Vector{Int}(undef, ◺(npar))
    idx  = 1
    @inbounds for j in 1:npar, i in 1:j
        arr1[idx] = i
        arr2[idx] = j
        idx      += 1
    end
    return zip(arr1, arr2)
end

function MOI.eval_hessian_lagrangian(
    m   :: WSVarLmmModel,
    H   :: AbstractVector{T},
    par :: AbstractVector{T},
    σ   :: T,
    μ   :: AbstractVector{T}
    ) where {T<:BlasReal}
    q, l = m.q, m.l
    # refresh obj, gradient, and hessian
    optimpar_to_modelpar!(m, par)
    nlsv_obj!(m, true, true, false)
    # Hττ
    idx = 1
    @inbounds for j in 1:l, i in 1:j
        H[idx] = m.Hττ[i, j]
        idx   += 1
    end
    j = 1 # index columns of HτLγ and HLγLγ
    @inbounds for j2 in 1:q, j1 in j2:q # traverse lower triangular of Lγ[j1, j2]
        # HτLγ
        for i in 1:l # i index rows of HτLγ
            H[idx] = m.HτLγ[i, j]
            j1 == j2 && (H[idx] *= m.Lγ[j1, j2])
            idx += 1
        end
        # HLγLγ
        i = 1 # index rows of HLγLγ
        for i2 in 1:q, i1 in i2:q # traverse lower triangular of Lγ[i1, i2]
            i > j && break # skip lower triangular of HLγLγ
            H[idx] = m.HLγLγ[i, j]
            # different diagonal entries of Lγ
            i1 == i2 && (H[idx] *= m.Lγ[i1, i2])
            j1 == j2 && (H[idx] *= m.Lγ[j1, j2])
            # same diagonal entry of Lγ
            i1 == i2 == j1 == j2 && (H[idx] += m.∇Lγ[j1, j2] * m.Lγ[j1, j2])
            idx += 1
            i   += 1
        end
        j += 1
    end
    lmul!(σ, H)
end


# This file is a part of Julia. License is MIT: https://julialang.org/license

# Conversion/Promotion

"""
    Date(dt::DateTime) -> Date

Converts a `DateTime` to a `Date`. The hour, minute, second, and millisecond parts of
the `DateTime` are truncated, so only the year, month and day parts are used in
construction.
"""
Date(dt::TimeType) = convert(Date, dt)

"""
    DateTime(dt::Date) -> DateTime

Converts a `Date` to a `DateTime`. The hour, minute, second, and millisecond parts of
the new `DateTime` are assumed to be zero.
"""
DateTime(dt::TimeType) = convert(DateTime, dt)

"""
    Time(dt::DateTime) -> Time

Converts a `DateTime` to a `Time`. The hour, minute, second, and millisecond parts of
the `DateTime` are used to create the new `Time`. Microsecond and nanoseconds are zero by default.
"""
Time(dt::DateTime) = convert(Time, dt)

Base.convert(::Type{DateTime}, dt::Date) = DateTime(UTM(value(dt) * 86400000))
Base.convert(::Type{Date}, dt::DateTime) = Date(UTD(days(dt)))
Base.convert(::Type{Time}, dt::DateTime) = Time(Nanosecond((value(dt) % 86400000) * 1000000))

Base.convert(::Type{DateTime},x::Millisecond)  = DateTime(Dates.UTInstant(x))  # Converts Rata Die milliseconds to a DateTime
Base.convert(::Type{Millisecond},dt::DateTime) = Millisecond(value(dt))        # Converts DateTime to Rata Die milliseconds
Base.convert(::Type{Date},x::Day)  = Date(Dates.UTInstant(x))  # Converts Rata Die days to a Date
Base.convert(::Type{Day},dt::Date) = Day(value(dt))            # Converts Date to Rata Die days

### External Conversions
const UNIXEPOCH = value(DateTime(1970)) #Rata Die milliseconds for 1970-01-01T00:00:00

"""
    unix2datetime(x) -> DateTime

Takes the number of seconds since unix epoch `1970-01-01T00:00:00` and converts to the
corresponding `DateTime`.
"""
function unix2datetime(x)
    rata = UNIXEPOCH + round(Int64, Int64(1000) * x)
    return DateTime(UTM(rata))
end
"""
    datetime2unix(dt::DateTime) -> Float64

Takes the given `DateTime` and returns the number of seconds
since the unix epoch `1970-01-01T00:00:00` as a `Float64`.
"""
datetime2unix(dt::DateTime) = (value(dt) - UNIXEPOCH) / 1000.0

"""
    now() -> DateTime

Returns a `DateTime` corresponding to the user's system time including the system timezone
locale.
"""
function now()
    tv = Libc.TimeVal()
    tm = Libc.TmStruct(tv.sec)
    return DateTime(tm.year + 1900, tm.month + 1, tm.mday, tm.hour, tm.min, tm.sec, div(tv.usec, 1000))
end

"""
    today() -> Date

Returns the date portion of `now()`.
"""
today() = Date(now())

"""
    now(::Type{UTC}) -> DateTime

Returns a `DateTime` corresponding to the user's system time as UTC/GMT.
"""
now(::Type{UTC}) = unix2datetime(time())

"""
    rata2datetime(days) -> DateTime

Takes the number of Rata Die days since epoch `0000-12-31T00:00:00` and returns the
corresponding `DateTime`.
"""
rata2datetime(days) = DateTime(yearmonthday(days)...)

"""
    datetime2rata(dt::TimeType) -> Int64

Returns the number of Rata Die days since epoch from the given `Date` or `DateTime`.
"""
datetime2rata(dt::TimeType) = days(dt)

# Julian conversions
const JULIANEPOCH = value(DateTime(-4713, 11, 24, 12))

"""
    julian2datetime(julian_days) -> DateTime

Takes the number of Julian calendar days since epoch `-4713-11-24T12:00:00` and returns the
corresponding `DateTime`.
"""
function julian2datetime(f)
    rata = JULIANEPOCH + round(Int64, Int64(86400000) * f)
    return DateTime(UTM(rata))
end

"""
    datetime2julian(dt::DateTime) -> Float64

Takes the given `DateTime` and returns the number of Julian calendar days since the julian
epoch `-4713-11-24T12:00:00` as a `Float64`.
"""
datetime2julian(dt::DateTime) = (value(dt) - JULIANEPOCH) / 86400000.0


using Flux
using Random

mutable struct ImageDQNAgent{M, TN, O, LU, AP<:AbstractValuePolicy, Φ, ERP, IMB} <: AbstractAgent
    model::M
    target_network::TN
    opt::O
    lu::LU
    ap::AP
    er::ExperienceReplay{ERP}
    image_buffer::IMB
    γ::Float32
    batch_size::Int64
    tn_counter_init::Int64
    target_network_counter::Int64
    action::Int64
    prev_s::Φ
end

ImageDQNAgent(model, target_network, opt, lu, ap, size_buffer, im, γ, batch_size, tn_counter_init, s) =
    ImageDQNAgent(model,
             target_network,
             opt,
             lu,
             ap,
             ExperienceReplay(size_buffer,
                              (Int64, Int64, Int64, Float32, Bool),
                              (:s, :a, :sp, :r, :t)),
             ImageBuffer(size_buffer+2, im, Array{Float32, 3}),
             γ,
             batch_size,
             tn_counter_init,
             tn_counter_init,
             0,
             0)

function JuliaRL.start!(agent::ImageDQNAgent, env_s_tp1, rng::AbstractRNG; kwargs...)
    # Start an Episode
    agent.action = sample(agent.ap,
                          agent.model(env_s_tp1),
                          rng)
    agent.prev_s = add!(agent.imb, env_s_tp1)
    return agent.action
end

function JuliaRL.step!(agent::ImageDQNAgent, env_s_tp1, r, terminal, rng::AbstractRNG; kwargs...)

    img_idx = add!(agent.imb, env_s_tp1)

    add!(agent.er, (agent.prev_s, agent.action, img_idx, r, terminal))

    if size(agent.er)[1] > 1000
        e = sample(agent.er, agent.batch_size; rng=rng)
        update_params!(agent, e)
    end

    agent.prev_s .= env_s_tp1
    agent.action = sample(agent.ap,
                          agent.model(agent.prev_s),
                          rng)

    return agent.action
end

function update_params!(agent::ImageDQNAgent, e)

    if agent.tn_counter_init > 0

        update!(agent.model, agent.lu, agent.opt, agent.image_buffer[e.s], e.a, agent.image_buffer[e.sp], e.r, e.t, agent.target_network)

        if agent.target_network_counter == 1
            agent.target_network_counter = agent.tn_counter_init
            agent.target_network = deepcopy(mapleaves(
                Flux.Tracker.data,
                agent.model))
        else
            agent.target_network_counter -= 1
        end
    else
        update!(agent.model, agent.lu, agent.opt, agent.image_buffer[e.s], e.a, agent.image_buffer[e.sp], e.r, e.t)
    end
    return nothing

end


#=
Copyright (c) 2018-2022 Chris Coey, Lea Kapelevich, and contributors

This Julia package Hypatia.jl is released under the MIT license; see LICENSE
file in the root directory or at https://github.com/chriscoey/Hypatia.jl
=#

insts = OrderedDict()
insts["minimal"] = [
    ((1, 3, 2, false, true, false, false, false),),
    ((1, 3, 2, false, false, true, false, false),),
    ((1, 3, 2, true, false, false, true, false),),
    ((1, 3, 2, true, false, true, false, false),),
    ((1, 3, 2, true, false, false, true, true),),
    ]
# insts["fast"] = [] # TODO
insts["various"] = [
    ((5, 20, 20, false, true, false, false, false),),
    ((5, 20, 20, false, false, true, false, false),),
    ((5, 20, 20, true, false, false, true, false),),
    ((5, 20, 20, true, false, true, false, false),),
    ((5, 20, 20, true, false, false, true, true),),
    ((10, 30, 20, false, true, false, false, false),),
    ((10, 30, 20, false, false, true, false, false),),
    ((10, 30, 20, true, false, false, true, false),),
    ((10, 30, 20, true, false, true, false, false),),
    ((10, 30, 20, true, false, false, true, true),),
    ((8, 50, 30, false, true, false, false, false),),
    ((8, 50, 30, false, false, true, false, false),),
    ((8, 50, 30, true, false, false, true, false),),
    ((8, 50, 30, true, false, true, false, false),),
    ((8, 50, 30, true, false, false, true, true),),
    ]
return (SparseLMIJuMP, insts)


__precompile__()
module OttawaSlabCell
using FilmDataStructures, ResponseModels,eps_gold,eps_InAsSbP,eps_InAs_ntype
export uOttawaSlabs
# Code for initializing uOttawa slab structure
include("uOttawaSlabStruct.jl")
# Creates uOttawaSlabs(tEmit::Float64, tBck::Float64, divNCell::Int, divPCell::Int)::Tuple{lyrDsc, Array{Int64,2}}
precompile(uOttawaSlabs, (Float64, Float64, Int, Int))
end


"""
    exampleobject(nrow, ncol)

Create an example `SummarizedExperiment` object with the specified number of rows and columns.
This is to be used to improve the succinctness of examples and tests.

# Examples
```jldoctest
julia> using SummarizedExperiments

julia> x = exampleobject(20, 10)
20x10 SummarizedExperiment
  assays(3): foo bar whee
  rownames: Gene1 Gene2 ... Gene19 Gene20
  rowdata(2): name Type
  colnames: Patient1 Patient2 ... Patient9 Patient10
  coldata(3): name Treatment Response
  metadata(1): version
```
"""
function exampleobject(nrow::Int64, ncol::Int64)
    first_assay = Array{Int64, 2}(undef, (nrow, ncol))
    for i in 1:length(first_assay)
        first_assay[i] = round(rand() * 100)
    end

    second_assay = Array{Float64, 2}(undef, (nrow, ncol))
    for i in 1:length(second_assay)
        second_assay[i] = rand()
    end

    third_assay = Array{Float32, 3}(undef, (nrow, ncol, 2))
    for i in 1:length(third_assay)
        third_assay[i] = -rand()
    end

    assays = OrderedDict{String, AbstractArray}(
        "foo" => first_assay,
        "bar" => second_assay,
        "whee" => third_assay
    )

    tchoices = ["normal", "drug1", "drug2"]
    coldata = DataFrame(
        "name" => [ "Patient" * string(i) for i in 1:ncol ],
        "Treatment" => [ tchoices[Int(floor(rand() * length(tchoices) + 1))] for i in 1:ncol ],
        "Response" => [ rand() for i in 1:ncol ]
    )

    gchoices = ["Protein-coding", "Pseudogene"]
    rowdata = DataFrame(
        "name" => [ "Gene" * string(i) for i in 1:nrow ],
        "Type" => [ gchoices[Int(floor(rand() * length(gchoices) + 1))] for i in 1:nrow ]
    )

    metadata = Dict{String,Any}("version" => "1.1.0")
    return SummarizedExperiment(assays, rowdata, coldata, metadata)
end


struct Continuity3dModel{M} <: AbstractOceanModel
    ocean::M
    function Continuity3dModel(ocean::M) where {M}
        return new{M}(ocean)
    end
end
vars_state(cm::Continuity3dModel, ::Prognostic, FT) =
    vars_state(cm.ocean, Prognostic(), FT)

# Continuity3dModel is used to compute the horizontal divergence of u

vars_state(cm::Continuity3dModel, ::Auxiliary, T) = @vars()
vars_state(cm::Continuity3dModel, ::Gradient, T) = @vars()
vars_state(cm::Continuity3dModel, ::GradientFlux, T) = @vars()
vars_state(cm::Continuity3dModel, ::UpwardIntegrals, T) = @vars()
init_state_auxiliary!(cm::Continuity3dModel, _...) = nothing
init_state_prognostic!(cm::Continuity3dModel, _...) = nothing
@inline flux_second_order!(cm::Continuity3dModel, _...) = nothing
@inline source!(cm::Continuity3dModel, _...) = nothing
@inline update_penalty!(::RusanovNumericalFlux, ::Continuity3dModel, _...) =
    nothing

# This allows the balance law framework to compute the horizontal gradient of u
# (which will be stored back in the field θ)
@inline function flux_first_order!(
    m::Continuity3dModel,
    flux::Grad,
    state::Vars,
    aux::Vars,
    t::Real,
    direction,
)
    @inbounds begin
        u = state.u # Horizontal components of velocity
        v = @SVector [u[1], u[2], -0]

        # ∇ • (v)
        # Just using θ to store w = ∇h • u
        flux.θ += v
    end

    return nothing
end

# This is zero because when taking the horizontal gradient we're piggy-backing
# on θ and want to ensure we do not use it's jump
@inline wavespeed(cm::Continuity3dModel, n⁻, _...) = -zero(eltype(n⁻))

boundary_state!(
    ::CentralNumericalFluxSecondOrder,
    cm::Continuity3dModel,
    _...,
) = nothing

"""
    boundary_state!(nf, ::Continuity3dModel, Q⁺, A⁺, Q⁻, A⁻, bctype)

applies boundary conditions for the hyperbolic fluxes
dispatches to a function in OceanBoundaryConditions.jl based on bytype defined by a problem such as SimpleBoxProblem.jl
"""

@inline function boundary_state!(nf, cm::Continuity3dModel, args...)
    # hack for handling multiple boundaries for now
    # will fix with a future update
    boundary_conditions = (
        cm.ocean.problem.boundary_conditions[1],
        cm.ocean.problem.boundary_conditions[1],
        cm.ocean.problem.boundary_conditions[1],
    )
    return ocean_boundary_state!(nf, boundary_conditions, cm, args...)
end
#=
@inline function boundary_state!(
    ::NumericalFluxFirstOrder,
    ::Continuity3dModel,
    Q⁺,
    A⁺,
    n⁻,
    Q⁻,
    A⁻,
    bctype,
    t,
    args...,
)
    Q⁺.u = Q⁻.u

    return nothing
end
=#


using Test

@testset "CustomMAC" begin
    @test mac(n"Si", n"O K-L3", :Henke1974) == 8790.0
    @test mac(n"U", n"B K-L2", :Henke1982) == 9020.0
    @test mac(n"Ta", n"N K-L3", :Bastin1988) == 15000.0
    @test ismissing(mac(n"W", n"F K-L3", :Bastin1988))

    cust1 = getcustommacs(:Bastin1988, true)
    cust2 = getcustommacs(:Bastin1988, false)

    @test !isnothing(findfirst(x -> x == (n"Si", n"N K-L3", 17170.0), cust1))
    @test !isnothing(findfirst(x -> x == (n"Si", n"N K-L2", 17170.0), cust1))
    @test isnothing(findfirst(x -> x == (n"Si", n"N K-L2", 17170.0), cust2))
    @test !isnothing(findfirst(x -> x == (n"Zr", n"N K-L3", 24000.0), cust1))

    clear_user_macs!()
    for ( el, cxr, mac) in getcustommacs(:Bastin1988, true)
        set_user_mac!(el, cxr, mac)
    end
    @test mac(n"Si", n"N K-L3") == 17170.0
    @test mac(n"Zr", n"N K-L3") == 24000.0
    @test mac(n"Zr", n"O K-L3") == mac(n"Zr", n"O K-L3", FFASTDB)

    for ( el, cxr, mac) in getcustommacs(:Bastin1989, true)
        set_user_mac!(el, cxr, mac)
    end
    @test mac(n"Zr", n"O K-L3") == 16200.0
    @test mac(n"Zr", n"O K-L2") == 16200.0
    @test mac(n"Pb", n"O K-L3") == 11000.0
    delete_user_mac!(n"Ga", n"O K-L2")
    @test mac(n"Ga", n"O K-L2") == mac(n"Ga", n"O K-L2", FFASTDB)

    clear_user_macs!()
    @test mac(n"Zr", n"O K-L3") == mac(n"Zr", n"O K-L3", FFASTDB)
end


# Declarations mirroring the C declarations in xsum.h

struct xsum_small_accumulator
    chunk::NTuple{67,Int64}
    inf::Int64
    nan::Int64
    adds_until_propagate::Cint
end

struct xsum_large_accumulator
    chunk::NTuple{4096,UInt64}
    count::NTuple{4096,Int16}
    chunks_used::NTuple{64,UInt64}
    used_used::UInt64
    sacc::xsum_small_accumulator
end

# The following low-level functions do little checking of their arguments, e.g. they don't check
# that the array strides == 1 or that the dot product is of vectors with equal lengths.
# They do check that length(array) ≤ typemax(Cint), however.
for kind in ("xsum_small_", "xsum_large_")
    T = Symbol(kind, "accumulator")
    @eval begin
        $(Symbol(kind,"init"))(acc) =
            ccall(($(QuoteNode(Symbol(kind,"init"))),libxsum), Cvoid, (Ref{$T},), acc)
        $(Symbol(kind,"addv"))(acc, a::StridedVector{Float64}) =
            ccall(($(QuoteNode(Symbol(kind,"addv"))),libxsum), Cvoid, (Ref{$T},Ptr{Float64},Cint), acc, a, length(a))
        $(Symbol(kind,"add_sqnorm"))(acc, a::StridedVector{Float64}) =
            ccall(($(QuoteNode(Symbol(kind,"add_sqnorm"))),libxsum), Cvoid, (Ref{$T},Ptr{Float64},Cint), acc, a, length(a))
        $(Symbol(kind,"add_dot"))(acc, a::StridedVector{Float64}, b::StridedVector{Float64}) =
            ccall(($(QuoteNode(Symbol(kind,"add_dot"))),libxsum), Cvoid, (Ref{$T},Ptr{Float64},Ptr{Float64},Cint), acc, a, b, length(a))
        $(Symbol(kind,"round"))(acc) =
            ccall(($(QuoteNode(Symbol(kind,"round"))),libxsum), Cdouble, (Ref{$T},), acc)
        $(Symbol(kind,"chunks_used"))(acc) =
            ccall(($(QuoteNode(Symbol(kind,"chunks_used"))),libxsum), Cint, (Ref{$T},), acc)
    end
end

xsum_small_add1(acc, x::Real) =
    ccall((:xsum_small_add1,libxsum), Cvoid, (Ref{xsum_small_accumulator}, Cdouble), acc, x)


# define new aggregator algorithms here and specify their properties

# Direct type methods

"""
Gillespie's Direct method. `ConstantRateJump` rates and affects are stored in
tuples. Fastest for a small (total) number of `ConstantRateJump`s or
`MassActionJump`s (~10). For larger numbers of possible jumps use other
methods.

Gillespie, Daniel T. (1976). A General Method for Numerically Simulating the
Stochastic Time Evolution of Coupled Chemical Reactions. Journal of
Computational Physics. 22 (4): 403–434. doi:10.1016/0021-9991(76)90041-3.
"""
struct Direct <: AbstractAggregatorAlgorithm end

"""
Gillespie's Direct method. `ConstantRateJump` rates and affects are stored via
`FunctionWrappers`, which is more performant than `Direct` for very large
numbers of `ConstantRateJump`s. However, for such large numbers of jump
different classes of aggregators are usually much more performant (i.e.
`SortingDirect`, `DirectCR`, `RSSA` or `RSSACR`).

Gillespie, Daniel T. (1976). A General Method for Numerically Simulating the
Stochastic Time Evolution of Coupled Chemical Reactions. Journal of
Computational Physics. 22 (4): 403–434. doi:10.1016/0021-9991(76)90041-3.
"""
struct DirectFW <: AbstractAggregatorAlgorithm end

"""
The Composition-Rejection Direct method. Performs best relative to other methods
for systems with large numbers of jumps with special structure (for example a
linear chain of reactions, or jumps corresponding to particles hopping on a grid
or graph).

- A. Slepoy, A.P. Thompson and S.J. Plimpton, A constant-time kinetic Monte
  Carlo algorithm for simulation of large biochemical reaction networks, Journal
  of Chemical Physics, 128 (20), 205101 (2008). doi:10.1063/1.2919546

- S. Mauch and M. Stalzer, Efficient formulations for exact stochastic
  simulation of chemical systems, ACM Transactions on Computational Biology and
  Bioinformatics, 8 (1), 27-35 (2010). doi:10.1109/TCBB.2009.47
"""
struct DirectCR <: AbstractAggregatorAlgorithm end

"""
The Sorting Direct method. Often the fastest algorithm for smaller to moderate
sized systems (tens of jumps), or systems where a few jumps occur much more
frequently than others.

J. M. McCollum, G. D. Peterson, C. D. Cox, M. L. Simpson and N. F. Samatova, The
  sorting direct method for stochastic simulation of biochemical systems with
  varying reaction execution behavior, Computational Biology and Chemistry, 30
  (1), 39049 (2006). doi:10.1016/j.compbiolchem.2005.10.007
"""
struct SortingDirect <: AbstractAggregatorAlgorithm end

"""
The Rejection SSA method. One of the best methods for systems with hundreds to
many thousands of jumps (along with `RSSACR`) and sparse dependency graphs.

- V. H. Thanh, C. Priami and R. Zunino, Efficient rejection-based simulation of
  biochemical reactions with stochastic noise and delays, Journal of Chemical
  Physics, 141 (13), 134116 (2014). doi:10.1063/1.4896985

- V. H. Thanh, R. Zunino and C. Priami, On the rejection-based algorithm for
  simulation and analysis of large-scale reaction networks, Journal of Chemical
  Physics, 142 (24), 244106 (2015). doi:10.1063/1.4922923
"""
struct RSSA <: AbstractAggregatorAlgorithm end

"""
The Rejection SSA Composition-Rejection method. Often the best performer for
systems with tens of thousands of jumps and sparse depedency graphs.

V. H. Thanh, R. Zunino, and C. Priami, Efficient Constant-Time Complexity
Algorithm for Stochastic Simulation of Large Reaction Networks, IEEE/ACM
Transactions on Computational Biology and Bioinformatics, Vol. 14, No. 3,
657-667 (2017).
"""
struct RSSACR <: AbstractAggregatorAlgorithm end

"""
A rejection-based direct method.
"""
struct RDirect <: AbstractAggregatorAlgorithm end

# NRM-based methods

"""
Gillespie's First Reaction Method. Should not be used for practical applications
due to slow performance relative to all other methods.

Gillespie, Daniel T. (1976). A General Method for Numerically Simulating the
Stochastic Time Evolution of Coupled Chemical Reactions. Journal of
Computational Physics. 22 (4): 403–434. doi:10.1016/0021-9991(76)90041-3.
"""
struct FRM <: AbstractAggregatorAlgorithm end

"""
Gillespie's First Reaction Method with `FunctionWrappers` for handling
`ConstantRateJump`s. Should not be used for practical applications
due to slow performance relative to all other methods.

Gillespie, Daniel T. (1976). A General Method for Numerically Simulating the
Stochastic Time Evolution of Coupled Chemical Reactions. Journal of
Computational Physics. 22 (4): 403–434. doi:10.1016/0021-9991(76)90041-3.
"""
struct FRMFW <: AbstractAggregatorAlgorithm end

"""
The Next Reaction Method. Can significantly outperform Direct for systems with
large numbers of jumps and sparse dependency graphs, but is usually slower than
one of `DirectCR`, `RSSA`, or `RSSACR` for such systems.

M. A. Gibson and J. Bruck, Efficient exact stochastic simulation of chemical
systems with many species and many channels, Journal of Physical Chemistry A,
104 (9), 1876-1889 (2000). doi:10.1021/jp993732q
"""
struct NRM <: AbstractAggregatorAlgorithm end

# spatial methods

"""
The Next Subvolume Method for spatial jump process simulations. Usually slower
than `DirectCRDirect`. Uses an indexed priority queue tree structure to
determine where on the grid/graph the next jump occurs, and then the `Direct`
method to determine which jump at the given location occurs.

Elf, Johan and Ehrenberg, M, Spontaneous separation of bi-stable biochemical
systems into spatial domains of opposite phases,Systems Biology, 2004 vol. 1(2)
pp. 230-236. doi:10.1049/sb:20045021
"""
struct NSM <: AbstractAggregatorAlgorithm end

"""
The Direct Composition-Rejection Direct method. Uses the `DirectCR` method to
determine where on the grid/graph a jump occurs, and the `Direct` method to
determine which jump occurs at the sampled location.

Constant-complexity stochastic simulation algorithm with optimal binning, Kevin
R. Sanft and Hans G. Othmer, Journal of Chemical Physics 143, 074108 (2015);
doi: 10.1063/1.4928635
"""
struct DirectCRDirect <: AbstractAggregatorAlgorithm end

const JUMP_AGGREGATORS = (Direct(),DirectFW(),DirectCR(),SortingDirect(),RSSA(),FRM(),FRMFW(),NRM(),RSSACR(),RDirect())

# For JumpProblem construction without an aggregator
struct NullAggregator <: AbstractAggregatorAlgorithm end

# true if aggregator requires a jump dependency graph
needs_depgraph(aggregator::AbstractAggregatorAlgorithm) = false
needs_depgraph(aggregator::DirectCR) = true
needs_depgraph(aggregator::SortingDirect) = true
needs_depgraph(aggregator::NRM) = true
needs_depgraph(aggregator::RDirect) = true

# true if aggregator requires a map from solution variable to dependent jumps.
# It is implicitly assumed these aggregators also require the reverse map, from
# jumps to variables they update.
needs_vartojumps_map(aggregator::AbstractAggregatorAlgorithm) = false
needs_vartojumps_map(aggregator::RSSA) = true
needs_vartojumps_map(aggregator::RSSACR) = true

is_spatial(aggregator::AbstractAggregatorAlgorithm) = false
is_spatial(aggregator::NSM) = true
is_spatial(aggregator::DirectCRDirect) = true


export
    SimulationOptions,
    build_environment,
    build_uavs,
    get_joint_actions,
    load_policy,
    ACTIONS

#=
Functions for building the environment from options and default values.
=#

### Simulation settings and constants ###

# These constants do not go in the options class because they
# should essentially never change
const RANGES = [0, 25, 50, 100, 150, 200, 250, 300, 350, 400, 450, 500, 550, 600, 650, 700, 750, 800, 900, 1000, 1200, 1400, 1600, 1800, 2000, 2500, 3000, 4000, 6000, 8000, 10000, 12800]
const ACTIONS = deg2rad([-12, -6, 0, 6, 12, rad2deg(-1)])
const THETA_DIM = 37
const THETAS = collect(linspace(-pi, pi, THETA_DIM))
const BEARING_DIM = 37
const BEARINGS = collect(linspace(0, 2 * pi, BEARING_DIM))
const SPEED_DIM = 3
const SPEEDS = collect(linspace(35, 45, SPEED_DIM))
const RESPONDINGS = [0, 1]

### Helper functions for building the environment ###

# type for default simulation settings
type SimulationOptions
    # coordination table policy
    coordination_table_one_policy::Bool
    coordination_policy_filepath::String

    # ownship options
    ownship_policy_filepath::String
    is_ownship_policy_joint::Bool
    policy_name::String
    ownship_strategy::Function # {greedy_strategy, coordinated_strategy}
    ownship_action_selection::String # {best*, worst*, average*}
    ownship_range_sensor_std::Float64
    ownship_theta_sensor_std::Float64
    ownship_psi_sensor_std::Float64
    ownship_velocity_sensor_std::Float64
    use_ownship_sensor_noise::Bool

    # intruder options
    intruder_policy_filepath::String
    is_intruder_policy_joint::Bool
    intruder_strategy::Function # {greedy_strategy, coordinated_strategy}
    intruder_action_selection::String # {best*, worst*, average*}
    intruder_range_sensor_std::Float64
    intruder_theta_sensor_std::Float64
    intruder_psi_sensor_std::Float64
    intruder_velocity_sensor_std::Float64
    use_intruder_sensor_noise::Bool

    # environment options
    coordination_message_scheme::Int64
    start_state_generator::Function
    use_noise::Bool
    random_seed::Int64
    gauss_sigma::Float64
    gauss_mu::Float64
    max_steps::Int64
    min_start_range_ratio::Float64
    max_start_range_ratio::Float64
    range_std::Float64
    initial_heading_offset::Float64
    initial_min_separation::Float64
    delayed_pilot_response::Bool

    # dynamics options
    omega::Float64
    dt::Int64
    dt_step::Float64
    pid_b_std::Float64
    pid_b_std_coc::Float64

    function SimulationOptions()
        # optimal policy used for coordination table
        coordination_table_one_policy = true ## currently this is whether or not solve real time...
        table_filepath = joinpath(Pkg.dir("Coord"), "data", "qvalue_tables", "optimal.jld")
        coordination_policy_filepath = table_filepath

        # ownship options
        ownship_policy_filepath = table_filepath
        is_ownship_policy_joint = true # {false = non-joint, true = joint}

        policy_name = "solQ"
        ownship_strategy = greedy_strategy # {greedy_strategy, coordinated_strategy}
        ownship_action_selection = "average" # {best, worst, average}

        # intruder options
        intruder_policy_filepath = table_filepath
        is_intruder_policy_joint = true # {false = non-joint, true = joint}

        intruder_strategy = greedy_strategy
        intruder_action_selection = "average" # {best, worst, average}
        if ownship_strategy == greedy_strategy
            intruder_strategy = greedy_strategy
        end

        # environment options
        # coordination messaging scheme selection
        coordination_message_scheme = 0 # {0, 1, 2, 3, 4}

        # generators: {start_towards_origin_encounter, start_head_on_encounter}
        start_state_generator = start_towards_origin_encounter
        use_noise = true

        # for running on sherlock with multiple subtasks
        # get the id of this subtask to use as the seed
        random_seed = 1

        gauss_sigma = deg2rad(2)
        gauss_mu = 0
        max_steps = 20
        min_start_range_ratio = 5 / 64 # gives min range of 2000
        max_start_range_ratio = 5 / 32 # gives max range of 2000
        range_std = 100.
        initial_heading_offset = pi / 32
        initial_min_separation = 153.
        delayed_pilot_response = false

        # ownship and intruder sensor noise defaults
        ownship_range_sensor_std = 15. # meters
        ownship_theta_sensor_std = deg2rad(1.) # rad
        ownship_psi_sensor_std = deg2rad(1.) # rad
        ownship_velocity_sensor_std = 1. # m/s
        use_ownship_sensor_noise = true
        intruder_range_sensor_std = 15. # meters
        intruder_theta_sensor_std = deg2rad(1.) # rad
        intruder_psi_sensor_std = deg2rad(1.) # rad
        intruder_velocity_sensor_std = 1. # m/s
        use_intruder_sensor_noise = true

        # dynamics options
        omega = .2
        dt = 5
        dt_step = 1.
        pid_b_std = deg2rad(3)
        pid_b_std_coc = deg2rad(5)

        # build the default options
        return new(coordination_table_one_policy, coordination_policy_filepath, ownship_policy_filepath,
            is_ownship_policy_joint, policy_name, ownship_strategy, ownship_action_selection,
            ownship_range_sensor_std, ownship_theta_sensor_std, ownship_psi_sensor_std,
            ownship_velocity_sensor_std, use_ownship_sensor_noise, intruder_policy_filepath,
            is_intruder_policy_joint, intruder_strategy, intruder_action_selection,
            intruder_range_sensor_std, intruder_theta_sensor_std, intruder_psi_sensor_std,
            intruder_velocity_sensor_std, use_intruder_sensor_noise, coordination_message_scheme,
            start_state_generator, use_noise, random_seed, gauss_sigma, gauss_mu, max_steps,
            min_start_range_ratio, max_start_range_ratio, range_std, initial_heading_offset,
            initial_min_separation, delayed_pilot_response, omega, dt, dt_step, pid_b_std, pid_b_std_coc)
    end
end

#=
Description:
Returns an enviornment object

Return Value:
- environment: an environment object
=#
function build_environment()
    # should instead gather options from cmd line
    options = SimulationOptions()

    # build uavs and dynamics
    uavs = build_uavs(options)
    dynamics = build_dynamics(options)

    # build actual environment
    max_steps = options.max_steps
    use_noise = options.use_noise
    gauss_sigma = options.gauss_sigma
    gauss_mu = options.gauss_mu
    coordination_message_scheme = options.coordination_message_scheme
    start_state_generator = options.start_state_generator
    min_start_range_ratio = options.min_start_range_ratio
    max_start_range_ratio = options.max_start_range_ratio
    range_std = options.range_std
    initial_heading_offset = options.initial_heading_offset
    initial_min_separation = options.initial_min_separation
    delayed_pilot_response = options.delayed_pilot_response

    rng = MersenneTwister(options.random_seed)
    state_generator_rng = MersenneTwister(options.random_seed)
    e = Environment(max_steps, use_noise, uavs, gauss_sigma,
                    gauss_mu, dynamics, coordination_message_scheme, start_state_generator,
                    min_start_range_ratio, max_start_range_ratio, range_std,
                    initial_heading_offset, initial_min_separation,
                    delayed_pilot_response, rng, state_generator_rng)
    return e
end

#=
Description:
Make joint actions given list of actions

Parameters:
- actions: list of float valued actions

Return Value:
- array of actions such that each row corresponds to a uav
    and col to a joint action index, with the cell value being the action
=#
function get_joint_actions(actions)
    num_actions = length(actions)
    joint_actions = zeros(2, num_actions^2)
    i_actions = 1
    for ia2 = 1:num_actions
        for ia1 = 1:num_actions
            joint_actions[1, i_actions] = actions[ia1]
            joint_actions[2, i_actions] = actions[ia2]
            i_actions += 1
        end
    end
    return joint_actions
end

#=
Description:
Load in a policy from a filepath

Parameters:
- filepath: filepath to policy
- options: options type with info to build policy

Return Value:
- policy_table type object
=#
function load_policy(filepath::String, is_policy_joint::Bool, options::SimulationOptions)
    if is_policy_joint # joint policy
        q_table = load(filepath, options.policy_name)
        q_grid = RectangleGrid(RANGES, THETAS, BEARINGS, SPEEDS, SPEEDS, RESPONDINGS, RESPONDINGS)
        num_q_table_states = reduce(*, q_grid.cut_counts) + 1
        msg = "the q_table (rows = $(size(q_table, 1))) does not match the
            q_grid (length = $num_q_table_states), which was created using constant values"
        @assert num_q_table_states == size(q_table, 1) msg
        policy = PolicyTable(q_table, q_grid, num_q_table_states)
        return policy
    else # non joint policy
        q_table = load(filepath, options.policy_name)
        q_grid = RectangleGrid(RANGES, THETAS, BEARINGS, SPEEDS, SPEEDS, RESPONDINGS)
        num_q_table_states = reduce(*, q_grid.cut_counts) + 1
        msg = "the q_table (rows = $(size(q_table, 1))) does not match the
            q_grid (length = $num_q_table_states), which was created using constant values"
        @assert num_q_table_states == size(q_table, 1) msg
        policy = PolicyTable(q_table, q_grid, num_q_table_states)
        return policy
    end
end

#=
Description:
Returns a list of uav objects

Return Value:
- uavs: list of uav objects
=#
function build_uavs(options::SimulationOptions)
    # rng - this is for the sensor noise
    m = MersenneTwister(options.random_seed)

    # alias this constant value
    actions = ACTIONS

    # load ownship policy
    policy = load_policy(options.ownship_policy_filepath, options.is_ownship_policy_joint, options)

    # determine the joint actions
    joint_actions = get_joint_actions(actions)

    # build the coordination table
    if options.coordination_table_one_policy
        optimal_policy = load(options.coordination_policy_filepath, options.policy_name)
        c_table = make_coordination_table(optimal_policy, joint_actions)
        c_grid = RectangleGrid(RANGES, THETAS, BEARINGS, SPEEDS, SPEEDS, RESPONDINGS, RESPONDINGS)
        num_coord_states = reduce(*, c_grid.cut_counts)
        coord = CoordinationTable(c_table, c_grid, num_coord_states)
    else
        c_table = load(options.coordination_policy_filepath, options.policy_name)
        c_grid = RectangleGrid(RANGES, THETAS, BEARINGS, SPEEDS, SPEEDS, RESPONDINGS, RESPONDINGS)
        num_coord_states = reduce(*, c_grid.cut_counts)
        coord = CoordinationTable(c_table, c_grid, num_coord_states)
    end

 # create the ownship
    ownship_id = 1
    ownship_uav = UAV(ownship_id, actions, joint_actions, policy, coord, options.is_ownship_policy_joint,
                    options.ownship_strategy, options.ownship_action_selection, m,
                    options.ownship_range_sensor_std, options.ownship_theta_sensor_std,
                    options.ownship_psi_sensor_std, options.ownship_velocity_sensor_std,
                    options.use_ownship_sensor_noise)

    # create the intruder
    if options.intruder_policy_filepath == options.ownship_policy_filepath
        intruder_policy = policy
    else
        intruder_policy = load_policy(options.intruder_policy_filepath, options.is_intruder_policy_joint, options)
    end

    intruder_id = 2
    strategy = options.intruder_strategy
    intruder_uav = UAV(intruder_id, actions, joint_actions, intruder_policy, coord,
                    options.is_intruder_policy_joint, strategy, options.intruder_action_selection,
                    m, options.intruder_range_sensor_std, options.intruder_theta_sensor_std,
                    options.intruder_psi_sensor_std, options.intruder_velocity_sensor_std,
                    options.use_intruder_sensor_noise)

    return [ownship_uav, intruder_uav]
end

#=
Description:
Returns a dynamics object

Return Value:
- dynamics: a dynamics object
=#
function build_dynamics(options::SimulationOptions)
    m = MersenneTwister(options.random_seed)
    omega = options.omega
    use_noise = options.use_noise
    dt = options.dt
    dt_step = options.dt_step
    pid_b_std = options.pid_b_std
    pid_b_std_coc = options.pid_b_std_coc
    d = Dynamics(dt, dt_step, omega, pid_b_std, pid_b_std_coc, use_noise, m)
    return d
end


# A code to save the modeling data in the MATLAB structure to be used for EEG
# analysis
using SurNoR_2020
using DataFrames
using MAT

PyPlot.svg(true)
rcParams = PyPlot.PyDict(PyPlot.matplotlib."rcParams")
rcParams["svg.fonttype"] = "none"
rcParams["pdf.fonttype"] = 42

for opt_i = 3
    # Training Algorithms
    Agents_SurNoR   = SurNoR_2020.Func_SurNoR_Overall_train_opt_ind_i(opt_i)
    Agents_Tot_Alg = cat(Agents_SurNoR, dims=4)
    Path_save = "data/ModelData_for_EEG/"

    # MATLAB Version of struct
    for Sub = 1:12
        for Block = 1:2
            for Epi = 1:5
                if Epi==1

                    global Agents_Epi_MATLAB = SurNoR_2020.Str_Agent_MAT(Agent=deepcopy(Agents_SurNoR[Sub,Block,Epi]))
                else
                    Agents_Epi_MATLAB = cat(Agents_Epi_MATLAB,
                                            SurNoR_2020.Str_Agent_MAT(Agent=deepcopy(Agents_SurNoR[Sub,Block,Epi])),
                                            dims=3)
                end
            end
            if Block==1
                global Agents_Block_MATLAB = Agents_Epi_MATLAB
            else
                Agents_Block_MATLAB = cat(Agents_Block_MATLAB,
                                          Agents_Epi_MATLAB, dims=2)
            end
        end
        if Sub==1
            global Agents_Sub_MATLAB = Agents_Block_MATLAB
        else
            Agents_Sub_MATLAB = cat(Agents_Sub_MATLAB,
                                    Agents_Block_MATLAB, dims=1)
        end
    end

    Path = string(Path_save,"Trained_Data", string(opt_i) ,".mat" )
    file = matopen(Path, "w")
    write(file, "Data", Agents_Sub_MATLAB)
    close(file)
end


@testset "SubSetSimulation" begin
    proposal = Normal()
    subset = SubSetSimulation(2000, 0.2, 10, proposal)

    @test isa(subset, SubSetSimulation)
    @test subset.n == 2000
    @test subset.target == 0.2
    @test subset.levels == 10
    @test subset.proposal == proposal

    @test_throws ErrorException("proposal must be a symmetric distribution") SubSetSimulation(
        2000, 0.2, 10, Exponential()
    )
end


    using Revise, Parameters, Logging, Test
    using CUDA
    includet("C:\\GitHub\\GitHub\\NuclearMedEval\\test\\includeAllUseFullForTest.jl")
    using ..CUDAGpuUtils ,..IterationUtils,..ReductionUtils , ..MemoryUtils,..CUDAAtomicUtils, ..PrepareArrtoBool, ..BitWiseUtils
    using ..ResultListUtils, ..MetadataAnalyzePass,..MetaDataUtils,..WorkQueueUtils,..ProcessMainDataVerB,..HFUtils, ..ScanForDuplicates, ..Housdorff

    mainArrDims= (512,512,500)

    arrGold = zeros(Int32,mainArrDims);
    arrSegm = zeros(Int32,mainArrDims);

        # so we should get all mins as 2 and all maxes as 5

        arrGold[9,2,5]= 2
        arrGold[3,6,4]= 2
        arrGold[3,3,4]= 2

        arrGold[4,4,4]= 2
        arrSegm[4,4,4]= 2

        arrSegm[3,18,4]= 2
        arrSegm[5,2,11]= 2

        gold3d = CuArray(arrGold)
        segm3d= CuArray(arrSegm)
        numberToLooFor =Int32(2)
        robustnessPercent = 0.9

        boolKernelArgs, mainKernelArgs,threadsBoolKern,blocksBoolKern ,threadsMainKern,blocksMainKern ,shmemSizeBool,shmemSizeMain=    preparehousedorfKernel(gold3d,segm3d,robustnessPercent,numberToLooFor)
        mainArrDims,dataBdim,metaData,metaDataDims,reducedGoldA,reducedSegmA,loopXinPlane,loopYinPlane,minxRes,maxxRes,minyRes,maxyRes,minzRes,maxzRes,fn,fp ,numberToLooFor,inBlockLoopXZIterWithPadding,shmemblockDataLoop,shmemblockDataLenght,loopAXFixed,loopBXfixed,loopAYFixed,loopBYfixed,loopAZFixed,loopBZfixed,loopdataDimMainX,loopdataDimMainY,loopdataDimMainZ,inBlockLoopX,inBlockLoopY,inBlockLoopZ,metaDataLength,loopMeta,loopWarpMeta,clearIterResShmemLoop,clearIterSourceShmemLoop,resShmemTotalLength,sourceShmemTotalLength=boolKernelArgs

        arrGold[dataBdim[1]*3+1,dataBdim[2]*4+1,dataBdim[3]*5+1]= 2
        arrSegm[dataBdim[1]*3+1,dataBdim[2]*4+1,dataBdim[3]*5+1]= 2

        golddd = CuArray(arrGold);
        segmmm= CuArray(arrSegm);

        # using CUDA
        # ff = CUDA.zeros(1)
        # function gggg(ff)
        # #   shmemblockData = @cuDynamicSharedMem(Float32,(dataBdim[1], dataBdim[2]))
        # #   shmemblockData = @cuDynamicSharedMem(Float32,(1,2))
        #    shmemblockData = CuDynamicSharedArray(Int, 6)
        #     return
        # end
        # @cuda threads=(32,2) blocks=2 gggg(ff)
        # ff[1]

        metaDataDims
        metaDataLength
        metaDataLength/blocksBoolKern
        loopMeta
        dataBdim
        Int64(blocksBoolKern)
        function locForKernel(goldGPU,segmGPU,mainArrDims,dataBdim,metaData,metaDataDims,reducedGoldA,reducedSegmA,loopXinPlane,loopYinPlane,minxRes,maxxRes,minyRes,maxyRes,minzRes,maxzRes,fn,fp ,numberToLooFor,inBlockLoopXZIterWithPadding,shmemblockDataLoop,shmemblockDataLenght,loopAXFixed,loopBXfixed,loopAYFixed,loopBYfixed,loopAZFixed,loopBZfixed,loopdataDimMainX,loopdataDimMainY,loopdataDimMainZ,inBlockLoopX,inBlockLoopY,inBlockLoopZ,metaDataLength,loopMeta,loopWarpMeta,clearIterResShmemLoop,clearIterSourceShmemLoop,resShmemTotalLength,sourceShmemTotalLength)

          @getBoolCubeKernel()

        return
    end
    minxRes[1]

        @cuda shmem=shmemSizeBool threads=threadsBoolKern blocks=blocksBoolKern locForKernel(golddd,segmmm,mainArrDims,dataBdim,metaData,metaDataDims,reducedGoldA,reducedSegmA,loopXinPlane,loopYinPlane,minxRes,maxxRes,minyRes,maxyRes,minzRes,maxzRes,fn,fp ,numberToLooFor,inBlockLoopXZIterWithPadding,shmemblockDataLoop,shmemblockDataLenght,loopAXFixed,loopBXfixed,loopAYFixed,loopBYfixed,loopAZFixed,loopBZfixed,loopdataDimMainX,loopdataDimMainY,loopdataDimMainZ,inBlockLoopX,inBlockLoopY,inBlockLoopZ,metaDataLength,loopMeta,loopWarpMeta,clearIterResShmemLoop,clearIterSourceShmemLoop,resShmemTotalLength,sourceShmemTotalLength)
          Int64( maxzRes[1])

        @test Int64(fn[])==3
        @test Int64(fp[])==2

       Int64(minxRes[1])
         Int64( maxxRes[])
         Int64(minyRes[])
          Int64(maxyRes[])
          Int64( minzRes[])
          Int64( maxzRes[])

          @test  Int64(minxRes[])==1
          @test  Int64( maxxRes[])==4
          @test  Int64(minyRes[])==1
          @test   Int64(maxyRes[])==5
          @test  Int64( minzRes[])==1
          @test  Int64( maxzRes[])==6

        #   mainArrDims
        # @test Int64(minxRes[])==2
        # @test Int64( maxxRes[])==5
        # @test Int64(minyRes[])==1
        # @test  Int64(maxyRes[])==9
        # @test  Int64( minzRes[])==2
        # @test  Int64( maxzRes[])==6
        # arrGold[32*3+1,32*4+1,32*5+1]
        # arrSegm[32*3+1,32*4+1,32*5+1]

Int64(sum(reducedGoldA))

Int64(reducedGoldA[4,4,1])
Int64(reducedSegmA[4,4,1])

Int64(reducedGoldA[3,18,1])

        @test  isBit1AtPos(reducedGoldA[9,2,1],5)==true

        @test  isBit1AtPos(reducedGoldA[3,3,1],4)==true
        @test  isBit1AtPos(reducedGoldA[3,3,1],4)==true
        @test  isBit1AtPos(reducedGoldA[4,4,1],4)==true
        @test  isBit1AtPos(reducedGoldA[5,5,1],5)==false
        @test  isBit1AtPos(reducedGoldA[9,2,1],5)==true
        @test  isBit1AtPos(reducedSegmA[3,18,1],4)==true
        @test  isBit1AtPos(reducedGoldA[3,18,1],4)==false

        @test  isBit1AtPos(reducedSegmA[dataBdim[1]*3+1,dataBdim[2]*4+1,6],1)==true

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

 using Revise, Parameters, Logging, Test
 using CUDA
 includet("C:\\GitHub\\GitHub\\NuclearMedEval\\test\\includeAllUseFullForTest.jl")
 using ..CUDAGpuUtils ,..IterationUtils,..ReductionUtils , ..MemoryUtils,..CUDAAtomicUtils, ..PrepareArrtoBool, ..BitWiseUtils
 using ..ResultListUtils, ..MetadataAnalyzePass,..MetaDataUtils,..WorkQueueUtils,..ProcessMainDataVerB,..HFUtils, ..ScanForDuplicates, ..Housdorff

 mainArrDims= (60,60,60);

 mainArrCPU= zeros(UInt8,mainArrDims);
 refArrCPU = zeros(UInt8,mainArrDims);
 ##### we will create two planes 20 units apart from each
mainArrCPU[10:50,10:50,10].= 1;
refArrCPU[10:50,10:50,30].= 1;

     gold3d = CuArray(mainArrCPU)
     segm3d= CuArray(refArrCPU)
     numberToLooFor =Int32(1)
     robustnessPercent = 0.9

     boolKernelArgs, mainKernelArgs,threadsBoolKern,blocksBoolKern ,threadsMainKern,blocksMainKern ,shmemSizeBool,shmemSizeMain=    preparehousedorfKernel(gold3d,segm3d,robustnessPercent,numberToLooFor)
     mainArrDims,dataBdim,metaData,metaDataDims,reducedGoldA,reducedSegmA,loopXinPlane,loopYinPlane,minxRes,maxxRes,minyRes,maxyRes,minzRes,maxzRes,fn,fp ,numberToLooFor,inBlockLoopXZIterWithPadding,shmemblockDataLoop,shmemblockDataLenght,loopAXFixed,loopBXfixed,loopAYFixed,loopBYfixed,loopAZFixed,loopBZfixed,loopdataDimMainX,loopdataDimMainY,loopdataDimMainZ,inBlockLoopX,inBlockLoopY,inBlockLoopZ,metaDataLength,loopMeta,loopWarpMeta,clearIterResShmemLoop,clearIterSourceShmemLoop,resShmemTotalLength,sourceShmemTotalLength=boolKernelArgs

     golddd = CuArray(mainArrCPU);
     segmmm= CuArray(refArrCPU);

     @cuda shmem=shmemSizeBool threads=threadsBoolKern blocks=blocksBoolKern boolKernelLoad(golddd,segmmm,boolKernelArgs...)
     Int64(sum(mainArrCPU))
     Int64(fn[])
     @test Int64(fn[])==Int64(sum(mainArrCPU))


function get_acceleration(obj::PropulsionModel, arg0::FieldSpacecraftState, arg1::FieldAttitude, arg2::Vector{RealFieldElement})
    return jcall(obj, "getAcceleration", FieldVector3D, (FieldSpacecraftState, FieldAttitude, Vector{RealFieldElement}), arg0, arg1, arg2)
end

function get_acceleration(obj::PropulsionModel, arg0::SpacecraftState, arg1::Attitude, arg2::Vector{jdouble})
    return jcall(obj, "getAcceleration", Vector3D, (SpacecraftState, Attitude, Vector{jdouble}), arg0, arg1, arg2)
end

function get_mass_derivatives(obj::PropulsionModel, arg0::FieldSpacecraftState, arg1::Vector{RealFieldElement})
    return jcall(obj, "getMassDerivatives", RealFieldElement, (FieldSpacecraftState, Vector{RealFieldElement}), arg0, arg1)
end

function get_mass_derivatives(obj::PropulsionModel, arg0::SpacecraftState, arg1::Vector{jdouble})
    return jcall(obj, "getMassDerivatives", jdouble, (SpacecraftState, Vector{jdouble}), arg0, arg1)
end

function get_name(obj::PropulsionModel)
    return jcall(obj, "getName", JString, ())
end


# strongly inspired by https://github.com/s-macke/VoxelSpace
module VoxelSpace
using MiniFB, ColorTypes, PNGFiles, Printf
using Base: unsafe_trunc

include("map_utils.jl")
include("render.jl")
include("input.jl")

const VIEW_DIST = 1024
const WIN_WIDTH = 1200
const WIN_HEIGHT = 900
const MAP_WIDTH = 1024
const MAP_HEIGHT = 1024
const ZSCALE = 1024f0/2
const HORIZON = 1024f0/2

function run(map="C1W")
    # Load map and create framebuffer
    datac, datah = read_map(map)
    fbuffer = Matrix{RGB24}(undef, WIN_WIDTH, WIN_HEIGHT)
    hbuffer = Vector{Int}(undef, WIN_WIDTH)

    # Create window
    window = mfb_open("VoxelSpace.jl", WIN_WIDTH, WIN_HEIGHT)

    # Set up input callbacks
    key_cb_fun_c = @cfunction(key_cb_fun, Cvoid, (Ptr{Cvoid}, mfb_key, mfb_key_mod, Bool))
    mouse_cb_fun_c = @cfunction(mouse_cb_fun, Cvoid, (Ptr{Cvoid}, Int32, Int32))
    mfb_set_keyboard_callback(window, key_cb_fun_c)
    mfb_set_mouse_move_callback(window, mouse_cb_fun_c)

    # Set up FPS counter
    t1 = time_ns()
    @printf("FPS: %3.0f",0.0)

    # Initial position and orientation
    px = 512f0
    py = 512f0
    pz = 80f0
    θ  = 0f0

    while mfb_wait_sync(window)
        # Compute movement direction and apply to pos
        dx,dy,dz,dθ = compute_movement!(θ)
        px += dx
        py += dy
        pz += dz
        θ  += dθ

        # Render viewport
        render!(fbuffer,hbuffer,datac,datah,(px,py,pz),θ)

        # Compute FPS
        t2 = time_ns()
        @printf("\b\b\b%3.0f", 1E9/(t2-t1))
        t1 = t2

        # Update framebuffer
        state = mfb_update(window, fbuffer)
        state == MiniFB.STATE_OK || break
    end

    mfb_close(window)
    println()

    return nothing
end

end # module


mutable struct IPM_Ensemble
	path::String #ensemble models and popped-out posterior samples serialised here
	models::Vector{Model_Record} #ensemble keeps paths to serialised models and their likelihood tuples rather than keeping the models in memory

	contour::AbstractFloat#current log_Li[end]
	log_Li::Vector{AbstractFloat} #likelihood of lowest-ranked model at iterate i
	log_Xi::Vector{AbstractFloat} #amt of prior mass included in ensemble contour at Li
	log_wi::Vector{AbstractFloat} #width of prior mass covered in this iterate
	log_Liwi::Vector{AbstractFloat} #evidentiary weight of the iterate
	log_Zi::Vector{AbstractFloat} #ensemble evidence
	Hi::Vector{AbstractFloat} #ensemble information

	obs_array::Matrix{Integer} #observations Txo
	obs_lengths::Vector{Integer}

	source_priors::AbstractVector{<:Union{<:AbstractVector{<:Dirichlet{<:AbstractFloat}},<:Bool}} #source pwm priors
	mix_prior::Tuple{BitMatrix,AbstractFloat} #prior on %age of observations that any given source contributes to

	bg_scores::Matrix{AbstractFloat} #precalculated background HMM scores, same dims as obs

	sample_posterior::Bool
	posterior_samples::Vector{Model_Record} #list of posterior sample records

	model_counter::Integer

	naive_lh::AbstractFloat #the likelihood of the background model without any sources
end

####IPM_Ensemble FUNCTIONS
IPM_Ensemble(path::String, no_models::Integer, source_priors::AbstractVector{<:Union{<:AbstractVector{<:Dirichlet{<:AbstractFloat}},<:Bool}}, mix_prior::Tuple{BitMatrix,<:AbstractFloat}, bg_scores::AbstractMatrix{<:AbstractFloat}, obs::AbstractArray{<:Integer}, source_length_limits; posterior_switch::Bool=false) =
IPM_Ensemble(
	path,
	assemble_IPMs(path, no_models, source_priors, mix_prior, bg_scores, obs, source_length_limits)...,
	[-Inf], #L0 = 0
	[0], #X0 = 1
	[-Inf], #w0 = 0
	[-Inf], #Liwi0 = 0
	[-Inf], #Z0 = 0
	[0], #H0 = 0,
	obs,
	[findfirst(iszero,obs[:,o])-1 for o in 1:size(obs)[2]],
	source_priors,
	mix_prior,
	bg_scores, #precalculated background score
	posterior_switch,
	Vector{String}(),
	no_models+1,
	IPM_likelihood(init_logPWM_sources(source_priors, source_length_limits), obs, [findfirst(iszero,obs[:,o])-1 for o in 1:size(obs)[2]], bg_scores, falses(size(obs)[2],length(source_priors))))

IPM_Ensemble(worker_pool::AbstractVector{<:Integer}, path::String, no_models::Integer, source_priors::AbstractVector{<:Union{<:AbstractVector{<:Dirichlet{<:AbstractFloat}},<:Bool}}, mix_prior::Tuple{BitMatrix,<:AbstractFloat}, bg_scores::AbstractMatrix{<:AbstractFloat}, obs::Array{<:Integer}, source_length_limits; posterior_switch::Bool=false) =
IPM_Ensemble(
	path,
	distributed_IPM_assembly(worker_pool, path, no_models, source_priors, mix_prior, bg_scores, obs, source_length_limits)...,
	[-Inf], #L0 = 0
	[0], #X0 = 1
	[-Inf], #w0 = 0
	[-Inf], #Liwi0 = 0
	[-Inf], #Z0 = 0
	[0], #H0 = 0,
	obs,
	[findfirst(iszero,obs[:,o])-1 for o in 1:size(obs)[2]],
	source_priors,
	mix_prior,
	bg_scores, #precalculated background score
	posterior_switch,
	Vector{String}(),
	no_models+1,
	IPM_likelihood(init_logPWM_sources(source_priors, source_length_limits), obs, [findfirst(iszero,obs[:,o])-1 for o in 1:size(obs)[2]], bg_scores, falses(size(obs)[2],length(source_priors))))

function assemble_IPMs(path::String, no_models::Integer, source_priors::AbstractVector{<:Union{<:AbstractVector{<:Dirichlet{<:AbstractFloat}},<:Bool}}, mix_prior::Tuple{BitMatrix,<:AbstractFloat}, bg_scores::AbstractArray{<:AbstractFloat}, obs::AbstractArray{<:Integer}, source_length_limits::UnitRange{<:Integer})
	ensemble_records = Vector{Model_Record}()
	!isdir(path) && mkpath(path)

	@assert size(obs)[2]==size(bg_scores)[2]

	@showprogress 1 "Assembling IPM ensemble..." for model_no in 1:no_models
		model_path = string(path,'/',model_no)
		if !isfile(model_path)
			model = ICA_PWM_Model(string(model_no), source_priors, mix_prior, bg_scores, obs, source_length_limits)
			serialize(model_path, model) #save the model to the ensemble directory
			push!(ensemble_records, Model_Record(model_path,model.log_Li))
		else #interrupted assembly pick up from where we left off
			model = deserialize(model_path)
			push!(ensemble_records, Model_Record(model_path,model.log_Li))
		end
	end

	return ensemble_records, minimum([record.log_Li for record in ensemble_records])
end

function distributed_IPM_assembly(worker_pool::Vector{Int64}, path::String, no_models::Integer, source_priors::AbstractVector{<:Union{<:AbstractVector{<:Dirichlet{<:AbstractFloat}},<:Bool}}, mix_prior::Tuple{BitMatrix,<:AbstractFloat}, bg_scores::AbstractArray{<:AbstractFloat}, obs::AbstractArray{<:Integer}, source_length_limits::UnitRange{<:Integer})
	ensemble_records = Vector{Model_Record}()
	!isdir(path) && mkpath(path)

	@assert size(obs)[2]==size(bg_scores)[2]

    model_chan= RemoteChannel(()->Channel{ICA_PWM_Model}(length(worker_pool)))
    job_chan = RemoteChannel(()->Channel{Union{Tuple,Nothing}}(1))
	put!(job_chan,(source_priors, mix_prior, bg_scores, obs, source_length_limits))

	sequence_workers(worker_pool, worker_assemble, job_chan, model_chan)

	assembly_progress=Progress(no_models, desc="Assembling IPM ensemble...")

	model_counter=check_assembly!(ensemble_records, path, no_models, assembly_progress)

	while model_counter <= no_models
		wait(model_chan)
		candidate=take!(model_chan)
		model = ICA_PWM_Model(string(model_counter), candidate.origin, candidate.sources, candidate.source_length_limits, candidate.mix_matrix, candidate.log_Li)
		model_path=string(path,'/',model_counter)
		serialize(model_path,model)
		push!(ensemble_records, Model_Record(model_path,model.log_Li))
		model_counter+=1
		next!(assembly_progress)
	end

	take!(job_chan),put!(job_chan,nothing)

	return ensemble_records, minimum([record.log_Li for record in ensemble_records])
end
				function check_assembly!(ensemble_records::AbstractVector{<:Model_Record}, path::String, no_models::Integer, assembly_progress::Progress)
					counter=1
					while counter <= no_models
						model_path=string(path,'/',counter)
						if isfile(model_path)
							model=deserialize(model_path)
							push!(ensemble_records, Model_Record(model_path,model.log_Li))
							counter+=1
							next!(assembly_progress)
						else
							return counter
						end
					end
					return counter
				end

				function worker_assemble(job_chan::RemoteChannel, models_chan::RemoteChannel, comms_chan::RemoteChannel)
					put!(comms_chan,myid())
					wait(job_chan)
					params=fetch(job_chan)
					while !(fetch(job_chan) === nothing)
						model=ICA_PWM_Model(string(myid()),params...)
						put!(models_chan,model)
					end
				end

function Base.show(io::IO, e::IPM_Ensemble; nsrc=0, progress=false)
	livec=[model.log_Li for model in e.models]
	maxLH=maximum(livec)
	printstyled(io, "ICA PWM Model Ensemble @ $(e.path)\n", bold=true)
	msg = @sprintf "Contour: %3.6e MaxLH:%3.3e Max/Naive:%3.3e log Evidence:%3.6e" e.contour maxLH (maxLH-e.naive_lh) e.log_Zi[end]
	println(io, msg)
	hist=UnicodePlots.histogram(livec, title="Ensemble Likelihood Distribution")
	show(io, hist)
	println()
	progress && return(nrows(hist.graphics)+6)
end


using ServiceMeshLearning, Flux

hyperparams = Dict(pairs((;
    # Set params
    tag = "test_simple",
    env = :single,

    # Optimized params
    lr_alpha = 5.0e-5,
    target_entropy = 7.5,
    frames = 2,
    γ = 0.93,
    τ = 0.01,
    start_steps = 50000,
    update_after = 1000,
    replay_size = 500000,
    hidden_units_policy = 20,
    hidden_units_value = 120,
    hidden_layers_policy = 2,
    hidden_layers_value = 4,
    batch_size = 100,
    update_freq = 8,
    actfun = elu,
    reward_scaling = 0.5,
    closing_cost = 0.0,

    # Always
    logging = true,
    verbose = true,
    seed = 24837293,
    timesteps = 20_000_000,
    start_policy = :RandomAgent,
)))

# Run RL + baselines locally with logging and different seed

for agent in [:SAC, :SimpleAgent, :K8sAgent, :OracleAgent]
    run_experiment(;
        hyperparams...,
        alg=agent,
    )
end