"""
A file with structs and functions for sampling hops and updating hopping rates
"""

# TODO to simplify the design can do two things:
# 1. make an abstract type HopRatesGridOptim (grid-optimized) and dispatch on it. Or maybe dispatch on the type of hop_const_cumulative_sums instead (Array{F} vs Array{Vector{F}})
# 2. make an abstract type HopRatesGraphDsi, which has hop rates uniform wrt neighbors, and treat things not of that type as having hop_const_cumulative_sums field

### spatial hop rates ###
abstract type AbstractHopRates end

HopRates(hopping_constants::Matrix{F}, spatial_system) where F <: Number = HopRatesGraphDsi(hopping_constants)
HopRates(hopping_constants::Matrix{F}, grid::Union{CartesianGridRej, CartesianGridIter}) where F <: Number = HopRatesGraphDsi(hopping_constants)

HopRates(hopping_constants::Matrix{Vector{F}}, spatial_system) where F <: Number = HopRatesGraphDsij(hopping_constants)
HopRates(hopping_constants::Matrix{Vector{F}}, grid::Union{CartesianGridRej, CartesianGridIter}) where F <: Number = HopRatesGridDsij(hopping_constants, grid)

HopRates(p::Pair{SpecHop, SiteHop}, spatial_system) where {SpecHop <: Vector{F}, SiteHop <: Vector{Vector{F}}} where F <: Number = HopRatesGraphDsLij(p...)
HopRates(p::Pair{SpecHop, SiteHop}, grid::Union{CartesianGridRej, CartesianGridIter}) where {SpecHop <: Vector{F}, SiteHop <: Vector{Vector{F}}} where F <: Number = HopRatesGridDsLij(p..., grid)

HopRates(p::Pair{SpecHop, SiteHop}, spatial_system) where {SpecHop <: Matrix{F}, SiteHop <: Vector{Vector{F}}} where F <: Number = HopRatesGraphDsiLij(p...)
HopRates(p::Pair{SpecHop, SiteHop}, grid::Union{CartesianGridRej, CartesianGridIter}) where {SpecHop <: Matrix{F}, SiteHop <: Vector{Vector{F}}} where F <: Number = HopRatesGridDsiLij(p..., grid)

"""
    update_hop_rates!(hop_rates::AbstractHopRates, species::AbstractArray, u, site, spatial_system)

update rates of all specs in species at site
"""
function update_hop_rates!(hop_rates::AbstractHopRates, species::AbstractArray, u, site, spatial_system)
    @inbounds for spec in species
        update_hop_rate!(hop_rates, spec, u, site, spatial_system)
    end
end

"""
    update_hop_rate!(hop_rates::HopRatesGraphDsi, species, u, site, spatial_system)

update rates of single species at site
"""
function update_hop_rate!(hop_rates::AbstractHopRates, species, u, site, spatial_system)
    rates = hop_rates.rates
    @inbounds old_rate = rates[species, site]
    @inbounds rates[species, site] = evalhoprate(hop_rates, u, species, site, spatial_system)
    @inbounds hop_rates.sum_rates[site] += rates[species, site] - old_rate
    old_rate
end

"""
    total_site_hop_rate(hop_rates::AbstractHopRates, site)

return total hopping rate out of site
"""
total_site_hop_rate(hop_rates::AbstractHopRates, site) = @inbounds hop_rates.sum_rates[site]

"""
    reset!(hop_rates::AbstractHopRates)

make all rates zero
"""
function reset!(hop_rates::AbstractHopRates)
    hop_rates.rates .= zero(eltype(hop_rates.rates))
    hop_rates.sum_rates .= zero(eltype(hop_rates.sum_rates))
    nothing
end

"""
    sample_species(hop_rates::AbstractHopRates, site, rng)

sample species to hop from site
"""
sample_species(hop_rates::AbstractHopRates, site, rng) = @inbounds linear_search((@view hop_rates.rates[:,site]), rand(rng) * total_site_hop_rate(hop_rates, site))

"""
    sample_hop_at_site(hop_rates::AbstractHopRates, site, rng, spatial_system)

sample a reaction at site, return (species, target_site)
"""
function sample_hop_at_site(hop_rates::AbstractHopRates, site, rng, spatial_system)
    species = sample_species(hop_rates, site, rng)
    return species, sample_target_site(hop_rates, site, species, rng, spatial_system)
end

############## hopping rates of form D_{s,i} ################
struct HopRatesGraphDsi{F} <: AbstractHopRates
    "hopping_constants[i,j] is the hop constant of species i at site j"
    hopping_constants::Matrix{F}

    "rates[i,j] is total hopping rate of species i at site j"
    rates::Matrix{F}

    "sum_rates[j] is the sum of hopping rates at site j"
    sum_rates::Vector{F}
end

function Base.show(io::IO, ::MIME"text/plain", hop_rates::HopRatesGraphDsi)
    num_specs, num_sites = size(hop_rates.rates)
    println(io, "HopRates with $num_specs species and $num_sites sites. \nHopping constants of form D_{s,i} where s is species, and i is source.")
end

"""
    HopRatesGraphDsi(hopping_constants::Matrix{F}) where F <: Number

initializes HopRatesGraphDsi with zero rates
"""
function HopRatesGraphDsi(hopping_constants::Matrix{F}) where F <: Number
    rates = zeros(F, size(hopping_constants))
    HopRatesGraphDsi{F}(hopping_constants, rates, zeros(F, size(rates, 2)))
end

sample_target_site(hop_rates::HopRatesGraphDsi, site, species, rng, spatial_system) = rand_nbr(rng, spatial_system, site)

"""
return hopping rate of species at site
"""
function evalhoprate(hop_rates::HopRatesGraphDsi, u, species, site, spatial_system)
    @inbounds u[species,site]*hop_rates.hopping_constants[species,site]*outdegree(spatial_system, site)
end

############## hopping rates of form D_{s,i,j} ################
struct HopRatesGraphDsij{F} <: AbstractHopRates
    "hop_const_cumulative_sums[s,i] is the vector of cumulative sums of hopping constants of species s at site i"
    hop_const_cumulative_sums::Matrix{Vector{F}}

    "rates[s,i] is the total hopping rate of species s at site i"
    rates::Matrix{F}

    "sum_rates[i] is the total hopping rate out of site i"
    sum_rates::Vector{F}
end

function Base.show(io::IO, ::MIME"text/plain", hop_rates::HopRatesGraphDsij)
    num_specs, num_sites = size(hop_rates.rates)
    println(io, "HopRates with $num_specs species and $num_sites sites. \nHopping constants of form D_{s,i,j} where s is species, i is source and j is destination.")
end

"""
    HopRatesGraphDsij(hopping_constants::Matrix{Vector{F}}) where F <: Number

initializes HopRates with zero rates
"""
function HopRatesGraphDsij(hopping_constants::Matrix{Vector{F}}; do_cumsum = true) where F <: Number
    do_cumsum && (hopping_constants = map(cumsum, hopping_constants))
    rates = zeros(F, size(hopping_constants))
    sum_rates = zeros(F, size(rates, 2))
    HopRatesGraphDsij{F}(hopping_constants, rates, sum_rates)
end

function sample_target_site(hop_rates::HopRatesGraphDsij, site, species, rng, spatial_system)
    @inbounds cumulative_hop_constants = hop_rates.hop_const_cumulative_sums[species, site]
    @inbounds n = searchsortedfirst(cumulative_hop_constants, rand(rng) * cumulative_hop_constants[end])
    return nth_nbr(spatial_system, site, n)
end

function evalhoprate(hop_rates::HopRatesGraphDsij, u, species, site, spatial_system)
    @inbounds u[species, site] * hop_rates.hop_const_cumulative_sums[species, site][end]
end

#################  hopping rates of form L_{s,i,j} optimized for cartesian grid  ######################
"""
Analogue of HopRatesGraphDsij, optimized for CartesianGrid
"""
struct HopRatesGridDsij{F} <: AbstractHopRates
    "hop_const_cumulative_sums[:,s,i] is the vector of cumulative sums of hopping constants of species s at site i"
    hop_const_cumulative_sums::Array{F, 3}

    "rates[s,i] is the total hopping rate of species s at site i"
    rates::Matrix{F}

    "sum_rates[i] is the total hopping rate out of site i"
    sum_rates::Vector{F}
end

function Base.show(io::IO, ::MIME"text/plain", hop_rates::HopRatesGridDsij)
    num_specs, num_sites = size(hop_rates.rates)
    println(io, "HopRates with $num_specs species and $num_sites sites, optimized for CartesianGrid. \nHopping constants of form L_{s,i,j} where s is species, i is source and j is destination.")
end

"""
    HopRatesGridDsij(hopping_constants::Matrix{Vector{F}}) where F <: Number

initializes HopRates with zero rates
"""
function HopRatesGridDsij(hopping_constants::Array{F, 3}; do_cumsum = true) where F <: Number
    do_cumsum && (hopping_constants = mapslices(cumsum, hopping_constants, dims = 1))
    rates = zeros(F, size(hopping_constants)[2:3])
    sum_rates = zeros(F, size(rates, 2))
    HopRatesGridDsij{F}(hopping_constants, rates, sum_rates)
end

function HopRatesGridDsij(hopping_constants::Matrix{Vector{F}}, grid) where F <: Number
    new_hopping_constants = Array{F, 3}(undef, 2*dimension(grid), size(hopping_constants)...)
    for ci in CartesianIndices(hopping_constants)
        species, site = Tuple(ci)
        nb_constants = @view new_hopping_constants[:, species, site]
        pad_hop_vec!(nb_constants, grid, site, hopping_constants[ci])
        cumsum!(nb_constants, nb_constants)
    end
    HopRatesGridDsij(new_hopping_constants, do_cumsum = false)
end

function sample_target_site(hop_rates::HopRatesGridDsij, site, species, rng, grid)
    @inbounds cumulative_hop_constants = @view hop_rates.hop_const_cumulative_sums[:,species, site]
    @inbounds n = searchsortedfirst(cumulative_hop_constants, rand(rng) * cumulative_hop_constants[end])
    return nth_potential_nbr(grid, site, n)
end

function evalhoprate(hop_rates::HopRatesGridDsij, u, species, site, spatial_system)
    @inbounds u[species, site] * hop_rates.hop_const_cumulative_sums[end, species, site]
end

############## hopping rates of form D_s * L_{i,j} ################
struct HopRatesGraphDsLij{F} <: AbstractHopRates
    "hopping constants of species -- D_s"
    species_hop_constants::Vector{F}

    "nbs_cumulative[i] is the vector of cumulative sums of hopping constants from site i to its neighbors"
    hop_const_cumulative_sums::Vector{Vector{F}}

    "rates[s,i] is the total hopping rate of species s at site i"
    rates::Matrix{F}

    "sum_rates[i] is the total hopping rate out of site i"
    sum_rates::Vector{F}
end

function Base.show(io::IO, ::MIME"text/plain", hop_rates::HopRatesGraphDsLij)
    num_specs, num_sites = length(hop_rates.species_hop_constants), length(hop_rates.hop_const_cumulative_sums)
    println(io, "HopRates with $num_specs species and $num_sites sites. \nHopping constants of form D_s * L_{i,j} where s is species, i is source and j is destination.")
end

"""
    HopRatesGraphDsLij(species_hop_constants::Vector{F}, site_hop_constants::Vector{Vector{F}}) where F <: Number

initializes HopRates with zero rates
"""
function HopRatesGraphDsLij(species_hop_constants::Vector{F}, site_hop_constants::Vector{Vector{F}}; do_cumsum = true) where F <: Number
    do_cumsum && (site_hop_constants = map(cumsum, site_hop_constants))
    rates = zeros(F, length(species_hop_constants), length(site_hop_constants))
    sum_rates = zeros(size(rates, 2))
    HopRatesGraphDsLij{F}(species_hop_constants, site_hop_constants, rates, sum_rates)
end

function sample_target_site(hop_rates::HopRatesGraphDsLij, site, species, rng, spatial_system)
    @inbounds cumulative_hop_constants = hop_rates.hop_const_cumulative_sums[site]
    @inbounds n = searchsortedfirst(cumulative_hop_constants, rand(rng) * cumulative_hop_constants[end])
    return nth_nbr(spatial_system, site, n)
end

function evalhoprate(hop_rates::HopRatesGraphDsLij, u, species, site, spatial_system)
    @inbounds u[species, site] * hop_rates.species_hop_constants[species] * hop_rates.hop_const_cumulative_sums[site][end]
end

############## hopping rates of form D_s * L_{i,j} optimized for cartesian grid ################
struct HopRatesGridDsLij{F} <: AbstractHopRates
    "hopping constants of species -- D_s"
    species_hop_constants::Vector{F}

    "nbs_cumulative[i] is the vector of cumulative sums of hopping constants from site i to its neighbors"
    hop_const_cumulative_sums::Matrix{F}

    "rates[s,i] is the total hopping rate of species s at site i"
    rates::Matrix{F}

    "sum_rates[i] is the total hopping rate out of site i"
    sum_rates::Vector{F}
end

function Base.show(io::IO, ::MIME"text/plain", hop_rates::HopRatesGridDsLij)
    num_specs, num_sites = length(hop_rates.species_hop_constants), size(hop_rates.hop_const_cumulative_sums, 2)
    println(io, "HopRates with $num_specs species and $num_sites sites, optimized for CartesianGrid. \nHopping constants of form D_s * L_{i,j} where s is species, i is source and j is destination.")
end

function HopRatesGridDsLij(species_hop_constants::Vector{F}, site_hop_constants::Matrix{F}; do_cumsum = true) where F <: Number
    do_cumsum && (site_hop_constants = mapslices(cumsum, site_hop_constants, dims = 1))
    rates = zeros(F, length(species_hop_constants), size(site_hop_constants, 2))
    sum_rates = zeros(size(rates, 2))
    HopRatesGridDsLij{F}(species_hop_constants, site_hop_constants, rates, sum_rates)
end

function HopRatesGridDsLij(species_hop_constants::Vector{F}, site_hop_constants::Vector{Vector{F}}, grid) where F <: Number
    new_hopping_constants = Matrix{F}(undef, 2*dimension(grid), length(site_hop_constants))
    for site in 1:length(site_hop_constants)
        nb_constants = @view new_hopping_constants[:, site]
        pad_hop_vec!(nb_constants, grid, site, site_hop_constants[site])
        cumsum!(nb_constants, nb_constants)
    end
    HopRatesGridDsLij(species_hop_constants, new_hopping_constants, do_cumsum = false)
end

function sample_target_site(hop_rates::HopRatesGridDsLij, site, species, rng, grid)
    @inbounds cumulative_hop_constants = @view hop_rates.hop_const_cumulative_sums[:, site]
    @inbounds n = searchsortedfirst(cumulative_hop_constants, rand(rng) * cumulative_hop_constants[end])
    return nth_potential_nbr(grid, site, n)
end

function evalhoprate(hop_rates::HopRatesGridDsLij, u, species, site, grid)
    @inbounds u[species, site] * hop_rates.species_hop_constants[species] * hop_rates.hop_const_cumulative_sums[end, site]
end

############## hopping rates of form D_{s,i} * L_{i,j} ################
struct HopRatesGraphDsiLij{F} <: AbstractHopRates
    "hopping constants of species -- D_s"
    species_hop_constants::Matrix{F}

    "nbs_cumulative[i] is the vector of cumulative sums of hopping constants from site i to its neighbors"
    hop_const_cumulative_sums::Vector{Vector{F}}

    "rates[s,i] is the total hopping rate of species s at site i"
    rates::Matrix{F}

    "sum_rates[i] is the total hopping rate out of site i"
    sum_rates::Vector{F}
end

function Base.show(io::IO, ::MIME"text/plain", hop_rates::HopRatesGraphDsiLij)
    num_specs, num_sites = size(hop_rates.species_hop_constants)
    println(io, "HopRates with $num_specs species and $num_sites sites. \nHopping constants of form D_{s,i} * L_{i,j} where s is species, i is source and j is destination.")
end

function HopRatesGraphDsiLij(species_hop_constants::Matrix{F}, site_hop_constants::Vector{Vector{F}}; do_cumsum = true) where F <: Number
    @assert size(species_hop_constants, 2) == length(site_hop_constants)
    do_cumsum && (site_hop_constants = map(cumsum, site_hop_constants))
    rates = zeros(F, length(species_hop_constants), length(site_hop_constants))
    sum_rates = zeros(size(rates, 2))
    HopRatesGraphDsiLij{F}(species_hop_constants, site_hop_constants, rates, sum_rates)
end

function sample_target_site(hop_rates::HopRatesGraphDsiLij, site, species, rng, spatial_system)
    @inbounds cumulative_hop_constants = hop_rates.hop_const_cumulative_sums[site]
    @inbounds n = searchsortedfirst(cumulative_hop_constants, rand(rng) * cumulative_hop_constants[end])
    return nth_nbr(spatial_system, site, n)
end

function evalhoprate(hop_rates::HopRatesGraphDsiLij, u, species, site, spatial_system)
    @inbounds u[species, site] * hop_rates.species_hop_constants[species, site] * hop_rates.hop_const_cumulative_sums[site][end]
end

############## hopping rates of form D_{s,i} * L_{i,j} optimized for cartesian grid ################
struct HopRatesGridDsiLij{F} <: AbstractHopRates
    "hopping constants of species -- D_s"
    species_hop_constants::Matrix{F}

    "nbs_cumulative[i] is the vector of cumulative sums of hopping constants from site i to its neighbors"
    hop_const_cumulative_sums::Matrix{F}

    "rates[s,i] is the total hopping rate of species s at site i"
    rates::Matrix{F}

    "sum_rates[i] is the total hopping rate out of site i"
    sum_rates::Vector{F}
end

function Base.show(io::IO, ::MIME"text/plain", hop_rates::HopRatesGridDsiLij)
    num_specs, num_sites = length(hop_rates.species_hop_constants), size(hop_rates.hop_const_cumulative_sums, 2)
    println(io, "HopRates with $num_specs species and $num_sites sites, optimized for CartesianGrid. \nHopping constants of form D_{s,i} * L_{i,j} where s is species, i is source and j is destination.")
end

function HopRatesGridDsiLij(species_hop_constants::Matrix{F}, site_hop_constants::Matrix{F}; do_cumsum = true) where F <: Number
    @assert size(species_hop_constants, 2) == size(site_hop_constants, 2)
    do_cumsum && (site_hop_constants = mapslices(cumsum, site_hop_constants, dims = 1))
    rates = zeros(F, size(species_hop_constants))
    sum_rates = zeros(size(rates, 2))
    HopRatesGridDsiLij{F}(species_hop_constants, site_hop_constants, rates, sum_rates)
end

function HopRatesGridDsiLij(species_hop_constants::Matrix{F}, site_hop_constants::Vector{Vector{F}}, grid) where F <: Number
    new_hopping_constants = Matrix{F}(undef, 2*dimension(grid), length(site_hop_constants))
    for site in 1:length(site_hop_constants)
        nb_constants = @view new_hopping_constants[:, site]
        pad_hop_vec!(nb_constants, grid, site, site_hop_constants[site])
        cumsum!(nb_constants, nb_constants)
    end
    HopRatesGridDsiLij(species_hop_constants, new_hopping_constants, do_cumsum = false)
end

function sample_target_site(hop_rates::HopRatesGridDsiLij, site, species, rng, grid)
    @inbounds cumulative_hop_constants = @view hop_rates.hop_const_cumulative_sums[:, site]
    @inbounds n = searchsortedfirst(cumulative_hop_constants, rand(rng) * cumulative_hop_constants[end])
    return nth_potential_nbr(grid, site, n)
end

function evalhoprate(hop_rates::HopRatesGridDsiLij, u, species, site, grid)
    @inbounds u[species, site] * hop_rates.species_hop_constants[species, site] * hop_rates.hop_const_cumulative_sums[end, site]
end


module MonteCarloTools

using LinearAlgebra
using QuantumModels
using ExactDiagonalizationTools

include("mctools.jl")
include("detmatrix.jl")
include("gutzwillerslater.jl")
include("vmc.jl")

export runVMC

end


function J2DifferentialEffect(arg0::Orbit, arg1::Orbit, arg2::jboolean, arg3::UnnormalizedSphericalHarmonicsProvider)
    return J2DifferentialEffect((Orbit, Orbit, jboolean, UnnormalizedSphericalHarmonicsProvider), arg0, arg1, arg2, arg3)
end

function J2DifferentialEffect(arg0::Orbit, arg1::Orbit, arg2::jboolean, arg3::jdouble, arg4::jdouble, arg5::jdouble)
    return J2DifferentialEffect((Orbit, Orbit, jboolean, jdouble, jdouble, jdouble), arg0, arg1, arg2, arg3, arg4, arg5)
end

function J2DifferentialEffect(arg0::SpacecraftState, arg1::AdapterPropagator_DifferentialEffect, arg2::jboolean, arg3::UnnormalizedSphericalHarmonicsProvider)
    return J2DifferentialEffect((SpacecraftState, AdapterPropagator_DifferentialEffect, jboolean, UnnormalizedSphericalHarmonicsProvider), arg0, arg1, arg2, arg3)
end

function J2DifferentialEffect(arg0::SpacecraftState, arg1::AdapterPropagator_DifferentialEffect, arg2::jboolean, arg3::jdouble, arg4::jdouble, arg5::jdouble)
    return J2DifferentialEffect((SpacecraftState, AdapterPropagator_DifferentialEffect, jboolean, jdouble, jdouble, jdouble), arg0, arg1, arg2, arg3, arg4, arg5)
end

function apply(obj::J2DifferentialEffect, arg0::Orbit)
    return jcall(obj, "apply", Orbit, (Orbit,), arg0)
end

function apply(obj::J2DifferentialEffect, arg0::SpacecraftState)
    return jcall(obj, "apply", SpacecraftState, (SpacecraftState,), arg0)
end


"""
    Frames

Stores the actual computed frames and the user input
which can be `:same` or `Rel(10)`.
The `frames` are computed in `javis`.
"""
mutable struct Frames{T}
    frames::Union{Nothing,UnitRange}
    user::T
end

Base.convert(::Type{Frames}, x::Union{Symbol,Rel}) = Frames(nothing, x)
Base.convert(::Type{Frames}, x::UnitRange) = Frames(x, x)

"""
    set_frames!(a::AbstractAction, last_frames::UnitRange)

Compute the frames based on a.frames and `last_frames`.
Save the result in `a.frames.frames` which can be accessed via [`get_frames`](@ref).
"""
function set_frames!(a::AbstractAction, last_frames::UnitRange)
    frames = a.frames.user
    a.frames.frames = get_frames(frames, last_frames)
end

"""
    get_frames(a::AbstractAction)

Return `a.frames.frames` which holds the computed frames for the AbstractAction `a`.
"""
get_frames(a::AbstractAction) = a.frames.frames

"""
    get_frames(frames::Symbol, last_frames::UnitRange)

Get the frames based on a symbol (currently only `same`) and the `last_frames`.
Throw `ArgumentError` if symbol is unknown
"""
function get_frames(frames::Symbol, last_frames::UnitRange)
    if frames === :same
        return last_frames
    else
        throw(ArgumentError("Currently the only symbol supported for defining frames is `:same`"))
    end
end

"""
    get_frames(frames::Rel, last_frames::UnitRange)

Return the frames based on a relative frames [`Rel`](@ref) object and the `last_frames`.
"""
function get_frames(frames::Rel, last_frames::UnitRange)
    start_frame = last(last_frames) + first(frames.rel)
    last_frame = last(last_frames) + last(frames.rel)
    return start_frame:last_frame
end


function findtracks(A,t,p,l);
    #FINDTRACKS creates track associative array.

        # Find docs that have person
        DocIDwPer = getrow(A[:,p])

        # Find docs that have person and location.
        DocIDwPerLoc = getrow(A[DocIDwPer,l])

        # Find docs that have person, location and time.
        DocIDwPerLocTime = getrow(A[DocIDwPerLoc,t])

        # Limit to these documents.
        AA = A[DocIDwPerLocTime,:]

        # Get person sub array.
        Aper = AA[DocIDwPerLocTime,p]
        TrackPer,DocAper = find(Aper')

        # Get location sub array.
        Aloc = AA[DocIDwPerLocTime,l]
        EntAloc,DocAloc = find(Aloc')

        # Get Single location per document- order by actors
        uLocDocs = unique(DocAloc)
        uLocDocIdx = indexin(uLocDocs,DocAloc) # get first index of unique docs (correct for getting highest)
        uDocLocs = EntAloc[uLocDocIdx] # single locations per document
        uLocDocIdxinDocAper = indexin(DocAper,uLocDocs) # locations of unique loc docs in per doc list
        TrackLoc = uDocLocs[uLocDocIdxinDocAper]

        # Get time sub array.
        Atime = AA[DocIDwPerLocTime,t]
        EntAtime,DocAtime = find(Atime')

        # Get Single time per document- order by actors
        uTimeDocs = unique(DocAtime)
        uTimeDocIdx = indexin(uTimeDocs,DocAtime) # get first index of unique docs (correct for getting highest)
        uDocTimes = EntAtime[uTimeDocIdx] # single times per document
        uTimeDocIdxinDocAper = indexin(DocAper,uTimeDocs) # locations of unique time docs in per doc list
        TrackTime = uDocTimes[uTimeDocIdxinDocAper]

        Tr = Assoc(TrackTime,TrackPer,TrackLoc);

    end


@testset "harmonic transformer model" begin
    @testset "YN - yn - 0" begin
        path  = joinpath(_HPM.BASE_DIR,"test/data/matpower/case_xfmr_YNyn0.m")
        data  = _PMs.parse_file(path)
        hdata = _HPM.replicate(data)

        run_hopf_iv(hdata, _PMs.IVRPowerModel, solver)
    end

    @testset "D - yn - 11" begin
        path  = joinpath(_HPM.BASE_DIR,"test/data/matpower/case_xfmr_Dyn11.m")
        data  = _PMs.parse_file(path)
        hdata = _HPM.replicate(data)

        run_hopf_iv(hdata, _PMs.IVRPowerModel, solver)
    end

    @testset "Y - d - 11" begin
        path  = joinpath(_HPM.BASE_DIR,"test/data/matpower/case_xfmr_Yd11.m")
        data  = _PMs.parse_file(path)
        hdata = _HPM.replicate(data)

        run_hopf_iv(hdata, _PMs.IVRPowerModel, solver)
    end

    @testset "YN - d - 11" begin
        path  = joinpath(_HPM.BASE_DIR,"test/data/matpower/case_xfmr_YNd11.m")
        data  = _PMs.parse_file(path)
        hdata = _HPM.replicate(data)

        run_hopf_iv(hdata, _PMs.IVRPowerModel, solver)
    end
end


module Utilities

export ComponentValue, ComponentType, Component, ComponentChild, Enumerated,
        FlagList, RangeValue, Size, Count, Normalized, Color, FontSize, Font,
        Side, AxisRef, XAnchor, YAnchor, Margin, Title, ColorScale,
        ColorScaleSet, Orientation, Easing, Ordering, Transition, NameLength,
        HoverLabel, Calendar, TickMode, Ticks, TickFormatStops, ShowTickPrefix,
        ShowTickSuffix, ExponentFormat, ShowExponent, Dash, default

abstract type ComponentValue end
abstract type ComponentType end
abstract type Component <: ComponentType end
abstract type ComponentChild <: ComponentType end

const default = Dict{String,Component}()

diff(ref::ComponentValue,obj::ComponentValue) = diff(ref.value,obj.value)

function diff(ref,obj)
    ref != obj && return obj
    return
end

struct Enumerated{E,T} <: ComponentValue
    value::Union{String,Bool}

    function Enumerated{E,T}(v) where {E,T}
        Symbol(v) in T || v in T ||
        error("$(E): \"$(v)\" not found in $(tuple([isbits(t) ? t : string(t) for t in T]...))")
        return new{E,T}(v)
    end
end

struct FlagList{S,T} <: ComponentValue
    value::Union{String,Bool}

    function FlagList{S,T}(s::Union{String,Bool}) where {S,T}
        flags = isa(s,String) ? Symbol.(split(s,'+')) : Bool[s]
        length(flags) == 1 && (Symbol(flags[1]) in T || flags[1] in T) && return new(s)
        for flag in flags
            Symbol(flag) in S || flag in S || error("$(flag) not found in $(S)")
        end
        return new(s)
    end
end

struct RangeValue{N,R,T} <: ComponentValue
    value::N

    function RangeValue{N,R,T}(v::N) where {N,R,T}
        errmsg = "$(v) is outside the range $(T[1] ? "[" : "(")$(R[1]),$(R[2])$(T[2] ? "]" : ")")"
        (v < R[1]) && error(errmsg)
        (v > R[2]) && error(errmsg)
        (!T[1] && v == R[1]) && error(errmsg)
        (!T[2] && v == R[2]) && error(errmsg)
        return new(v)
    end
end
RangeValue(range,value) = RangeValue{typeof(value),range,(true,true)}(value)

const Size{S} = RangeValue{S,(0,Inf),(true,false)}

const Count = RangeValue{Int,(1,Inf),(true,false)}

const Normalized = RangeValue{Float64,(0,1),(true,true)}

struct Color <: ComponentValue
    value::String
end

const FontSize = RangeValue{Int,(1,Inf),(true,false)}

struct Font <: ComponentType
    family::Union{String,Vector{String}}
    size::Union{FontSize,Vector{FontSize}}
    color::Union{Color,Vector{Color}}

    function Font(;
        family::Union{String,Vector{String}}="",
        size::Union{Int,Vector{Int}}=6,
        color::Union{String,Vector{String}}="#000")
        return new(family,FontSize.(size),Color.(color))
    end
end

const Side = Enumerated{:Side,(:left,:right,:top,:bottom)}

const AxisRef = Enumerated{:AxisRef,(:container,:paper)}

const XAnchor = Enumerated{:XAnchor,(:auto,:left,:center,:right)}

const YAnchor = Enumerated{:YAnchor,(:auto,:top,:middle,:bottom)}

struct Margin <: ComponentType
    t::Size{Int}
    r::Size{Int}
    b::Size{Int}
    l::Size{Int}
    pad::Size{Int}
    autoexpand::Bool

    function Margin(;
        t::Int = 0,
        r::Int = 0,
        b::Int = 0,
        l::Int = 0,
        pad::Int = 0,
        autoexpand::Bool = true)
        return new(
            Size{Int}(t),
            Size{Int}(r),
            Size{Int}(b),
            Size{Int}(l),
            Size{Int}(pad),
            autoexpand)
    end
end

struct Title <: ComponentType
    text::String
    font::Font
    side::Side
    xref::AxisRef
    yref::AxisRef
    x::Normalized
    y::Normalized
    xanchor::XAnchor
    yanchor::YAnchor
    pad::Margin

    function Title(;
        text::String = "",
        font::Font = Font(),
        side::String = "top",
        xref::String = "container",
        yref::String = "container",
        x = .5,
        y = .5,
        xanchor::String = "auto",
        yanchor::String = "auto",
        pad = Margin())
        return new(
            text,
            font,
            Side(side),
            AxisRef(xref),
            AxisRef(yref),
            Normalized(Float64(x)),
            Normalized(Float64(y)),
            XAnchor(xanchor),
            YAnchor(yanchor),
            pad)
    end
end
Title(t::Title) = t

struct ColorScale <: ComponentType
    colors::Vector{Tuple{Normalized,Color}}
end

colorscalesets = (:Greys,:YlGnBu,:Greens,:YlOrRd,:Bluered,:RdBu,:Reds,:Blues,:Picnic,:Rainbow,:Portland,:Jet,:Hot,:Blackbody,:Earth,:Electric,:Viridis,:Cividis)
const ColorScaleSet = Enumerated{:colorscaleset,colorscalesets}

const Orientation = Enumerated{:Orientation,(:v,:h)}

easings = Symbol.(("linear","quad","cubic","sin","exp","circle","elastic","back","bounce","linear-in","quad-in","cubic-in","sin-in","exp-in","circle-in","elastic-in","back-in","bounce-in","linear-out","quad-out","cubic-out","sin-out","exp-out","circle-out","elastic-out","back-out","bounce-out","linear-in-out","quad-in-out","cubic-in-out","sin-in-out","exp-in-out","circle-in-out","elastic-in-out","back-in-out","bounce-in-out"))
const Easing = Enumerated{:Easing,easings}

const Ordering = Enumerated{:Ordering,Symbol.(("layout first","traces first"))}

struct Transition
    duration::Size{Int}
    easing::Easing
    ordering::Ordering

    function Transition(;
        duration::Int = 500,
        easing::String = "cubic-in-out",
        ordering::String = "layout first")
        return new(
            Size{Int}(duration),
            Easing(easing),
            Ordering(ordering)
        )
    end
end

const NameLength = RangeValue{Int,(-1,Inf),(true,false)}

struct HoverLabel <: ComponentType
    bgcolor::Union{Color,Vector{Color}}
    bordercolor::Union{Color,Vector{Color}}
    font::Font
    namelength::Union{NameLength,Vector{NameLength}}

    function HoverLabel(;
        bgcolor::Union{String,Vector{String}} = "#fff",
        bordercolor::Union{String,Vector{String}} = "#fff",
        font = Font(),
        namelength::Union{Int,Vector{Int}} = -1)
        return new(
            Color.(bgcolor),
            Color.(bordercolor),
            font,
            NameLength.(namelength))
    end
end

const Calendar = Enumerated{:Calendar,(:gregorian,:chinese,:coptic,:discworld,:ethiopian,:hebrew,:islamic,:mayan,:nanakshahi,:nepali,:persian,:jalali,:taiwan,:thai,:ummalqura)}

const TickMode = Enumerated{:TickMode,(:auto,:linear,:array)}

const TickLocation = Enumerated{:TickLocation,(:outside,:inside,Symbol(""))}

const TicksOn = Enumerated{:TicksOn,(:labels,:boundaries)}

const Mirror = Enumerated{:Mirror,(true,:ticks,false,:all,:allticks)}

struct TickFormatStops <: ComponentType
    enable::Bool
    dtickrange::Vector{NamedTuple{(:min,:max),NTuple{2,Float64}}}
    value::String

    function TickFormatStops(;
        enable::Bool = true,
        dtickrange = Vector{NamedTuple{(:min,:max),NTuple{2,Float64}}}(),
        value::String = "")
        return new(enable,dtickrange,value)
    end
end

const ShowTickPrefix = Enumerated{:ShowTickPrefix,(:all,:first,:last,:none)}

const ShowTickSuffix = Enumerated{:ShowTickSuffix,(:all,:first,:last,:none)}

const ExponentFormat = Enumerated{:ExponentFormat,(:none,:e,:E,:power,:SI,:B)}

const ShowExponent = Enumerated{:ShowExponent,(:all,:first,:last,:none)}

struct Ticks <: ComponentChild
    tickmode::TickMode
    nticks::Size{Int}
    tick0::Any
    dtick::Any
    tickvals::Vector{Float64}
    ticktext::Vector{String}
    ticks::TickLocation
    tickson::TicksOn
    mirror::Mirror
    ticklen::Size{Int}
    tickwidth::Size{Int}
    tickcolor::Color
    showticklabels::Bool
    automargin::Bool
    tickfont::Font
    tickangle::Float64
    tickprefix::String
    showtickprefix::ShowTickPrefix
    ticksuffix::String
    showticksuffix::ShowTickSuffix
    showexponent::ShowExponent
    exponentformat::ExponentFormat
    separatethousands::Bool
    tickformat::String
    tickformatstops::TickFormatStops

    function Ticks(;
        tickmode::String = "auto",
        nticks::Int = 0,
        tick0::Any = "",
        dtick::Any = "",
        tickvals::Vector{Float64} = Float64[],
        ticktext::Vector{String} = String[],
        ticks::String = "inside",
        tickson::String = "labels",
        mirror::Union{String,Bool} = false,
        ticklen::Int = 5,
        tickwidth::Int = 1,
        tickcolor::String = "#444",
        showticklabels::Bool = true,
        automargin::Bool = false,
        tickfont::Font = Font(),
        tickangle::Float64 = 0.,
        tickprefix::String = "",
        showtickprefix::String = "all",
        ticksuffix::String = "",
        showticksuffix::String = "all",
        showexponent::String = "all",
        exponentformat::String = "B",
        separatethousands::Bool = false,
        tickformat::String = "",
        tickformatstops::TickFormatStops = TickFormatStops())
        return new(
            TickMode(tickmode),
            Size{Int}(nticks),
            tick0,
            dtick,
            tickvals,
            ticktext,
            TickLocation(ticks),
            TicksOn(tickson),
            Mirror(mirror),
            Size{Int}(ticklen),
            Size{Int}(tickwidth),
            Color(tickcolor),
            showticklabels,
            automargin,
            tickfont,
            tickangle,
            tickprefix,
            ShowTickPrefix(showtickprefix),
            ticksuffix,
            ShowTickSuffix(showticksuffix),
            ShowExponent(showexponent),
            ExponentFormat(exponentformat),
            separatethousands,
            tickformat,
            tickformatstops)
    end
end

const Dash = Enumerated{:Dash,(:solid,:dot,:dash,:longdash,:dashdot,:longdashdot)}

end


#####
##### The LESbrary (so to speak)
#####

# mutable struct CaseMetaData{F, W, C, T}
#     filename :: F
#     with_wind_stress :: W
#     with_coriolis :: C
#     targets :: T
# end

# function CaseMetaData(; filename, with_wind_stress=true, with_coriolis=true, first=nothing, last=nothing, targets)
#    #targets = ...
#    metadata = CaseMetaData(filename, with_wind_stress, with_coriolis, targets)
#    return metadata
# end

# function stride_targets(metadata; stride=1)
#    first = metadata.targets[1]
#    last = metadata.targets[end]
#    new_targets = first:stride:last
#    return CaseMetaData(metadata.filename, metadata.with_wind_stress, metadata.with_coriolis, new_targets)
# end

# https://engaging-web.mit.edu/~alir/lesbrary/4DaySuite/
function FourDaySuite(directory)
   return OrderedDict(
                    "4d_free_convection" => (
                        filename = joinpath(directory, "free_convection/instantaneous_statistics.jld2"),
                        stressed = false,
                        rotating = true,
                           first = 13, # cut out the first 2 hours
                            last = nothing),

                    "4d_strong_wind" => (
                        filename = joinpath(directory, "strong_wind/instantaneous_statistics.jld2"),
                        stressed = true,
                        rotating = true,
                           first = 13,
                            last = nothing),

                    "4d_strong_wind_no_rotation" => (
                        filename = joinpath(directory, "strong_wind_no_rotation/instantaneous_statistics.jld2"),
                        stressed = true,
                        rotating = false,
                           first = 13,
                            last = nothing),

                     # This file is currently corrupted (on tartarus? on engaging?) and may need to be regenerated
                     # "4d_strong_wind_weak_cooling" => (
                     #    filename = joinpath(directory, "strong_wind_weak_cooling/instantaneous_statistics.jld2"),
                     #    stressed = true,
                     #    rotating = true,
                     #       first = 13,
                     #        last = nothing),

                     "4d_weak_wind_strong_cooling" => (
                        filename = joinpath(directory, "weak_wind_strong_cooling/instantaneous_statistics.jld2"),
                        stressed = true,
                        rotating = true,
                           first = 13,
                            last = nothing),
                 )
end

# https://engaging-web.mit.edu/~alir/lesbrary/2DaySuite/
function TwoDaySuite(directory)
      return OrderedDict(
                   "2d_free_convection" => (
                       filename = joinpath(directory, "free_convection/instantaneous_statistics.jld2"),
                       stressed = false,
                       rotating = true,
                          first = 13, # cut out the first 2 hours
                           last = nothing),

                   "2d_strong_wind" => (
                       filename = joinpath(directory, "strong_wind/instantaneous_statistics.jld2"),
                       stressed = true,
                       rotating = true,
                          first = 13,
                           last = nothing),

                   "2d_strong_wind_no_rotation" => (
                       filename = joinpath(directory, "strong_wind_no_rotation/instantaneous_statistics.jld2"),
                       stressed = true,
                       rotating = false,
                          first = 13,
                           last = nothing),

                    "2d_strong_wind_weak_cooling" => (
                       filename = joinpath(directory, "strong_wind_weak_cooling/instantaneous_statistics.jld2"),
                       stressed = true,
                       rotating = true,
                          first = 13,
                           last = nothing),

                    "2d_weak_wind_strong_cooling" => (
                       filename = joinpath(directory, "weak_wind_strong_cooling/instantaneous_statistics.jld2"),
                       stressed = true,
                       rotating = true,
                          first = 13,
                           last = nothing),
                )
end

# https://engaging-web.mit.edu/~alir/lesbrary/6DaySuite/
function SixDaySuite(directory)
   return OrderedDict(
                 "6d_free_convection" => (
                     filename = joinpath(directory, "free_convection/instantaneous_statistics.jld2"),
                     stressed = false,
                     rotating = true,
                        first = 13, # cut out the first 2 hours
                         last = nothing),

                 "6d_strong_wind" => (
                     filename = joinpath(directory, "strong_wind/instantaneous_statistics.jld2"),
                     stressed = true,
                     rotating = true,
                        first = 13,
                         last = nothing),

                 "6d_strong_wind_no_rotation" => (
                     filename = joinpath(directory, "strong_wind_no_rotation/instantaneous_statistics.jld2"),
                     stressed = true,
                     rotating = false,
                        first = 13,
                         last = nothing),

                  "6d_strong_wind_weak_cooling" => (
                     filename = joinpath(directory, "strong_wind_weak_cooling/instantaneous_statistics.jld2"),
                     stressed = true,
                     rotating = true,
                        first = 13,
                         last = nothing),

                  "6d_weak_wind_strong_cooling" => (
                     filename = joinpath(directory, "weak_wind_strong_cooling/instantaneous_statistics.jld2"),
                     stressed = true,
                     rotating = true,
                        first = 13,
                         last = nothing),
              )
end

function GeneralStrat(directory)
   return OrderedDict(
                   "general_strat_4" => (
                       filename = joinpath(directory, "general_strat_4/instantaneous_statistics.jld2"),
                       stressed = false,
                       rotating = true,
                          first = 37, # cut out the first 6 hours
                           last = 288), # 2 days -- mixed layer depth reaches about 75 meters

                   "general_strat_8" => (
                       filename = joinpath(directory, "general_strat_8/instantaneous_statistics.jld2"),
                       stressed = false,
                       rotating = true,
                          first = 13,
                           last = 648), # 4 days -- mixed layer depth reaches about 75 meters

                   "general_strat_16" => (
                       filename = joinpath(directory, "general_strat_16/instantaneous_statistics.jld2"),
                       stressed = false,
                       rotating = true,
                          first = 13,
                           last = nothing),

                    "general_strat_32" => (
                       filename = joinpath(directory, "general_strat_32/instantaneous_statistics.jld2"),
                       stressed = false,
                       rotating = true,
                          first = 13,
                           last = nothing),
                )
end


#######################################
#
# rba file
#
#######################################

"""
Read rba from MAT file
"""
function read_rba_mat(mat_path)
    # Define variables here so that they can be accessed within the scope of try constructs
    modulation_frequency  = NaN
    stimulation_amplitude = NaN
    stimulation_side      = ""
    participant_name      = ""
    carrier_frequency     = NaN

    # Old RBA format
    try
        rba = matread(mat_path)

        modulation_frequency = rba["properties"]["stimulation_properties"]["stimulus_1"]["rounded_modulation_frequency"]
        carrier_frequency    = rba["properties"]["stimulation_properties"]["stimulus_1"]["rounded_carrier_frequency"]

        @info("Imported matching .mat file in old format")
    catch
        # nothing
    end

    # New RBA format
    try
        rba = matopen(mat_path)
        mat = read(rba, "properties")

        modulation_frequency1 = mat["stimulation_properties"]["stimulus_1"]["rounded_modulation_frequency"]
        modulation_frequency2 = mat["stimulation_properties"]["stimulus_2"]["rounded_modulation_frequency"]

        carrier_frequency = mat["stimulation_properties"]["stimulus_1"]["rounded_carrier_frequency"]

        if modulation_frequency1 != modulation_frequency2
            err("Different modulation frequency in each stimulus. Taking stimulus 1")
        end
        modulation_frequency  = modulation_frequency1

        stimulus_amplitude1 = mat["stimulation_properties"]["stimulus_1"]["amplitude"]
        stimulus_amplitude2 = mat["stimulation_properties"]["stimulus_2"]["amplitude"]
        if stimulus_amplitude1 == stimulus_amplitude2
            stimulation_side = "Bilateral"
        else
            stimulation_side = stimulus_amplitude1 > stimulus_amplitude2 ? "Left" : "Right"
        end
        stimulation_amplitude = max(stimulus_amplitude1, stimulus_amplitude2)

        participant_name    = mat["metas"]["subject"]

        @info("Imported matching .mat file in new format")
    catch
        # nothing
    end

    if modulation_frequency == NaN && stimulation_side == NaN && participant_name == NaN
        @warn("Reading of .mat file failed")
    end

    @debug("Frequency: $modulation_frequency Side: $stimulation_side Name: $participant_name Carrier: $carrier_frequency")

    return modulation_frequency, stimulation_side, participant_name, stimulation_amplitude, carrier_frequency
end


module MuseInference

using AbstractDifferentiation
using Base.Iterators: peel, repeated
using ComponentArrays
using CovarianceEstimation
using Dates
using Distributed
using Distributions
using FileIO
using FiniteDifferences
using ForwardDiff
using LinearAlgebra
using Markdown
using NamedTupleTools
using Optim
using ProgressMeter
using Random
using Requires
using Setfield
using Statistics
using UnPack

export MuseProblem, MuseResult, muse, muse!, get_J!, get_H!

include("util.jl")
include("ad.jl")
include("progress.jl")
include("muse.jl")
include("manual.jl")
@init @require Turing="fce5fe82-541a-59a6-adf8-730c64b5f9a0" include("turing.jl")

end


module ImageStitch

using FileIO, Images
#using OEMTIFF #zero array
using TiffImages
using ImageView
using NRRD

using RegisterQD, CoordinateTransformations
using OffsetArrays
using LinearAlgebra
using StatsBase
using JLD2
using Statistics #for median
using Unitful

export checkimgs, merged_axes, merge_images, getOffsetOrigin
# Write your package code here.
#import Unitful: s

## Check image
function checkimgs(img1, img2)
  imgsub0 = similar(img1)
    for i in 1:size(img1, 3)
      if i%2 != 0
        imgsub0[:,:,i] = img2[:,:,i]
      else
        imgsub0[:,:,i] = img1[:,:,i-1]
      end
    end
  return(imgsub0)
end

## Merge image and save data into the hard drive
function merged_axes(img1, imgw)
  axs =
      min(minimum(axes(imgw, 1)), minimum(axes(img1, 1))):max(maximum(axes(imgw, 1)), maximum(axes(img1, 1))),
      min(minimum(axes(imgw, 2)), minimum(axes(img1, 2))):max(maximum(axes(imgw, 2)), maximum(axes(img1, 2))),
      min(minimum(axes(imgw, 3)), minimum(axes(img1, 3))):max(maximum(axes(imgw, 3)), maximum(axes(img1, 3)))
  axs
end

function merge_images(fn, img1, imgw)
  axs = merged_axes(img1, imgw)
  ss = open(fn, "w+")
  for i in axs[3]
    for j in axs[2]
      for k in axs[1]
        if i in axes(imgw,3) && j in axes(imgw,2) && k in axes(imgw,1) && imgw[k,j,i] != 0
          write(ss, UInt16(imgw[k,j,i].*typemax(UInt16)))
        elseif i in axes(img1,3) && j in axes(img1,2) && k in axes(img1,1)
          write(ss, UInt16(img1[k,j,i].*typemax(UInt16)))
        else
          write(ss, zero(UInt16))
        end
      end
    end
  end
  close(ss)
end

## calculate the origin for image stitch
function getOffsetOrigin(img, overlap_coords)
  xOrigin = -(length(axes(img,1))-length(overlap_coords[1])-1)
  yOrigin = -(length(axes(img,2))-length(overlap_coords[2])-1)
  zOrigin = -(length(axes(img,3))-length(overlap_coords[3])-1)
  return(xOrigin, yOrigin, zOrigin)
end

end


# -----------------------------------------------------------------------------
# Implementation of Atomistic MolecularDynamicsResult API
# -----------------------------------------------------------------------------

"""
    NBSResult <: MolecularDynamicsResult

The result generated from running an `NBSimulator`.

**Field descriptions**
- `result::SimulationResult` the standard simulation result from `NBodySimulator`
"""
struct NBSResult <: MolecularDynamicsResult
    result::SimulationResult
end

function get_system(result::NBSResult, t::Integer = 0)
    sr = result.result
    simulation = sr.simulation
    system = simulation.system
    time = sr.solution.t[t > 0 ? t : end]
    positions = get_position(sr, time)
    velocities = get_velocity(sr, time)
    masses = get_masses(system)
    symbols = get_atomic_symbols(system)
    numbers = get_atomic_numbers(system)
    particles = [ElementMassBody(SVector{3}(positions[:, i]), SVector{3}(velocities[:, i]), masses[i], symbols[i], numbers[i]) for i ∈ 1:length(system.bodies)]
    DynamicSystem(FlexibleSystem(particles, simulation.boundary_conditions), time * TIME_UNIT)
end

function get_time_range(result::NBSResult)
    result.result.solution.t * TIME_UNIT
end

function temperature(result::NBSResult, t::Integer = 0)
    time = result.result.solution.t[t > 0 ? t : end]
    NBodySimulator.temperature(result.result, time) * TEMPERATURE_UNIT
end

function reference_temperature(result::NBSResult)
    thermostat = result.result.simulation.thermostat
    thermostat isa NullThermostat ? missing : thermostat.T * TEMPERATURE_UNIT
end

function kinetic_energy(result::NBSResult, t::Integer = 0)
    time = result.result.solution.t[t > 0 ? t : end]
    NBodySimulator.kinetic_energy(result.result, time) * ENERGY_UNIT
end

function potential_energy(result::NBSResult, t::Integer = 0)
    potentials = result.result.simulation.system.potentials
    # https://github.com/SciML/NBodySimulator.jl/issues/44
    if :custom ∈ keys(potentials)
        return InteratomicPotentials.potential_energy(get_system(result, t), potentials[:custom].potential) * ENERGY_UNIT
    end
    time = result.result.solution.t[t > 0 ? t : end]
    NBodySimulator.potential_energy(result.result, time) * ENERGY_UNIT
end

function rdf(result::NBSResult, sample_fraction::Float64 = 1.0)
    @assert 0 < sample_fraction ≤ 1
    sr = result.result
    n = length(sr.simulation.system.bodies)
    pbc = sr.simulation.boundary_conditions
    trange = sr.solution.t[end-floor(Int, length(sr.solution.t) * sample_fraction)+1:end]

    maxbin = 1000
    dr = pbc.L / 2 / maxbin
    hist = zeros(maxbin)
    for t ∈ trange
        cc = get_position(sr, t)
        for i ∈ 1:n
            ri = SVector{3}(cc[:, i])
            for j ∈ i+1:n
                rj = SVector{3}(cc[:, j])
                # TODO: potential for major performance improvements here
                (rij, r, r2) = NBodySimulator.get_interparticle_distance(ri, rj, pbc)
                if r2 < (0.5 * pbc.L)^2
                    bin = ceil(Int, r / dr)
                    if bin > 1 && bin <= maxbin
                        hist[bin] += 2
                    end
                end
            end
        end
    end

    c = 4 / 3 * π * n / pbc.L^3
    gr = zeros(maxbin)
    rs = zeros(maxbin)
    tlen = length(trange)
    for bin ∈ 1:maxbin
        rlower = (bin - 1) * dr
        rupper = rlower + dr
        nideal = c * (rupper^3 - rlower^3)
        gr[bin] = (hist[bin] / (tlen * n)) / nideal
        rs[bin] = rlower + dr / 2
    end

    rs, gr
end


function writeapp(app,filename)

appname = 0;
app.flag = [app.tdep app.wave app.linearproblem app.debugmode app.matvecorder app.GMRESortho app.preconditioner app.precMatrixType app.NLMatrixType app.runmode app.tdfunc app.source app.modelnumber app.flag];
app.problem = [app.hybrid appname app.temporalscheme app.torder app.nstage app.convStabMethod app.diffStabMethod app.rotatingFrame app.viscosityModel app.SGSmodel app.ALE app.AV app.linearsolver app.NLiter app.linearsolveriter app.GMRESrestart app.RBdim app.saveSolFreq app.saveSolOpt app.timestepOffset app.stgNmode app.saveSolBouFreq app.ibs app.dae_steps app.problem];
app.factor = [app.time app.dae_alpha app.dae_beta app.dae_gamma app.dae_epsilon app.factor app.factor];
app.solversparam = [app.NLtol app.linearsolvertol app.matvectol app.NLparam app.solversparam];

ndims = zeros(40,1);
ndims[1] = app.mpiprocs;  # number of processors
ndims[2] = app.nd;
ndims[3] = 0;
ndims[4] = 0;
ndims[5] = 0;
ndims[6] = app.nc;
ndims[7] = app.ncu;
ndims[8] = app.ncq;
ndims[9] = app.ncp;
ndims[10] = app.nco;
ndims[11] = app.nch;
ndims[12] = app.ncx;
ndims[13] = app.nce;
ndims[14] = app.ncw;

nsize = zeros(20,1);
nsize[1] = length(ndims[:]);
nsize[2] = length(app.flag[:]);  # length of flag
nsize[3] = length(app.problem[:]); # length of physics
nsize[4] = length(app.externalparam[:]); # boundary data
nsize[5] = length(app.dt[:]); # number of time steps
nsize[6] = length(app.factor[:]); # length of factor
nsize[7] = length(app.physicsparam[:]); # number of physical parameters
nsize[8] = length(app.solversparam[:]); # number of solver parameters
nsize[9] = length(app.tau[:]); # number of solver parameters
nsize[10] = length(app.stgdata[:]);
nsize[11] = length(app.stgparam[:]);
nsize[12] = length(app.stgib[:]);
nsize[13] = length(app.vindx[:]);
nsize[14] = length(app.dae_dt[:]);

# app.nsize = nsize;
# app.ndims = ndims;

print("Writing app into file...\n");
fileID = open(filename,"w");

write(fileID,Float64(length(nsize[:])));
write(fileID,Float64.(nsize[:]));

if nsize[1]>0
    write(fileID,Float64.(ndims[:]));
end
if nsize[2]>0
    write(fileID,Float64.(app.flag[:]));
end
if nsize[3]>0
    write(fileID,Float64.(app.problem[:]));
end
if nsize[4]>0
    write(fileID,app.externalparam[:]);
end
if nsize[5]>0
    write(fileID,app.dt[:]);
end
if nsize[6]>0
    write(fileID,app.factor[:]);
end
if nsize[7]>0
    write(fileID,app.physicsparam[:]);
end
if nsize[8]>0
    write(fileID,app.solversparam[:]);
end
if nsize[9]>0
    write(fileID,app.tau[:]);
end
if nsize[10]>0
    write(fileID,app.stgdata[:]);
end
if nsize[11]>0
    write(fileID,app.stgparam[:]);
end
if nsize[12]>0
    write(fileID,app.stgib[:]);
end
if nsize[13]>0
    write(fileID,app.vindx[:].-1);
end
if nsize[14]>0
    write(fileID,app.dae_dt[:]);
end

close(fileID);

return app;

end


#TODO fully implement this!!!
# TODO: vectorize this
# TODO: make vectorized version
"""
    Will add a Met to a given idx. See Met struct for details
"""
function set_rxn!(metnet::MetNet, rxni::Int, rxn::Rxn)
    metnet.rxns[rxni] = rxn.id
    metnet.lb[rxni] = rxn.lb
    metnet.ub[rxni] = rxn.ub
    metnet.c[rxni] = rxn.c
    for (met, s) in zip(rxn.mets, rxn.S)
        meti = metindex(metnet, met)
        metnet.S[meti, rxni] = s
    end
    return metnet
end

set_rxn!(metnet::MetNet, rxn::Rxn) =
    set_rxn!(metnet, findempty(metnet, :rxns; check = true), rxn)

set_rxn!(metnet::MetNet, rxni::Int, id::String) = set_rxn!(metnet, rxni, Rxn(id))
set_rxn!(metnet::MetNet, id::String) =
    set_rxn!(metnet, findempty(metnet, :rxns; check = true), id)


using Pkg.Artifacts
const artifact_toml = find_artifacts_toml(@__DIR__)

# Kernels to be downloaded lazily
const SPICE_url_base = "https://naif.jpl.nasa.gov/pub/naif/generic_kernels"
const all_kernels = Dict(
    "meta"    => ["pck/pck00010.tpc",
                  "pck/gm_de431.tpc",
                  "lsk/latest_leapseconds.tls"],  # Required for all other kernels
    "default" => ["spk/planets/de430.bsp"],
    "mars"    => ["spk/satellites/mar097.bsp"],
    "jupiter" => ["spk/satellites/jup310.bsp"],
    "saturn"  => ["spk/satellites/sat427.bsp"],
    "uranus"  => ["spk/satellites/ura111.bsp"],
    "neptune" => ["spk/satellites/nep095.bsp"],
    "pluto"   => ["spk/satellites/plu055.bsp"]
)
for (system, kernels) in all_kernels
    artifact_name = "$(system)_kernels"
    artifact = artifact_hash(artifact_name, artifact_toml)
    if isnothing(artifact)
        # Only download the kernel if it doesn't exist in Artifacts.toml. Otherwise, we want to
        # download it lazily in SpiceUtils.
        download_info = [(join((SPICE_url_base, kernel), "/"), "") for kernel in kernels]
        artifact = create_artifact() do artifact_dir
            for (download_URL, _) in download_info
                @info "Downloading $(basename(download_URL))..."
                download(download_URL, joinpath(artifact_dir, basename(download_URL)))
            end
        end
        if Sys.islinux()
            bind_artifact!(artifact_toml, artifact_name, artifact; download_info, lazy=true, force=true)
        else
            @warn "Could not bind artifact, need to run on Linux so the hashes are correct!"
        end
    end
end


import Base.InteractiveUtils
@inline uses_affinemap(::Type{T}) where {T} = supports_affine(T) || supports_affineview(T)

@inline supports_eager(::Type) = true
@inline supports_affine(::Type{<:AffineOperation}) = true
@inline supports_affine(::Type) = false
@inline supports_affineview(::Type{T}) where {T} = supports_affine(T)
@inline supports_permute(::Type) = false
@inline supports_view(::Type) = false
@inline supports_stepview(::Type) = false
@inline supports_lazy(::Type{T}) where {T} = supports_affine(T) || supports_affineview(T) || supports_stepview(T) || supports_view(T) || supports_permute(T)

@inline supports_eager(A)      = supports_eager(typeof(A))
@inline supports_affine(A)     = supports_affine(typeof(A))
@inline supports_affineview(A) = supports_affineview(typeof(A))
@inline supports_permute(A)    = supports_permute(typeof(A))
@inline supports_view(A)       = supports_view(typeof(A))
@inline supports_stepview(A)   = supports_stepview(typeof(A))
@inline supports_lazy(A)       = supports_lazy(typeof(A))

# --------------------------------------------------------------------

"""
    prepareaffine(img)

Make sure `img` is either a `InvWarpedView` or a `SubArray` of
one. If that is already the case, `img` will be returned as is.
Otherwise `invwarpedview` will be called using a `Flat()`
extrapolation scheme.

Doing this will tell subsequent operations that they should also
participate as affine operations (i.e. use `AffineMap` if they
can).
"""
prepareaffine(img) = invwarpedview(img, toaffinemap(NoOp(), img), Flat())
prepareaffine(img::AbstractExtrapolation) = invwarpedview(img, toaffinemap(NoOp(), img))
@inline prepareaffine(img::SubArray{T,N,<:InvWarpedView}) where {T,N} = img
@inline prepareaffine(img::InvWarpedView) = img
prepareaffine(imgs::Tuple) = map(prepareaffine, imgs)

# currently unused
@inline preparelazy(img) = img

# --------------------------------------------------------------------
# Operation and AffineOperation fallbacks

# discard param if operations doesn't need it
@inline randparam(op::Operation, img) = nothing
@inline toaffinemap(op::Operation, img, param) = toaffinemap(op, img)

for FUN in (:applyeager, :applylazy, :applypermute,
            :applyaffine, :applyaffineview,
            :applyaffine_common, :applyaffineview_common,
            :applyview, :applystepview)
    @eval begin
        function ($FUN)(op::Operation, imgs::Tuple)
            param = randparam(op, imgs)
            map(img -> ($FUN)(op, img, param), imgs)
        end
        @inline function ($FUN)(op::Operation, img::AbstractArray)
            ($FUN)(op, img, randparam(op, img))
        end
    end
end

function applyeager(op::Operation, img::AbstractArray, param)
    maybe_copy(applylazy(op, img, param))
end

function applyaffineview(op::Operation, img::AbstractArray, param)
    wv = applyaffine(op, img, param)
    direct_view(wv, indices(wv))
end

function applyaffine(op::AffineOperation, img::AbstractArray, param)
    invwarpedview(img, toaffinemap(op, img, param))
end

# Allow affine operations to omit specifying a custom
# "applylazy". On the other hand this also makes sure that a
# custom implementation of "applylazy" is preferred over
# "applylazy_fallback" which by default just calls "applyaffine".
function applylazy(op::AffineOperation, img::AbstractArray, param)
    _applylazy(op, img, param)
end

# The purpose of having a separate "_applylazy" is to not
# force "applylazy" implementations to specify the type of "img",
# when typeof(img) <: AbstractArray.
function _applylazy(op::AffineOperation, img::InvWarpedView, param)
    applyaffine(op, img, param)
end

function _applylazy(op::AffineOperation, img::SubArray{T,N,<:InvWarpedView}, param) where {T,N}
    applyaffine(op, img, param)
end

function _applylazy(op::AffineOperation, img, param)
    applylazy_fallback(op, img, param)
end

# Defining "applylazy_fallback" instead of "applylazy" will
# make sure that the custom implementation is only used if
# "img" is not already an "InvWarpedView", in which case
# "applyaffine" would be called instead of "applylazy_fallback".
function applylazy_fallback(op::AffineOperation, img, param)
    applyaffine(op, prepareaffine(img), param)
end

# --------------------------------------------------------------------
# Since an "AffineMap" is strongly typed, the resulting type of
# "toaffinemap" can be very different from one operation to the
# next. To allow for type-stable stochastic branches (i.e. "Either")
# we need a way to force a common "AffineMap" type using only
# "SArray" internally (i.e. no "RotMatrix" or other special types).

@generated function toaffinemap_common(op::AffineOperation, img::AbstractArray{T,N}, param) where {T,N}
    quote
        tfm = toaffinemap(op, img, param)
        AffineMap(SMatrix(tfm.linear), SVector(tfm.translation))::AffineMap{SArray{Tuple{$N,$N},Float64,$N,$(N*N)},SVector{$N,Float64}}
    end
end

function applyaffine_common(op::AffineOperation, img::AbstractArray, param)
    invwarpedview(img, toaffinemap_common(op, img, param))
end

function applyaffineview_common(op::AffineOperation, img::AbstractArray, param)
    wv = applyaffine_common(op, img, param)
    direct_view(wv, indices(wv))
end

# We trust that non-affine operations use SArray-only AffineMap.
function applyaffineview_common(op::Operation, img::AbstractArray, param)
    applyaffineview(op, img, param)
end

# --------------------------------------------------------------------
# Functions to unroll sequences of Operation. These are called by
# the pipeline logic to prefer specific behaviour.

# "applyaffine" or "applyaffineview" is accepted
@inline @generated function unroll_applyaffine(op::Operation, img)
    if supports_affine(op)
        :(applyaffine(op, img))
    elseif supports_affineview(op)
        :(applyaffineview(op, img))
    else
        :(throw(MethodError(unroll_applyaffine, (op, img))))
    end
end

@inline unroll_applylazy(op::Operation, img) = applylazy(op, img)

# Recursively unroll the sequentially applied operations.
# This will result in a function that is equivalent to hardcoded
# chaining of operations.
for KIND in (:affine, :lazy)
    FUN = Symbol(:unroll_apply, KIND)
    PRE = Symbol(:prepare, KIND)
    @eval begin
        function ($FUN)(operations::NTuple{N,Operation}, img) where N
            ($FUN)(first(operations), Base.tail(operations), ($PRE)(img))
        end

        @inline function ($FUN)(head::Operation, tail::Tuple, img)
            ($FUN)(first(tail), Base.tail(tail), ($FUN)(head, img))
        end

        @inline function ($FUN)(head::Operation, tail::Tuple{}, img)
            ($FUN)(head, img)
        end
    end
end


module Visualization

export
    Widget, # from InteractBase
    PushRecoveryVisualizer

using MeshCat
using MeshCatMechanisms
using CoordinateTransformations
using NamedColors

import Observables
import InteractBase
import WebIO

using InteractBase: button, slider, vbox, pad, px, Widget, style, container
using Observables: Observable, observe
using ..QPWalkingControl: PushApplier, icp
using RigidBodyDynamics: FreeVector3D, MechanismState
using MeshCat: Visualizer
using GeometryTypes: HyperSphere, Point

function InteractBase.Widget(controller::PushApplier; max_force::Number, max_Δt::Number)
    pushbutton = button("Apply push")
    angleslider = slider(range(Float64(-π), stop=Float64(π), length=51), label="θ")
    magnitudeslider = slider(range(0.0, stop=max_force, length=51), label="mag")
    timeslider = slider(range(0, stop=max_Δt, length=51), label="Δt")

    let controller = controller
        Observables.on(pushbutton) do _
            controller.new_push = true
        end
        Observables.onany(angleslider, magnitudeslider) do θ, magnitude
            frame = controller.jacobian.frame
            s, c = sincos(θ)
            controller.force = FreeVector3D(frame, magnitude * c, magnitude * s, 0.0)
        end
        Observables.on(timeslider) do Δt
            controller.Δt = Δt
        end
    end

    angleslider[] = 0
    magnitudeslider[] = max_force / 2
    timeslider[] = max_Δt

    style(container(
        pad(10px, style(vbox(pushbutton, angleslider, magnitudeslider, timeslider), :fontFamily => "sans-serif", :fontSize => "10pt")),
    ), :width => "350px")
end

struct PushRecoveryVisualizer
    mvis::MechanismVisualizer
    icpvis::Visualizer
end

function PushRecoveryVisualizer(mvis::MechanismVisualizer)
    state = MeshCatMechanisms.state(mvis)
    vis = MeshCatMechanisms.visualizer(mvis)
    icpvis = vis[:icp]
    setobject!(icpvis, HyperSphere(Point(0., 0., 0.), 0.015), MeshPhongMaterial(color=colorant"greyish blue"))
    PushRecoveryVisualizer(mvis, icpvis)
end

function Base.copyto!(vis::PushRecoveryVisualizer, state::Union{MechanismState, AbstractVector})
    mvis = vis.mvis
    copyto!(mvis, state)
    settransform!(vis.icpvis, Translation(icp(MeshCatMechanisms.state(mvis)).v))
end

Base.wait(vis::PushRecoveryVisualizer) = Base.wait(vis.mvis)

MeshCatMechanisms.visualizer(vis::PushRecoveryVisualizer) = MeshCatMechanisms.visualizer(vis.mvis)

end # module


import Base: copy, copyto!, ==
import Random: rand, seed!
import RandomNumbers: gen_seed, seed_type

@inline rand(r::MT19937, ::Type{UInt32}) = mt_get(r)

seed!(r::MT19937, seed::Integer) = mt_set!(r, seed % UInt32)
seed!(r::MT19937, seed::NTuple{N, UInt32}=gen_seed(UInt32, N)) = mt_set!(r, seed)

seed_type(::Type{MT19937}) = NTuple{N, UInt32}

function copyto!(dest::MT19937, src::MT19937)
    copyto!(dest.mt, src.mt)
    dest.mti = src.mti
    dest
end

copy(src::MT19937) = MT19937(copy(src.mt), src.mti)

==(r1::MT19937, r2::MT19937) = r1.mt == r2.mt && r1.mti == r2.mti


using Test
using BenchmarkTools
using ComplexMixtures, PDBTools
const CM = ComplexMixtures

@testset "Allocations" begin

  dir="./data/NAMD"
  atoms = readPDB("$dir/structure.pdb")

  options = Options(lastframe=1,nthreads=1,silent=true,seed=321,StableRNG=true)
  t_options = @benchmark Options(lastframe=1,seed=321,StableRNG=true,nthreads=1,silent=true) samples=1 evals=1
  @test t_options.allocs == 0

  protein = Selection(select(atoms,"protein"),nmols=1)
  t_selection1 = @benchmark Selection(select($atoms,"protein"),nmols=1) samples=1 evals=1
  @test t_selection1.allocs == 62082

  tmao = Selection(select(atoms,"resname TMAO"),natomspermol=14)
  t_selection2 = @benchmark Selection(select($atoms,"resname TMAO"),natomspermol=14) samples=1 evals=1
  @test t_selection2.allocs == 186144

  trajfile = "$dir/trajectory.dcd" # because of the interpolation of @benchmark
  traj = Trajectory(trajfile,protein,tmao)
  t_trajectory = @benchmark Trajectory($trajfile,$protein,$tmao) samples=1 evals=1
  @test t_trajectory.allocs == 839

  samples = CM.Samples(md=(traj.solvent.nmols-1)/2,random=options.n_random_samples)
  t_samples = @benchmark CM.Samples(md=($traj.solvent.nmols-1)/2,random=$options.n_random_samples) samples=1 evals=1
  @test t_samples.allocs == 0

  R = Result(traj,options)
  t_result = @benchmark Result($traj,$options) samples=1 evals=1
  @test t_result.allocs == 66

  framedata = CM.FrameData(traj,R)
  t_framedata = @benchmark CM.FrameData($traj,$R) samples=1 evals=1
  @test t_framedata.allocs == 206

  CM.nextframe!(traj)
  t_nextframe = @benchmark CM.nextframe!($traj) samples=1 evals=1
  @test t_nextframe.allocs == 34

  RNG = CM.init_random(options)
  t_RNG = @benchmark CM.init_random($options) samples=1 evals=1
  @test t_RNG.allocs == 2

  CM.mddf_frame!(1,framedata,options,RNG,R)
  t_mddf_frame = @benchmark CM.mddf_frame!(1,$framedata,$options,$RNG,$R) samples=1 evals=1
  @test t_mddf_frame.allocs == 0

end


"The bedGraph record."
struct Record{T<:Real}
    chrom::String
    first::Int
    last::Int
    value::T

    Record{T}(chrom::String, first::Int, last::Int, value::T) where {T} = new{T}(chrom, first, last, value)
    Record{T}(chrom, first, last, value) where {T} = new{T}(_string(chrom), _int(first), _int(last), _real(value))
end

function Record(chrom, first, last, value)
    v = _real(value)
    return Record{typeof(v)}(chrom, first, last, v)
end

function Record(data) :: Record
    return data
end

function Record{T}(data) :: Record{T} where T
    return data
end

function Base.convert(::Type{NamedTuple{names}}, record::Record) where {names}
    return NamedTuple{names}((
        record.chrom,
        record.first,
        record.last,
        record.value,
    ))
end

Base.convert(::Type{NamedTuple}, record::Record) = convert(NamedTuple{(:chrom, :first, :last, :value)}, record)

function Base.convert(::Type{Record}, nt::NamedTuple{(:chrom, :first, :last, :value)})
    return Record(nt.chrom, nt.first, nt.last, nt.value)
end

function Base.convert(::Type{Record}, str::AbstractString)
    data = _split_line(str)
    return Record(data[1], data[2], data[3], data[4])
end

function Base.convert(::Type{Record{T}}, str::AbstractString) where T
    data = _split_line(str)
    return Record{T}(data[1], data[2], data[3], data[4])
end

function Base.convert(::Type{Vector{String}}, record::Record)
    return String[record.chrom, string(record.first), string(record.last), string(record.value)]
end

function Base.convert(::Type{String}, record::Record, delim::String = " ")

    strs = [
        record.chrom,
        string(record.first),
        string(record.last),
        string(record.value)
    ]

    return join(strs, delim)
end

"Access [`Record`](@ref)'s chrom field."
function chrom(record::Record)
    return record.chrom
end

"Access [`Record`](@ref)'s left position."
function Base.first(record::Record)
    return record.first
end

"Access [`Record`](@ref)'s last position."
function Base.last(record::Record)
    return record.last
end

"Access [`Record`](@ref)'s value."
function value(record::Record)
    return record.value
end

function Base.:(==)(a::Record, b::Record)
    return a.chrom == b.chrom &&
           a.first == b.first &&
           a.last  == b.last &&
           a.value == b.value
end

function Base.isless(a::Record, b::Record, chrom_isless::Function=isless)

    if chrom_isless(a.chrom, b.chrom)
        return true
    end

    if a.first < b.first
        return true
    end

    if a.last < b.last
        return true
    end

    if a.value < b.value
        return true
    end

    return false
end

## Internal helper functions.
function _split_line(line::AbstractString)
    return split(line, r"\s", limit=5, keepempty=false) #Note: may return 5 elements.
end

function _int(pos)
    return pos
end

function _string(str)
    return str
end

function _real(x)
    return x
end

function _int(str::AbstractString)
    return parse(Float64, str) #Note: conversion to Int is completed when the Record is constructed.
end

function _real(str::AbstractString)
    return parse(Float64, str)
end


# This file is part of Kpax3. License is MIT.

abstract type State end

mutable struct AminoAcidState <: State
  R::Vector{Int}
  C::Matrix{UInt8}

  emptycluster::BitArray{1}
  cl::Vector{Int}
  k::Int

  v::Vector{Int}
  n1s::Matrix{Float64}
  unit::Vector{Vector{Int}}

  logpR::Float64
  logpC::Vector{Float64}
  loglik::Float64
  logpp::Float64
end

function AminoAcidState(data::Matrix{UInt8},
                        partition::Vector{Int},
                        priorR::PriorRowPartition,
                        priorC::AminoAcidPriorCol,
                        settings::KSettings)
  (m, n) = size(data)

  R = normalizepartition(partition, n)
  k = maximum(R)

  maxclust = min(n, max(k, settings.maxclust))

  C = zeros(UInt8, maxclust, m)

  emptycluster = trues(maxclust)
  cl = zeros(Int, maxclust)

  v = zeros(Int, maxclust)
  n1s = zeros(Float64, maxclust, m)
  unit = Vector{Int}[zeros(Int, settings.maxunit) for g in 1:maxclust]

  g = 0
  for a in 1:n
    g = R[a]

    if emptycluster[g]
      emptycluster[g] = false
    end

    if v[g] == length(unit[g])
      resize!(unit[g], min(n, v[g] + settings.maxunit))
    end

    v[g] += 1
    unit[g][v[g]] = a

    for b in 1:m
      n1s[g, b] += float(data[b, a])
    end
  end

  # we do it here and not in the previous loop because we want them sorted
  # it is not really necessary but they will be sorted during the simulation
  # anyway (we will loop on emptycluster from now on)
  i = 1
  for a in 1:length(emptycluster)
    if !emptycluster[a]
      cl[i] = a
      i += 1
    end
  end

  logpR = logdPriorRow(n, k, v, priorR)

  logpC = zeros(Float64, 2)
  computelocalmode!(v, n1s, C, cl, k, logpC, priorC)

  loglik = loglikelihood(C, cl, k, v, n1s, priorC)

  logpp = logpR + logpC[1] + loglik

  AminoAcidState(R, C, emptycluster, cl, k, v, n1s, unit, logpR, logpC, loglik,
                 logpp)
end

function copystate(x::AminoAcidState)
  R = copy(x.R)
  C = copy(x.C)

  emptycluster = copy(x.emptycluster)
  cl = copy(x.cl)
  k = x.k
  v = copy(x.v)
  n1s = copy(x.n1s)

  unit = Array{Vector{Int}}(undef, length(x.unit))
  for l in 1:length(x.unit)
    unit[l] = copy(x.unit[l])
  end

  logpR = x.logpR
  logpC = copy(x.logpC)
  loglik = x.loglik
  logpp = x.logpp

  AminoAcidState(R, C, emptycluster, cl, k, v, n1s, unit, logpR, logpC, loglik,
                 logpp)
end

function copystate!(dest::AminoAcidState,
                    src::AminoAcidState)
  resizestate!(dest, src.cl[src.k])

  if length(dest.R) == length(src.R)
    copyto!(dest.R, src.R)
  else
    dest.R = copy(src.R)
  end

  if size(dest.C, 2) != size(src.C, 2)
    dest.C = zeros(UInt8, size(dest.C, 1), size(src.C, 2))
    dest.n1s = zeros(Float64, size(dest.n1s, 1), size(src.n1s, 2))
  end

  fill!(dest.emptycluster, true)

  g = 0
  for l in 1:src.k
    g = src.cl[l]

    dest.C[g, 1] = src.C[g, 1]
    dest.emptycluster[g] = false
    dest.cl[l] = src.cl[l]
    dest.v[g] = src.v[g]
    dest.n1s[g, 1] = src.n1s[g, 1]

    if length(dest.unit[g]) < src.v[g]
      resize!(dest.unit[g], src.v[g])
    end

    copyto!(dest.unit[g], 1, src.unit[g], 1, src.v[g])
  end

  for b in 2:size(dest.C, 2)
    for l in 1:src.k
      g = src.cl[l]

      dest.C[g, b] = src.C[g, b]
      dest.n1s[g, b] = src.n1s[g, b]
    end
  end

  dest.k = src.k

  dest.logpR = src.logpR

  dest.logpC[1] = src.logpC[1]
  dest.logpC[2] = src.logpC[2]

  dest.loglik = src.loglik

  dest.logpp = src.logpp

  nothing
end

function initializestate(data::Matrix{UInt8},
                         D::Matrix{Float64},
                         kset::UnitRange{Int},
                         priorR::PriorRowPartition,
                         priorC::PriorColPartition,
                         settings::KSettings)
  n = size(data, 2)

  R = ones(Int, n)

  s = AminoAcidState(data, R, priorR, priorC, settings)

  t1 = copystate(s)
  t2 = copystate(s)

  # TODO: remove the hack once kmedoids is fixed

  if settings.verbose
    Printf.@printf("Log-posterior (plus a constant) for one cluster: %.4f.\n",
                   s.logpp)
    Printf.@printf("Now scanning %d to %d clusters.\n", kset[1], kset[end])
  end

  niter = 0
  for k in kset
    if settings.verbose && (k % 10 == 0)
      Printf.@printf("Total number of clusters = %d.\n", k)
    end

    fill!(R, 0)
    try
      copyto!(R, Clustering.kmedoids(D, k).assignments)
    catch
      StatsBase.sample!(1:k, R, replace=true)
      R[StatsBase.sample(1:n, k, replace=false)] = collect(1:k)
    end

    updatestate!(t1, data, R, priorR, priorC, settings)

    niter = 0
    while niter < 10
      try
        copyto!(R, Clustering.kmedoids(D, k).assignments)
      catch
        StatsBase.sample!(1:k, R, replace=true)
        R[StatsBase.sample(1:n, k, replace=false)] = collect(1:k)
      end

      updatestate!(t2, data, R, priorR, priorC, settings)

      if t2.logpp > t1.logpp
        copystate!(t1, t2)
      end

      niter += 1
    end

    if t1.logpp > s.logpp
      copystate!(s, t1)

      if settings.verbose
        Printf.@printf("Found a better solution! ")
        Printf.@printf("Log-posterior (plus a constant) for %d clusters: %.4f.\n",
                       k, s.logpp)
      end
    end
  end

  s
end

function updatestate!(state::AminoAcidState,
                      data::Matrix{UInt8},
                      R::Vector{Int},
                      priorR::PriorRowPartition,
                      priorC::AminoAcidPriorCol,
                      settings::KSettings)
  (m, n) = size(data)

  copyto!(state.R, normalizepartition(R, n))
  k = maximum(state.R)

  resizestate!(state, k, settings)

  state.k = k

  fill!(state.emptycluster, true)
  fill!(state.v, 0)
  fill!(state.n1s, 0.0)

  g = 0
  for a in 1:n
    g = state.R[a]

    if state.emptycluster[g]
      state.emptycluster[g] = false
    end

    if state.v[g] == length(state.unit[g])
      resize!(state.unit[g], min(n, state.v[g] + settings.maxunit))
    end

    state.v[g] += 1
    state.unit[g][state.v[g]] = a

    for b in 1:m
      state.n1s[g, b] += float(data[b, a])
    end
  end

  i = 1
  for a in 1:length(state.emptycluster)
    if !state.emptycluster[a]
      state.cl[i] = a
      i += 1
    end
  end

  state.logpR = logdPriorRow(n, state.k, state.v, priorR)

  computelocalmode!(state.v, state.n1s, state.C, state.cl, state.k,
                    state.logpC, priorC)

  state.loglik = loglikelihood(state.C, state.cl, state.k, state.v, state.n1s,
                               priorC)

  state.logpp = state.logpR + state.logpC[1] + state.loglik

  nothing
end

function resizestate!(state::AminoAcidState,
                      k::Int,
                      settings::KSettings)
  oldlen = size(state.C, 1)

  if oldlen < k
    m = size(state.C, 2)
    n = length(state.R)

    # we don't want to allocate new resources too often, so allocate double the
    # previous size. this should guarantee a logarithmic number of allocations
    newlen = min(n, max(k, 2 * oldlen))

    resize!(state.emptycluster, newlen)
    resize!(state.cl, newlen)
    resize!(state.v, newlen)
    resize!(state.unit, newlen)

    for l in (oldlen + 1):newlen
      state.emptycluster[l] = true
      state.cl[l] = 0
      state.v[l] = 0
      state.unit[l] = zeros(Int, settings.maxunit)
    end

    C = zeros(UInt8, newlen, m)
    n1s = zeros(Float64, newlen, m)

    for b in 1:m
      for l in 1:state.k
        g = state.cl[l]
        C[g, b] = state.C[g, b]
        n1s[g, b] = state.n1s[g, b]
      end
    end

    state.C = C
    state.n1s = n1s
  end

  nothing
end

function resizestate!(state::AminoAcidState,
                      k::Int)
  oldlen = size(state.C, 1)

  if oldlen < k
    m = size(state.C, 2)
    n = length(state.R)

    # we don't want to allocate new resources too often, so allocate double the
    # previous size. this should guarantee a logarithmic number of allocations
    newlen = min(n, max(k, 2 * oldlen))

    resize!(state.emptycluster, newlen)
    resize!(state.cl, newlen)
    resize!(state.v, newlen)
    resize!(state.unit, newlen)

    for l in (oldlen + 1):newlen
      state.emptycluster[l] = true
      state.cl[l] = 0
      state.v[l] = 0
      state.unit[l] = zeros(Int, 1)
    end

    C = zeros(UInt8, newlen, m)
    n1s = zeros(Float64, newlen, m)

    for b in 1:m
      for l in 1:state.k
        g = state.cl[l]
        C[g, b] = state.C[g, b]
        n1s[g, b] = state.n1s[g, b]
      end
    end

    state.C = C
    state.n1s = n1s
  end

  nothing
end

function optimumstate(x::AminoAcidData,
                      partition::String,
                      settings::KSettings)
  R = normalizepartition(partition, x.id)

  priorR = EwensPitman(settings.α, settings.θ)
  priorC = AminoAcidPriorCol(x.data, settings.γ, settings.r)

  AminoAcidState(x.data, R, priorR, priorC, settings)
end


function partitionnementBiObjectif(instance)
    V, A, P = AEPbiFiltrage(instance)
    graphiqueBis = plt.figure()
    plotPointsNonDomines(V, A, graphiqueBis)
    choix = choixBinaire("\n --> Souhaitez-vous calculer les points non-dominés supportés (o/n) ? ")
    if choix == "o"
        calculPointsSupportes(V, A, graphiqueBis)
    end
    solution = choixSolution(V, A)
    close(graphiqueBis)
    backtracking(instance, V, A, P, solution)
end

function backtracking(instance, V, A, P, solution)
    operationMinIJ = Array{Int64,2}(undef, length(instance.solution.ordreVisiteCamion), length(instance.solution.ordreVisiteCamion))
    for i in 1:length(instance.solution.ordreVisiteCamion)
        for j in 1:length(instance.solution.ordreVisiteCamion)
            operationMinIJ[i,j] = -100
        end
    end

    temp = []
    pred = P[end][solution]
    push!(temp, length(instance.solution.ordreVisiteCamion))
    push!(temp, pred[1])
    operationMinIJ[temp[end], temp[end-1]] = pred[3]
    while pred[1] != 1
        pred = P[pred[1]][pred[2]]
        push!(temp, pred[1])
        operationMinIJ[temp[end], temp[end-1]] = pred[3]
    end
    temp = reverse(temp)

    ordreVisiteCamion = []
    operationsDrone = []
    push!(ordreVisiteCamion, 1)
    for i in 1:length(temp)-1
        operation = operationMinIJ[temp[i],temp[i+1]]
        if operation == -1 # Trajet direct
            noeudCombine = instance.solution.ordreVisiteCamion[temp[i+1]]
            push!(ordreVisiteCamion, noeudCombine)
        else # Opération drone
            noeudDrone = instance.solution.ordreVisiteCamion[operation]
            noeudDepart = instance.solution.ordreVisiteCamion[temp[i]]
            noeudFin = instance.solution.ordreVisiteCamion[temp[i+1]]
            push!(operationsDrone, [noeudDepart, noeudDrone, noeudFin])
            for l in temp[i]+1:temp[i+1]-1
                noeudActu = instance.solution.ordreVisiteCamion[l]
                if noeudActu != noeudDrone # Noeud camion
                    push!(ordreVisiteCamion, noeudActu)
                end
            end
            push!(ordreVisiteCamion, noeudFin)
        end
    end

    instance.solution.tempsParcours = V[end][solution]
    instance.solution.tempsAttente = A[end][solution]
    instance.solution.ordreVisiteCamion = ordreVisiteCamion
    instance.solution.operationsDrone = operationsDrone
    instance.codeResolution = 5
end

function choixSolution(V, A) # On choisit la solution à afficher (avec laquelle on effectue le backtracking)
    println("\n Liste des points non-dominés :")
    println(" ------------------------------")
    for i in 1:length(V[end])
        println(" --> $i) ($(V[end][i]),$(A[end][i]))")
        text(V[end][i]+1, A[end][i]+1, i)
    end
    choix = choixListe(length(V[end]))
    return parse(Int64, choix)
end

function plotPointsNonDomines(V, A, graphiqueBis)
    ax = graphiqueBis.gca()
    ax.plot(V[end], A[end], marker="x", linestyle="", color="C0", label="Non-dominated points")
    title(L"#$Y_N =$ "*string(length(V[end])))
    grid(true)
    xlabel(L"$z_1$")
    ylabel(L"$z_2$")
    legend(loc="best")
end

function calculPointsSupportes(V, A, graphiqueBis)
    ax = graphiqueBis.gca()
    listeDesPoints = [tuple(V[end][i],A[end][i]) for i in 1:length(V[end])]
    hull = convex_hull_graham_scan!(listeDesPoints) # On doit enlever les points non-supportés qui sont sur les bords de l'enveloppe convexe

    # On cherche les deux points extrêmes
    maxY = -Inf32
    maxIndice = -1
    minY = Inf32
    minIndice = -1
    for i in 1:length(hull)
        if hull[i][2] < minY
            minY = hull[i][2]
            minIndice = i
        end
        if hull[i][2] > maxY
            maxY = hull[i][2]
            maxIndice = i
        end
    end

    # On cherche les points qui sont au-dessus (et donc à enlever)
    vecteurMinMax = [hull[minIndice][1] - hull[maxIndice][1], hull[minIndice][2] - hull[maxIndice][2]]
    hullNew = [hull[minIndice]]
    for i in 1:length(hull)
        if i != minIndice && i != maxIndice
            vecteur = [hull[i][1] - hull[maxIndice][1], hull[i][2] - hull[maxIndice][2]]
            produitVectoriel = (vecteurMinMax[1]*vecteur[2])-(vecteurMinMax[2]*vecteur[1])
            if produitVectoriel < 0 # ok : point supporté
                push!(hullNew, hull[i])
            end
        end
    end
    push!(hullNew, hull[maxIndice])
    hull = hullNew

    YSN = length(hull)
    ax.plot(hull[1][1], hull[1][2], marker="x", color="red", linestyle="", label="Supported non-dominated points")
    for i in 1:length(hull)
        #println(cvh_graham[i])
        ax.plot(hull[i][1], hull[i][2], marker="x", color="red", linestyle="")
        if i != length(hull)
            ax.plot([hull[i][1],hull[i+1][1]], [hull[i][2],hull[i+1][2]], color="red", linestyle="--", linewidth=0.7)
        end
    end

    YN = length(V[end])
    YSN = length(hull)
    YNN = YN - YSN
    title(L"#$Y_{N} =$ "*string(YN)*L"; #$Y_{SN} =$ "*string(YSN)*L"; #$Y_{NN} =$ "*string(YNN))
    legend(loc="best")
end

#################### Bi-objectif ####################

function tempsOperationBi(instance, i,j,k)
    solutionTSP = instance.solution.ordreVisiteCamion # Solution TSP : ordre de visite  du camion
    # Debug
    #println(" Calcul du temps de l'opération : (i=$i, j=$j, k=$k)")
    #println(" Solution TSP : R_tsp = $solutionTSP")
    #print(" r_$i = $(solutionTSP[i]);")
    #print(" r_$j = $(solutionTSP[j]);")
    #if k != -1
    #    println(" r_$k = $(solutionTSP[k])")
    #end
    # Calcul distance parcourue par le drone
    if k == -1 # Pas de noeud drone
        distanceDrone = 0
        ri = solutionTSP[i]
        rj = solutionTSP[j]
        distanceCamion = instance.D[ri,rj]
        tempsAttente = 0.0
    else
        ri = solutionTSP[i]
        rj = solutionTSP[j]
        rk = solutionTSP[k]
        distanceDrone = instance.D[ri,rk] + instance.D[rk,rj]
        distanceCamion = 0
        for l in i:(k-2)
            r1 = solutionTSP[l]
            r2 = solutionTSP[l+1]
            distanceCamion += instance.D[r1,r2]
        end
        rkAvant = solutionTSP[k-1]
        rkApres = solutionTSP[k+1]
        distanceCamion += instance.D[rkAvant,rkApres]
        for l in (k+1):(j-1)
            r1 = solutionTSP[l]
            r2 = solutionTSP[l+1]
            distanceCamion += instance.D[r1,r2]
        end
        tempsAttente = abs((distanceCamion/instance.vitesseCamion) - (distanceDrone/instance.vitesseDrone))
    end
    tempsParcours = max(distanceCamion/instance.vitesseCamion, distanceDrone/instance.vitesseDrone)
    return tempsParcours, tempsAttente
    #println(" Distance parcourue par le drone : $distanceDrone")
    #println(" Distance parcourue par le camion : $distanceCamion")
end

function calculMeilleuresOperationsBi(instance)
    solutionTSP = instance.solution.ordreVisiteCamion
    tempsMinIJ = Array{Vector{Float64},2}(undef, length(solutionTSP), length(solutionTSP))
    attenteIJ = Array{Vector{Float64},2}(undef, length(solutionTSP), length(solutionTSP))
    operationMinIJ = Array{Vector{Int64},2}(undef, length(solutionTSP), length(solutionTSP))
    for i in 1:length(solutionTSP)
        for j in 1:length(solutionTSP)
            tempsMinIJ[i,j] = Float64[]
            attenteIJ[i,j] = Float64[]
            operationMinIJ[i,j] = Int64[]
        end
    end
    for i in 1:length(solutionTSP)-1
        for j in i+1:length(solutionTSP)
            if j == i+1 # k = -1
                Tij, Aij = tempsOperationBi(instance, i, j, -1)
                operation = -1
                push!(tempsMinIJ[i, j], Tij)
                push!(attenteIJ[i, j], Aij)
                push!(operationMinIJ[i, j], operation)
            else
                # k > 0
                for k in (i+1):(j-1)
                    Tij, Aij = tempsOperationBi(instance, i, j, k)
                    operation = k
                    push!(tempsMinIJ[i, j], Tij)
                    push!(attenteIJ[i, j], Aij)
                    push!(operationMinIJ[i, j], operation)
                end
            end
        end
    end
    return tempsMinIJ, operationMinIJ, attenteIJ
end

########## Filtrage pendant (iterative dominance filtering) ##########

function AEPbiFiltrage(instance)
    tempsMinIJ, operationMinIJ, attenteIJ = calculMeilleuresOperationsBi(instance)
    tempsMinIJ, attenteIJ, operationMinIJ = filtrageOperations(tempsMinIJ, operationMinIJ, attenteIJ, instance.solution.ordreVisiteCamion) # Filtrage par dominance des opérations T(i,j)

    V = []
    A = []
    P = []
    for i in 1:length(instance.solution.ordreVisiteCamion)
        push!(V, Float64[])
        push!(A, Float64[])
        push!(P, [])
    end

    push!(V[1], 0.0)
    push!(A[1], 0.0)
    push!(P[1], [-1,-1,-1]) # A changer
    for i in 2:length(V)
        #println(" $(i-1) : $(length(V[i-1]))")
        for k in 1:(i-1)
            for y in 1:length(V[k])
                for z in 1:length(tempsMinIJ[k,i])
                    temps = V[k][y] + (tempsMinIJ[k,i])[z]
                    attente = A[k][y] + (attenteIJ[k,i])[z]
                    push!(V[i], temps)
                    push!(A[i], attente)
                    push!(P[i], [k,y, trunc(Int64, (operationMinIJ[k,i])[z]) ]) # On vient de V[k,m] et on a pris l'opération T(k,i,n)
                end
            end
        end
        Vtemp = []
        Atemp = []
        Ptemp = []
        for a in 1:length(V[i]) # On veut tester si le point i est dominé par un point j
            estDomine = false
            for b in 1:length(V[i])
                if b != a
                    if round(V[i][b], digits=5) == round(V[i][a], digits=5) && round(A[i][b], digits=5) == round(A[i][a], digits=5)
                        println(" Egalité V[i]")
                    end
                    if round(V[i][b], digits=5) <= round(V[i][a], digits=5) && round(A[i][b], digits=5) <= round(A[i][a], digits=5)
                        estDomine = true
                        break
                    end
                end
            end
            if estDomine == false
                push!(Vtemp, V[i][a])
                push!(Atemp, A[i][a])
                push!(Ptemp, P[i][a])
            end
        end
        V[i] = Vtemp
        A[i] = Atemp
        P[i] = Ptemp
    end
    #println(" V = $V")
    #println(" A = $A")
    #println(" P = $P")
    return V, A, P
end

########## Filtrage par dominance des T(i,j) ##########

function filtrageOperations(tempsMinIJ, operationMinIJ, attenteIJ, solutionTSP)
    tempsMinIJBis = Array{Vector{Float64},2}(undef, length(solutionTSP), length(solutionTSP))
    attenteIJBis = Array{Vector{Float64},2}(undef, length(solutionTSP), length(solutionTSP))
    operationMinIJBis = Array{Vector{Float64},2}(undef, length(solutionTSP), length(solutionTSP))
    for i in 1:length(solutionTSP)
        for j in 1:length(solutionTSP)
            tempsMinIJBis[i,j] = Float64[]
            attenteIJBis[i,j] = Float64[]
            operationMinIJBis[i,j] = Int64[]
        end
    end

    for i in 1:length(solutionTSP)
        for j in i+1:length(solutionTSP)
            for l1 in 1:length(tempsMinIJ[i,j]) # Pour toutes les opérations dans T(i,j)
                estDomine = false # Test du point (i,j,l1)
                for l2 in 1:length(tempsMinIJ[i,j]) # Pour toutes les opérations dans T(i,j)
                    if l1 != l2
                        if round(tempsMinIJ[i,j][l2], digits=5) == round(tempsMinIJ[i,j][l1], digits=5) && round(attenteIJ[i,j][l2], digits=5) == round(attenteIJ[i,j][l1], digits=5) # Cas d'égalité
                            println(" Egalité operation")
                        end
                        if round(tempsMinIJ[i,j][l2], digits=5) <= round(tempsMinIJ[i,j][l1], digits=5) && round(attenteIJ[i,j][l2], digits=5) <= round(attenteIJ[i,j][l1], digits=5)
                            estDomine = true
                            break
                        end
                    end
                end
                if estDomine == false
                    push!(tempsMinIJBis[i,j], tempsMinIJ[i,j][l1])
                    push!(attenteIJBis[i,j], attenteIJ[i,j][l1])
                    push!(operationMinIJBis[i,j], operationMinIJ[i,j][l1])
                end
            end
        end
    end
    return tempsMinIJBis, attenteIJBis, operationMinIJBis
end

########## Graham ##########

function orient(p,q,r)
    pr_x = r[1] - p[1]
    pr_y = r[2] - p[2]
    rq_x = q[1] - r[1]
    rq_y = q[2] - r[2]
    sign(pr_x * rq_y - pr_y * rq_x)
end

function convex_hull_graham_scan!(points)
    @assert length(points) >= 3 "Too few points"
    sort!(points)  # in lexical order

    function build_convex_hull(points)
        n = length(points)
        hull = points[1:2]
        sizehint!(hull, n)

        for i in 3:n
            while true
                if orient(hull[end - 1], hull[end], points[i]) > 0
                    break
                else
                    pop!(hull)
                    if length(hull) < 2
                        break
                    end
                end
            end
            push!(hull, points[i])
        end

        hull
    end

    upper_hull = build_convex_hull(points)
    pop!(upper_hull)
    lower_hull = build_convex_hull(@view points[end:-1:1])
    pop!(lower_hull)

    [upper_hull; lower_hull]
end

function convex_hull_length(hull)
    hull2 = circshift(hull, -1)
    map((x, y) -> hypot((x .- y)...), hull2, hull) |> sum
end


""" Verify that the data is self-consistent and consistent with the likelihood """
function check_data!(X::AbstractArray{T1,N1},y::AbstractArray{T2,N2},likelihood::Likelihood) where {T1<:Real,T2,N1,N2}
    @assert (size(y,1)==size(X,1)) "There is not the same number of samples in X and y";
    @assert N1 <= 2 "The input matrix X can only be a vector or a matrix"
    if N1 == 1 #Reshape a X vector as a matrix
        X = reshape(X,length(X),1);
    end
    y,nLatent,likelihood = treat_labels!(y,likelihood)
    return X,y,nLatent,likelihood
end

""" Verify that the likelihood and inference are compatible (are implemented) """
function check_implementation(model::Symbol,likelihood::L,inference::I) where {I<:Inference,L<:Likelihood}
    if isa(likelihood,GaussianLikelihood)
        if model == :GP && inference isa Analytic
            return true
        elseif model == :SVGP && inference isa AnalyticVI
            return true
        else
            return false
        end
    elseif likelihood isa StudentTLikelihood
        if inference isa AnalyticVI
            return true
        elseif model == :VGP && inference isa GibbsSampling
            return true
        else
            return false
        end
    elseif likelihood isa LaplaceLikelihood
        if inference isa AnalyticVI
            return true
        else
            return false
        end
    elseif likelihood isa LogisticLikelihood
        if inference isa AnalyticVI
            return true
        elseif model == :VGP && inference isa GibbsSampling
            return true
        else
            return false
        end
    elseif likelihood isa BayesianSVM
        if inference isa AnalyticVI
            return true
        else
            return false
        end
    elseif likelihood isa SoftMaxLikelihood
        if inference isa MCIntegrationVI
            return true
        else
            return false
        end
    elseif likelihood isa LogisticSoftMaxLikelihood
        if inference isa AnalyticVI
            return true
        elseif model == :VGP && inference isa GibbsSampling
            return true
        elseif inference isa MCIntegrationVI
            return true
        else
            return false
        end
    elseif likelihood isa PoissonLikelihood
        if inference isa AnalyticVI
            return true
        else
            return false
        end
    else
        return false
    end
end