"""
    Calculators

This module exists to bridge the gap between the `Models` and the `Dynamics`.

Here, we provide functions and types for evaluating and storing quantities obtained from the
`Models`.
In addition any further manipulation of those quantities, such as computing eigenvalues,
is included here.

This module is largely needed to facilitate integration of both ring polymer and
classical dynamics to allow using the same models and functions for both.
Specific ring polymer types are provided that have the extra fields and methods needed
to evaluate the quantities for each bead.
"""
module Calculators

using LinearAlgebra: LinearAlgebra, Hermitian, I, Eigen, tr
using StaticArrays: SMatrix, SVector
using RingPolymerArrays: get_centroid!

using NQCModels: NQCModels, Model, nstates
using NQCModels.AdiabaticModels: AdiabaticModel
using NQCModels.DiabaticModels: DiabaticModel, DiabaticFrictionModel
using NQCModels.FrictionModels: AdiabaticFrictionModel

using NQCDynamics: ndofs

"""
    AbstractCalculator{M<:Model}

Top-level type for all calculators.

Each concrete calculator contains the `Model` and the fields to store the quantities
obtained from the model.
"""
abstract type AbstractCalculator{T,M<:Model} end
abstract type AbstractAdiabaticCalculator{T,M<:AdiabaticModel} <: AbstractCalculator{T,M} end
abstract type AbstractDiabaticCalculator{T,M<:Union{DiabaticFrictionModel,DiabaticModel}} <: AbstractCalculator{T,M} end
abstract type AbstractStaticDiabaticCalculator{T,M} <: AbstractDiabaticCalculator{T,M} end
abstract type AbstractFrictionCalculator{T,M<:AdiabaticFrictionModel} <: AbstractCalculator{T,M} end

NQCModels.nstates(calc::AbstractCalculator) = NQCModels.nstates(calc.model)
NQCModels.eachstate(calc::AbstractCalculator) = NQCModels.eachstate(calc.model)
NQCModels.nelectrons(calc::AbstractCalculator) = NQCModels.nelectrons(calc.model)
NQCModels.eachelectron(calc::AbstractCalculator) = NQCModels.eachelectron(calc.model)
NQCModels.mobileatoms(calc::AbstractCalculator) = NQCModels.mobileatoms(calc.model, size(calc.derivative, 2))
NQCModels.dofs(calc::AbstractCalculator) = NQCModels.dofs(calc.model)

Base.eltype(::AbstractCalculator{T}) where {T} = T

mutable struct DependentField{V,R}
    value::V
    position::R
end
needsupdate(field::DependentField, r) = field.position != r

function Base.getproperty(calc::AbstractCalculator, name::Symbol)
    field = getfield(calc, name)
    if field isa DependentField
        return field.value
    else
        return field
    end
end

function Base.setproperty!(calc::AbstractCalculator, name::Symbol, x)
    field = getfield(calc, name)
    if field isa DependentField
        setfield!(field, :value, x)
    else
        setfield!(calc, name, x)
    end
end

"""
Each of the quantities specified here has functions:
`get_quantity(calculator, r)`
`evaluate_quantity!(calculator, r)!`

The user should access only the former.
This will ensure quantities are correctly evaluated and cached accordingly.

The latter is called by the former and is where the details required to calculate the quantity are found.
"""
const quantities = [
    :potential,
    :derivative,
    :eigen,
    :adiabatic_derivative,
    :nonadiabatic_coupling,

    :traceless_potential,
    :V̄,
    :traceless_derivative,
    :D̄,
    :traceless_adiabatic_derivative,

    :centroid,
    :centroid_potential,
    :centroid_derivative,
    :centroid_eigen,
    :centroid_adiabatic_derivative,
    :centroid_nonadiabatic_coupling,

    :friction,
]

for quantity in quantities
    get_quantity = Symbol(:get_, quantity)
    evaluate_quantity! = Symbol(:evaluate_, quantity, :!)
    field = Expr(:call, :getfield, :calculator, QuoteNode(quantity))

    @eval function $(get_quantity)(calculator, r)

        if needsupdate($field, r)
            copyto!($field.position, r)
            $(evaluate_quantity!)(calculator, r)
        end
        return $field.value
    end
end

struct AdiabaticCalculator{T,M} <: AbstractAdiabaticCalculator{T,M}
    model::M
    potential::DependentField{T,Matrix{T}}
    derivative::DependentField{Matrix{T},Matrix{T}}
    stats::Dict{Symbol,Int}
    function AdiabaticCalculator{T}(model::M, atoms::Integer) where {T,M<:Model}
        potential = zero(T)
        derivative = zeros(T, ndofs(model), atoms)
        position = fill(NaN, ndofs(model), atoms)

        stats = Dict{Symbol,Int}(
            :potential=>0,
            :derivative=>0,
        )

        new{T,M}(
            model,
            DependentField(potential, copy(position)),
            DependentField(derivative, copy(position)),
            stats
        )
    end
end

struct RingPolymerAdiabaticCalculator{T,M} <: AbstractAdiabaticCalculator{T,M}
    model::M
    potential::DependentField{Vector{T},Array{T,3}}
    derivative::DependentField{Array{T,3},Array{T,3}}
    stats::Dict{Symbol,Int}
    function RingPolymerAdiabaticCalculator{T}(model::M, atoms::Integer, beads::Integer) where {T,M<:Model}
        potential = zeros(T, beads)
        derivative = zeros(T, ndofs(model), atoms, beads)
        position = fill(NaN, ndofs(model), atoms, beads)

        stats = Dict{Symbol,Int}(
            :potential=>0,
            :derivative=>0,
        )

        new{T,M}(
            model,
            DependentField(potential, copy(position)),
            DependentField(derivative, copy(position)),
            stats
        )
    end
end

struct DiabaticCalculator{T,M,S,L} <: AbstractDiabaticCalculator{T,M}
    model::M
    potential::DependentField{Hermitian{T,SMatrix{S,S,T,L}},Matrix{T}}
    derivative::DependentField{Matrix{Hermitian{T,SMatrix{S,S,T,L}}},Matrix{T}}
    eigen::DependentField{LinearAlgebra.Eigen{T,T,SMatrix{S,S,T,L},SVector{S,T}},Matrix{T}}
    adiabatic_derivative::DependentField{Matrix{SMatrix{S,S,T,L}},Matrix{T}}
    nonadiabatic_coupling::DependentField{Matrix{SMatrix{S,S,T,L}},Matrix{T}}
    stats::Dict{Symbol,Int}
    function DiabaticCalculator{T}(model::M, atoms::Integer) where {T,M<:Model}
        n = nstates(model)
        mat = NQCModels.DiabaticModels.matrix_template(model, T)
        vec = NQCModels.DiabaticModels.vector_template(model, T)

        potential = Hermitian(zero(mat))
        derivative = [Hermitian(zero(mat)) for _=1:ndofs(model), _=1:atoms]
        eigen = Eigen(zero(vec), zero(mat) + I)
        adiabatic_derivative = [zero(mat) for _ in CartesianIndices(derivative)]
        nonadiabatic_coupling = [zero(mat) for _ in CartesianIndices(derivative)]
        position = fill(NaN, ndofs(model), atoms)

        stats = Dict{Symbol,Int}(
            :potential=>0,
            :derivative=>0,
            :eigen=>0,
            :adiabatic_derivative=>0,
            :nonadiabatic_coupling=>0
        )

        new{T,M,n,n^2}(
            model,
            DependentField(potential, copy(position)),
            DependentField(derivative, copy(position)),
            DependentField(eigen, copy(position)),
            DependentField(adiabatic_derivative, copy(position)),
            DependentField(nonadiabatic_coupling, copy(position)),
            stats
        )
    end
end

struct RingPolymerDiabaticCalculator{T,M,S,L} <: AbstractDiabaticCalculator{T,M}
    model::M
    potential::DependentField{Vector{Hermitian{T,SMatrix{S,S,T,L}}},Array{T,3}}
    derivative::DependentField{Array{Hermitian{T,SMatrix{S,S,T,L}},3},Array{T,3}}
    eigen::DependentField{Vector{LinearAlgebra.Eigen{T,T,SMatrix{S,S,T,L},SVector{S,T}}},Array{T,3}}
    adiabatic_derivative::DependentField{Array{SMatrix{S,S,T,L},3},Array{T,3}}
    nonadiabatic_coupling::DependentField{Array{SMatrix{S,S,T,L},3},Array{T,3}}

    traceless_potential::DependentField{Vector{Hermitian{T,SMatrix{S,S,T,L}}},Array{T,3}}
    V̄::DependentField{Vector{T},Array{T,3}}
    traceless_derivative::DependentField{Array{Hermitian{T,SMatrix{S,S,T,L}},3},Array{T,3}}
    D̄::DependentField{Array{T,3},Array{T,3}}
    traceless_adiabatic_derivative::DependentField{Array{SMatrix{S,S,T,L},3},Array{T,3}}

    centroid::DependentField{Matrix{T},Array{T,3}}
    centroid_potential::DependentField{Hermitian{T,SMatrix{S,S,T,L}},Array{T,3}}
    centroid_derivative::DependentField{Matrix{Hermitian{T,SMatrix{S,S,T,L}}},Array{T,3}}
    centroid_eigen::DependentField{LinearAlgebra.Eigen{T,T,SMatrix{S,S,T,L},SVector{S,T}},Array{T,3}}
    centroid_adiabatic_derivative::DependentField{Matrix{SMatrix{S,S,T,L}},Array{T,3}}
    centroid_nonadiabatic_coupling::DependentField{Matrix{SMatrix{S,S,T,L}},Array{T,3}}

    stats::Dict{Symbol,Int}
    function RingPolymerDiabaticCalculator{T}(model::M, atoms::Integer, beads::Integer) where {T,M<:Model}
        n = nstates(model)
        mat = NQCModels.DiabaticModels.matrix_template(model, T)
        vec = NQCModels.DiabaticModels.vector_template(model, T)

        potential = [Hermitian(zero(mat)) for _=1:beads]
        derivative = [Hermitian(zero(mat)) for _=1:ndofs(model), _=1:atoms, _=1:beads]
        adiabatic_derivative = [zero(mat) for _=1:ndofs(model), _=1:atoms, _=1:beads]
        eigen = [Eigen(zero(vec), zero(mat) + I) for _=1:beads]
        nonadiabatic_coupling = [zero(mat) for _=1:ndofs(model), _=1:atoms, _=1:beads]

        traceless_potential = [Hermitian(zero(mat)) for _=1:beads]
        V̄ = zeros(beads)
        traceless_derivative = [Hermitian(zero(mat)) for _=1:ndofs(model), _=1:atoms, _=1:beads]
        D̄ = zeros(ndofs(model), atoms, beads)
        traceless_adiabatic_derivative = [zero(mat) for _=1:ndofs(model), _=1:atoms, _=1:beads]

        centroid = zeros(T, ndofs(model), atoms)
        centroid_potential = Hermitian(zero(mat))
        centroid_derivative = [Hermitian(zero(mat)) for _=1:ndofs(model), _=1:atoms]
        centroid_eigen = Eigen(zero(vec), zero(mat) + I)
        centroid_adiabatic_derivative = [zero(mat) for _ in CartesianIndices(centroid_derivative)]
        centroid_nonadiabatic_coupling = [zero(mat) for _ in CartesianIndices(centroid_derivative)]

        position = fill(NaN, ndofs(model), atoms, beads)

        stats = Dict{Symbol,Int}(
            :potential=>0,
            :derivative=>0,
            :eigen=>0,
            :adiabatic_derivative=>0,
            :nonadiabatic_coupling=>0,
            :traceless_potential=>0,
            :V̄=>0,
            :traceless_derivative=>0,
            :D̄=>0,
            :traceless_adiabatic_derivative=>0,
            :centroid=>0,
            :centroid_potential=>0,
            :centroid_derivative=>0,
            :centroid_eigen=>0,
            :centroid_adiabatic_derivative=>0,
            :centroid_nonadiabatic_coupling=>0,
        )

        new{T,M,n,n^2}(
            model,
            DependentField(potential, copy(position)),
            DependentField(derivative, copy(position)),
            DependentField(eigen, copy(position)),
            DependentField(adiabatic_derivative, copy(position)),
            DependentField(nonadiabatic_coupling, copy(position)),

            DependentField(traceless_potential, copy(position)),
            DependentField(V̄, copy(position)),
            DependentField(traceless_derivative, copy(position)),
            DependentField(D̄, copy(position)),
            DependentField(traceless_adiabatic_derivative, copy(position)),

            DependentField(centroid, copy(position)),
            DependentField(centroid_potential, copy(position)),
            DependentField(centroid_derivative, copy(position)),
            DependentField(centroid_eigen, copy(position)),
            DependentField(centroid_adiabatic_derivative, copy(position)),
            DependentField(centroid_nonadiabatic_coupling, copy(position)),
            stats
        )
    end
end

function Calculator(model::DiabaticModel, atoms::Integer, t::Type{T}) where {T}
    DiabaticCalculator{t}(model, atoms)
end
function Calculator(model::AdiabaticModel, atoms::Integer, t::Type{T}) where {T}
    AdiabaticCalculator{t}(model, atoms)
end
function Calculator(model::DiabaticModel, atoms::Integer, beads::Integer, t::Type{T}) where {T}
    RingPolymerDiabaticCalculator{t}(model, atoms, beads)
end
function Calculator(model::AdiabaticModel, atoms::Integer, beads::Integer, t::Type{T}) where {T}
    RingPolymerAdiabaticCalculator{t}(model, atoms, beads)
end

function evaluate_potential!(calc::AbstractCalculator, R)
    calc.stats[:potential] += 1
    calc.potential = NQCModels.potential(calc.model, R)
    return nothing
end

function evaluate_potential!(calc::AbstractCalculator, R::AbstractArray{T,3}) where {T}
    calc.stats[:potential] += 1
    @views for i in axes(R, 3)
        calc.potential[i] = NQCModels.potential(calc.model, R[:,:,i])
    end
    return nothing
end

function evaluate_V̄!(calc::RingPolymerDiabaticCalculator, r)
    calc.stats[:V̄] += 1
    potential = get_potential(calc, r)
    for i in 1:length(calc.V̄)
        calc.V̄[i] = tr(potential[i]) / nstates(calc.model)
    end
end

function evaluate_traceless_potential!(calc::RingPolymerDiabaticCalculator, r)
    calc.stats[:traceless_potential] += 1
    n = nstates(calc.model)
    potential = get_potential(calc, r)
    V̄ = get_V̄(calc, r)
    for I in eachindex(potential)
        calc.traceless_potential[I] = Hermitian(SMatrix{n,n}(
            i != j ? potential[I][j,i] : potential[I][j,i] - V̄[I] for j=1:n, i=1:n
        ))
    end
end

function evaluate_centroid!(calc::AbstractCalculator, r::AbstractArray{T,3}) where {T}
    calc.stats[:centroid] += 1
    get_centroid!(calc.centroid, r)
end

function evaluate_centroid_potential!(calc::AbstractCalculator, r::AbstractArray{T,3}) where {T}
    calc.stats[:centroid_potential] += 1
    centroid = get_centroid(calc, r)
    calc.centroid_potential = NQCModels.potential(calc.model, centroid)
end

function evaluate_derivative!(calc::AbstractCalculator, R)
    calc.stats[:derivative] += 1
    NQCModels.derivative!(calc.model, calc.derivative, R)
end

function evaluate_derivative!(calc::AbstractCalculator, R::AbstractArray{T,3}) where {T}
    calc.stats[:derivative] += 1
    @views for i in axes(R, 3)
        NQCModels.derivative!(calc.model, calc.derivative[:,:,i], R[:,:,i])
    end
end

function evaluate_D̄!(calc::RingPolymerDiabaticCalculator, r)
    calc.stats[:D̄] += 1
    derivative = get_derivative(calc, r)
    for I in eachindex(derivative)
        calc.D̄[I] = tr(derivative[I]) / nstates(calc.model)
    end
end

function evaluate_traceless_derivative!(calc::RingPolymerDiabaticCalculator, r)
    calc.stats[:traceless_derivative] += 1
    n = nstates(calc.model)
    derivative = get_derivative(calc, r)
    D̄ = get_D̄(calc, r)
    for I in eachindex(derivative)
        calc.traceless_derivative[I] = Hermitian(SMatrix{n,n}(
            i != j ? derivative[I][j,i] : derivative[I][j,i] - D̄[I] for j=1:n, i=1:n
        ))
    end
end

function evaluate_traceless_adiabatic_derivative!(calc::RingPolymerDiabaticCalculator, r)
    calc.stats[:traceless_adiabatic_derivative] += 1
    n = nstates(calc.model)
    adiabatic_derivative = get_adiabatic_derivative(calc, r)
    D̄ = get_D̄(calc, r)
    for I in eachindex(D̄)
        calc.traceless_adiabatic_derivative[I] = SMatrix{n,n}(
            i != j ? adiabatic_derivative[I][j,i] : adiabatic_derivative[I][j,i] - D̄[I] for j=1:n, i=1:n
        )
    end
end

function evaluate_centroid_derivative!(calc::AbstractCalculator, r::AbstractArray{T,3}) where {T}
    calc.stats[:centroid_derivative] += 1
    centroid = get_centroid(calc, r)
    NQCModels.derivative!(calc.model, calc.centroid_derivative, centroid)
end

function evaluate_eigen!(calc::AbstractDiabaticCalculator, r)
    calc.stats[:eigen] += 1
    potential = get_potential(calc, r)
    eig = LinearAlgebra.eigen(potential)
    corrected_vectors = correct_phase(eig.vectors, calc.eigen.vectors)
    calc.eigen = Eigen(eig.values, corrected_vectors)
    return nothing
end

function evaluate_centroid_eigen!(calc::RingPolymerDiabaticCalculator, r)
    calc.stats[:centroid_eigen] += 1
    potential = get_centroid_potential(calc, r)
    eig = LinearAlgebra.eigen(potential)
    corrected_vectors = correct_phase(eig.vectors, calc.centroid_eigen.vectors)
    calc.centroid_eigen = Eigen(eig.values, corrected_vectors)
    return nothing
end

function evaluate_eigen!(calc::RingPolymerDiabaticCalculator, r)
    calc.stats[:eigen] += 1
    potential = get_potential(calc, r)
    for i=1:length(potential)
        eig = LinearAlgebra.eigen(potential[i])
        corrected_vectors = correct_phase(eig.vectors, calc.eigen[i].vectors)
        calc.eigen[i] = Eigen(eig.values, corrected_vectors)
    end
    return nothing
end

function correct_phase(new_vectors::SMatrix, old_vectors::SMatrix)
    n = size(new_vectors, 1)
    vect = SVector{n}(sign(LinearAlgebra.dot(new_vectors[:,i], old_vectors[:,i])) for i=1:n)
    return new_vectors .* vect'
end

function evaluate_adiabatic_derivative!(calc::AbstractDiabaticCalculator, r)
    calc.stats[:adiabatic_derivative] += 1
    U = get_eigen(calc, r).vectors
    diabatic_derivative = get_derivative(calc, r)
    for I in eachindex(diabatic_derivative)
        calc.adiabatic_derivative[I] = U' * diabatic_derivative[I] * U
    end
end

function evaluate_centroid_adiabatic_derivative!(calc::RingPolymerDiabaticCalculator, r)
    calc.stats[:centroid_adiabatic_derivative] += 1
    centroid_derivative = get_centroid_derivative(calc, r)
    centroid_eigen = get_centroid_eigen(calc, r)
    for I in eachindex(centroid_derivative)
        calc.centroid_adiabatic_derivative[I] = centroid_eigen.vectors' * centroid_derivative[I] * centroid_eigen.vectors
    end
end

function evaluate_adiabatic_derivative!(calc::RingPolymerDiabaticCalculator, r)
    calc.stats[:adiabatic_derivative] += 1
    derivative = get_derivative(calc, r)
    eigen = get_eigen(calc, r)
    for i in axes(derivative, 3) # Beads
        for j in axes(derivative, 2) # Atoms
            for k in axes(derivative, 1) # DoFs
                calc.adiabatic_derivative[k,j,i] = eigen[i].vectors' * derivative[k,j,i] * eigen[i].vectors
            end
        end
    end
end

function evaluate_nonadiabatic_coupling!(calc::AbstractDiabaticCalculator, r)
    calc.stats[:nonadiabatic_coupling] += 1
    adiabatic_derivative = get_adiabatic_derivative(calc, r)
    eigen = get_eigen(calc, r)
    for I in eachindex(calc.adiabatic_derivative)
        calc.nonadiabatic_coupling[I] = evaluate_nonadiabatic_coupling(adiabatic_derivative[I], eigen.values)
    end
end

function evaluate_centroid_nonadiabatic_coupling!(calc::RingPolymerDiabaticCalculator, r)
    calc.stats[:centroid_nonadiabatic_coupling] += 1
    centroid_adiabatic_derivative = get_centroid_adiabatic_derivative(calc, r)
    centroid_eigen = get_centroid_eigen(calc, r)
    for I in eachindex(centroid_adiabatic_derivative)
        calc.centroid_nonadiabatic_coupling[I] = evaluate_nonadiabatic_coupling(centroid_adiabatic_derivative[I], centroid_eigen.values)
    end
end

function evaluate_nonadiabatic_coupling!(calc::RingPolymerDiabaticCalculator, r)
    calc.stats[:nonadiabatic_coupling] += 1
    adiabatic_derivative = get_adiabatic_derivative(calc, r)
    eigen = get_eigen(calc, r)
    for i in eachindex(eigen) # Beads
        for I in CartesianIndices(size(adiabatic_derivative)[1:2])
            calc.nonadiabatic_coupling[I,i] = evaluate_nonadiabatic_coupling(adiabatic_derivative[I,i], eigen[i].values)
        end
    end
end

"""
# References

- HammesSchifferTully_JChemPhys_101_4657_1994 Eq. (32)
- SubotnikBellonzi_AnnuRevPhyschem_67_387_2016, section 2.3
"""
function evaluate_nonadiabatic_coupling(adiabatic_derivative::SMatrix, eigenvalues::SVector)
    n = length(eigenvalues)
    SMatrix{n,n}(
        (i != j ? adiabatic_derivative[j,i] / (eigenvalues[i] - eigenvalues[j]) : 0
        for j=1:n, i=1:n))
end

include("friction.jl")
include("large_diabatic.jl")

"""
Evaluates all electronic properties for the current position `r`.
# Properties evaluated:
- Diabatic potential
- Diabatic derivative
- Eigenvalues and eigenvectors
- Adiabatic derivative
- Nonadiabatic coupling

This should no longer be used, instead access the quantities directly with `get_quantity(calc, r)`.
"""
function update_electronics!(calculator::AbstractDiabaticCalculator, r::AbstractArray)
    get_nonadiabatic_coupling(calculator, r)
    return nothing
end

function update_electronics!(calculator::RingPolymerDiabaticCalculator, r::AbstractArray{T,3}) where {T}
    get_nonadiabatic_coupling(calculator, r)
    update_centroid_electronics!(calculator, r)
    return nothing
end

function update_centroid_electronics!(calculator::RingPolymerDiabaticCalculator, r::AbstractArray{T,3}) where {T}
    get_centroid_nonadiabatic_coupling(calculator, r)
    return nothing
end

end # module


function mol_volume_impl(mt::SingleΦT,model::SatLiquidModel,t)
    return mol_volume_impl(QuickStates.sat_t(),model,t)
end

function mol_volume_impl(mt::MultiΦT,model::SatLiquidModel,t,x)
    return mol_volume_impl(QuickStates.sat_tx(),model,t,x)
end

function mol_volume(model::SatLiquidModel,st::ThermodynamicState,unit=u"m^3/mol",mw=nothing)
    return mol_volume(state_type(st),model,st,unit)
end

function mass_volume(model::SatLiquidModel,st::ThermodynamicState,unit=u"m^3/kg",mw=nothing)
    return mol_volume(state_type(st),model,st,unit)
end

function total_volume(model::SatLiquidModel,st::ThermodynamicState,unit=u"m^3",mw=nothing)
    return total_volume(state_type(st),model,st,unit)
end

function mol_density(model::SatLiquidModel,st::ThermodynamicState,unit=u"mol/m^3",mw=nothing)
    return mol_density(state_type(st),model,st,unit)
end

function mass_density(model::SatLiquidModel,st::ThermodynamicState,unit=u"kg/m^3",mw=nothing)
    return mass_density(state_type(st),model,st,unit)
end

function mol_volume(mt::SingleSatT,model::SatLiquidModel,st::ThermodynamicState,unit,mw)
    t = temperature(FromState(),st)
    res = mol_volume_impl(mt,model,t)
    return convert_unit(u"m^3/mol",unit,res)
end

function mol_volume(mt::MultiSatT,model::SatLiquidModel,st::ThermodynamicState,unit,mw)
    t = temperature(FromState(),st)
    x = mol_fraction(FromState(),st,nothing,mw)
    res = mol_volume_impl(mt,model,t,x)
    return convert_unit(u"m^3/mol",unit,res)
end

function mass_volume(mt::SingleSatT,model::SatLiquidModel,st::ThermodynamicState,unit,mw)
    t = temperature(FromState(),st)
    res = mol_volume_impl(mt,model,t)
    return convert_unit(u"m^3/kg",unit,res)
end

function mass_volume(mt::MultiSatT,model::SatLiquidModel,st::ThermodynamicState,unit,mw)
    t = temperature(FromState(),st)
    x = mol_fraction(FromState(),st,nothing,mw)
    res = mol_volume_impl(mt,model,t,x)/molar_mass(FromState(),st,u"kg/mol",mw)
    return convert_unit(u"m^3/kg",unit,)
end

function total_volume(mt::SingleSatT,model::SatLiquidModel,st::ThermodynamicState,unit,mw)
    t = temperature(FromState(),st)
    res = mol_volume_impl(mt,model,t)
    return convert_unit(u"m^3",unit,res)
end

function total_volume(mt::MultiSatT,model::SatLiquidModel,st::ThermodynamicState,unit,mw)
    t = temperature(FromState(),st)
    x = mol_fraction(FromState(),st,nothing,mw)
    res = mol_volume_impl(mt,model,t,x)*moles(FromState(),st,u"mol",mw)
    return convert_unit(u"m^3",unit,)
end

function mol_density(mt::SingleSatT,model::SatLiquidModel,st::ThermodynamicState,unit,mw)
    t = temperature(FromState(),st)
    res = mol_volume_impl(mt,model,t)
    return convert_unit(u"mol/m^3",unit,one(res)/res)
end

function mol_density(mt::MultiSatT,model::SatLiquidModel,st::ThermodynamicState,unit,mw)
    t = temperature(FromState(),st)
    x = mol_fraction(FromState(),st,nothing,mw)
    res = mol_volume_impl(mt,model,t,x)
    return convert_unit(u"mol/m^3",unit,one(res)/res)
end

function mass_density(mt::SingleSatT,model::SatLiquidModel,st::ThermodynamicState,unit,mw)
    t = temperature(FromState(),st)
    res = mol_volume_impl(mt,model,t)
    return convert_unit(u"kg/m^3",unit,one(res)/res)
end

function mass_density(mt::MultiSatT,model::SatLiquidModel,st::ThermodynamicState,unit,mw)
    t = temperature(FromState(),st)
    x = mol_fraction(FromState(),st,nothing,mw)
    res = mol_volume_impl(mt,model,t,x)/molar_mass(FromState(),st,u"kg/mol",mw)
    return convert_unit(u"kg/m^3",unit,one(res)/res)
end


export SpfPosteriorMcmcOutput, SpfPosteriorModeConditionalMcmcOutput, SpfPosteriorDispersionConditionalMcmcOutput
export SisPosteriorMcmcOutput, SisPosteriorModeConditionalMcmcOutput, SisPosteriorDispersionConditionalMcmcOutput
export SimPosteriorModeConditionalMcmcOutput, SimPosteriorDispersionConditionalMcmcOutput

# ==========
#    SPF
# ==========

struct SpfPosteriorMcmcOutput{T<:Union{Int, String}}
    I_sample::Vector{Path{T}}
    γ_sample::Vector{Float64}
    log_post::Dict # Had to do dict since might want different output when (Iᵐ, γ) updated jointly
    dist::Metric
    I_prior::SPF{T}
    γ_prior::ContinuousUnivariateDistribution
    data::Vector{Path{T}}
    performance_measures::Dict
end

struct SpfPosteriorModeConditionalMcmcOutput{T<:Union{Int, String}}
    γ_fixed::Float64
    I_sample::Vector{Path{T}}
    dist::Metric
    I_prior::SPF{T}
    data::Vector{Path{T}}
    performance_measures::Dict
end

struct SpfPosteriorDispersionConditionalMcmcOutput{T<:Union{Int, String}}
    I_fixed::Path{T}
    γ_sample::Vector{Float64}
    γ_prior::ContinuousUnivariateDistribution
    data::Vector{Path{T}}
    performance_measures::Dict
end

function Base.show(io::IO, output::T) where {T<:SpfPosteriorMcmcOutput}
    title = "MCMC Sample for SPF Posterior"
    n = length(title)
    println(io, title)
    println(io, "-"^n)
    for (key, value) in output.performance_measures
        println(io, key, ": ", value)
    end
end

function Base.show(io::IO, output::T) where {T<:SpfPosteriorModeConditionalMcmcOutput}
    title = "MCMC Sample for SPF Posterior (Mode Conditional)"
    n = length(title)
    println(io, title)
    println(io, "-"^n)
    for (key, value) in output.performance_measures
        println(io, key, ": ", value)
    end
end

function Base.show(io::IO, output::T) where {T<:SpfPosteriorDispersionConditionalMcmcOutput}
    title = "MCMC Sample for SPF Posterior (Dispersion Conditional)"
    n = length(title)
    println(io, title)
    println(io, "-"^n)
    for (key, value) in output.performance_measures
        println(io, key, ": ", value)
    end
end

struct SisPosteriorMcmcOutput
    S_sample::Vector{InteractionSequence{Int}}
    γ_sample::Vector{Float64}
    posterior::SisPosterior
    suff_stat_trace::Vector{Float64}
    performace_measures::Dict
end

struct SisPosteriorModeConditionalMcmcOutput
    γ_fixed::Float64
    S_sample::Vector{Vector{Path{Int}}}
    posterior::SisPosterior
    suff_stat_trace::Vector{Float64}
    performance_measures::Dict
end

struct SisPosteriorDispersionConditionalMcmcOutput
    S_fixed::Vector{Path{Int}}
    γ_sample::Vector{Float64}
    posterior::SisPosterior
    performance_measures::Dict
end

# ==========
#    SIM
# ==========

struct SimPosteriorMcmcOutput
    S_sample::Vector{InteractionSequence{Int}}
    γ_sample::Vector{Float64}
    posterior::SimPosterior
    suff_stat_trace::Vector{Float64}
    performace_measures::Dict
end

struct SimPosteriorModeConditionalMcmcOutput
    γ_fixed::Float64
    S_sample::Vector{Vector{Path{Int}}}
    posterior::SimPosterior
    suff_stat_trace::Vector{Float64}
    performance_measures::Dict
end

struct SimPosteriorDispersionConditionalMcmcOutput
    S_fixed::Vector{Path{Int}}
    γ_sample::Vector{Float64}
    γ_prior::ContinuousUnivariateDistribution
    data::Vector{Vector{Path{Int}}}
    performance_measures::Dict
end


import ForwardDiff

using Celeste: Model, SensitiveFloats, DeterministicVI
using Distributions
import DeterministicVI: BvnComponent, GalaxyCacheComponent
import DeterministicVI: eval_bvn_pdf!, get_bvn_derivs!, transform_bvn_derivs!
using DerivativeTestUtils

function test_elbo()
    blob, ea, body = gen_two_body_dataset()
    keep_pixels = 10:11
    trim_tiles!(ea.images, keep_pixels)

    # vp_vec is a vector of the parameters from all the active sources.
    function wrap_elbo{NumType <: Number}(vp_vec::Vector{NumType})
        ea_local = unwrap_vp_vector(vp_vec, ea)
        elbo = DeterministicVI.elbo(ea_local, calculate_derivs=false)
        elbo.v[1]
    end

    ea.active_sources = [1]
    vp_vec = wrap_vp_vector(ea, true)
    elbo_1 = DeterministicVI.elbo(ea)
    test_with_autodiff(wrap_elbo, vp_vec, elbo_1)
    #test_elbo_mp(ea, elbo_1)

    ea.active_sources = [2]
    vp_vec = wrap_vp_vector(ea, true)
    elbo_2 = DeterministicVI.elbo(ea)
    test_with_autodiff(wrap_elbo, vp_vec, elbo_2)
    #test_elbo_mp(ea, elbo_2)

    ea.active_sources = [1, 2]
    vp_vec = wrap_vp_vector(ea, true)
    elbo_12 = DeterministicVI.elbo(ea)
    test_with_autodiff(wrap_elbo, vp_vec, elbo_12)

    P = length(CanonicalParams)
    @test size(elbo_1.d) == size(elbo_2.d) == (P, 1)
    @test size(elbo_12.d) == (length(CanonicalParams), 2)

    @test size(elbo_1.h) == size(elbo_2.h) == (P, P)
    @test size(elbo_12.h) == size(elbo_12.h) == (2 * P, 2 * P)
end

function test_real_image()
    # TODO: replace this with stamp tests having non-trivial WCS transforms.
    # TODO: streamline the creation of small real images.

    run, camcol, field = (3900, 6, 269)

    images = SDSSIO.load_field_images(RunCamcolField(run, camcol, field), datadir)
    tiled_images = TiledImage[TiledImage(img) for img in images]
    dir = joinpath(datadir, "$run/$camcol/$field")
    fname = @sprintf "%s/photoObj-%06d-%d-%04d.fits" dir run camcol field
    catalog = SDSSIO.read_photoobj_celeste(fname)

    # Pick an object.
    objid = "1237662226208063499"
    objids = [ce.objid for ce in catalog]
    sa = findfirst(objids, objid)
    neighbors = Infer.find_neighbors([sa], catalog, tiled_images)[1]

    cat_local = vcat(catalog[sa], catalog[neighbors])
    vp = Vector{Float64}[init_source(ce) for ce in cat_local]
    patches, tile_source_map = Infer.get_tile_source_map(tiled_images, cat_local)
    ea = ElboArgs(tiled_images, vp, tile_source_map, patches, [1])
    Infer.fit_object_psfs!(ea, ea.active_sources)
    Infer.trim_source_tiles!(ea)

    elbo = DeterministicVI.elbo(ea)

    function wrap_elbo{NumType <: Number}(vs1::Vector{NumType})
        ea_local = forward_diff_model_params(NumType, ea)
        ea_local.vp[1][:] = vs1
        local_elbo = DeterministicVI.elbo(ea_local, calculate_derivs=false)
        local_elbo.v[1]
    end

    test_with_autodiff(wrap_elbo, ea.vp[1], elbo)
end

function test_transform_sensitive_float()
	blob, ea, body = gen_two_body_dataset();

	# Only keep a few pixels to make the autodiff results faster.
  keep_pixels = 10:11
	for b = 1:ea.N
	  pixels1 = ea.images[b].tiles[1,1].pixels
      h_width, w_width = size(pixels1)
	  pixels1[setdiff(1:h_width, keep_pixels), :] = NaN;
	  pixels1[:, setdiff(1:w_width, keep_pixels)] = NaN;
	end

	function wrap_elbo{NumType <: Number}(vp_free_vec::Vector{NumType})
		vp_free_array = reshape(vp_free_vec, length(UnconstrainedParams), length(ea.active_sources))
		vp_free = Vector{NumType}[ zeros(NumType, length(UnconstrainedParams)) for
		                           sa in ea.active_sources ];
		Transform.array_to_free_vp!(vp_free_array, vp_free, Int[])
		ea_local = forward_diff_model_params(NumType, ea);
		transform.to_vp!(vp_free, ea_local.vp)
		elbo = DeterministicVI.elbo(ea_local, calculate_derivs=false)
		elbo.v[1]
	end

	transform = Transform.get_mp_transform(ea.vp, ea.active_sources, loc_width=1.0);
	elbo = DeterministicVI.elbo(ea);
	elbo_trans = transform.transform_sensitive_float(elbo, ea.vp, ea.active_sources);

	free_vp_vec = reduce(vcat, transform.from_vp(ea.vp));
	ad_grad = ForwardDiff.gradient(wrap_elbo, free_vp_vec);
	ad_hess = ForwardDiff.hessian(wrap_elbo, free_vp_vec);

	@test_approx_eq ad_grad reduce(vcat, elbo_trans.d)
	@test_approx_eq ad_hess elbo_trans.h

  # Test with a subset of sources.
	ea.active_sources = [2]
	transform = Transform.get_mp_transform(ea.vp, ea.active_sources, loc_width=1.0);
	elbo = DeterministicVI.elbo(ea);
	elbo_trans = transform.transform_sensitive_float(elbo, ea.vp, ea.active_sources);

	free_vp_vec = reduce(vcat, transform.from_vp(ea.vp));
	ad_grad = ForwardDiff.gradient(wrap_elbo, free_vp_vec);
	ad_hess = ForwardDiff.hessian(wrap_elbo, free_vp_vec);

	@test_approx_eq ad_grad reduce(vcat, elbo_trans.d)
	@test_approx_eq ad_hess elbo_trans.h
end

# ForwardDiff 0.2's compilation time is very slow, so only run these tests
# if explicitly requested.

@time test_elbo()
@time test_real_image()
@time test_transform_sensitive_float()


using Revise
using MccDaqHats
using FFTW
using DelimitedFiles
using Infiltrator
includet(joinpath(@__DIR__, "scan_utils.jl"))

#Constants
CURSOR_BACK_2 = "\x1b[2D"
ERASE_TO_END_OF_LINE = "\x1b[0K"

"""
MCC 172 Functions Demonstrated:
    mcc172_iepe_config_write
    mcc172_a_in_clock_config_write
    mcc172_a_in_clock_config_read
    mcc172_a_in_scan_start
    mcc172_a_in_scan_read_numpy
    mcc172_a_in_scan_stop

Purpose:
    Perform finite acquisitions on both channels, calculate the FFTs, and
    display peak information.

Description:
    Acquires blocks of analog input data for both channels then performs
    FFT calculations to determine the frequency content. The highest
    frequency peak is detected and displayed, along with harmonics. The
    time and frequency data are saved to CSV files.

    This example requires the NumPy library.

"""
function fft_scan()

    channels = [0, 1]
    channel_mask = chan_list_to_mask(channels)

    samples_per_channel = 12800
    scan_rate = 51200.0
    options = :OPTS_DEFAULT

    try
        # Select an MCC 172 HAT device to use.
        hat = select_hat_devices(:MCC_172, 1)
        address = hat[1].address

        println("\nSelected MCC 172 HAT device at address $address")

        # Turn on IEPE supply?
        iepe_enable = get_iepe()

        mcc172_open(address)
        for channel in channels
            mcc172_iepe_config_write(address, channel, iepe_enable)
        end

        # Configure the clock and wait for sync to complete.
        mcc172_a_in_clock_config_write(address, :SOURCE_LOCAL, scan_rate)

        synced = false
        actual_scan_rate = 0.0
        while !synced
            (_source_type, actual_scan_rate, synced) = mcc172_a_in_clock_config_read(address)
            if !synced
                sleep(0.005)
            end
        end

        println("\nMCC 172 Multi channel FFT example")
        println("    Functions demonstrated:")
        println("         mcc172_iepe_config_write")
        println("         mcc172_a_in_clock_config_write")
        println("         mcc172_a_in_clock_config_read")
        println("         mcc172_a_in_scan_start")
        println("         mcc172_a_in_scan_read_numpy")
        println("         mcc172_a_in_scan_stop")
        println("         mcc172_a_in_scan_cleanup")
        println("    IEPE power: $(iepe_enable ? "on" : "off")")
        println("    Channels: $(join(channels, ", "))")
        println("    Requested scan rate: $scan_rate")
        println("    Actual scan rate: $(round(actual_scan_rate, digits=3))")
        println("    Samples per channel $(round(samples_per_channel, digits=3))")
        println("    Options: $(options)")

        try
            println("\nPress 'Enter' to continue")
            readline()
        catch
            error("^C to end program")
        end

        # Configure and start the scan.
        channel_mask = chan_list_to_mask(channels)
        mcc172_a_in_scan_start(address, channel_mask, samples_per_channel, options)

        try
            read_and_display_data(address, samples_per_channel, num_channels)

        catch e # KeyboardInterrupt
            if isa(e, InterruptException)
                # Clear the "^C" from the display.
                println("$CURSOR_BACK_2, $ERASE_TO_END_OF_LINE, Aborted\n")
            else
                println("\n $e")
            end
            mcc172_a_in_scan_stop(address)
        end

        mcc172_a_in_scan_cleanup(address)

    catch err  # (HatError, ValueError) as err
        println("\n $err")
    end
    return nothing
end

""" Hann window function. """
function window(index, max_index)
    return 0.5 - 0.5*cos(2π*index / max_index)
end

""" Hann window compensation factor. """
function window_compensation()
    return 2.0
end

""" Calculate a real-only FFT, returning the spectrum in dBFS. """
function calculate_real_fft(data)
    n_samples = length(data)
    in_data = Vector{Float64}(undef, n_samples)

    # Apply the window and normalize the time data.
    info = mcc172_info()
    max_v = info.AI_MAX_RANGE
    for i in 1:n_samples
        in_data[i] = window(i, n_samples) * data[i] / max_v
    end

    # Perform the FFT.
    out = fft(in_data)

    # Convert the complex results to real and convert to dBFS.
    spectrum = Vector{Float64}(undef, length(out))

    real_part = real(out)
    imag_part = imag(out)
    for i in 1:length(out)
        if i == 0
            # Don"t multiply DC value times 2.
            spectrum[i] = 20*log10(window_compensation() *
                sqrt(real_part[i] ^ 2 + imag_part[i] ^ 2) / n_samples)
        else
            spectrum[i] = 20*log10(window_compensation() * 2 *
                sqrt(real_part[i] ^ 2 + imag_part[i] ^ 2) / n_samples)
        end
    end

    return spectrum
end

"""
Interpolate between the bins of an FFT peak to find a more accurate
frequency.  bin1 is the FFT value at the detected peak, bin0 and bin2
are the values from adjacent bins above and below the peak. Returns
the offset value from the index of bin1.
"""
function quadratic_interpolate(bin0, bin1, bin2)
    y_1 = abs(bin0)
    y_2 = abs(bin1)
    y_3 = abs(bin2)

    result = (y_3 - y_1) / (2 * (2 * y_2 - y_1 - y_3))

    return result
end

"""
Return a given number with the appropriate suffix as a string.
"""
function add_suffix(index::Integer)
    if index == 1
        suffix = "st"
    elseif index == 2
        suffix = "nd"
    elseif index == 3
        suffix = "rd"
    else
        suffix = "th"
    end

    return string(i) * suffix
end

"""
Wait for all of the scan data, perform an FFT, find the peak frequency,
and display the frequency information.  Only supports 1 channel in the data.

Args:
hat (mcc172): The mcc172 HAT device object.
samples_per_channel: The number of samples to read for each channel.

Returns:
Nothing
"""
function read_and_display_data(address, channels, samples_per_channel, scan_rate)

    timeout = 5.0

    # Wait for all the data, and read it as an array.
    read_result = mcc172_a_in_scan_read(address, samples_per_channel, timeout)

    # Separate the data by channel (deinterleave)
    read_data = deinterleave(read_result, length(channels))

    for channel in channels
        print("===== Channel:    $channel\n")

        # Calculate the FFT
        spectrum = calculate_real_fft(read_data[:, channel])

        # Calculate dBFS and find peak.
        f_i = 0.0
        peak_index = 0
        peak_val = -1000.0

        # Save data to CSV file
        logname = "fft_scan_" * string(i) * ".csv"
        io = open(logname, "w")
        writedlm("Time data (V), Frequency (Hz), Spectrum (dBFS)\n")

        for (i, spec_val) in enumerate(spectrum)
            # Find the peak value and index.
            if spec_val > peak_val
                peak_val = spec_val
                peak_index = i
            end

            # Save to the CSV file.
            writedlm(read_data[i], f_i, spec_val)

            f_i += scan_rate / samples_per_channel
        end

        close(io)

        # Interpolate for a more precise peak frequency.
        peak_offset = quadratic_interpolate(
            spectrum[peak_index - 1], spectrum[peak_index], spectrum[peak_index + 1])
        peak_freq = ((peak_index + peak_offset) * scan_rate /
                     samples_per_channel)
        println("Peak: $(round(peak_val, digits = 2)) dBFS at $(round(peak_freq, digitis = 2)) Hz")

        # Find and display harmonic levels.
        i = 2
        h_freq = 0
        nyquist = scan_rate / 2.0
        while (i < 8) && (h_freq <= nyquist)
            # Stop when frequency exceeds Nyquist rate or at the 8th harmonic.
            h_freq = peak_freq * i
            if h_freq <= nyquist
                h_index = floor(h_freq * samples_per_channel / scan_rate + 0.5)
                h_val = spectrum[h_index]
                print(" $(add_suffix(i)) harmonic: $h_val dBFS at $h_freq Hz")
            end
            i += 1
        end

        println("Data and FFT saved in $logname")
    end
end


module execute1d
using LPDM
using POMDPs, Parameters, StaticArrays, D3Trees, Distributions, SparseArrays
using Combinatorics
using StatsFuns
using Printf
using Dates

include("LightDarkLPDM.jl")
using LightDarkPOMDPs

include("1D/LPDMBounds1d.jl")

struct LPDMTest
    solver_mode::Symbol
    action_mode::Symbol
    obs_mode::Symbol
    reward_mode::Symbol
end

struct LPDMScenario
    s0::LD1State
end

# reward function options - :quadratic or :fixed
function batch_execute(;n::Int64=1, debug::Int64=1, reward_mode=:quadratic)
    test=Array{LPDMTest}(undef,0)

    push!(test, LPDMTest(:lpdm, :standard, :discrete, reward_mode))
    push!(test, LPDMTest(:lpdm, :standard, :continuous, reward_mode))
    # push!(test, LPDMTest(:lpdm, :adapt, reward_mode)) # simulated annealing
    # push!(test, LPDMTest(:despot, :small, reward_mode))
    # push!(test, LPDMTest(:despot, :large, reward_mode))
    # push!(test, LPDMTest(:lpdm_bv, :bv, reward_mode)) # blind value

    scen=Array{LPDMScenario}(undef,0)
    push!(scen, LPDMScenario(LD1State(-2*π)))
    # push!(scen, LPDMScenario(LD1State(π/2)))
    # push!(scen, LPDMScenario(LD1State(3/2*π)))
    # push!(scen, LPDMScenario(LD1State(2*π)))

    # Dummy execution, just to make sure all the code is compiled and loaded,
    # to improve uniformity of subsequent executions.
    # execute(solv_mode = :lpdm, action_space_type = :adapt, n_sims = 1, s0 = LD1State(π), output = 0)

    # General solver parameters
    steps::Int64                = -1
    time_per_move::Float64      = 1.0
    search_depth::Int64         = 30
    n_particles::Int64          = 50
    max_trials::Int64           = -1

    f = open("results_" * Dates.format(now(),"yyyy-mm-dd_HH_MM") * ".txt", "w")

    Printf.@printf(f,"GENERAL SOLVER PARAMETERS\n")
    Printf.@printf(f,"\treward function:\t\t\t%s\n", reward_func)
    Printf.@printf(f,"\tsteps:\t\t\t%d\n", steps)
    Printf.@printf(f,"\ttime per move:\t\t\t%f\n", time_per_move)
    Printf.@printf(f,"\tsearch depth:\t\t\t%d\n", search_depth)
    Printf.@printf(f,"\tN particles:\t\t\t%d\n", n_particles)
    Printf.@printf(f,"\tmax trials:\t\t\t%d\n\n", max_trials)

    for i in 1:length(scen)
        # write(f,"SCENARIO $i, s0 = $(scen[i].s0)\n\n")
        if debug >= 0
            println("")
            println("SCENARIO $i, s0 = $(scen[i].s0)")
            println("------------------------")
        end

        Printf.@printf(f,"SCENARIO %d, s0 = %f\n", i, scen[i].s0)
        Printf.@printf(f,"=====================================================================\n")
        # Printf.@printf(f,"mode\t\tact. space\t\tsteps(std)\t\treward(std)\n")
        Printf.@printf(f,"SOLVER\t\tACT. MODE\t\tOBS. MODE\t\tSTEPS (STD)\t\t\tREWARD (STD)\n")
        Printf.@printf(f,"=====================================================================\n")
        for t in test
            if debug >= 0
                println("mode: $(t.solver_mode), actions: $(t.action_mode), observations: $(t.obs_mode), rewards: $(t.reward_mode)")
            end
            steps, steps_std, reward, reward_std =
                        execute(solver_mode       = t.solver_mode,
                                action_mode       = t.action_mode,
                                obs_mode          = t.obs_mode,
                                reward_mode       = t.reward_mode,
                                n_sims            = n,
                                s0                = scen[i].s0,
                                output            = debug)

            Printf.@printf(f,"%s\t\t%s\t\t%s\t\t\t%05.2f (%06.2f)\t\t%06.2f (%06.2f)\n",
                            string(t.solver_mode), string(t.action_mode), string(t.obs_mode), steps, steps_std, reward, reward_std)
        end
        Printf.@printf(f,"=====================================================================\n")
        Printf.@printf(f,"%d tests per scenario\n\n", n)
    end
    close(f)
end

function execute(;vis::Vector{Int64}=Int64[],
                solver_mode::Symbol=:despot,
                action_mode::Symbol=:standard,
                obs_mode::Symbol=:discrete,
                reward_mode = :quadratic,
                n_sims::Int64=1,
                steps::Int64=-1,
                s0::LD1State=LD1State(1.9),
                output::Int64=1)#n_sims::Int64 = 100)

    if solver_mode == :despot
        p = LightDark1DDespot(action_mode = action_mode,
                            obs_mode = obs_mode,
                            reward_mode = reward_mode)
    elseif solver_mode == :lpdm
        p = LightDark1DLpdm(action_mode = action_mode,
                            obs_mode = obs_mode,
                            reward_mode = reward_mode)
    end

    sim_rewards = Vector{Float64}(undef,n_sims)
    sim_steps   = Vector{Int64}(undef,n_sims)

    for sim in 1:n_sims
        # world_seed  ::UInt32   = convert(UInt32, 42)

        world_rng = RNGVector(1, UInt32(sim))
        LPDM.set!(world_rng, 1)

        # NOTE: restrict initial state to positive numbers only, for now
        # s::LD1State                  = LD1State(π);    # initial state
        step_rewards::Array{Float64}     = Vector{Float64}(undef,0)
        # println("$(supertype(LightDark1DDespot)))")

        # NOTE: Original version
        # solver = LPDM.LPDMSolver{LD1State, LD1Action, LD1Obs, LDBounds1d{LD1State, LD1Action, LD1Obs}, RNGVector}(
        #                                                                     # rng = sim_rng,
        #                                                                     debug = output,
        #                                                                     time_per_move = 1.0,  #sec
        #                                                                     # time_per_move = 1.0,  #sec
        #                                                                     sim_len = -1,
        #                                                                     search_depth = 50,
        #                                                                     n_particles = 20,
        #                                                                     # seed = UInt32(5),
        #                                                                     seed = UInt32(2*sim+1),
        #                                                                     # max_trials = 10)
        #                                                                     max_trials = -1,
        #                                                                     mode = solv_mode)

        # NOTE: Test version
        solver = LPDM.LPDMSolver{LD1State, LD1Action, LD1Obs, LDBounds1d{LD1State, LD1Action, LD1Obs}, RNGVector}(
                                                                            # rng = sim_rng,
                                                                            debug = output,
                                                                            time_per_move = -1.0,  #sec
                                                                            # time_per_move = 1.0,  #sec
                                                                            sim_len = steps,
                                                                            search_depth = 50,
                                                                            n_particles = 20,
                                                                            # seed = UInt32(2),
                                                                            seed = UInt32(2*sim+1),
                                                                            # max_trials = 10)
                                                                            max_trials = 1,
                                                                            solver_mode = solver_mode,
                                                                            action_mode = action_mode,
                                                                            obs_mode    = obs_mode)

    #---------------------------------------------------------------------------------
        # Belief
        bu = LPDMBeliefUpdater(p,
                               n_particles = solver.config.n_particles,
                               seed = UInt32(3*sim+1));  # initialize belief updater
        initial_states = state_distribution(p, s0, solver.config, world_rng)     # create initial  distribution
        current_belief = LPDM.create_belief(bu)                       # allocate an updated belief object
        LPDM.initialize_belief(bu, initial_states, current_belief)    # initialize belief
        # show(current_belief)
        updated_belief = LPDM.create_belief(bu)
    #---------------------------------------------------------------------------------

        policy::LPDMPolicy = POMDPs.solve(solver, p)
        s = s0
        step::Int64 = 0
        r::Float64 = 0.0

        # output >= 1 && println("=============== SIMULATION # $sim ================")
        if output >= 0
            println("")
            println("*** SIM $sim (solver mode: $solver_mode, action mode: $action_mode, obs mode: $obs_mode, s0: $s0)***")
            println("")
        end

        val, run_time, bytes, gctime, memallocs =
        @timed while !isterminal(p, s) && (solver.config.sim_len == -1 || step < solver.config.sim_len)
            step += 1

            if output >= 1
                println("")
                println("=============== Step $step ================")
                show(current_belief)
            end

            a = POMDPs.action(policy, current_belief)
            if output >= 1
                println("s: $s")
                println("a: $a")
            end

            s, o, r = POMDPs.generate_sor(p, s, a, world_rng)
            push!(step_rewards, r)
            if output >= 1
                println("s': $s")
                println("o': $o")
                println("r: $r")
                println("=======================================")
            end

            # error("$(@which(POMDPs.update(bu, current_belief, a, o, updated_belief)))")
            # update belief
            POMDPs.update(bu, current_belief, a, o, updated_belief)
            current_belief = deepcopy(updated_belief)
            # show(updated_belief) #NOTE: don't show for now

            if LPDM.isterminal(p, current_belief.particles)
                if output >= 1
                    println("Terminal belief. Execution completed.")
                    show(current_belief)
                end
                break
            end

            if step ∈ vis
                t = LPDM.d3tree(solver,
                                detect_repeat=false,
                                title="Step $step",
                                init_expand=10)
                # # show(t)
                inchrome(t)
                # blink(t)
            end
            if output >= 1
                solver_mode == :lpdm && println("root actions: $(solver.root.action_space)")
            end
        end

        sim_steps[sim] = step
        sim_rewards[sim] = sum(step_rewards)
        # # Compute discounted reward
        # discounted_reward::Float64 = 0.0
        # multiplier::Float64 = 1.0
        # for r in rewards
        #     discounted_reward += multiplier * r
        #     multiplier *= p.discount
        # end
        # println("Discounted reward: $discounted_reward")
    end
    # println("SIM_STEPS: $sim_steps, SIM_REWARDS: $sim_rewards")
    return mean(sim_steps), std(sim_steps), mean(sim_rewards), std(sim_rewards)
    # return t
end

function test_move()
    p = LightDarkPOMDPs.LightDark1DDespot()

    coordinates = Vector{Tuple{Float64,Float64}}()
    push!(coordinates,(5.0,0.0))
    push!(coordinates,(-5.0,0.0))
    push!(coordinates,(3.14,4.50))

    for c in coordinates
        r,a = move(p,c[1],c[2])
        println("x1=$(c[1]), x2=$(c[2]), r=$r, a=$a")
    end
end

end #module


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

"Represents a Basic Block, in the DomTree"
struct DomTreeNode
    # How deep we are in the DomTree
    level::Int
    # The BB indices in the CFG for all Basic Blocks we immediately dominate
    children::Vector{Int}
end
DomTreeNode() = DomTreeNode(1, Vector{Int}())

"Data structure that encodes which basic block dominates which."
struct DomTree
    # Which basic block immediately dominates each basic block (ordered by BB indices)
    # Note: this is the inverse of the nodes, children field
    idoms::Vector{Int}

    # The nodes in the tree (ordered by BB indices)
    nodes::Vector{DomTreeNode}
end

"""
    Checks if bb1 dominates bb2.
    bb1 and bb2 are indexes into the CFG blocks.
    bb1 dominates bb2 if the only way to enter bb2 is via bb1.
    (Other blocks may be in between, e.g bb1->bbX->bb2).
"""
function dominates(domtree::DomTree, bb1::Int, bb2::Int)
    bb1 == bb2 && return true
    target_level = domtree.nodes[bb1].level
    source_level = domtree.nodes[bb2].level
    source_level < target_level && return false
    for _ in (source_level - 1):-1:target_level
        bb2 = domtree.idoms[bb2]
    end
    return bb1 == bb2
end

bb_unreachable(domtree::DomTree, bb::Int) = bb != 1 && domtree.nodes[bb].level == 1

function update_level!(domtree::Vector{DomTreeNode}, node::Int, level::Int)
    worklist = Tuple{Int, Int}[(node, level)]
    while !isempty(worklist)
        (node, level) = pop!(worklist)
        domtree[node] = DomTreeNode(level, domtree[node].children)
        foreach(domtree[node].children) do child
            push!(worklist, (child, level+1))
        end
    end
end

"Iterable data structure that walks though all dominated blocks"
struct DominatedBlocks
    domtree::DomTree
    worklist::Vector{Int}
end

"Returns an iterator that walks through all blocks dominated by the basic block at index `root`"
function dominated(domtree::DomTree, root::Int)
    doms = DominatedBlocks(domtree, Vector{Int}())
    push!(doms.worklist, root)
    doms
end

function iterate(doms::DominatedBlocks, state::Nothing=nothing)
    isempty(doms.worklist) && return nothing
    bb = pop!(doms.worklist)
    for dominated in doms.domtree.nodes[bb].children
        push!(doms.worklist, dominated)
    end
    return (bb, nothing)
end

function naive_idoms(cfg::CFG)
    nblocks = length(cfg.blocks)
    # The extra +1 helps us detect unreachable blocks below
    dom_all = BitSet(1:nblocks+1)
    dominators = BitSet[n == 1 ? BitSet(1) : copy(dom_all) for n = 1:nblocks]
    changed = true
    while changed
        changed = false
        for n = 2:nblocks
            if isempty(cfg.blocks[n].preds)
                continue
            end
            firstp, rest = Iterators.peel(Iterators.filter(p->p != 0, cfg.blocks[n].preds))
            new_doms = copy(dominators[firstp])
            for p in rest
                intersect!(new_doms, dominators[p])
            end
            push!(new_doms, n)
            changed = changed || (new_doms != dominators[n])
            dominators[n] = new_doms
        end
    end
    # Compute idoms
    idoms = fill(0, nblocks)
    for i = 2:nblocks
        if dominators[i] == dom_all
            idoms[i] = 0
            continue
        end
        doms = collect(dominators[i])
        for dom in doms
            i == dom && continue
            hasany = false
            for p in doms
                if p !== i && p !== dom && dom in dominators[p]
                    hasany = true; break
                end
            end
            hasany && continue
            idoms[i] = dom
        end
    end
    idoms
end

# Construct Dom Tree
function construct_domtree(cfg::CFG)
    idoms = SNCA(cfg)
    # Compute children
    nblocks = length(cfg.blocks)
    domtree = DomTreeNode[DomTreeNode() for _ = 1:nblocks]
    for (idx, idom) in Iterators.enumerate(idoms)
        (idx == 1 || idom == 0) && continue
        push!(domtree[idom].children, idx)
    end
    # Recursively set level
    update_level!(domtree, 1, 1)
    DomTree(idoms, domtree)
end

#================================ [SNCA] ======================================#
#
#   This section implements the Semi-NCA (SNCA) dominator tree construction from
#   described in Georgiadis' PhD thesis [LG05], which itself is a simplification
#   of the Simple Lenguare-Tarjan (SLT) algorithm [LG79]. This algorithm matches
#   the algorithm choice in LLVM and seems to be a sweet spot in implementation
#   simplicity and efficiency.
#
#   [LG05]  Linear-Time Algorithms for Dominators and Related Problems
#           Loukas Georgiadis, Princeton University, November 2005, pp. 21-23:
#           ftp://ftp.cs.princeton.edu/reports/2005/737.pdf
#
#   [LT79]  A fast algorithm for finding dominators in a flowgraph
#           Thomas Lengauer, Robert Endre Tarjan, July 1979, ACM TOPLAS 1-1
#           http://www.dtic.mil/dtic/tr/fulltext/u2/a054144.pdf
#
begin
    # We could make these real structs, but probably not worth the extra
    # overhead. Still, give them names for documentary purposes.
    const BBNumber = UInt
    const DFSNumber = UInt

    """
    Keeps the per-BB state of the Semi NCA algorithm. In the original
    formulation, there are three separate length `n` arrays, `label`, `semi` and
    `ancestor`. Instead, for efficiency, we use one array in a array-of-structs
    style setup.
    """
    struct Node
        semi::DFSNumber
        label::DFSNumber
    end

    struct DFSTree
        # Maps DFS number to BB number
        numbering::Vector{BBNumber}
        # Maps BB number to DFS number
        reverse::Vector{DFSNumber}
        # Records parent relationships in the DFS tree (DFS number -> DFS number)
        # Storing it this way saves a few lookups in the snca_compress! algorithm
        parents::Vector{DFSNumber}
    end
    length(D::DFSTree) = length(D.numbering)
    preorder(D::DFSTree) = OneTo(length(D))
    _drop(xs::AbstractUnitRange, n::Integer) = (first(xs)+n):last(xs)

    function DFSTree(nblocks::Int)
        DFSTree(
            Vector{BBNumber}(undef, nblocks),
            zeros(DFSNumber, nblocks),
            Vector{DFSNumber}(undef, nblocks))
    end

    function DFS(cfg::CFG, current_node::BBNumber)::DFSTree
        dfs = DFSTree(length(cfg.blocks))
        # TODO: We could reuse the storage in DFSTree for our worklist. We're
        # guaranteed for the worklist to be smaller than the remaining space in
        # DFSTree
        worklist = Tuple{DFSNumber, BBNumber}[(0, current_node)]
        dfs_num = 1
        parent = 0
        while !isempty(worklist)
            (parent, current_node) = pop!(worklist)
            dfs.reverse[current_node] != 0 && continue
            dfs.reverse[current_node] = dfs_num
            dfs.numbering[dfs_num] = current_node
            dfs.parents[dfs_num] = parent
            for succ in cfg.blocks[current_node].succs
                push!(worklist, (dfs_num, succ))
            end
            dfs_num += 1
        end
        # If all blocks are reachable, this is a no-op, otherwise,
        # we shrink these arrays.
        resize!(dfs.numbering, dfs_num - 1)
        resize!(dfs.parents, dfs_num - 1)
        dfs
    end

    """
    Matches the snca_compress algorithm in Figure 2.8 of [LG05], with the
    modification suggested in the paper to use `last_linked` to determine
    whether an ancestor has been processed rather than storing `0` in the
    ancestor array.
    """
    function snca_compress!(state::Vector{Node}, ancestors::Vector{DFSNumber},
                            v::DFSNumber, last_linked::DFSNumber)
        u = ancestors[v]
        @assert u < v
        if u >= last_linked
            snca_compress!(state, ancestors, u, last_linked)
            if state[u].label < state[v].label
                state[v] = Node(state[v].semi, state[u].label)
            end
            ancestors[v] = ancestors[u]
        end
        nothing
    end

    function snca_compress_worklist!(
            state::Vector{Node}, ancestors::Vector{DFSNumber},
            v::DFSNumber, last_linked::DFSNumber)
        # TODO: There is a smarter way to do this
        u = ancestors[v]
        worklist = Tuple{DFSNumber, DFSNumber}[(u,v)]
        @assert u < v
        while !isempty(worklist)
            u, v = last(worklist)
            if u >= last_linked
                if ancestors[u] >= last_linked
                    push!(worklist, (ancestors[u], u))
                    continue
                end
                if state[u].label < state[v].label
                    state[v] = Node(state[v].semi, state[u].label)
                end
                ancestors[v] = ancestors[u]
            end
            pop!(worklist)
        end
    end

    """
        SNCA(cfg::CFG)

    Determines a map from basic blocks to the block which immediately dominate them.
    Expressed as indexes into `cfg.blocks`.

    The main Semi-NCA algrithm. Matches Figure 2.8 in [LG05].
    Note that the pseudocode in [LG05] is not entirely accurate.
    The best way to understand what's happening is to read [LT79], then the
    description of SLT in in [LG05] (warning: inconsistent notation), then
    the description of Semi-NCA.
    """
    function SNCA(cfg::CFG)
        D = DFS(cfg, BBNumber(1))
        # `label` is initialized to the identity mapping (though
        # the paper doesn't make that clear). The rational for this is Lemma
        # 2.4 in [LG05] (i.e. Theorem 4 in ). Note however, that we don't
        # ever look at `semi` until it is fully initialized, so we could leave
        # it uninitialized here if we wanted to.
        state = Node[ Node(typemax(DFSNumber), w) for w in preorder(D) ]
        # Initialize idoms to parents. Note that while idoms are eventually
        # BB indexed, we keep it DFS indexed until a final post-processing
        # pass to avoid extra memory references during the O(N^2) phase below.
        idoms_dfs = copy(D.parents)
        # We abuse the parents array as the ancestors array.
        # Semi-NCA does not look at the parents array at all.
        # SLT would, but never simultaneously, so we could still
        # do this.
        ancestors = D.parents
        for w::DFSNumber ∈ reverse(_drop(preorder(D), 1))
            # LLVM initializes this to the parent, the paper initializes this to
            # `w`, but it doesn't really matter (the parent is a predecessor,
            # so at worst we'll discover it below). Save a memory reference here.
            semi_w = typemax(DFSNumber)
            for v ∈ cfg.blocks[D.numbering[w]].preds
                # For the purpose of the domtree, ignore virtual predecessors
                # into catch blocks.
                v == 0 && continue
                vdfs = D.reverse[v]
                # Ignore unreachable predecessors
                vdfs == 0 && continue
                last_linked = DFSNumber(w + 1)
                # N.B.: This conditional is missing from the psuedocode
                # in figure 2.8 of [LG05]. It corresponds to the
                # `ancestor[v] != 0` check in the `eval` implementation in
                # figure 2.6
                if vdfs >= last_linked
                    # For performance, if the number of ancestors is small
                    # avoid the extra allocation of the worklist.
                    if length(ancestors) <= 32
                        snca_compress!(state, ancestors, vdfs, last_linked)
                    else
                        snca_compress_worklist!(state, ancestors, vdfs, last_linked)
                    end
                end
                semi_w = min(semi_w, state[vdfs].label)
            end
            state[w] = Node(semi_w, semi_w)
        end
        for v ∈ _drop(preorder(D), 1)
            idom = idoms_dfs[v]
            vsemi = state[v].semi
            while idom > vsemi
                idom = idoms_dfs[idom]
            end
            idoms_dfs[v] = idom
        end
        # Reexpress the idom relationship in BB indexing
        idoms_bb = Int[ (i == 1 || D.reverse[i] == 0) ? 0 : D.numbering[idoms_dfs[D.reverse[i]]] for i = 1:length(cfg.blocks) ]
        idoms_bb
    end
end


export get_config
import Cambrian.get_config

"overrides Cambrian.get_config(::Dict), converts function names to functions and tracks arity"
function get_config(config::Dict)
    # parse all function names, assign to function value
    two_arity = falses(length(config["functions"]))
    functions = Array{Function}(undef, length(config["functions"]))
    for i in eachindex(config["functions"])
        fname = config["functions"][i]
        if CGPFunctions.arity[fname] == 2
            two_arity[i] = true
        end
        functions[i] = eval(Meta.parse(string("CGPFunctions.", fname)))
    end
    config["two_arity"] = two_arity
    config["functions"] = functions
    (; (Symbol(k)=>v for (k, v) in config)...)
end


# Ocean put and get calls

"""
function preOcean(csolver)
    - saves and regrids the `BoundaryEnergyFlux` couplerfield (i.e. regridded `mAtmos.state.F_ρθ_accum[mAtmos.boundary]`) to `F_ρθ_prescribed[mOcean.boundary]` on the `mOcean` grid
    csolver::CplSolver
"""
function preOcean(csolver)
    mOcean = csolver.component_list.domainOcean.component_model
    mAtmos = csolver.component_list.domainAtmos.component_model

    mOcean.odesolver.rhs!.state_auxiliary.F_ρe_prescribed[mOcean.boundary] .=
        coupler_get(csolver.coupler, :BoundaryEnergyFlux, mOcean.grid, DateTime(0), u"J")
end

"""
function postOcean(csolver)
    - updates couplerfield `OceanSurfaceTemerature` with `mOcean.state.T_sfc[mOcean.boundary]` regridded to the coupler grid, and updates the coupler time
    csolver::CplSolver
"""
function postOcean(csolver)
    mOcean = csolver.component_list.domainOcean.component_model
    mAtmos = csolver.component_list.domainAtmos.component_model
    coupler_put!(
        csolver.coupler,
        :SeaSurfaceTemerature,
        mOcean.state.T_sfc[mOcean.boundary],
        mOcean.grid.numerical,
        DateTime(0),
        u"K",
    )
end

# Atmos put and get calls
"""
function preAtmos(csolver)
    - saves and regrids the `SeaSurfaceTemerature` couplerfield (i.e. regridded `mOcean.state.T_sfc[mOcean.boundary]``) on coupler grid to `mAtmos.aux.T_sfc[mAtmos.boundary]` on the `mAtmos` grid
    - resets the accumulated flux `F_ρe_accum` in `mAtmos` to 0
    csolver::CplSolver
"""
function preAtmos(csolver)
    mOcean = csolver.component_list.domainOcean.component_model
    mAtmos = csolver.component_list.domainAtmos.component_model
    # Set boundary T_sfc used in atmos at the start of the coupling cycle.
    mAtmos.odesolver.rhs!.state_auxiliary.T_sfc[mAtmos.boundary] .=
        coupler_get(csolver.coupler, :SeaSurfaceTemerature, mAtmos.grid.numerical, DateTime(0), u"K")
    # Set atmos boundary flux accumulator to 0.
    mAtmos.state.F_ρe_accum .= 0

    # get indicies of the required variables from the `mAtmos` MPIStateArray
    idx = varsindex(vars(mAtmos.state), :ρe)[1]
    idx_rho = varsindex(vars(mAtmos.state), :ρ)[1]

    # Calculate, print and log domain-averaged energy
    FT = eltype(mAtmos.state)

    model_type = cpl_solver.component_list.domainAtmos.component_model.model[1] # TODO: bl should be more accessible
    p = model_type.model.physics.parameters # maybe the parameter list could be saved in the coupler? (this would requre all the compute kernels below to import the coupler)
    nel = mAtmos.grid.resolution.elements.vertical
    po = mAtmos.grid.resolution.polynomial_order.vertical

    Earth_sfc_area = FT(4π * 6371000.0^2)
    ocean_layer_volume = sum((mOcean.state.weights[:, :, mOcean.state.realelems])[mOcean.boundary])
    E_Ocean = weightedsum(mOcean.state, 1) .* p.ρ_o .* p.c_o .* p.h_o ./ ocean_layer_volume # J / m^2
    E_Atmos = weightedsum(mAtmos.state, idx) ./ Earth_sfc_area  # J / m^2

    @info(
        "preatmos",
        time = string(csolver.t) * "/" * string(mAtmos.time.finish),
        total_energyA = E_Ocean,
        total_energyB = E_Atmos,
        total_energy = E_Atmos + E_Ocean,
        Ocean_T_sfc = maximum(mOcean.state.T_sfc[mOcean.boundary]),
        atmos_ρe_sfc = maximum(mAtmos.state.ρe[mAtmos.boundary]),
    )

    isnothing(csolver.fluxlog) ? nothing : csolver.fluxlog.A[csolver.steps] = E_Ocean
    isnothing(csolver.fluxlog) ? nothing : csolver.fluxlog.B[csolver.steps] = E_Atmos
end

"""
function postAtmos(csolver)
    - updates couplerfield `BoundaryEnergyFlux` with mAtmos.state.F_ρe_accum[mAtmos.boundary] regridded to the coupler grid, and updates the coupler time
    csolver::CplSolver
"""
function postAtmos(csolver)
    mOcean = csolver.component_list.domainOcean.component_model
    mAtmos = csolver.component_list.domainAtmos.component_model

    # Pass atmos exports to "coupler" namespace
    # 1. Save mean θ flux at the Atmos boundary during the coupling period

    @info(
        "postatmos",
        time = string(csolver.t) * "/" * string(mAtmos.time.finish),
        Ocean_T_sfc = maximum(mOcean.state.T_sfc[mOcean.boundary]),
        atmos_ρe_sfc = maximum(mAtmos.state.ρe[mAtmos.boundary]),
    )

    coupler_put!(
        csolver.coupler,
        :BoundaryEnergyFlux,
        mAtmos.state.F_ρe_accum[mAtmos.boundary] ./ csolver.dt,
        mAtmos.grid.numerical,
        DateTime(0),
        u"J",
    )
end


#InspectDR: Operation on datasets
#-------------------------------------------------------------------------------

#==Types declarations
===============================================================================#

#Input dataset data:
#TODO: make x an immutable (read-only) type so data cannot be changed once identified.
mutable struct IDataset{DATAF1} #DATAF1::Bool: Data is function of 1 argument (optimization possible)
	x::Vector
	y::Vector
end

#==Accessors
===============================================================================#

Point(ds::IDataset, i::Int) = Point2D(ds.x[i], ds.y[i])
Point(ds::Vector{Point2D}, i::Int) = ds[i]

#TODO: Specialize for sorted datasets?
function getextents(ds::IDataset)
	(xmin, xmax) = extrema_nan(ds.x)
	(ymin, ymax) = extrema_nan(ds.y)
	return PExtents2D(xmin, xmax, ymin, ymax)
end

#==Pre-processing display data
===============================================================================#

#_reduce: Obtain reduced dataset by limiting to the extents & max resolution:
#-------------------------------------------------------------------------------
#Reduce data without dropping points (Just transfering all data for now)
#    xres_max: Max number of x-points in data window.
#returns: Vector{Point2D}
#TODO: clip data beyond extents.
#WARNING: not clipping might cause display issues when applying the transform
function _reduce_nodrop(input::IDataset, xext::PExtents1D, xres_max::Integer)
	x = input.x; y = input.y
	n_ds = length(x) #numer of points of input dataset
	result = Array{Point2D}(undef, n_ds)
	for i in 1:n_ds
		result[i] = Point2D(x[i], y[i])
	end
	return result
end

#Optimized for functions of 1 argument (F1-acceleration):
#    xres_max: Max number of x-points in data window.
#returns: Vector{Point2D}
function _reduce_droppts(input::IDataset{true}, xext::PExtents1D, xres_max::Integer)
	xres = (xext.max - xext.min)/ xres_max #WANTCONST
	min_lookahead = 3 #WANTCONST number of xres windows to potentially collapse
	thresh_xres = min_lookahead*xres #WANTCONST maximum x-distance to look ahead for reduction
	x = input.x; y = input.y
	n_ds = length(x) #numer of points of input dataset
	#FIXME: improve algorithm
	#Algorithm succeptible to cumulative errors (esp w large x-values & large xres_max):
	sz = min(n_ds, xres_max)+1+2*min_lookahead #Allow extra points in case
	result = Array{Point2D}(undef, sz)
	n = 0 #Number of points in reduced dataset
	i = 1 #Index into input dataset

	if length(x) != length(y)
		error("x & y - vector length mismatch.")
	elseif length(x) < 1
		resize!(result, 0)
		return result
	end

	#Discard data before visible extents:
	while i < n_ds
		if x[i] > xext.min; break; end
		i+=1
	end

	i = max(i-1, 1)
	prevx = x[i]; prevy = y[i]
	n+=1
	result[n] = Point2D(prevx, prevy)
	i+=1

	while i <= n_ds
		if prevx >= xext.max; break; end
		xthresh = min(prevx+thresh_xres, xext.max)
		ilahead_start = i+(min_lookahead-1) #If we consolidate, must at least reach here
		ilahead = min(ilahead_start, n_ds)
		while ilahead < n_ds
			if x[ilahead] >= xthresh; break; end
			ilahead += 1
		end

		#@assert(ilahead<=nds)

		if ilahead > ilahead_start #More data points than xres allows
			#TODO: make this a function??
			#final point of the data-reduced window:
			ifin = ilahead-1 #ilahead is now beyond minimum xres window
			pfin = Point2D(x[ifin], y[ifin])

			#Find smallest/largest y-value (before/after xres window):
			(ymin, ymax) = minmax(prevy, pfin.y)

			#Add points representing "internal limits"
			#(if they exceed external):
			ywnd = view(y, i:ifin) #y-window ~within 1 xres.
			i_ymin = i-1+argmin(ywnd); i_ymax = i-1+argmax(ywnd)
			ymin_wnd = y[i_ymin]; ymax_wnd = y[i_ymax]
			xint = [x[i_ymin], x[i_ymax]]
			yint = [ymin_wnd, ymax_wnd]
			ysel = Bool[false, false]
			ysel[1] = ymin_wnd < ymin #Don't display twice
			ysel[2] = ymax_wnd > ymax #Don't display twice
#ysel = Bool[true, true] #Debug: add points no matter what
#@show ilahead-i
			#offset: Decide whether to display min or max first - improves look of minres "glitch"
			offset = prevy < pfin.y ? 0 : 1 #Depending if trend is increasing or decreasing
			for j in (offset .+ (1:2))
				idx = 1+(j&0x1)
				if !ysel[idx]; continue; end #Only add points if desired
				n+=1;
				result[n] = Point2D(xint[idx],yint[idx])
			end
			#Done adding points

			prevx = pfin.x; prevy = pfin.y
			n+=1;
			result[n] = Point2D(prevx, prevy)
			i = ilahead
		else #Plot actual data points
			while i <= ilahead
				n += 1
				prevx = x[i]
				prevy = y[i]
				result[n] = Point2D(prevx, prevy)
				i += 1
			end
		end
	end

	resize!(result, n)
	return result
end

#Call generic algorithm when not dealing with functions of 1 argument:
_reduce(input::IDataset, xext::PExtents1D, xres_max::Integer) =
	_reduce_nodrop(input, xext, xres_max)

#Call data-reduction algorithm when dealing with functions of 1 argument:
function _reduce(input::IDataset{true}, xext::PExtents1D, xres_max::Integer)
	try
		return _reduce_droppts(input, xext, xres_max)
	catch
		msg = "Failed to perform data reduction on function of 1 argument.\n"
		msg *= "Algorithm needs to be made more robust.\n"
		msg *= "Defaulting to slower, naive solution (drawing all points).\n"
		@warn(msg)

		return _reduce_nodrop(input, xext, xres_max)
	end
end

#Last line


using TimetableSolver
using JuMP
using ConstraintSolver
using DataStructures: OrderedDict

# TODO: add return type for functions

const CS = ConstraintSolver

"""
    VarRefOrExpr

Represents either a model variable reference or expression reference.

Alias for union of VariableRef and AffExpr.
"""
VarRefOrExpr = Union{VariableRef,AffExpr}

#=
Each variable is a string of the form:
    `{div}_{row}_{period}_{type}`,
where div is the str representation of division,
row is the number of row, period is the number of period,
and type is either 'subject' or 'teacher'.

Each domain value represents a subject or teacher, so its value
should be a str of the form: `{subject}` or `{teacher}`.

However, all the strings are hidden and ints are used for solving directly.
Mappings are generated and stored in VariableData struct.
=#

"""
    VariableData

Immutable struct to store data related to divisions, variables and values,
and their mappings from ints to strings, because the solver only works on ints.

Fields:
- `divvarval_maps::OrderedDict{Symbol,OrderedDict}`: Keys are a symbol `:subject` or `:teacher`.
  Values are map of ints (div) to 2-tuple of vecs of ints (var, val).
- `divvarval_mapstrs::OrderedDict{Symbol,OrderedDict}`: Same as `div_var_val_map` but with strings instead of ints.

- `alldivs::Vector{Int}`: Vector of all divs as ints.
- `allvars::Vector{Int}`: Vector of all vars as ints. (no `:subject` `:teacher` bifurcation)
- `allvals::Vector{Int}`: Vector of all vals as ints. (no `:subject` `:teacher` bifurcation)

- `divmap::OrderedDict{Int,String}`: Map from div ints to strings.
- `varmap::OrderedDict{Int,String}`: Map from variable ints to strings.
- `valmap::OrderedDict{Int,String}`: Map from values ints to strings.
- `inversedivmap::OrderedDict{String,Int}`: Inverted `divmap`.
- `inversevarmap::OrderedDict{String,Int}`: Inverted `varmap`.
- `inversevalmap::OrderedDict{String,Int}`: Inverted `valmap`.

- `schedule::Schedule`: Schedule type instance.

"""
struct VariableData
    divvarval_maps::OrderedDict{Symbol,OrderedDict}
    divvarval_mapstrs::OrderedDict{Symbol,OrderedDict}

    alldivs::Vector{Int}
    allvars::Vector{Int}
    allvals::Vector{Int}

    divmap::OrderedDict{Int,String}
    varmap::OrderedDict{Int,String}
    valmap::OrderedDict{Int,String}
    inversedivmap::OrderedDict{String,Int}
    inversevarmap::OrderedDict{String,Int}
    inversevalmap::OrderedDict{String,Int}

    schedule::Schedule
end

"""
    VariableData(schedule::Schedule)::VariableData

Return `VariableData` instance with all the fields filled based on the schedule.
"""
function VariableData(schedule::Schedule)::VariableData
    # get divvarval_mapstrs
    divvarval_mapstrs = get_divvarval_mapstrs(schedule)

    # extract str representations of all data
    alldiv_strs = Vector{String}([])
    allvar_strs = OrderedDict{Symbol,Vector{String}}(
        :subject => [],
        :teacher => []
    )
    allval_strs = OrderedDict{Symbol,Vector{String}}(
        :subject => [],
        :teacher => []
    )
    # symbol is `:subject` or `:teacher`
    # symbol_map is div_var_val_map_strs but for either subject or teacher only
    for (symbol, symbolmap) in divvarval_mapstrs
        for (div_str, varval_strs) in symbolmap
            push!(alldiv_strs, div_str)
            for var_str in varval_strs[1]
                # push the var to the proper key in all_var_strs
                push!(allvar_strs[symbol], var_str)
            end
            for val_str in varval_strs[2]
                # push the val to the proper key in all_var_strs
                push!(allval_strs[symbol], val_str)
            end
        end
    end

    # remove duplicate values from allval_strs
    allval_strs[:subject] = unique(allval_strs[:subject])
    allval_strs[:teacher] = unique(allval_strs[:teacher])

    # remove duplicate values from alldiv_strs
    alldiv_strs = unique(alldiv_strs)

    # define int representations of all data
    # and int to str mappings
    alldivs = Vector{Int}(1:length(alldiv_strs))
    divmap::OrderedDict{Int,String} = OrderedDict(
        enumerate(alldiv_strs)
    )
    inversedivmap::OrderedDict{String,Int} = OrderedDict(v => k for (k, v) in divmap)

    allvars = Vector{Int}(1:(length(allvar_strs[:subject]))+length(allvar_strs[:teacher]))
    # nested list comp
    # return var_str for var_str in symbol_var_strs which are in all_var_strs
    varmap::OrderedDict{Int,String} = OrderedDict(
        enumerate([v for (_, var_strs) in allvar_strs for v in var_strs])
    )
    inversevarmap::OrderedDict{String,Int} = OrderedDict(v => k for (k, v) in varmap)

    allvals = Vector{Int}(1:(length(allval_strs[:subject]))+length(allval_strs[:teacher]))
    # nested list comp (same as for `var_map`)
    valmap::OrderedDict{Int,String} = OrderedDict(
        enumerate([v for (_, val_strs) in allval_strs for v in val_strs])
    )
    inversevalmap::OrderedDict{String,Int} = OrderedDict(v => k for (k, v) in valmap)

    # dict to store ordereddict of all vars and vals mapped to divs
    # dict keys are :subject and :teacher
    divvarval_maps::OrderedDict{Symbol,OrderedDict} = OrderedDict(
        :subject => OrderedDict(),
        :teacher => OrderedDict()
    )

    # generate values for div_var_val_map
    # symbol is `:subject` or `:teacher`
    # symbol_map is divvarval_mapstrs but for either subject or teacher only
    for (symbol, symbolmap) in divvarval_mapstrs
        for (i, (div_str, (var_strs, val_strs))) in enumerate(symbolmap)
            # int representing current div
            div::Int = inversedivmap[div_str]

            # ints representing current vars
            vars = Vector{Int}([inversevarmap[v] for v in var_strs])
            vals = Vector{Int}([inversevalmap[v] for v in val_strs])

            # add key value pair to div_var_val_map
            divvarval_maps[symbol][div] = (vars, vals)
        end
    end

    return VariableData(
        divvarval_maps, divvarval_mapstrs,
        alldivs, allvars, allvals,
        divmap, varmap, valmap, inversedivmap, inversevarmap, inversevalmap,
        schedule
    )
end

"""
    get_divvarval_mapstrs(schedule::Schedule)::OrderedDict{Symbol,OrderedDict}

Return `OrderedDict` mapping symbol (either `:subject` or `:teacher`), to an `OrderedDict`.

The first-order values (`OrderedDict`), are mappings from division string, to a 2-tuple of variables and values.
"variables" and "values" are vectors of variable strings and value strings.
"""
function get_divvarval_mapstrs(schedule::Schedule)::OrderedDict{Symbol,OrderedDict}
    # dict to store ordereddict of all vars and vals mapped to divs
    # dict keys are :subject and :teacher
    divvarval_maps::OrderedDict{
        Symbol,OrderedDict{String,Tuple{Vector{String},Vector{String}}}
    } = OrderedDict(
        :subject => OrderedDict(),
        :teacher => OrderedDict()
    )

    # fill divvarval_maps with data
    for (div, tt) in schedule.data
        subjectvals::Vector{String} = []
        teachervals::Vector{String} = []

        # add dom vals for both teachers and subjects
        for (subject, teachers) in tt.subjectteachers
            push!(subjectvals, subject)
            for t in teachers
                push!(teachervals, t)
            end
        end

        subjectvars::Vector{String} = []
        teachervars::Vector{String} = []

        for (i, row) in enumerate(tt.data)
            for j in (1:length(row))
                # add subject variable
                subject = "$(div)_$(i)_$(j)_subject"
                push!(subjectvars, subject)
                # add teacher variable
                teacher = "$(div)_$(i)_$(j)_teacher"
                push!(teachervars, teacher)
            end
        end

        # only use unique vals
        subjectvals = unique(subjectvals)
        teachervals = unique(teachervals)

        # create a key for every div, in teacher and subject
        divvarval_maps[:subject][div] = (subjectvars, subjectvals)
        divvarval_maps[:teacher][div] = (teachervars, teachervals)
    end

    return divvarval_maps
end

# TODO: add kwargs arg to this and other function to pass kwargs to Model()
"""
    get_model(vardata::VariableData, timeout::Float64=Inf, all_solutions::Bool=false)::Tuple{Model,OrderedDict}

Return a JuMP model created using `vardata`, and an OrderedDict of the variables defined in the model.

# Arguments
- `vardata::VariableData`: VariableData object representing the variables and constraints to be added to the model.
- `timeout::Float64=Inf`: Float representing the maximum time to run the model.
- `all_solutions::Bool=false`: Bool representing whether to return all solutions or just the first.

# Notes
- Create a model with the correct variables, and adds constraints to the model.
- The variables defined are anonymous, so they are stored in an OrderedDict mapped to their int representation.
"""
function get_model(vardata::VariableData, timeout::Float64 = Inf, all_solutions::Bool = false
)::Tuple{Model,OrderedDict}
    # create a constraint solver model and set ConstraintSolver as the optimizer
    m = Model(optimizer_with_attributes(CS.Optimizer,
        # logging vector containing :Info and/or :Table or empty
        "logging" => [],

        # whether to return all possible solutions
        "all_solutions" => all_solutions,

        # seed for random operations given for reproducibility
        "seed" => 0,
        # default to :Auto
        # "traverse_strategy" => :BFS,
        # "traverse_strategy"=> :DFS,
        # "traverse_strategy"=> :DBFS,

        # default :Auto
        # "branch_split"=>:Smallest,
        # "branch_split"=>:Biggest,
        # "branch_split" => :InHalf,

        "branch_strategy" => :IMPS, # default

        # simplify alldifferent, etc. constraints
        "simplify" => true, # default

        # max seconds before exiting with status `MOI.TIME_LIMIT`
        "time_limit" => timeout,

        # backtracking options
        "backtrack" => true, # default
        "backtrack_sorting" => true, # default

        # "lp_optimizer" => cbc_optimizer,
        # "lp_optimizer" => glpk_optimizer,
        # "lp_optimizer" => ipopt_optimizer,
    ))

    # define variables
    modelvars = define_variables!(m, vardata)

    # define constraints
    define_constraints!(m, modelvars, vardata)

    return m, modelvars
end

"""
    define_variables!(m::Model, vardata::VariableData)::OrderedDict{Int,VarRefOrExpr}

Define variables in the model `m` using the `vardata` object.
Return an OrderedDict mapping the int representation of the variable to the variable reference/expression reference.

# Notes
- Modify the model `m` in place.
- See [`VarRefOrExpr`](@ref).
"""
function define_variables!(
    m::Model, vardata::VariableData)::OrderedDict{Int,VarRefOrExpr}
    divvarval_maps = vardata.divvarval_maps

    # mapping variable int repr to its model variableref
    modelvars = OrderedDict{Int,VarRefOrExpr}()

    for (symbol, symbolmap) in divvarval_maps
        for (div, (vars, vals)) in symbolmap
            # define a set of vars at a time
            currentmodelvars = @variable(m, [vars], CS.Integers(vals))
            # allocate them to the value of var in modelvars one by one
            for v in vars
                modelvars[v] = currentmodelvars[v]
            end
        end
    end

    return modelvars
end

"""
    define_subjectconstraints!(m::Model, modelvars::OrderedDict, vardata::VariableData)::Nothing

Define subject constraints in the model `m` using the `vardata` object and the model variables `modelvars`.

# Notes
- Modify the model `m` in place.
"""
function define_subjectconstraints!(m::Model, modelvars::OrderedDict, vardata::VariableData)::Nothing
    schedule = vardata.schedule
    divvarval_maps = vardata.divvarval_maps
    inversedivmap = vardata.inversedivmap
    inversevalmap = vardata.inversevalmap
    for (div_str, tt) in schedule.data
        # define terms
        div = inversedivmap[div_str]
        counts = OrderedDict{Int,Int}(
            inversevalmap[subject_str] => count
            for (subject_str, count) in tt.subjectcounts
        )
        vars = divvarval_maps[:subject][div][1]
        vals = collect(Int, keys(counts))
        # tl;dr (var::Int = 2) -> (var::Binary[] = [0, 1, 0]) index shows the value
        # create a matrix of binary variables:
        # `vars` as indices and `vals` as number of values
        # so instead of setting an int as value of each var
        # each var has `vals` corresponding binary variables
        binvars = @variable(m, [vars, vals], Bin)
        # add constraint to only choose one val
        @constraint(m, [i = vars], sum(binvars[i, :]) == 1)
        # add constraint to only choose that val `counts[val]` times
        @constraint(m, [j = vals], sum(binvars[:, j]) == counts[j])
        # array of expressions that actually store the value of var as int
        exprs = @expression(m, [i = vars], sum(j * binvars[i, j] for j in vals))
        # replace the corresponding vars in modelvars with exprs
        for var in vars
            # delete the var from the model
            delete(m, modelvars[var])
            # set the var to an expr defined above
            modelvars[var] = exprs[var]
        end
    end
end

"""
    define_subjectteacherconstraints!(m::Model, modelvars::OrderedDict, vardata::VariableData)::Nothing

Define subject-teacher constraints in the model `m` using the `vardata` object and the model variables `modelvars`.

# Notes
- Modify the model `m` in place.
"""
function define_subjectteacherconstraints!(m::Model, modelvars::OrderedDict, vardata::VariableData)::Nothing
    divvarval_maps = vardata.divvarval_maps
    inversevalmap = vardata.inversevalmap
    schedule = vardata.schedule

    # mapping div to corresponding subjectteachers
    subjectteacher_maps = OrderedDict{Int,OrderedDict}()
    for (div_str, tt) in schedule.data
        div = vardata.inversedivmap[div_str]
        # subject teacher mappings as ints
        subjectteachers = OrderedDict{Int,Vector{Int}}()
        # get the subjectteacher value for this timetable
        subjectteacher_strs = tt.subjectteachers
        # convert strs to ints, and add to `subjectteachers`
        for (subject, teachers) in subjectteacher_strs
            subject_int = inversevalmap[subject]
            teacher_ints = [inversevalmap[t] for t in teachers]
            subjectteachers[subject_int] = teacher_ints
        end
        subjectteacher_maps[div] = subjectteachers

    end

    # iterate over every set of subject variables (by div)
    for (div, (subjectvars, _)) in divvarval_maps[:subject]
        # generate pairs of all possible subject teacher combinations
        # and store them in a matrix
        numpairs = sum(length.(values(subjectteacher_maps[div])))
        subjectteacherpairs = Matrix{Int}(undef, numpairs, 2)
        i = 0
        for (subject, teachers) in subjectteacher_maps[div]
            for teacher in teachers
                i += 1
                subjectteacherpairs[i, 1:2] = [subject, teacher]
            end
        end
        # add using CS.TableSet(table)
        teachervars = divvarval_maps[:teacher][div][1]
        for (svar, tvar) in zip(subjectvars, teachervars)
            # model vars corresponding to this pair of subject and teacher vars
            (modelsvar, modeltvar) = modelvars[svar], modelvars[tvar]
            @constraint(m, [modelsvar, modeltvar] in CS.TableSet(subjectteacherpairs))
        end
    end

    return nothing
end

"""
    define_teacherconstraints!(m::Model, modelvars::OrderedDict, vardata::VariableData)::Nothing

Define teacher constraints in the model `m` using the `vardata` object and the model variables `modelvars`.

# Notes
- Modify the model `m` in place.
"""
function define_teacherconstraints!(m::Model, modelvars::OrderedDict, vardata::VariableData)::Nothing
    # variables with same row,period should have different teachers
    divvarval_maps = vardata.divvarval_maps

    # following block of code creates a list of vars
    # grouped by the same row and period
    i = 0
    while true
        i += 1
        # store var or 0 if var not found
        rowperiodvars = [get(vars, i, 0) for (_, (vars, _)) in divvarval_maps[:teacher]]
        # remove zeroes from rowperiodvars
        filter!(!iszero, rowperiodvars)
        # if no vars were pushed to list,
        # then we have exceeded the number of rows and periods
        if length(rowperiodvars) == 0
            break
        end
        # now get model vars from these vars
        rowperiodmodelvars = [modelvars[x] for x in rowperiodvars]
        # add alldifferent constraint
        @constraint(m, rowperiodmodelvars in CS.AllDifferent())
    end

    return nothing
end

"""
    define_constraints!(m::Model, modelvars::OrderedDict, vardata::VariableData)::Nothing

Define all constraints in the model `m` using the `vardata` object and the model variables `modelvars`.

# Notes
- Modify the model `m` in place.
- Call [`define_subjectconstraints!`](@ref), [`define_subjectteacherconstraints!`](@ref)
  and [`define_teacherconstraints!`](@ref) internally, in the respective order.
"""
function define_constraints!(m::Model, modelvars::OrderedDict, vardata::VariableData)::Nothing
    define_subjectconstraints!(m, modelvars, vardata)

    define_subjectteacherconstraints!(m, modelvars, vardata)

    define_teacherconstraints!(m, modelvars, vardata)

    return nothing
end

"""
    get_solution(m::Model, modelvars::OrderedDict, vardata::VariableData)::Tuple{OrderedDict{Int,Int},MOI.TerminationStatusCode}

Return a solution and status code after solving the given model `m` and using model variables `modelvars` and variable data `vardata` to access the solution values.

# Notes
- Model is optimally solved only if the status code is [`MOI.OPTIMAL`](https://jump.dev/JuMP.jl/stable/moi/reference/models/#MathOptInterface.TerminationStatusCode).
- Solution is a mapping from variable to its value (both as ints).
"""
function get_solution(
    m::Model, modelvars::OrderedDict, vardata::VariableData)::Tuple{OrderedDict{Int,Int},MOI.TerminationStatusCode}
    # solve model
    optimize!(m)

    # convert model values to var=>sol mapping
    solution = OrderedDict{Int,Int}()

    # model status
    status = JuMP.termination_status(m)

    if status != MOI.OPTIMAL
        error("Model unsolved!\nstatus = $status")
    end

    for (symbol, symbolmap) in vardata.divvarval_maps
        for (div, (vars, vals)) in symbolmap
            for var in vars
                # get the corresponding var defined in the model
                mvar = modelvars[var]
                # get its value
                mval = Int(JuMP.value(mvar))
                # add var val mapping to solution
                solution[var] = mval
            end
        end
    end

    # return solution
    return solution, status
end

"""
    convertsolution(rawsolution::OrderedDict{Int,Int}, vardata::VariableData)::OrderedDict{String,String}

Return a solution in the form of a mapping from variable to its value (both as strings).

# Notes
- Convert an `Int=>Int` solution to a `String=>String` one using the `inverse*` mappings defined in `vardata`.
"""
function convertsolution(
    rawsolution::OrderedDict{Int,Int}, vardata::VariableData)::OrderedDict{String,String}
    solution = OrderedDict{String,String}()
    # iterate over every var and val in solution
    for (var, val) in rawsolution
        # get var val strs using mappings
        var_str = vardata.varmap[var]
        val_str = vardata.valmap[val]
        # add it to solution
        solution[var_str] = val_str
    end
    return solution
end

"""
    applysolution!(schedule::Schedule, solution::OrderedDict{String,String})::Nothing

Apply a solution (of strings) to the given schedule.

# Notes
- Modify the schedule in place.
"""
function applysolution!(schedule::Schedule, solution::OrderedDict{String,String})::Nothing
    # iterate over every var and val in solution
    for (var, val) in solution
        # modify schedule with var and val
        modify!(schedule, var, val)
    end
end

"""
    solve!(schedule::Schedule, timeout::Float64 = Inf, all_solutions::Bool = false)::Tuple{OrderedDict{String,String},MOI.TerminationStatusCode}

Solve a given schedule in place and return a `String=>String` solution and status code.

# Arguments
- `schedule::Schedule`: Schedule object representing the problem.
- `timeout::Float64=Inf`: Float representing the maximum time to run the model.
- `all_solutions::Bool=false`: Bool representing whether to return all solutions or just the first.

# Notes
- Schedule is optimally solved only if the status code is [`MOI.OPTIMAL`](https://jump.dev/JuMP.jl/stable/moi/reference/models/#MathOptInterface.TerminationStatusCode).
"""
function solve!(
    schedule::Schedule, timeout::Float64 = Inf, all_solutions::Bool = false
)::Tuple{OrderedDict{String,String},MOI.TerminationStatusCode}
    # get vardata
    vardata = VariableData(schedule)

    # get model and model vars
    m, modelvars = get_model(vardata, timeout, all_solutions)

    # get solution
    rawsolution, status = get_solution(m, modelvars, vardata)

    # convert solution from ints to strs
    solution = convertsolution(rawsolution, vardata)

    # apply solution
    applysolution!(schedule, solution)

    # return solution
    return solution, status
end


# Atomic weights taken from:

# Atomic weights of the elements 2013 (IUPAC Technical  Report)
# http://doi.org/10.1515/pac-2015-0305

# The "Conventional" weights are used where a range is provided, shown in comments.

# hydrogen 1
@mol H 1.008 # [1.0078, 1.0082]
# helium 2
@mol He 4.0026
# lithium 3
@mol Li 6.94 # [6.938, 6.997]
# beryllium 4
@mol Be 9.0122
# boron 5
@mol B 10.81 # [10.806, 10.821]
# carbon 6
@mol C 12.011 # [12.009, 12.012]
# nitrogen 7
@mol N 14.007 # [14.006, 14.008]
# oxygen 8
@mol O 15.999 # [15.999, 16.000]
# fluorine 9
@mol F 18.998
# neon 10
@mol Ne 20.180
# sodium 11
@mol Na 22.990
# magnesium 12
@mol Mg 24.305 # [24.304, 24.307]
# aluminium (aluminum) 13
@mol Al 26.982
# silicon 14
@mol Si 28.085 # [28.084, 28.086]
# phosphorus 15
@mol P 30.974
# sulfur 16
@mol S 32.06 # [32.059, 32.076]
# chlorine 17
@mol Cl 35.45 # [35.446, 35.457]
# argon 18
@mol Ar 39.948
# potassium 19
@mol K 39.098
# calcium 20
@mol Ca 40.078 # (4)
# scandium 21
@mol Sc 44.956
# titanium 22
@mol Ti 47.867
# vanadium 23
@mol V 50.942
# chromium 24
@mol Cr 51.996
# manganese 25
@mol Mn 54.938
# iron 26
@mol Fe 55.845 # (2)
# cobalt 27
@mol Co 58.933
# nickel 28
@mol Ni 58.693
# copper 29
@mol Cu 63.546 # (3)
# zinc 30
@mol Zn 65.38 # (2)
# gallium 31
@mol Ga 69.723
# germanium 32
@mol Ge 72.630 # (8)
# arsenic 33
@mol As 74.922
# selenium 34
@mol Se 78.971 # (8)
# bromine 35
@mol Br 79.904 # [79.901, 79.907]
# krypton 36
@mol Kr 83.798 # (2)
# rubidium 37
@mol Rb 85.468
# strontium 38
@mol Sr 87.62
# yttrium 39
@mol Y 88.906
# zirconium 40
@mol Zr 91.224 # (2)
# niobium 41
@mol Nb 92.906
# molybdenum 42
@mol Mo 95.95
# technetium* 43
# @mol Tc –
# ruthenium 44
@mol Ru 101.07 # (2)
# rhodium 45
@mol Rh 102.91
# palladium 46
@mol Pd 106.42
# silver 47
@mol Ag 107.87
# cadmium 48
@mol Cd 112.41
# indium 49
@mol In 114.82
# tin 50
@mol Sn 118.71
# antimony 51
@mol Sb 121.76
# tellurium 52
@mol Te 127.60 # (3)
# iodine 53
@mol I 126.90
# xenon 54
@mol Xe 131.29
# caesium (cesium) 55
@mol Cs 132.91
# barium 56
@mol Ba 137.33
# lanthanum 57
@mol La 138.91
# cerium 58
@mol Ce 140.12
# praseodymium 59
@mol Pr 140.91
# neodymium 60
@mol Nd 144.24
# promethium* 61
# @mol Pm –
# samarium 62
@mol Sm 150.36 # (2)
# europium 63
@mol Eu 151.96
# gadolinium 64
@mol Gd 157.25 # (3)
# terbium 65
@mol Tb 158.93
# dysprosium 66
@mol Dy 162.50
# holmium 67
@mol Ho 164.93
# erbium 68
@mol Er 167.26
# thulium 69
@mol Tm 168.93
# ytterbium 70
@mol Yb 173.05
# lutetium 71
@mol Lu 174.97
# hafnium 72
@mol Hf 178.49 # (2)
# tantalum 73
@mol Ta 180.95
# tungsten 74
@mol W 183.84
# rhenium 75
@mol Re 186.21
# osmium 76
@mol Os 190.23 # (3)
# iridium 77
@mol Ir 192.22
# platinum 78
@mol Pt 195.08
# gold 79
@mol Au 196.97
# mercury 80
@mol Hg 200.59
# thallium 81
@mol Tl 204.38 # [204.38, 204.39]
# lead 82
@mol Pb 207.2
# bismuth* 83
@mol Bi 208.98