"""
`module  JAC.ReducedDensityMatrix`
    ... a submodel of JAC that contains all methods for computing reduced density matrices, natural orbitals, radial density
        distributions and related information for some level(s).
"""
module ReducedDensityMatrix

    using Printf, ..Basics, ..Defaults, ..InteractionStrength, ..ManyElectron, ..Nuclear, ..Radial, ..RadialIntegrals,
                  ..SpinAngular, ..TableStrings

    """
    `struct  ReducedDensityMatrix.Settings  <:  AbstractPropertySettings`
        ... defines a type for providing and parameters for reduced density matrices and other parameters.

        + calc1pRDM                ::Bool             ... True if 1-particle RDM need to be calculated, and false otherwise.
        + calc2pRDM                ::Bool             ... True if 2-particle RDM need to be calculated, and false otherwise.
        + calcCumulant             ::Bool             ... True if the cumulant need to be calculated, and false otherwise.
        + calcNatural              ::Bool             ... True if natural orbitals need to be calculated, and false otherwise.
        + calcDensity              ::Bool             ... True if the radial density need to be calculated, and false otherwise.
        + calcIpq                  ::Bool             ... True if orbital interaction I_pq need to be calculated, and false otherwise.
        + printBefore              ::Bool             ... True if a list of selected levels is printed before the actual computations start.
        + levelSelection           ::LevelSelection   ... Specifies the selected levels, if any.
    """
    struct Settings  <:  AbstractPropertySettings
        calc1pRDM                  ::Bool
        calc2pRDM                  ::Bool
        calcCumulant               ::Bool
        calcNatural                ::Bool
        calcDensity                ::Bool
        calcIpq                    ::Bool
        printBefore                ::Bool
        levelSelection             ::LevelSelection
    end

    """
    `ReducedDensityMatrix.Settings(; calc1pRDM::Bool=true,` calc2pRDM::Bool=false, calcCumulant::Bool=false, calcNatural::Bool=false,
                                     calcDensity::Bool=false, calcIpq::Bool=false, printBefore::Bool=true,
                                     levelSelection::LevelSelection=LevelSelection())
        ... keyword constructor to overwrite selected value of reduced-density matrix computations.
    """
    function Settings(; calc1pRDM::Bool=true, calc2pRDM::Bool=false, calcCumulant::Bool=false, calcNatural::Bool=false,
                        calcDensity::Bool=false, calcIpq::Bool=false, printBefore::Bool=true,
                        levelSelection::LevelSelection=LevelSelection())
        Settings(calc1pRDM, calc2pRDM, calcCumulant, calcNatural, calcDensity, calcIpq, printBefore, levelSelection)
    end

    # `Base.show(io::IO, settings::ReducedDensityMatrix.Settings)`
    #       ... prepares a proper printout of the variable settings::ReducedDensityMatrix.Settings.
    function Base.show(io::IO, settings::ReducedDensityMatrix.Settings)
        println(io, "calc1pRDM:                $(settings.calc1pRDM)  ")
        println(io, "calc2pRDM:                $(settings.calc2pRDM)  ")
        println(io, "calcCumulant:             $(settings.calcCumulant)  ")
        println(io, "calcNatural:              $(settings.calcNatural)  ")
        println(io, "calcDensity:              $(settings.calcDensity)  ")
        println(io, "calcIpq:                  $(settings.calcIpq)  ")
        println(io, "printBefore:              $(settings.printBefore)  ")
        println(io, "levelSelection:           $(settings.levelSelection)  ")
    end

    """
    `struct  ReducedDensityMatrix.Outcome`
        ... defines a type to keep the outcome of a reduced density-matrix and natural-orbital computation.

        + level                     ::Level                 ... Atomic level to which the outcome refers to.
        + reduced1p                 ::Array{Float64,2}      ... 1-particle RDM (ip,iq)
        + reduced2p                 ::Array{Float64,4}      ... 1-particle RDM (ip,iq,ir,is)
        + orbitalInteraction        ::Array{Float64,2}      ... orbital interaction I_pq = (ip,iq)
        + naturalOccupation         ::Array{Float64,1}      ... Occupation numbers of natural orbitals.
        + naturalSubshells          ::Array{Subshell,1}     ... List of natural orbitals (subshells).
        + naturalOrbitals           ::Dict{Subshell, Orbital} . Dictionary of (one-electron) natural orbitals.
        + electronDensity           ::Radial.Density        ... Radial density distribution.
    """
    struct Outcome
        level                       ::Level
        reduced1p                   ::Array{Float64,2}
        reduced2p                   ::Array{Float64,4}
        orbitalInteraction          ::Array{Float64,2}
        naturalOccupation           ::Array{Float64,1}
        naturalSubshells            ::Array{Subshell,1}
        naturalOrbitals             ::Dict{Subshell, Orbital}
        electronDensity             ::Radial.Density
    end

    """
    `ReducedDensityMatrix.Outcome()`  ... constructor for an `empty` instance of Hfs.Outcome for the computation of isotope-shift properties.
    """
    function Outcome()
        Outcome(Level(), zeros(1,1), zeros(1,1,1,1), zeros(1,1), Float64[], Subshell[], Dict{Subshell, Orbital}(), Radial.Density() )
    end

    # `Base.show(io::IO, outcome::ReducedDensityMatrix.Outcome)`  ... prepares a proper printout of the variable outcome::ReducedDensityMatrix.Outcome.
    function Base.show(io::IO, outcome::ReducedDensityMatrix.Outcome)
        println(io, "level:                   $(outcome.level)  ")
        println(io, "reduced1p:               $(outcome.reduced1p)  ")
        println(io, "reduced2p:               $(outcome.reduced2p)  ")
        println(io, "orbitalInteraction:      $(outcome.orbitalInteraction)  ")
        println(io, "naturalOccupation:       $(outcome.naturalOccupation)  ")
        println(io, "naturalSubshells:        $(outcome.naturalSubshells)  ")
        println(io, "naturalOrbitals:         $(outcome.naturalOrbitals)  ")
        println(io, "electronDensity:         $(outcome.electronDensity)  ")
    end

    """
    `ReducedDensityMatrix.computeProperties(outcome::ReducedDensityMatrix.Outcome, nm::Nuclear.Model, grid::Radial.Grid,
                                            settings::ReducedDensityMatrix.Settings)
        ... to compute all properties for a given level; an outcome::ReducedDensityMatrix.Outcome is returned for
            which the properties are now evaluated explicitly.
    """
    function  computeProperties(outcome::ReducedDensityMatrix.Outcome, nm::Nuclear.Model, grid::Radial.Grid,
                                settings::ReducedDensityMatrix.Settings)
        reduced1p = outcome.reduced1p;     naturalOccupation = outcome.naturalOccupation;
        reduced2p = outcome.reduced2p;     naturalSubshells  = outcome.naturalSubshells;     naturalOrbitals = outcome.naturalOrbitals;
        orbitalInteraction = outcome.orbitalInteraction;       electronDensity = outcome.electronDensity;

        if  settings.calc1pRDM      @warn("Option calc1pRDM not yet implemented.")
        end

        if  settings.calc2pRDM      @warn("Option calc2pRDM not yet implemented.")
        end

        if  settings.calcCumulant   println("Option calcCumulant not yet implemented.")
        end

        if  settings.calcNatural    println("Option calcNatural not yet implemented.")
        end

        if  settings.calcDensity    println("Option calcDensity not yet implemented.")
        end

        if  settings.calcIpq        println("Option calcIpq not yet implemented.")
        end

        newOutcome = ReducedDensityMatrix.Outcome( outcome.level, reduced1p, reduced2p, orbitalInteraction,
                                                   naturalOccupation, naturalSubshells, naturalOrbitals, electronDensity )
        return( newOutcome )
    end

    """
    `ReducedDensityMatrix.computeOutcomes(multiplet::Multiplet, nm::Nuclear.Model, grid::Radial.Grid,
                                          settings::ReducedDensityMatrix.Settings; output=true)`
        ... to compute (as selected) the 1-particle and 2-particle reduced density matrices, natural orbitals or other
            requested information for the levels of interest.
    """
    function computeOutcomes(multiplet::Multiplet, nm::Nuclear.Model, grid::Radial.Grid,
                             settings::ReducedDensityMatrix.Settings; output=true)
        println("")
        printstyled("ReducedDensityMatrix.computeOutcomes(): The computation of the reduced density matrices starts now ... \n", color=:light_green)
        printstyled("------------------------------------------------------------------------------------------------------ \n", color=:light_green)
        #
        outcomes = ReducedDensityMatrix.determineOutcomes(multiplet, settings)
        # Display all selected levels before the computations start
        if  settings.printBefore    ReducedDensityMatrix.displayOutcomes(outcomes, settings)    end
        # Calculate all requested properties
        newOutcomes = ReducedDensityMatrix.Outcome[]
        for  outcome in outcomes
            newOutcome = ReducedDensityMatrix.computeProperties(outcome, nm, grid, settings)
            push!( newOutcomes, newOutcome)
        end
        # Print all results to screen
        ReducedDensityMatrix.displayResults(stdout, newOutcomes, settings)
        printSummary, iostream = Defaults.getDefaults("summary flag/stream")
        if  printSummary    ReducedDensityMatrix.displayResults(iostream, newOutcomes, settings)   end
        #
        if    output    return( newOutcomes )
        else            return( nothing )
        end
    end

    """
    `ReducedDensityMatrix.determineOutcomes(multiplet::Multiplet, settings::ReducedDensityMatrix.Settings)`
        ... to determine a list of Outcomes's for the computation of reduced density matrices and other information
            for the given multiplet. It takes into account the particular selections and settings.
            An Array{ReducedDensityMatrix.Outcome,1} is returned. Apart from the level specification, all physical
            properties are set to zero during the initialization process.
    """
    function  determineOutcomes(multiplet::Multiplet, settings::ReducedDensityMatrix.Settings)
        outcomes = ReducedDensityMatrix.Outcome[]
        @show settings.levelSelection
        for  level  in  multiplet.levels
            if  Basics.selectLevel(level, settings.levelSelection)
                push!( outcomes, ReducedDensityMatrix.Outcome(level, zeros(1,1), zeros(1,1,1,1), zeros(1,1),
                                                              Float64[], Subshell[], Dict{Subshell, Orbital}(), Radial.Density()) )
            end
        end
        return( outcomes )
    end

    """
    `ReducedDensityMatrix.displayOutcomes(outcomes::Array{ReducedDensityMatrix.Outcome,1}, settings::ReducedDensityMatrix.Settings)`
        ... to display a list of levels that have been selected for the computations. A small neat table of all
            selected levels and their energies is printed but nothing is returned otherwise. Moreover, the
            selected properties and information is printed as well, though not yet calculated.
    """
    function  displayOutcomes(outcomes::Array{ReducedDensityMatrix.Outcome,1}, settings::ReducedDensityMatrix.Settings)
        nx = 43
        println(" ")
        println("  Selected levels for reduced density matrix & natural orbital calculations:")
        println(" ")
        println("  ", TableStrings.hLine(nx))
        sa = "  ";   sb = "  "
        sa = sa * TableStrings.center(10, "Level"; na=2);                             sb = sb * TableStrings.hBlank(12)
        sa = sa * TableStrings.center(10, "J^P";   na=4);                             sb = sb * TableStrings.hBlank(14)
        sa = sa * TableStrings.center(14, "Energy"; na=4);
        sb = sb * TableStrings.center(14, TableStrings.inUnits("energy"); na=4)
        println(sa);    println(sb);    println("  ", TableStrings.hLine(nx))
        #
        for  outcome in outcomes
            sa  = "  ";    sym = Basics.LevelSymmetry( outcome.level.J, outcome.level.parity)
            sa = sa * TableStrings.center(10, TableStrings.level(outcome.level.index); na=2)
            sa = sa * TableStrings.center(10, string(sym); na=4)
            sa = sa * @sprintf("%.8e", Defaults.convertUnits("energy: from atomic", outcome.level.energy)) * "    "
            println( sa )
        end
        println("  ", TableStrings.hLine(nx))
        println("\n  The following properties and information will be computed:\n")
        if  settings.calc1pRDM      println("  ++ compute 1-particle RDM.")                       end
        if  settings.calc2pRDM      println("  ++ compute 2-particle RDM.")                       end
        if  settings.calcCumulant   println("  ++ compute cumulant.")                             end
        if  settings.calcNatural    println("  ++ compute natural orbitals & occupation.")        end
        if  settings.calcDensity    println("  ++ compute radial density.")                       end
        if  settings.calcIpq        println("  ++ compute orbital interaction I(p,q).")           end
        #
        return( nothing )
    end

    """
    `ReducedDensityMatrix.displayResults(stream::IO, outcomes::Array{ReducedDensityMatrix.Outcome,1},
                                         settings::ReducedDensityMatrix.Settings)`
        ... to display the reduced density matrices, natural orbitals, etc. for the selected levels.
            A neat table is printed but nothing is returned otherwise.
    """
    function  displayResults(stream::IO, outcomes::Array{ReducedDensityMatrix.Outcome,1}, settings::ReducedDensityMatrix.Settings)
        nx = 133
        println(stream, " ")
        println(stream, "  ReducedDensityMatrix parameters and information:  !! Not yet worked out !! ")
        println(stream, " ")
        println(stream, "  ", TableStrings.hLine(nx))
        sa = "  ";   sb = "  ";   sc = "  "
        sa = sa * TableStrings.center(10, "Level"; na=2);                             sb = sb * TableStrings.hBlank(12)
        sa = sa * TableStrings.center(10, "J^P";   na=4);                             sb = sb * TableStrings.hBlank(14)
        sa = sa * TableStrings.center(14, "Energy"; na=5)
        sb = sb * TableStrings.center(14, TableStrings.inUnits("energy"); na=5);  sc = sc * TableStrings.hBlank(45)
        sa = sa * TableStrings.center(11, "K_nms"; na=4)
        sb = sb * TableStrings.center(11, "[GHz u]" ; na=4)
        sc = sc * TableStrings.center(11, "[a.u.]" ; na=4)
        println(stream, sa);    println(stream, sb);    ## println(stream, sc);
        println(stream, "  ", TableStrings.hLine(nx))
        #
        for  outcome in outcomes
            sa  = "  ";    sym = LevelSymmetry( outcome.level.J, outcome.level.parity)
            sa = sa * TableStrings.center(10, TableStrings.level(outcome.level.index); na=2)
            sa = sa * TableStrings.center(10, string(sym); na=4)
            sa = sa * @sprintf("%.8e", Defaults.convertUnits("energy: from atomic", outcome.level.energy))            * "    "
            println(stream, sa )
            ## sb = TableStrings.hBlank(47)
        end
        println(stream, "  ", TableStrings.hLine(nx), "\n\n")

        return( nothing )
    end

end # module


using Distributed, Pkg
Pkg.activate(joinpath(@__DIR__, ".."))
if haskey(ENV, "FOOD_CACHING_PARALLEL")
    nprocs() < Sys.CPU_THREADS && addprocs(20,
                                           exeflags = "--project",
                                           dir = joinpath(@__DIR__, ".."))
end
using FoodCachingModels
import FoodCachingFitting: resume_fit

resume_fit(ARGS[1])


# Object-based wrapper for level-2 BLIS routines.
#

macro blis_object_api_level2_mv(funcname)
    return quote
        @blis_ccall_object($funcname,
                           Cvoid,
                           α, xObj,
                           a, xObj,
                           x, xObj,
                           β, xObj,
                           y, xObj)
    end
end

@blis_object_api_level2_mv gemv
@blis_object_api_level2_mv hemv
@blis_object_api_level2_mv symv

macro blis_object_api_level2_r2(funcname)
    return quote
        @blis_ccall_object($funcname,
                           Cvoid,
                           α, xObj,
                           x, xObj,
                           y, xObj,
                           a, xObj)
    end
end

@blis_object_api_level2_r2 ger
@blis_object_api_level2_r2 her2
@blis_object_api_level2_r2 syr2

macro blis_object_api_level2_r(funcname)
    return quote
        @blis_ccall_object($funcname,
                           Cvoid,
                           α, xObj,
                           x, xObj,
                           a, xObj)
    end
end

@blis_object_api_level2_r her
@blis_object_api_level2_r syr

macro blis_object_api_level2_tr(funcname)
    return quote
        @blis_ccall_object($funcname,
                           Cvoid,
                           α, xObj,
                           a, xObj,
                           x, xObj)
    end
end

@blis_object_api_level2_tr trmv
@blis_object_api_level2_tr trsv


module CMBLensing

using Adapt
using Base.Broadcast: AbstractArrayStyle, ArrayStyle, Broadcasted,
    DefaultArrayStyle, preprocess_args, Style, result_style, Unknown
using Base.Iterators: flatten, product, repeated, cycle, countfrom, peel, partition
using Base.Threads
using Base: @kwdef, @propagate_inbounds, Bottom, OneTo, showarg, show_datatype,
    show_default, show_vector, typed_vcat, typename
using Combinatorics
using CompositeStructs
using DataStructures
using DelimitedFiles
using Distributed: pmap, nworkers, myid, workers, addprocs, @everywhere, remotecall_wait,
    @spawnat, pgenerate, procs, @fetchfrom, default_worker_pool, RemoteChannel
using EllipsisNotation
using FileIO
using FFTW
using InteractiveUtils
using IterTools: flagfirst
using JLD2
using JLD2: jldopen, JLDWriteSession
using KahanSummation
using Loess
using LinearAlgebra
using LinearAlgebra: diagzero, matprod, promote_op
using MacroTools: @capture, combinedef, isdef, isexpr, postwalk, prewalk, rmlines, splitdef
using Match
using Markdown
using Measurements
using Memoization
using NamedTupleTools: select, delete
using Pkg
using Printf
using ProgressMeter
using QuadGK
using Random
using Random: seed!, AbstractRNG
using Roots
using Requires
using Serialization
using Setfield
using SparseArrays
using StaticArrays: @SMatrix, @SVector, SMatrix, StaticArray, StaticArrayStyle,
    StaticMatrix, StaticVector, SVector, SArray, SizedArray
using Statistics
using StatsBase
using TimerOutputs: @timeit, get_defaulttimer, reset_timer!
using Tullio
using UnPack
using Zygote
using Zygote: unbroadcast, Numeric, @adjoint, @nograd
using Zygote.ChainRules: @thunk, NoTangent

import Adapt: adapt_structure
import Base: +, -, *, \, /, ^, ~, ≈, <, <=, |, &, ==, !,
    abs, adjoint, all, any, axes, broadcast, broadcastable, BroadcastStyle, conj, copy, convert,
    copy, copyto!, eltype, eps, fill!, getindex, getproperty, hash, hcat, hvcat, inv, isfinite,
    iterate, keys, lastindex, length, literal_pow, mapreduce, materialize!,
    materialize, one, permutedims, print_array, promote, promote_rule,
    promote_rule, promote_type, propertynames, real, setindex!, setproperty!, show,
    show_datatype, show_vector, similar, size, sqrt, string, sum, summary,
    transpose, zero
import Base.Broadcast: materialize, preprocess, broadcasted
import Zygote.ChainRules: rrule
import LinearAlgebra: checksquare, diag, dot, isnan, ldiv!, logdet, mul!, norm,
    pinv, StructuredMatrixStyle, structured_broadcast_alloc, tr
import Measurements: ±
import Statistics: std

export
    @⌛, @show⌛, @ismain, @namedtuple, @repeated, @unpack, @cpu!, @gpu!, @cu!, animate,
    argmaxf_lnP, AzFourier, BandPassOp, BaseDataSet, batch, batch_index, batch_length,
    BlockDiagEquiRect, beamCℓs, cache, CachedLenseFlow, camb, cov_to_Cℓ, cpu, Cℓ_2D,
    Cℓ_to_Cov, DataSet, DerivBasis, diag, Diagonal, DiagOp, dot, EBFourier, EBMap, expnorm,
    Field, FieldArray, fieldinfo, FieldMatrix, FieldOrOpArray, FieldOrOpMatrix, FieldOrOpRowVector,
    FieldOrOpVector, FieldRowVector, FieldTuple, FieldVector, FieldVector,
    firsthalf, fixed_white_noise, BlockDiagIEB, Fourier, FuncOp, get_max_lensing_step,
    get_Cℓ, get_Cℓ, get_Dℓ, get_αℓⁿCℓ, get_ρℓ, get_ℓ⁴Cℓ, gradhess, gradient, HighPass,
    IEBFourier, IEBMap, InterpolatedCℓs, IQUAzFourier, IQUFourier, IQUMap, kde,
    lasthalf, LazyBinaryOp, LenseBasis, LenseFlow, FieldOp, lnP, load_camb_Cℓs,
    load_chains, load_nolensing_sim, load_sim, LowPass, make_mask, Map, MAP_joint, MAP_marg,
    mean_std_and_errors, MidPass, mix, nan2zero, noiseCℓs,
    ParamDependentOp, pixwin, PowerLens, ProjLambert, ProjEquiRect, ProjHealpix, project,
    QUAzFourier, QUFourier, QUMap, resimulate!, resimulate, RK4Solver, sample_joint, shiftℓ,
    simulate, SymmetricFuncOp, symplectic_integrate, Taylens, toCℓ, toDℓ,
    ud_grade, unbatch, unmix, white_noise, Ð, Ł,
    ℓ², ℓ⁴, ∇, ∇², ∇ᵢ, ∇ⁱ

# bunch of sampling-related exports
export gibbs_initialize_f!, gibbs_initialize_ϕ!, gibbs_initialize_θ!,
    gibbs_sample_f!, gibbs_sample_ϕ!, gibbs_sample_slice_θ!,
    gibbs_mix!, gibbs_unmix!, gibbs_postprocess!,
    once_every, start_after_burnin, mass_matrix_ϕ, hmc_step

# util
include("util.jl")
include("util_fft.jl")
include("util_parallel.jl")
include("numerical_algorithms.jl")

# generic field / operator structure
include("generic.jl")
include("cls.jl")
include("field_tuples.jl")
include("field_vectors.jl")
include("base_fields.jl")
include("specialops.jl")
include("flowops.jl")
include("batched_reals.jl")

# lensing operators
include("lenseflow.jl")
include("powerlens.jl")

# field types
include("proj_cartesian.jl")
include("proj_lambert.jl")
include("proj_equirect.jl")
include("proj_healpix.jl")

# other field-specific stuff
include("masking.jl")
include("taylens.jl")
include("bilinearlens.jl")

# plotting
function animate end
@init @require PyPlot="d330b81b-6aea-500a-939a-2ce795aea3ee" include("plotting.jl")

# sampling and maximizing the posteriors
include("dataset.jl")
include("posterior.jl")
include("maximization.jl")
include("sampling.jl")
include("chains.jl")

# other estimates
include("quadratic_estimate.jl")

# AD
include("autodiff.jl")

# gpu
is_gpu_backed(x) = false
@init @require CUDA="052768ef-5323-5732-b1bb-66c8b64840ba" include("gpu.jl")

# misc init
# see https://github.com/timholy/ProgressMeter.jl/issues/71 and links therein
@init if ProgressMeter.@isdefined ijulia_behavior
    ProgressMeter.ijulia_behavior(:clear)
end

end


@memoize function version()
    buf = Vector{Cchar}(undef, NVML_SYSTEM_NVML_VERSION_BUFFER_SIZE)
    nvmlSystemGetNVMLVersion(pointer(buf), length(buf))

    # the version string is too long for Julia to handle, e.g. 11.450.36.06,
    # so split off the driver part into the build suffix
    ver = unsafe_string(pointer(buf))
    parts = parse.(Int, split(ver, '.'))
    return VersionNumber(parts[1], 0, 0, (), Tuple(parts[2:end]))
end

@memoize function driver_version()
    buf = Vector{Cchar}(undef, NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE)
    nvmlSystemGetDriverVersion(pointer(buf), length(buf))
    return VersionNumber(unsafe_string(pointer(buf)))
end

@memoize function cuda_driver_version()
    ref = Ref{Cint}()
    nvmlSystemGetCudaDriverVersion_v2(ref)
    major, ver = divrem(ref[], 1000)
    minor, patch = divrem(ver, 10)
    return VersionNumber(major, minor, patch)
end


import TracedRandom: Here

## Rewrite the following random primitives as calls to Gen distributions
const randprims = Set([
    rand, randn, randexp, randperm, randcycle, shuffle, sample,
    rand!, randn!, randexp!, randperm!, randcycle!, shuffle!, sample!
])
const primnames = Set(Symbol(fn) for fn in randprims)

## Manual addressing ##

# Forward manually addressed calls for each random primitive
for fn in randprims
    @eval @inline trace(options::Options, state, ::ManualAddress, ::typeof($fn), addr::Address, args...) =
        trace(options, state, addr::Address, $fn, args...)
    @eval @inline trace(options::Options, state, ::ManualAddress, ::typeof($fn), ::AbstractRNG, addr::Address, args...) =
        trace(options, state, addr::Address, $fn, rng, args...)
end

# Forward traced function calls that are manually addressed via rand
@inline trace(options::Options, state, ::ManualAddress, ::typeof(rand), addr::Address, fn::Function, args...) =
    trace(options, state, addr::Address, fn, args...)

# Splice function calls with `here` addresses into the current context
@inline trace(options::Options, state, ::ManualAddress, ::typeof(rand), ::Here, fn::Function, args...) =
    splice(options, state, fn, args...)

## Strip RNGs supplied to primitives

for fn in randprims
    @eval @inline trace(options::Options, state, addr::Address, ::typeof($fn), ::AbstractRNG, args...) =
        trace(options, state, addr::Address, $fn, args...)
end

## rand ##

# Default to uniformly sampling Float64 values from [0, 1]
@inline trace(options::Options, state, addr::Address, ::typeof(rand)) =
    trace(options, state, addr::Address, rand, Float64)
@inline trace(options::Options, state, addr::Address, ::typeof(rand), d::Integer, dims::Integer...) =
    trace(options, state, addr::Address, rand, Float64, d, dims...)

# rand for numeric types
@inline trace(::Options, state, addr::Address, ::typeof(rand), T::Type{<:Real}) =
    Gen.traceat(state, TypedScalarDistribution(T), (), addr)
@inline trace(::Options, state, addr::Address, ::typeof(rand), T::Type{<:Real}, ::Tuple{}) =
    Gen.traceat(state, TypedScalarDistribution(T), (), addr)
@inline trace(::Options, state, addr::Address, ::typeof(rand), T::Type{<:Real}, dims::Dims) =
    Gen.traceat(state, TypedArrayDistribution(T, dims), (), addr)
@inline trace(::Options, state, addr::Address, ::typeof(rand), T::Type{<:Real}, d::Integer, dims::Integer...) =
    Gen.traceat(state, TypedArrayDistribution(T, d, dims...), (), addr)

# rand for indexable collections
@inline trace(::Options, state, addr::Address, ::typeof(rand), c::T) where {T <: Indexable} =
    Gen.traceat(state, LabeledUniform{eltype(T)}(), (c,), addr)
@inline trace(::Options, state, addr::Address, ::typeof(rand), c::T, ::Tuple{}) where {T <: Indexable} =
    Gen.traceat(state, LabeledUniform{eltype(T)}(), (c,), addr)
@inline trace(::Options, state, addr::Address, ::typeof(rand), c::T, dims::Dims) where {T <: Indexable} =
    Gen.traceat(state, ArrayedDistribution(LabeledUniform{eltype(T)}(), dims), (c,), addr)
@inline trace(::Options, state, addr::Address, ::typeof(rand), c::T, d::Integer, dims::Integer...) where {T <: Indexable} =
    Gen.traceat(state, ArrayedDistribution(LabeledUniform{eltype(T)}(), d, dims...), (c,), addr)

# rand for set-like collections
@inline trace(::Options, state, addr::Address, ::typeof(rand), c::T) where {T <: Setlike} =
    Gen.traceat(state, SetUniform{eltype(T)}(), (c,), addr)
@inline trace(::Options, state, addr::Address, ::typeof(rand), c::T, ::Tuple{}) where {T <: Setlike} =
    Gen.traceat(state, SetUniform{eltype(T)}(), (c,), addr)
@inline trace(::Options, state, addr::Address, ::typeof(rand), c::T, dims::Dims) where {T <: Setlike} =
    Gen.traceat(state, ArrayedDistribution(SetUniform{eltype(T)}(), dims), (c,), addr)
@inline trace(::Options, state, addr::Address, ::typeof(rand), c::T, d::Integer, dims::Integer...) where {T <: Setlike} =
    Gen.traceat(state, ArrayedDistribution(SetUniform{eltype(T)}(), d, dims...), (c,), addr)

# rand for Distributions.jl distributions
@inline trace(::Options, state, addr::Address, ::typeof(rand), dist::D) where {D <: Distributions.Distribution} =
    Gen.traceat(state, WrappedDistribution(dist), params(dist), addr)
@inline trace(::Options, state, addr::Address, ::typeof(rand), dist::D, ::Tuple{}) where {D <: Distributions.Distribution} =
    Gen.traceat(state, WrappedDistribution(dist), params(dist), addr)
@inline trace(::Options, state, addr::Address, ::typeof(rand), dist::D, dims::Dims) where {D <: Distributions.Distribution} =
    Gen.traceat(state, ArrayedDistribution(WrappedDistribution(dist), dims), params(dist), addr)
@inline trace(::Options, state, addr::Address, ::typeof(rand), dist::D, d::Integer, dims::Integer...) where {D <: Distributions.Distribution} =
    Gen.traceat(state, ArrayedDistribution(WrappedDistribution(dist), d, dims...), params(dist), addr)

# rand! (simulate in-place allocation)
@inline trace(options::Options, state, addr::Address, ::typeof(rand!), A::AbstractArray{T}) where {T} =
    copyto!(A, trace(options, state, addr::Address, rand, T, size(A)))
@inline trace(options::Options, state, addr::Address, ::typeof(rand!), A::AbstractArray, S) =
    copyto!(A, trace(options, state, addr::Address, rand, S, size(A)))

## randn ##

# Default to sampling Float64 values
@inline trace(options::Options, state, addr::Address, ::typeof(randn)) =
    trace(options, state, addr::Address, randn, Float64)
@inline trace(options::Options, state, addr::Address, ::typeof(randn), dims::Dims) =
    trace(options, state, addr::Address, randn, Float64, dims)
@inline trace(options::Options, state, addr::Address, ::typeof(randn), d::Integer, dims::Integer...) =
    trace(options, state, addr::Address, randn, Float64, d, dims...)

# Forward sampling of all other Float types
@inline trace(::Options, state, addr::Address, ::typeof(randn), T::Type{<:AbstractFloat}, args...) =
    trace(options, state, addr::Address, randn, Float64, args...)

# Sampling of Float64 values
@inline trace(::Options, state, addr::Address, ::typeof(randn), T::Type{Float64}) =
    Gen.traceat(state, Gen.normal, (0, 1), addr)
@inline trace(::Options, state, addr::Address, ::typeof(randn), T::Type{Float64}, ::Tuple{}) =
    Gen.traceat(state, Gen.normal, (0, 1), addr)
@inline trace(::Options, state, addr::Address, ::typeof(randn), T::Type{Float64}, dims::Dims) =
    Gen.traceat(state, ArrayedDistribution(Gen.normal, dims), (0, 1), addr)
@inline trace(::Options, state, addr::Address, ::typeof(randn), T::Type{Float64}, d::Integer, dims::Integer...) =
    Gen.traceat(state, ArrayedDistribution(Gen.normal, d, dims...), (0, 1), addr)

# randn! (simulate in-place allocation)
@inline trace(options::Options, state, addr::Address, ::typeof(randn!), A::AbstractArray{T}) where {T} =
    copyto!(A, trace(options, state, addr::Address, randn, T, size(A)))

## randexp ##

# Default to sampling Float64 values
@inline trace(options::Options, state, addr::Address, ::typeof(randexp)) =
    trace(options, state, addr::Address, randexp, Float64)
@inline trace(options::Options, state, addr::Address, ::typeof(randexp), dims::Dims) =
    trace(options, state, addr::Address, randexp, Float64, dims)
@inline trace(options::Options, state, addr::Address, ::typeof(randexp), d::Integer, dims::Integer...) =
    trace(options, state, addr::Address, randexp, Float64, d, dims...)

# Forward sampling of all other Float types
@inline trace(::Options, state, addr::Address, ::typeof(randexp), T::Type{<:AbstractFloat}, args...) =
    trace(options, state, addr::Address, randexp, Float64, args...)

# Sampling of Float64 values
@inline trace(::Options, state, addr::Address, ::typeof(randexp), T::Type{Float64}) =
    Gen.traceat(state, Gen.exponential, (1,), addr)
@inline trace(::Options, state, addr::Address, ::typeof(randexp), T::Type{Float64}, ::Tuple{}) =
    Gen.traceat(state, Gen.exponential, (1,), addr)
@inline trace(::Options, state, addr::Address, ::typeof(randexp), T::Type{Float64}, dims::Dims) =
    Gen.traceat(state, ArrayedDistribution(Gen.exponential, dims), (1,), addr)
@inline trace(::Options, state, addr::Address, ::typeof(randexp), T::Type{Float64}, d::Integer, dims::Integer...) =
    Gen.traceat(state, ArrayedDistribution(Gen.exponential, d, dims...), (1,), addr)

# randexp! (simulate in-place allocation)
@inline trace(options::Options, state, addr::Address, ::typeof(randexp!), A::AbstractArray{T}) where {T} =
    copyto!(A, trace(options, state, addr::Address, randexp, T, size(A)))

## randperm, randcycle. shuffle ##

# Forward to corresponding Gen distributions
@inline trace(options::Options, state, addr::Address, ::typeof(randperm), n::Integer) =
    Gen.traceat(state, RandomPermutation(n), (), addr)
@inline trace(options::Options, state, addr::Address, ::typeof(randcycle), n::Integer) =
    Gen.traceat(state, RandomCycle(n), (), addr)
@inline trace(options::Options, state, addr::Address, ::typeof(shuffle), v::AbstractArray) =
    Gen.traceat(state, RandomShuffle(v), (), addr)

# Simulate in-place versions
@inline trace(options::Options, state, addr::Address, ::typeof(randperm!), A::AbstractArray{<:Integer}) =
    copyto!(A, trace(options, state, addr::Address, randperm, length(A)))
@inline trace(options::Options, state, addr::Address, ::typeof(randcycle!), A::AbstractArray{<:Integer}) =
    copyto!(A, trace(options, state, addr::Address, randcycle, length(A)))
@inline trace(options::Options, state, addr::Address, ::typeof(shuffle!), v::AbstractArray) =
    copyto!(v, trace(options, state, addr::Address, shuffle, copy(v)))

## sample ##

# Unweighted sample
@inline trace(::Options, state, addr::Address, ::typeof(sample), a::AbstractArray{T}) where {T} =
    Gen.traceat(state, LabeledUniform{T}(), (a,), addr)
@inline trace(::Options, state, addr::Address, ::typeof(sample), a::AbstractArray{T}, n::Integer) where {T} =
    Gen.traceat(state, ArrayedDistribution(LabeledUniform{T}(), n), (a,), addr)
@inline trace(::Options, state, addr::Address, ::typeof(sample), a::AbstractArray{T}, dims::Dims) where {T} =
    Gen.traceat(state, ArrayedDistribution(LabeledUniform{T}(), dims), (a,), addr)

# Weighted sample
@inline trace(::Options, state, addr::Address, ::typeof(sample), w::AbstractWeights) =
    Gen.traceat(state, Categorical(), (Vector(w) ./ sum(w)), addr)
@inline trace(::Options, state, addr::Address, ::typeof(sample), a::AbstractArray{T}, w::AbstractWeights) where {T} =
    Gen.traceat(state, LabeledCategorical{T}(), (a, w), addr)
@inline trace(::Options, state, addr::Address, ::typeof(sample), a::AbstractArray{T}, w::AbstractWeights, n::Integer) where {T} =
    Gen.traceat(state, ArrayedDistribution(LabeledCategorical{T}(), n), (a, w), addr)
@inline trace(::Options, state, addr::Address, ::typeof(sample), a::AbstractArray{T}, w::AbstractWeights, dims::Dims) where {T} =
    Gen.traceat(state, ArrayedDistribution(LabeledCategorical{T}(), dims), (a, w), addr)

# sample! (simulate in-place allocation)
@inline trace(options::Options, state, addr::Address, ::typeof(sample!), a::AbstractArray, x::AbstractArray) =
    copyto!(x, trace(options, state, addr::Address, sample, a, size(x)))
@inline trace(options::Options, state, addr::Address, ::typeof(sample!), a::AbstractArray, w::AbstractWeights, x::AbstractArray) =
    copyto!(x, trace(options, state, addr::Address, sample, a, w, size(x)))


# Downward initialization through a cascading model
# these functions need to be consolidated refactored and deprecated accordingly
# see #910

## ============================================================================================
## Down init sibling priority order
## ============================================================================================

"""
$SIGNATURES

Test waiting order between siblings for cascading downward tree initialization.

Notes
- State machine function 8j.

DevNotes
- FIXME, this guy is never called during WIP of 459 dwnMsg consolidation
- FIXME, something wrong with CSM sequencing, https://github.com/JuliaRobotics/IncrementalInference.jl/issues/602#issuecomment-682114232
- This might be replaced with 4-stroke tree-init.
"""
function dwnInitSiblingWaitOrder_StateMachine(csmc::CliqStateMachineContainer)

  prnt_ = getParent(csmc.tree, csmc.cliq)
  if 0 == length(prnt_)
    # go to 7
    return determineCliqNeedDownMsg_StateMachine
  end
  prnt = prnt_[1]

  opt = getSolverParams(csmc.cliqSubFg) # csmc.dfg

  # now get the newly computed message from the appropriate container
  # make sure this is a pull model #674 (pull msg from source/prnt)
  # FIXME must be consolidated as part of #459
  # dwinmsgs = getfetchCliqueInitMsgDown(getCliqueData(prnt), from=:dwnInitSiblingWaitOrder_StateMachine)
  dwinmsgs = fetchDwnMsgConsolidated(prnt)
  dwnkeys_ = collect(keys(dwinmsgs.belief))
  infocsm(csmc, "8j, dwnInitSiblingWaitOrder_StateMachine, dwinmsgs keys=$(dwnkeys_)")

  # add downward belief prop msgs
  msgfcts = addMsgFactors!(csmc.cliqSubFg, dwinmsgs, DownwardPass)
  ## update solveable dims with newly avaiable init information
  # determine if more info is needed for partial
  sdims = getCliqVariableMoreInitDims(csmc.cliqSubFg, csmc.cliq)
  infocsm(csmc, "8j, dwnInitSiblingWaitOrder_StateMachine, sdims=$(sdims)")
  updateCliqSolvableDims!(csmc.cliq, sdims, csmc.logger)

  _dbgCSMSaveSubFG(csmc, "fg_DWNCMN_8j")

  ## FIXME, if this new, and all sibling clique's solvableDim are 0, then go back to waitChangeOnParentCondition_StateMachine

  # go to 8o.i
  return testDirectDwnInit_StateMachine
end

"""
    $SIGNATURES

Can this clique initialize directly from available down message info?

Notes
- State machine function nr. 8o.i
- Assume must have parent since this only occurs after 8j.
"""
function testDirectDwnInit_StateMachine(csmc::CliqStateMachineContainer)

  prnt = getParent(csmc.tree, csmc.cliq)[1]
  dwinmsgs = fetchDwnMsgConsolidated(prnt)
  dwnkeys_ = collect(keys(dwinmsgs.belief))
  # NOTE, only use separators, not all parent variables (DF ???)
  # dwnkeys_ = lsf(csmc.cliqSubFg, tags=[:__DOWNWARD_COMMON__;]) .|> x->ls(csmc.cliqSubFg, x)[1]
  # @assert length(intersect(dwnkeys, dwnkeys_)) == length(dwnkeys) "split dwnkeys_ is not the same, $dwnkeys, and $dwnkeys_, separators: $(getCliqSeparatorVarIds(csmc.cliq))"

  # priorize solve order for mustinitdown with lowest dependency first
  # follow example from issue #344
  # mustwait = false
  if length(intersect(dwnkeys_, getCliqSeparatorVarIds(csmc.cliq))) == 0
    # can directly use DOWNWARD_COMMON
    infocsm(csmc, "8j, dwnInitSiblingWaitOrder_StateMachine, no can do, must wait for siblings to update parent first.")
    # mustwait = true

    # go to 8o.iv
    return sibsDwnPriorityInit_StateMachine
  end

  # go to 8o.ii
  return testDelayOrderDwnInit_StateMachine
end

"""
    $SIGNATURES

Return true if this clique's down init should be delayed on account of prioritization among sibling separators.

Notes
- State machine function nr. 8o.ii
- process described in issue #344

Dev Notes
- not priorizing order yet (TODO), just avoiding unsolvables at this time.
- Very closely related to 8o.iii -- refactor likely (NOTE).
- should precompute `allinters`.
- # FIXME ON FIRE, this is doing `getCliqueStatus` of siblings without following channels.
  - NOTE This is only used in old tree-init, not part of active regular solve code.
  - See #954
"""
function testDelayOrderDwnInit_StateMachine(csmc::CliqStateMachineContainer)

  prnt = getParent(csmc.tree, csmc.cliq)[1]
  dwinmsgs = fetchDwnMsgConsolidated(prnt)
  dwnkeys = collect(keys(dwinmsgs.belief))

  tree = csmc.tree
  cliq = csmc.cliq
  logger = csmc.logger

  # when is a cliq upsolved
  solvedstats = Symbol[:upsolved; :marginalized; :uprecycled]

  # safety net double check
  cliqst = getCliqueStatus(cliq)
  if cliqst in solvedstats
    infocsm(csmc, "getSiblingsDelayOrder -- clique status should not be here with a solved cliqst=$cliqst")
    # go to 8o.iii
    return testPartialNeedsDwnInit_StateMachine
  end

  # get siblings separators
  sibs = getCliqSiblings(tree, cliq, true)
  ids = map(s->s.index, sibs)
  len = length(sibs)
  sibidx = collect(1:len)[ids .== cliq.index][1]
  seps = getCliqSeparatorVarIds.(sibs)
  lielbls = setdiff(ids, cliq.index)
  # get intersect matrix of siblings (should be exactly the same across siblings' csm)
  allinters = Array{Int,2}(undef, len, len)
  dwninters = Vector{Int}(undef, len)
  infocsm(csmc, "getSiblingsDelayOrder -- number siblings=$(len), sibidx=$sibidx")

  # sum matrix with all "up solved" rows and columns eliminated
  fill!(allinters, 0)
  for i in 1:len
    for j in i:len
      if i != j
        allinters[i,j] = length(intersect(seps[i],seps[j]))
      end
    end
    dwninters[i] = length(intersect(seps[i], dwnkeys))
  end

  # sum "across/over" rows, then columns (i.e. visa versa "along" columns then rows)
  rows = sum(allinters, dims=1)
  cols = sum(allinters, dims=2)

  infocsm(csmc, "getSiblingsDelayOrder -- allinters=$(allinters), getSiblingsDelayOrder -- rows=$(rows), getSiblingsDelayOrder -- rows=$(cols)")

  # is this clique a non-zero row -- i.e. sum across columns? if not, no further special care needed
  if cols[sibidx] == 0
    infocsm(csmc, "getSiblingsDelayOrder -- cols[sibidx=$(sibidx))] == 0, no special care needed")
    # go to 8o.iii
    return testPartialNeedsDwnInit_StateMachine
  end

  # now determine if initializing from below or needdownmsg
  if cliqst in Symbol[:needdownmsg;]
    # be super careful about delay (true) vs pass (false) at this point -- might be partial too TODO
    # return true if delay beneficial to initialization accuracy

    # find which siblings this cliq epends on
    symm = allinters + allinters'
    maskcol = 0 .< symm[:,sibidx]
    # lenm = length(maskcol)
    stat = Vector{Symbol}(undef, len)
    stillbusymask = fill(false, len)

    # get each sibling status (entering atomic computation segment -- until wait command)
    stat .= getCliqueStatus.(sibs) #[maskcol]

    ## (long down chain case)
    # need different behaviour when all remaining siblings are blocking with :needdownmsg
    remainingmask = stat .== :needdownmsg
    if sum(remainingmask) == length(stat)
      infocsm(csmc, "getSiblingsDelayOrder -- all blocking: sum(remainingmask) == length(stat), stat=$stat")
      # pick sibling with most overlap in down msgs from parent
      # list of similar length siblings
      candidates = dwninters .== maximum(dwninters)
      if candidates[sibidx]
        # must also pick minimized intersect with other remaing siblings
        maxcan = collect(1:len)[candidates]
        infocsm(csmc, "getSiblingsDelayOrder -- candidates=$candidates, maxcan=$maxcan, rows=$rows")
        if rows[sibidx] == minimum(rows[maxcan])
          infocsm(csmc, "getSiblingsDelayOrder -- FORCE DOWN INIT SOLVE ON THIS CLIQUE: $(cliq.index), $(getLabel(cliq))")
          # go to 8o.iii
          return testPartialNeedsDwnInit_StateMachine
        end
      end
      infocsm(csmc, "getSiblingsDelayOrder -- not a max and should block")
      # go to 8o.iv
      return sibsDwnPriorityInit_StateMachine
    end

    # still busy solving on branches, so potential to delay
    for i in 1:len
      stillbusymask[i] = maskcol[i] && !(stat[i] in solvedstats)
    end
    infocsm(csmc, "getSiblingsDelayOrder -- busy solving: maskcol=$maskcol, stillbusy=$stillbusymask")

    # Too blunt -- should already have returned false by this point perhaps
    if 0 < sum(stillbusymask)
      # yes something to delay about
      infocsm(csmc, "getSiblingsDelayOrder -- yes delay, stat=$stat symm=$symm")

      # go to 8o.iv
      return sibsDwnPriorityInit_StateMachine
    end
  end

  infocsm(csmc, "getSiblingsDelayOrder -- default will not delay")
  # carry over default from partial init process

  # go to 8o.iii
  return testPartialNeedsDwnInit_StateMachine
end

"""
    $SIGNATURES

Return true if both, i.) this clique requires more downward information, ii.) more
downward message information could potentially become available.

Notes
- State machine function nr. 8o.iii
- Delay initialization to the last possible moment.

Dev Notes:
- # FIXME ON FIRE, this is doing `getCliqueStatus` of siblings without following channels.
  - NOTE This is only used in old tree-init, not part of active regular solve code.
  - See #954
- Determine clique truely isn't able to proceed any further:
  - should be as self reliant as possible (using clique's status as indicator)
  - OBSOLETE ??
    - change status to :mustinitdown if have only partial beliefs so far:
    - combination of status, while partials belief siblings are not :mustinitdown
"""
function testPartialNeedsDwnInit_StateMachine(csmc::CliqStateMachineContainer)
  #
  prnt = getParent(csmc.tree, csmc.cliq)[1]
  dwinmsgs = fetchDwnMsgConsolidated(prnt)

  tree = csmc.tree
  cliq = csmc.cliq
  logger = csmc.logger

  # which incoming messages are partials
  hasPartials = Dict{Symbol, Int}()
  for (sym, tmsg) in dwinmsgs.belief
    # assuming any down message per label that is not partial voids further partial consideration
    if sum(tmsg.inferdim) > 0
      if !haskey(hasPartials, sym)
        hasPartials[sym] = 0
      end
      hasPartials[sym] += 1
    end
  end
  partialKeys = collect(keys(hasPartials))

  ## determine who might be able to help init this cliq
  # check sibling separator sets against this clique's separator
  sibs = getCliqSiblings(tree, cliq)

  infocsm(csmc, "getCliqSiblingsPartialNeeds -- CHECK PARTIAL")
  # identify which cliques might have useful information
  localsep = getCliqSeparatorVarIds(cliq)
  seps = Dict{Int, Vector{Symbol}}()
  for si in sibs
    # @show getLabel(si)
    mighthave = intersect(getCliqSeparatorVarIds(si), localsep)
    if length(mighthave) > 0
      seps[si.index] = mighthave
      if getCliqueStatus(si) in [:initialized; :null; :needdownmsg]
        # partials treated special -- this is slightly hacky
        if length(intersect(localsep, partialKeys)) > 0 && length(mighthave) > 0
          # this sibling might have info to delay about
          setCliqueDrawColor!(cliq,"magenta")
          # go to 8o.iv
          return sibsDwnPriorityInit_StateMachine
        end
      end
    end
  end
  # determine if those cliques will / or will not be able to provide more info
  # when does clique change to :mustinitdown
  # default

  # go to 8e.ii.
  return tryDwnInitCliq_StateMachine
end

"""
$SIGNATURES

Return true there is no other sibling that will make progress.

Notes
- State machine function nr. 8o.iv
- Relies on sibling priority order with only one "currently best" option that will force progress in global upward inference.
- Return false if one of the siblings is still busy

DevNotes
- # FIXME ON FIRE, calling getCliqueStatus of siblings directly
  - Only used for older tree init, not part of regular up-down solving code
  - see #954
- Best is likely if parent actually made determination on which child to solve first in this case.
"""
function sibsDwnPriorityInit_StateMachine(csmc::CliqStateMachineContainer)

  tree = csmc.tree
  cliq = csmc.cliq
  prnt_ = getParent(tree, cliq)
  prnt = prnt_[1]

  dwinmsgs = fetchDwnMsgConsolidated(prnt)
  dwnkeys_ = collect(keys(dwinmsgs.belief))

  solvord = getCliqSiblingsPriorityInitOrder( csmc.tree, prnt, dwinmsgs, csmc.logger )

  # noOneElse = areSiblingsRemaingNeedDownOnly(csmc.tree, csmc.cliq)
  noOneElse = true
  stillbusylist = [:null; :initialized;]
  if 0 < length(prnt_)
    for si in getChildren(tree, prnt)
      # are any of the other siblings still busy?
      if si.index != cliq.index && getCliqueStatus(si) in stillbusylist
        noOneElse = false
      end
    end
  end
  # nope, everybody is waiting for something to change -- proceed with forcing a cliq solve

  infocsm(csmc, "8j, dwnInitSiblingWaitOrder_StateMachine, $(prnt.index), $noOneElse, solvord = $solvord")

  if csmc.cliq.index != solvord[1]
    # TODO, is this needed? fails hex if included with correction :needdowninit --> :needdownmsg
    # if dwinmsgs.status == :initialized && getCliqueStatus(csmc.cliq) == :needdownmsg  # used to be :needdowninit
    #   # go to 7e
    #   return slowWhileInit_StateMachine
    # end

    infocsm(csmc, "8j, dwnInitSiblingWaitOrder_StateMachine, must wait on change.")
    # remove all message factors
    # remove msg factors previously added
    fctstorm = deleteMsgFactors!(csmc.cliqSubFg, [:__DOWNWARD_COMMON__])
    infocsm(csmc, "8j, dwnInitSiblingWaitOrder_StateMachine, removing factors $fctstorm")

    # go to 8c
    return waitChangeOnParentCondition_StateMachine
  end

  # go to 8e.ii.
  return tryDwnInitCliq_StateMachine
end


###########################  Blotter

function show(io::IO, ta::Blotter)
  # variables
  nrow          = size(ta.values, 1)
  ncol          = size(ta.values, 2)
  intcatcher    = falses(ncol)
  for c in 1:ncol
      rowcheck =  trunc(ta.values[:,c]) - ta.values[:,c] .== 0
      if sum(rowcheck) == length(rowcheck)
          intcatcher[c] = true
      end
  end
  spacetime     = strwidth(string(ta.timestamp[1])) + 3
  firstcolwidth = strwidth(ta.colnames[1])
  colwidth      = Int[]
      for m in 1:ncol
          push!(colwidth, max(strwidth(ta.colnames[m]), strwidth(@sprintf("%.2f", maximum(ta.values[:,m])))))
      end

  # summary line
  println(io,"")
  #print(io,@sprintf("%dx%d %s %s to %s", nrow, ncol, typeof(ta), string(ta.timestamp[1]), string(ta.timestamp[nrow])))
  print(io,@sprintf("%dx%d Financial Blotter for Ticker *** %s ***, %s to %s",
                    nrow, ncol, string(ta.meta), string(ta.timestamp[1]), string(ta.timestamp[nrow])))
  println(io,"")

  # row label line
   print_with_color(:magenta, io, ^(" ", spacetime), ta.colnames[1], ^(" ", colwidth[1] + 2 -firstcolwidth))

   for p in 2:length(colwidth)
     print_with_color(:magenta, io, ta.colnames[p], ^(" ", colwidth[p] - strwidth(ta.colnames[p]) + 2))
   end
   println(io,"")

  # timestamp and values line
    if nrow > 7
        for i in 1:4
            print(io, ta.timestamp[i], " | ")
        for j in 1:ncol
            intcatcher[j] && ta.values[i,j] > 0?
            print_with_color(:green, io, rpad(iround(ta.values[i,j]), colwidth[j] + 2, " ")) :
            intcatcher[j] && ta.values[i,j] < 0 ?
            print_with_color(:red, io, rpad(iround(ta.values[i,j]), colwidth[j] + 2, " ")) :
            print_with_color(:blue, io, rpad(round(ta.values[i,j], 2), colwidth[j] + 2, " "))
        end
        println(io,"")
        end
        println(io,'\u22EE')
        for i in nrow-3:nrow
            print(io, ta.timestamp[i], " | ")
        for j in 1:ncol
            intcatcher[j] && ta.values[i,j] > 0?
            print_with_color(:green, io, rpad(iround(ta.values[i,j]), colwidth[j] + 2, " ")) :
            intcatcher[j] && ta.values[i,j] < 0 ?
            print_with_color(:red, io, rpad(iround(ta.values[i,j]), colwidth[j] + 2, " ")) :
            print_with_color(:blue, io, rpad(round(ta.values[i,j], 2), colwidth[j] + 2, " "))
        end
        println(io,"")
        end
    else
        for i in 1:nrow
            print(io, ta.timestamp[i], " | ")
        for j in 1:ncol
            intcatcher[j] ?
            print(io, rpad(iround(ta.values[i,j]), colwidth[j] + 2, " ")) :
            print(io, rpad(round(ta.values[i,j], 2), colwidth[j] + 2, " "))
        end
        println(io,"")
        end
    end
end

########################### Trade

# function show(io::IO, t::Trade)
#     print_with_color(:blue, io, @sprintf("%s:", string(t.start)))
#     print_with_color(:blue, io, @sprintf("%s ", string(t.finish)))
#     print_with_color(:yellow, io, @sprintf("%s ~ ", string(t.open)))
#     print_with_color(:yellow, io, @sprintf("%s ", string(t.close)))
#     if t.side == "long" && t.close - t.open> 0
#         print_with_color(:green, io, @sprintf("qty = "))
#     else
#         print_with_color(:red, io, @sprintf("qty = "))
#     end
#     print_with_color(:blue, io, @sprintf("%s", string(t.quantity)))
# end

# This code needs to be reduced by several fold, storing it here for now

###########################  Account

# function show(io::IO, ta::Account)
#   # variables
#   nrow          = size(ta.values, 1)
#   ncol          = size(ta.values, 2)
#   intcatcher    = falses(ncol)
#   for c in 1:ncol
#       rowcheck =  trunc(ta.values[:,c]) - ta.values[:,c] .== 0
#       if sum(rowcheck) == length(rowcheck)
#           intcatcher[c] = true
#       end
#   end
#   spacetime     = strwidth(string(ta.timestamp[1])) + 3
#   firstcolwidth = strwidth(ta.colnames[1])
#   colwidth      = Int[]
#       for m in 1:ncol
#           push!(colwidth, max(strwidth(ta.colnames[m]), strwidth(@sprintf("%.2f", maximum(ta.values[:,m])))))
#       end

#   # summary line
#   print(io,@sprintf("%dx%d %s %s to %s", nrow, ncol, typeof(ta), string(ta.timestamp[1]), string(ta.timestamp[nrow])))
#   println(io,"")
#   println(io,"")

#   # row label line
#    print_with_color(:magenta, io, ^(" ", spacetime), ta.colnames[1], ^(" ", colwidth[1] + 2 -firstcolwidth))

#    for p in 2:length(colwidth)
#      print_with_color(:magenta, io, ta.colnames[p], ^(" ", colwidth[p] - strwidth(ta.colnames[p]) + 2))
#    end
#    println(io,"")

#   # timestamp and values line
#     if nrow > 7
#         for i in 1:4
#             print(io, ta.timestamp[i], " | ")
#         for j in 1:ncol
#             intcatcher[j] && ta.values[i,j] > 0?
#             print_with_color(:green, io, rpad(iround(ta.values[i,j]), colwidth[j] + 2, " ")) :
#             intcatcher[j] && ta.values[i,j] < 0 ?
#             print_with_color(:red, io, rpad(iround(ta.values[i,j]), colwidth[j] + 2, " ")) :
#             print_with_color(:blue, io, rpad(round(ta.values[i,j], 2), colwidth[j] + 2, " "))
#         end
#         println(io,"")
#         end
#         println(io,'\u22EE')
#         for i in nrow-3:nrow
#             print(io, ta.timestamp[i], " | ")
#         for j in 1:ncol
#             intcatcher[j] && ta.values[i,j] > 0?
#             print_with_color(:green, io, rpad(iround(ta.values[i,j]), colwidth[j] + 2, " ")) :
#             intcatcher[j] && ta.values[i,j] < 0 ?
#             print_with_color(:red, io, rpad(iround(ta.values[i,j]), colwidth[j] + 2, " ")) :
#             print_with_color(:blue, io, rpad(round(ta.values[i,j], 2), colwidth[j] + 2, " "))
#         end
#         println(io,"")
#         end
#     else
#         for i in 1:nrow
#             print(io, ta.timestamp[i], " | ")
#         for j in 1:ncol
#             intcatcher[j] ?
#             print(io, rpad(iround(ta.values[i,j]), colwidth[j] + 2, " ")) :
#             print(io, rpad(round(ta.values[i,j], 2), colwidth[j] + 2, " "))
#         end
#         println(io,"")
#         end
#     end
# end

###########################  Portfolio

# function show(io::IO, ta::Portfolio)
#   # variables
#   nrow          = size(ta.values, 1)
#   ncol          = size(ta.values, 2)
#   intcatcher    = falses(ncol)
#   for c in 1:ncol
#       rowcheck =  trunc(ta.values[:,c]) - ta.values[:,c] .== 0
#       if sum(rowcheck) == length(rowcheck)
#           intcatcher[c] = true
#       end
#   end
#   spacetime     = strwidth(string(ta.timestamp[1])) + 3
#   firstcolwidth = strwidth(ta.colnames[1])
#   colwidth      = Int[]
#       for m in 1:ncol
#           push!(colwidth, max(strwidth(ta.colnames[m]), strwidth(@sprintf("%.2f", maximum(ta.values[:,m])))))
#       end

#   # summary line
#   print(io,@sprintf("%dx%d %s %s to %s", nrow, ncol, typeof(ta), string(ta.timestamp[1]), string(ta.timestamp[nrow])))
#   println(io,"")
#   println(io,"")

#   # row label line
#    print_with_color(:magenta, io, ^(" ", spacetime), ta.colnames[1], ^(" ", colwidth[1] + 2 -firstcolwidth))

#    for p in 2:length(colwidth)
#      print_with_color(:magenta, io, ta.colnames[p], ^(" ", colwidth[p] - strwidth(ta.colnames[p]) + 2))
#    end
#    println(io,"")

#   # timestamp and values line
#     if nrow > 7
#         for i in 1:4
#             print(io, ta.timestamp[i], " | ")
#         for j in 1:ncol
#             intcatcher[j] && ta.values[i,j] > 0?
#             print_with_color(:green, io, rpad(iround(ta.values[i,j]), colwidth[j] + 2, " ")) :
#             intcatcher[j] && ta.values[i,j] < 0 ?
#             print_with_color(:red, io, rpad(iround(ta.values[i,j]), colwidth[j] + 2, " ")) :
#             print_with_color(:blue, io, rpad(round(ta.values[i,j], 2), colwidth[j] + 2, " "))
#         end
#         println(io,"")
#         end
#         println(io,'\u22EE')
#         for i in nrow-3:nrow
#             print(io, ta.timestamp[i], " | ")
#         for j in 1:ncol
#             intcatcher[j] && ta.values[i,j] > 0?
#             print_with_color(:green, io, rpad(iround(ta.values[i,j]), colwidth[j] + 2, " ")) :
#             intcatcher[j] && ta.values[i,j] < 0 ?
#             print_with_color(:red, io, rpad(iround(ta.values[i,j]), colwidth[j] + 2, " ")) :
#             print_with_color(:blue, io, rpad(round(ta.values[i,j], 2), colwidth[j] + 2, " "))
#         end
#         println(io,"")
#         end
#     else
#         for i in 1:nrow
#             print(io, ta.timestamp[i], " | ")
#         for j in 1:ncol
#             intcatcher[j] ?
#             print(io, rpad(iround(ta.values[i,j]), colwidth[j] + 2, " ")) :
#             print(io, rpad(round(ta.values[i,j], 2), colwidth[j] + 2, " "))
#         end
#         println(io,"")
#         end
#     end
# end


"""
    RemainderModel(main::BalanceLaw, subcomponents::Tuple)

Compute the "remainder" contribution of the `main` model, after subtracting `subcomponents`.
"""
struct RemainderModel{M, S} <: BalanceLaw
    main::M
    subs::S
end

vars_state(rem::RemainderModel, FT) = vars_state(rem.main, FT)
vars_gradient(rem::RemainderModel, FT) = vars_gradient(rem.main, FT)
vars_diffusive(rem::RemainderModel, FT) = vars_diffusive(rem.main, FT)
vars_aux(rem::RemainderModel, FT) = vars_aux(rem.main, FT)
vars_integrals(rem::RemainderModel, FT) = vars_integrals(rem.main, FT)
vars_reverse_integrals(rem::RemainderModel, FT) = vars_integrals(rem.main, FT)
vars_gradient_laplacian(rem::RemainderModel, FT) =
    vars_gradient_laplacian(rem.main, FT)
vars_hyperdiffusive(rem::RemainderModel, FT) = vars_hyperdiffusive(rem.main, FT)

update_aux!(
    dg::DGModel,
    rem::RemainderModel,
    Q::MPIStateArray,
    t::Real,
    elems::UnitRange,
) = update_aux!(dg, rem.main, Q, t, elems)

integral_load_aux!(rem::RemainderModel, integ::Vars, state::Vars, aux::Vars) =
    integral_load_aux!(rem.main, integ, state, aux)

integral_set_aux!(rem::RemainderModel, aux::Vars, integ::Vars) =
    integral_set_aux!(rem.main, aux, integ)

reverse_integral_load_aux!(
    rem::RemainderModel,
    integ::Vars,
    state::Vars,
    aux::Vars,
) = reverse_integral_load_aux!(rem.main, integ, state, aux)

reverse_integral_set_aux!(rem::RemainderModel, aux::Vars, integ::Vars) =
    reverse_integral_set_aux!(rem.main, aux, integ)

function hyperdiffusive!(
    rem::RemainderModel,
    hyperdiffusive::Vars,
    hypertransform::Grad,
    state::Vars,
    aux::Vars,
    t::Real,
)
    hyperdiffusive!(rem.main, hyperdiffusive, hypertransform, state, aux, t)
end
function flux_diffusive!(
    rem::RemainderModel,
    flux::Grad,
    state::Vars,
    diffusive::Vars,
    hyperdiffusive::Vars,
    aux::Vars,
    t::Real,
)
    flux_diffusive!(rem.main, flux, state, diffusive, hyperdiffusive, aux, t)
end

gradvariables!(
    rem::RemainderModel,
    transform::Vars,
    state::Vars,
    aux::Vars,
    t::Real,
) = gradvariables!(rem.main, transform, state, aux, t)

diffusive!(
    rem::RemainderModel,
    diffusive::Vars,
    ∇transform::Grad,
    state::Vars,
    aux::Vars,
    t::Real,
) = diffusive!(rem.main, diffusive, ∇transform, state, aux, t)

function wavespeed(rem::RemainderModel, nM, state::Vars, aux::Vars, t::Real)
    ref = aux.ref_state
    return wavespeed(rem.main, nM, state, aux, t) -
           sum(sub -> wavespeed(sub, nM, state, aux, t), rem.subs)
end

boundary_state!(nf, rem::RemainderModel, x...) =
    boundary_state!(nf, rem.main, x...)
function normal_boundary_flux_diffusive!(
    nf,
    rem::RemainderModel,
    fluxᵀn::Vars{S},
    n⁻,
    state⁻,
    diff⁻,
    hyperdiff⁻,
    aux⁻,
    state⁺,
    diff⁺,
    hyperdiff⁺,
    aux⁺,
    bctype::Integer,
    t,
    args...,
) where {S}
    normal_boundary_flux_diffusive!(
        nf,
        rem.main,
        fluxᵀn,
        n⁻,
        state⁻,
        diff⁻,
        hyperdiff⁻,
        aux⁻,
        state⁺,
        diff⁺,
        hyperdiff⁺,
        aux⁺,
        bctype,
        t,
        args...,
    )
end

init_aux!(rem::RemainderModel, aux::Vars, geom::LocalGeometry) = nothing
init_state!(rem::RemainderModel, state::Vars, aux::Vars, coords, t) = nothing

function flux_nondiffusive!(
    rem::RemainderModel,
    flux::Grad,
    state::Vars,
    aux::Vars,
    t::Real,
)
    m = getfield(flux, :array)
    flux_nondiffusive!(rem.main, flux, state, aux, t)

    flux_s = similar(flux)
    m_s = getfield(flux_s, :array)

    for sub in rem.subs
        fill!(m_s, 0)
        flux_nondiffusive!(sub, flux_s, state, aux, t)
        m .-= m_s
    end
    nothing
end

function source!(
    rem::RemainderModel,
    source::Vars,
    state::Vars,
    diffusive::Vars,
    aux::Vars,
    t::Real,
    direction,
)
    m = getfield(source, :array)
    source!(rem.main, source, state, diffusive, aux, t, direction)

    source_s = similar(source)
    m_s = getfield(source_s, :array)

    for sub in rem.subs
        fill!(m_s, 0)
        source!(sub, source_s, state, diffusive, aux, t, direction)
        m .-= m_s
    end
    nothing
end


export fasta2matrix, letter2num, num2letter, matrix2fasta

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

"""
  function fasta2matrix(filename::AbstractString; max_gap_fraction = 1)

  Reads an aligned fasta file and returns a matrix with amino acids converted
  to numbers following the convention A --> 1, C --> 2, ... '-' --> 21.

    "filename": full path of the fasta file.
    "max_gap_fraction": to be removed.

"""

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

function fasta2matrix(filename::AbstractString; max_gap_fraction = 1)
    f = FastaReader(filename)

    max_gap_fraction = Float64(max_gap_fraction)

    # pass 1

    seqs = Int[]
    inds = Int[]
    fseqlen = 0

    for (name, seq) in f
        ngaps = 0
        if f.num_parsed == 1
            ls = length(seq)
            resize!(inds, ls)
            for i = 1:ls
                c = seq[i]
                if c != '.' && c == uppercase(c)
                    fseqlen += 1
                    inds[fseqlen] = i
                    c == '-' && (ngaps += 1)
                end
            end
        else
            ls = length(seq)
            ls == length(inds) || error("inputs are not aligned")
            tstfseqlen = 0
            for i = 1:ls
                c = seq[i]
                if c != '.' && c == uppercase(c)
                    tstfseqlen += 1
                    inds[tstfseqlen] == i || error("inconsistent inputs")
                    c == '-' && (ngaps += 1)
                end
            end
            tstfseqlen == fseqlen || error("inconsistent inputs")
        end
        ngaps / fseqlen <= max_gap_fraction && push!(seqs, f.num_parsed)
    end

    length(seqs) > 0 || error("Out of $(f.num_parsed) sequences, none passed the filter (max_gap_fraction=$max_gap_fraction)")

    # pass 2
    Z = Array{Int16}(undef, fseqlen, length(seqs))
    seqid = 1
    for (name, seq) in f
        seqs[end] < f.num_parsed && break
        seqs[seqid] == f.num_parsed || continue
        for i = 1:fseqlen
            c = seq[inds[i]]
            Z[i, seqid] = letter2num(c)
        end
        seqid += 1
    end
    @assert seqid == length(seqs) + 1

    close(f)

    size(Z)[2] == 1 && return Z[:, 1]
    return Int16.(Z')
end

##############################################################
"""
    letter2num(c::Union{Char,UInt8})

    Takes as input a char (representing an amino acid)
    and returns its conventional number representation.
    In this case with the convention "1 2 .. 21" <==> "A C .. -"
"""

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

let alphabet = [ 1, 21, 2, 3, 4, 5, 6, 7, 8, 21, 9, 10, 11, 12, 21, 13, 14, 15, 16, 17, 21, 18, 19, 21, 20]
               # A, B,  C, D, E, F, G, H, I, J,  K, L,  M,  N,  O,  P,  Q,  R,  S,  T,  U,  V,  W,  X,  Y
    global letter2num
    function letter2num(c::Union{Char,UInt8})
        i = UInt8(c) - 0x40
        1 <= i <= 25 && return alphabet[i]
        return 21
     end
end

##############################################################
"""
    num2letter(i::Integer)

    Takes as input an integer (representing an amino acid)
    and returns its conventional character representation.
    In this case with the convention "1 2 .. 21" <==> "A C .. -"
"""

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

let alphabet = ["A", "C", "D", "E", "F", "G", "H", "I",  "K", "L",  "M",  "N", "P",  "Q",  "R",
"S",  "T", "V",  "W",  "Y"]
    global num2letter
    function num2letter(i :: Integer)
        1 <= i <= 20 && return alphabet[i]
        return "-"
    end
end

##############################################################
"""
    vec2string(v::Array{:<Integer, 1})

    Takes as input a vector of integers (representing an amino acids)
    and returns a list of characters, the corresponding amino acids.
    In this case with the convention "1 2 .. 21" <==> "A C .. -"
"""

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

function vec2string(v::Array{<:Integer, 1})
    s = ""
    for i in v
        s = s*num2letter(i)
    end
    return s
end

##############################################################
"""
	string2vec(s::AbstractString)

    Takes as input a string (representing amino acids)
    and returns a vector of numbers, the corresponding number representation.
    In this case with the convention "1 2 .. 21" <==> "A C .. -"
"""

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

function string2vec(s::AbstractString)
    v = Vector{Int8}(undef, length(s))
    for (i, l) in enumerate(s)
        v[i] = letter2num(l)
    end
    return v
end

##############################################################
"""
	write_single_muts_MSA(path_wt::AbstractString, fitness_wt::Array{Union{Missing, Float64}},
			dir_out::AbstractString = ""; wt_name::AbstractString = "WT")
	write_single_muts_MSA(wt::Array{<:Integer, 1}, fitness_wt::Array{Union{Missing, Float64}},
			dir_out::AbstractString; wt_name::AbstractString = "WT")
	write_single_muts_MSA(path_wt::AbstractString, fitness_wt::Array{Any, 2},
			dir_out::AbstractString = ""; wt_name::AbstractString = "WT")
	write_single_muts_MSA(wt::Array{<:Integer, 1}, fitness_wt::Array{Any, 2},
			dir_out::AbstractString = ""; wt_name::AbstractString = "WT")

    USE:
    Given a sequence, it returns all single mutations in a MSA fashon.

    INPUT:
    "fitness_wt": if it is a 1d vector is must contain missings for missing mutations
    			  and should cover only the non gapped positions

    			  if it is 2d vector the first column must show the position, the second the mutation
    			  and the third the fitness value

    OUTPUT: writes a .fasta and a .fit file in the "dir_out" folder

"""

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

function write_single_muts_MSA(path_wt::AbstractString, fitness_wt::Array{Union{Missing, Float64}},
			dir_out::AbstractString = ""; wt_name::AbstractString = "WT")

	wt = fasta2matrix(path_wt)
	dir_out == "" && (  dir_out = (*).(split(path_wt, "/")[1:end-1])  )
	path_MSA_out = joinpath(dir_out, "single_muts_MSA_$(wt_name).fasta")
	path_fit_out = joinpath(dir_out, "single_muts_fitness_$(wt_name).fit")
	vec_fit = []
	k = 0
	FastaWriter(path_MSA_out, "w") do file
	    writeentry(file, "$k | $(wt_name)", vec2string(wt))
	    for i in [pos for (pos, amino) in enumerate(wt) if amino!= 21]
	        mutant = copy(wt)
	        for amino in 1:20
	            mutant[i] = amino
	            new_amino = num2letter(mutant[i])
	            old_amino = num2letter(wt[i])
	            if !ismissing(fitness_wt[(i-1)*20 + amino]) && (amino != wt[i])
	                k+=1
	                append!(vec_fit, fitness_wt[(i-1)*20 + amino])
	                writeentry(file, "$k | $(old_amino)$(i)$(new_amino)", vec2string(mutant))
	            end
	        end
	    end
	end

	pushfirst!(vec_fit, 0)
	writedlm(path_fit_out, vec_fit)
end

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

function write_single_muts_MSA(wt::Array{<:Integer, 1}, fitness_wt::Array{Union{Missing, Float64}},
			dir_out::AbstractString; wt_name::AbstractString = "WT")

	path_MSA_out = joinpath(dir_out, "single_muts_MSA_$(wt_name).fasta")
	path_fit_out = joinpath(dir_out, "single_muts_fitness_$(wt_name).fit")
	vec_fit = []
	k = 0
	FastaWriter(path_MSA_out, "w") do file
	    writeentry(file, "$k | $(wt_name)", vec2string(wt))
	    for i in [pos for (pos, amino) in enumerate(wt) if amino!= 21]
	        mutant = copy(wt)
	        for amino in 1:20
	            mutant[i] = amino
	            new_amino = num2letter(mutant[i])
	            old_amino = num2letter(wt[i])
	            if !ismissing(fitness_wt[(i-1)*20 + amino]) && (amino != wt[i])
	                k+=1
	                append!(vec_fit, fitness_wt[(i-1)*20 + amino])
	                writeentry(file, "$k | $(old_amino)$(i)$(new_amino)", vec2string(mutant))
	            end
	        end
	    end
	end

	pushfirst!(vec_fit, 0)
	writedlm(path_fit_out, vec_fit)
end

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

function write_single_muts_MSA(path_wt::AbstractString, fitness_wt::Array{Any, 2},
			dir_out::AbstractString = ""; wt_name::AbstractString = "WT")

	wt = fasta2matrix(path_wt)
	dir_out == "" && (  dir_out = (*).(split(path_wt, "/")[1:end-1])  )
	path_MSA_out = joinpath(dir_out, "single_muts_MSA_$(wt_name).fasta")
	path_fit_out = joinpath(dir_out, "single_muts_fitness_$(wt_name).fit")
	vec_fit = []

    l = 0
	FastaWriter(path_MSA_out, "w") do file
	    writeentry(file, "0 | $(wt_name)", vec2string(wt))
        for k in 1:size(fitness_wt, 1)

	        mutant = copy(wt)
            pos = fitness_wt[k, 1]

            old_amino_lettr = num2letter(wt[pos])
            new_amino_lettr = fitness_wt[k, 2][1]
            new_amino_num = letter2num(new_amino_lettr)

            mutant[pos] = new_amino_num

            if (new_amino_num != wt[pos])
                l +=1
                append!(vec_fit, fitness_wt[k, 3])
                writeentry(file, "$l | $(old_amino_lettr)$(pos)$(new_amino_lettr)", vec2string(mutant))
            end
	    end
	end

	pushfirst!(vec_fit, 0)
	writedlm(path_fit_out, vec_fit)
end

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

function write_single_muts_MSA(wt::Array{<:Integer, 1}, fitness_wt::Array{Any, 2},
			dir_out::AbstractString = ""; wt_name::AbstractString = "WT")

	dir_out == "" && (  dir_out = (*).(split(path_wt, "/")[1:end-1])  )
	path_MSA_out = joinpath(dir_out, "single_muts_MSA_$(wt_name).fasta")
	path_fit_out = joinpath(dir_out, "single_muts_fitness_$(wt_name).fit")
	vec_fit = []

    l = 0
	FastaWriter(path_MSA_out, "w") do file
	    writeentry(file, "0 | $(wt_name)", vec2string(wt))
        for k in 1:size(fitness_wt, 1)

	        mutant = copy(wt)
            pos = fitness_wt[k, 1]

            old_amino_lettr = num2letter(wt[pos])
            new_amino_lettr = fitness_wt[k, 2][1]
            new_amino_num = letter2num(new_amino_lettr)

            mutant[pos] = new_amino_num

            if (new_amino_num != wt[pos])
                l +=1
                append!(vec_fit, fitness_wt[k, 3])
                writeentry(file, "$l | $(old_amino_lettr)$(pos)$(new_amino_lettr)", vec2string(mutant))
            end
	    end
	end

	pushfirst!(vec_fit, 0)
	writedlm(path_fit_out, vec_fit)
end

##############################################################
"""
    function matrix2fasta(path::AbstractString, MSA, desc_name::AbstractString = "")

    (path, MSA) --> write MSA to path

    USE:
    Given a vector or matrix of numbers representing an amino acid, write it in letter format to "path".

    INPUT:
    "path": path of output
    "MSA": matrix or vector of numbers (1 to 21 in the convention used here)
    "desc_name": name of the sequence (works only if NOT an MSA)

    OUTPUT:

"""

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

function matrix2fasta(path::AbstractString, MSA, desc_name::AbstractString = "", start_index = 1)

    M = size(MSA, 1)

    single_seq_flag = (length(size(MSA)) == 1)

    if single_seq_flag == true
        FastaWriter(path, "w") do file
            desc_name == "" && (desc_name = "1")
            writeentry(file, desc_name, vec2string(MSA) )
        end
    else
	seq_index = start_index
        FastaWriter(path, "w") do file
            for i in 1:M
                seq = MSA[i, :]
                writeentry(file, "$(seq_index) $(desc_name)", vec2string(seq) )
		seq_index +=1
            end
        end
    end
end


"Single World Intervention Programs!"
module SWIP

import Distributions
import Cassette

# export ~, ⟂, ud

include("miniomega.jl")
include("intervene.jl")
# input("graph.jl")

end # module


using Test
using RadeonProRender, GeometryBasics, Colors
const RPR = RadeonProRender

context = RPR.Context(resource=RPR.RPR_CREATION_FLAGS_ENABLE_CPU)
scene = RPR.Scene(context)
matsys = RPR.MaterialSystem(context, 0)
camera = RPR.Camera(context)
lookat!(camera, Vec3f(8, 0, 5), Vec3f(2, 0, 2), Vec3f(0, 0, 1))
RPR.rprCameraSetFocalLength(camera, 45.0)
set!(scene, camera)
set!(context, scene)
env_light = RadeonProRender.EnvironmentLight(context)
image_path = RPR.assetpath("studio026.exr")
img = RPR.Image(context, image_path)
set!(env_light, img)
setintensityscale!(env_light, 1.1)
push!(scene, env_light)

light = RPR.PointLight(context)
transform!(light, Float32[1.0 0.0 0.0 2.0; 0.0 1.0 0.0 0.0; 0.0 0.0 1.0 5.0; 0.0 0.0 0.0 1.0])
f = 20
RPR.setradiantpower!(light, 500 / f, 641 / f, 630 / f)
push!(scene, light)

function add_shape!(scene, context, matsys, mesh; material=RPR.RPR_MATERIAL_NODE_DIFFUSE,
                    color=colorant"green", roughness=0.01)
    rpr_mesh = RadeonProRender.Shape(context, mesh)
    push!(scene, rpr_mesh)
    m = RPR.MaterialNode(matsys, material)
    set!(m, RPR.RPR_MATERIAL_INPUT_COLOR, color)
    set!(m, RPR.RPR_MATERIAL_INPUT_ROUGHNESS, roughness, roughness, roughness, roughness)
    set!(rpr_mesh, m)
    return rpr_mesh, m
end

mesh, mat = add_shape!(scene, context, matsys, Rect3f(Vec3f(-10, -10, -1), Vec3f(20, 20, 1));
                       color=colorant"white")
mesh, mat = add_shape!(scene, context, matsys, Rect3f(Vec3f(0, -10, 0), Vec3f(0.1, 20, 5));
                       color=colorant"white")
mesh, mat = add_shape!(scene, context, matsys, Rect3f(Vec3f(0, -2, 0), Vec3f(5, 0.1, 5));
                       color=colorant"white")
mesh, mat = add_shape!(scene, context, matsys, Rect3f(Vec3f(0, 2, 0), Vec3f(5, 0.1, 5));
                       color=colorant"white")

mesh, mat_sphere = add_shape!(scene, context, matsys, Tesselation(Sphere(Point3f(2, 0, 2), 1.0f0), 100);
                              material=RPR.RPR_MATERIAL_NODE_MICROFACET, roughness=0.2, color=colorant"red")
fb_size = (800, 600)
frame_buffer = RPR.FrameBuffer(context, RGBA, fb_size)
frame_bufferSolved = RPR.FrameBuffer(context, RGBA, fb_size)
set!(context, RPR.RPR_AOV_COLOR, frame_buffer)
set_standard_tonemapping!(context)
clear!(frame_buffer)
path = joinpath(@__DIR__, "test.png")
RPR.rprContextSetParameterByKey1u(context, RPR.RPR_CONTEXT_ITERATIONS, 1)
@time RPR.render(context)
RPR.rprContextResolveFrameBuffer(context, frame_buffer, frame_bufferSolved, false)
RPR.save(frame_bufferSolved, path)

@test isfile(path)


using Random

"""
    ContGridWorld

A Continuous grid world domain with generic dynamics to allow for different types of mazes.
    - state: [y, x]

# Args

"""
mutable struct ContGridWorld{F} <: AbstractEnvironment
    # State is stored as (y,x)
    state::Array{Float64, 1}
    walls::Array{Bool, 2}

    # goals stored as (y,x)
    goals::Array{Int, 2}
    reward_funcs::Dict{Int, F}

    max_action_noise::Float64
    drift_noise::Float64
    normalized::Bool
    collision::Bool

    edge_locs::Array{Array{Int64, 1}, 1}
    collision_check::Array{Bool, 1}
    new_state::Array{Float64, 1}
    ContGridWorld(walls, goals, rew_funcs::Dict{Int, F}, max_action_noise, drift_noise, normalized) where {F} =
        new{F}([0.0, 0.0], walls, goals, rew_funcs, max_action_noise, drift_noise, normalized, false, [[0,0] for i in 1:4], fill(false, 4), zeros(2))
end

module ContGridWorldParams
UP = 1
RIGHT = 2
DOWN = 3
LEFT = 4

ACTIONS = [UP, RIGHT, DOWN, LEFT]

const BODY_RADIUS = 0.2
const STEP = 0.5

const AGENT_BOUNDRIES = [
    [BODY_RADIUS, 0.0],
    [-BODY_RADIUS, 0.0],
    [0.0, BODY_RADIUS],
    [0.0, -BODY_RADIUS]
]
end

MinimalRLCore.is_terminal(env::ContGridWorld) = begin
    goal_id = which_goal(env)
    if goal_id ∈ keys(env.reward_funcs)
        is_terminal(env.reward_funcs[goal_id])
    else
        false
    end
end

MinimalRLCore.get_reward(env::ContGridWorld) = begin
    goal_id = which_goal(env)
    if goal_id ∈ keys(env.reward_funcs)
        get_reward(env.reward_funcs[goal_id])
    else
        0.0
    end
end
MinimalRLCore.get_state(env::ContGridWorld) = (env.normalized ? env.state./size(env) : env.state, env.collision)

random_state(env::ContGridWorld, rng) = [rand(rng)*size(env.walls)[1], rand(rng)*size(env.walls)[2]]
function random_start_state(env::ContGridWorld, rng)
    state = random_state(env, rng)
    while is_wall(env, state)
        state = random_state(env, rng)
    end
    return state
end
Base.size(env::ContGridWorld) = size(env.walls)
num_actions(env::ContGridWorld) = 4
get_states(env::ContGridWorld) = findall(x->x==false, env.walls)
MinimalRLCore.get_actions(env::ContGridWorld) = ContGridWorldParams.ACTIONS

valid_projected_state(env::ContGridWorld, prj) = begin
    gw_size = size(env.walls)
    prj[1] >= 1 && prj[2] >= 1 && prj[1] <= gw_size[1]  && prj[2] <= gw_size[2]
end

project(env::ContGridWorld, state) = [Int64(floor(state[1])) + 1, Int64(floor(state[2])) + 1]
project(env::ContGridWorld, state, loc) = begin; loc[1] = Int64(floor(state[1]) + 1); loc[2] = Int64(floor(state[2]) + 1); end;

function is_wall(env::ContGridWorld, state::Array{Float64, 1})
    prj = project(env, state)
    return is_wall(env, prj)
end

function is_wall(env::ContGridWorld, prj::Array{Int64, 1})
    gw_size = size(env.walls)
    if !valid_projected_state(env, prj)
        return false
    end
    env.walls[prj[1], prj[2]]
end

which_goal(env::ContGridWorld) = which_goal(env, env.state)

function which_goal(env::ContGridWorld, state::Array{Float64, 1})
    prj = project(env, state)
    if !valid_projected_state(env, prj)
        return -1
    end
    env.goals[prj[1], prj[2]]
end

function handle_collision(env::ContGridWorld, state, action)

    #Approximate by the square...
    CRP = ContGridWorldParams
    new_state = copy(state)
    # env.new_state .= state
    # new_state = copy(env.new_state)

    new_state[1] = clamp(new_state[1], CRP.BODY_RADIUS, size(env.walls)[1] - CRP.BODY_RADIUS)
    new_state[2] = clamp(new_state[2], CRP.BODY_RADIUS, size(env.walls)[2] - CRP.BODY_RADIUS)

    # Really basic collision detection for 2-d plane worlds.

    collided = new_state[1] != state[1] || new_state[2] != state[2]

    for i in 1:4
        project(env::ContGridWorld, new_state .+ CRP.AGENT_BOUNDRIES[i], env.edge_locs[i])
        env.collision_check[i] = is_wall(env, env.edge_locs[i])
    end

    collided = collided || any(env.collision_check)

    if env.collision_check[1] && env.collision_check[2]
        wall_piece = project(env, new_state)
        if action == CRP.DOWN
            new_state[1] = (wall_piece[1]) + CRP.BODY_RADIUS
        elseif action == CRP.UP
            new_state[1] = (wall_piece[1] - 1) - CRP.BODY_RADIUS
        end
    elseif env.collision_check[1]
        new_state[1] = (env.edge_locs[1][1] - 1) - CRP.BODY_RADIUS
    elseif env.collision_check[2]
        new_state[1] = (env.edge_locs[2][1]) + CRP.BODY_RADIUS
    end

    if env.collision_check[3] && env.collision_check[4]
        wall_piece = project(env, new_state)
        if action == CRP.LEFT
            new_state[2] = (wall_piece[2]) + CRP.BODY_RADIUS
        elseif action == CRP.RIGHT
            new_state[2] = (wall_piece[2] - 1) - CRP.BODY_RADIUS
        end
    elseif env.collision_check[3]
        new_state[2] = (env.edge_locs[3][2] - 1) - CRP.BODY_RADIUS
    elseif env.collision_check[4]
        new_state[2] = (env.edge_locs[4][2]) + CRP.BODY_RADIUS
    end

    new_state[1] = clamp(new_state[1], CRP.BODY_RADIUS, size(env.walls)[1] - CRP.BODY_RADIUS)
    new_state[2] = clamp(new_state[2], CRP.BODY_RADIUS, size(env.walls)[2] - CRP.BODY_RADIUS)

    return new_state, collided
end

function handle_collision!(env::ContGridWorld, action)
    env.state, collided = handle_collision(env, env.state, action)
    return collided
end

function MinimalRLCore.reset!(env::ContGridWorld, rng::Random.AbstractRNG=Random.GLOBAL_RNG)
    state = random_state(env, rng)

    while is_wall(env, state)
        state = random_state(env, rng)
    end
    env.state = state
    return state
end

function MinimalRLCore.reset!(env::ContGridWorld, state::AbstractArray)
    if is_wall(env, state)
        throw("Cannot reset environment to invalid state: $(state)")
    end

    env.state = state
    return state
end

function mini_step!(env::ContGridWorld, step, action)
    env.state .+= step
    collision = handle_collision!(env, action)
    return collision
end

function MinimalRLCore.environment_step!(env::ContGridWorld, action, rng=Random.GLOBAL_RNG)

    CRP = ContGridWorldParams
    next_step = zeros(2)

    if action == CRP.UP
        next_step[1] = (CRP.STEP - rand(rng)*env.max_action_noise)
        next_step[2] = randn(rng)*env.drift_noise
    elseif action == CRP.DOWN
        next_step[1] = -(CRP.STEP - rand(rng)*env.max_action_noise)
        next_step[2] = randn(rng)*env.drift_noise
    elseif action == CRP.RIGHT
        next_step[2] = CRP.STEP - rand(rng)*env.max_action_noise
        next_step[1] = randn(rng)*env.drift_noise
    elseif action == CRP.LEFT
        next_step[2] = -(CRP.STEP - rand(rng)*env.max_action_noise)
        next_step[1] = randn(rng)*env.drift_noise
    else
        throw("Wrong action")
    end

    # mini_physics simulation for 1 second (== 10 steps of 0.1 seconds)
    Δt=1.0
    τ=10
    next_step .*= Δt/τ
    collision = false
    for t in 1:τ
        collision = mini_step!(env, next_step, action)
        if collision
            break
        end
    end
    env.collision = collision
end

# function mini_step(env::ContGridWorld, state, step, action)

#     new_state = state .+ step
#     new_state, collision = handle_collision(env, new_state, action)
#     return new_state, collision
# end

# function _step(env::ContGridWorld, state, action; rng=Random.GLOBAL_RNG, kwargs...)

#     frp = FourRoomsContParams
#     next_step = zeros(2)

#     if action == frp.UP
#         next_step[1] = -(frp.STEP - rand(rng)*env.max_action_noise)
#         next_step[2] = randn(rng)*env.drift_noise
#     elseif action == frp.DOWN
#         next_step[1] = frp.STEP - rand(rng)*env.max_action_noise
#         next_step[2] = randn(rng)*env.drift_noise
#     elseif action == frp.RIGHT
#         next_step[2] = frp.STEP - rand(rng)*env.max_action_noise
#         next_step[1] = randn(rng)*env.drift_noise
#     elseif action == frp.LEFT
#         next_step[2] = -(frp.STEP - rand(rng)*env.max_action_noise)
#         next_step[1] = randn(rng)*env.drift_noise
#     else
#         throw("Wrong action")
#     end

#     Δt=1.0
#     τ=10
#     next_step .*= Δt/τ
#     collision = false
#     next_state = state[1]
#     for t in 1:τ
#         next_state, collision = mini_step(env, next_state, next_step, action)
#         if collision
#             break
#         end
#     end
#     # next_state = handle_collision(env, next_state, action)
#     return (next_state, collision), 0.0, false
# end


function setup_etc(backend::GPUBackend, layer::ChannelPoolingLayer, inputs, blobs)
  if isa(layer.pooling, Pooling.Max)
    masks = Array{CuPtr}(length(inputs))
    for i = 1:length(inputs)
      masks[i] = CUDA.cualloc(Csize_t, length(blobs[i]))
    end
    etc = masks
  elseif isa(layer.pooling, Pooling.Mean)
    integrals = Array{CuPtr}(length(inputs))
    for i = 1:length(inputs)
      integrals[i] = CUDA.cualloc(eltype(inputs[i]), prod(size(inputs[i])[1:end-1]))
    end
    etc = integrals
  else
    etc = nothing
  end
  return etc
end
function shutdown_etc(backend::GPUBackend, state::ChannelPoolingLayerState)
  if isa(state.layer.pooling, Pooling.Max)
    map(CUDA.free, state.etc)
  elseif isa(state.layer.pooling, Pooling.Mean)
    map(CUDA.free, state.etc)
  else
    error("Unknown pooling $(state.layer.pooling)")
  end
end

function forward(backend::GPUBackend, state::ChannelPoolingLayerState, inputs::Vector{Blob})
  forward(backend, state.layer.pooling, state, inputs)
end
function forward(backend::GPUBackend, pool::StdPoolingFunction,
    state::ChannelPoolingLayerState, inputs::Vector{Blob})

  for i = 1:length(inputs)
    input = inputs[i]
    output = state.blobs[i]

    if isa(pool, Pooling.Max)
      cuda_max_channel_pooling_forward(backend, input, output, state.etc[i], state.layer, state.op_dims[i])
    elseif isa(pool, Pooling.Mean)
      cuda_mean_channel_pooling_forward(backend, input, output, state.etc[i], state.layer, state.op_dims[i])
    else
      error("Pooling for $pool not implemented yet")
    end
  end
end

function backward(backend::GPUBackend, state::ChannelPoolingLayerState, inputs::Vector{Blob}, diffs::Vector{Blob})
  backward(backend, state.layer.pooling, state, inputs, diffs)
end

function backward(backend::GPUBackend, pool::StdPoolingFunction, state::ChannelPoolingLayerState,
    inputs::Vector{Blob}, diffs::Vector{Blob})

  for i = 1:length(inputs)
    diff = diffs[i]
    if !isa(diff, NullBlob)
      if isa(pool, Pooling.Max)
        cuda_max_channel_pooling_backward(backend, diff, state.blobs_diff[i], state.etc[i], state.layer, state.op_dims[i])
      elseif isa(pool, Pooling.Mean)
        cuda_mean_channel_pooling_backward(backend, diff, state.blobs_diff[i], state.layer, state.op_dims[i])
      else
        error("Pooling for $pool not implemented yet")
      end
    else
      continue # nothing to do if not propagating back
    end
  end
end

function cuda_mean_channel_pooling_forward{T}(backend::GPUBackend, input::CuTensorBlob{T},
    output::CuTensorBlob{T}, integral::CuPtr, layer, op_dim)

  spatial_dim_T, channels, num = split_dims(input, op_dim)
  pooled_chann = size(output, op_dim)
  one = convert(T, 1)
  neg_one = convert(T, -1)
  scale = convert(T, 1.0/layer.kernel)

  spatial_dim = spatial_dim_T * sizeof(T)
  fea_dim = spatial_dim * channels
  output_fea_dim = spatial_dim * pooled_chann

  for n = 1:num
    input_ptr = convert(Ptr{T}, input.ptr.p + fea_dim*(n-1))
    output_ptr = convert(Ptr{T}, output.ptr.p + output_fea_dim*(n-1))
    integral_ptr = convert(Ptr{T}, integral.p)

    # compute integral image
    CuBLAS.copy(backend.cublas_ctx, T, spatial_dim_T, input_ptr, 1, integral_ptr, 1)
    for c = 2:channels
      CuBLAS.copy(backend.cublas_ctx, T, spatial_dim_T, input_ptr + (c-1)*spatial_dim, 1,
          integral_ptr + (c-1)*spatial_dim, 1)
      CuBLAS.axpy(backend.cublas_ctx, spatial_dim_T, one, integral_ptr + (c-2)*spatial_dim, 1,
          integral_ptr + (c-1)*spatial_dim, 1)
    end

    for pc = 1:pooled_chann
      cstart = (pc-1)*layer.stride - layer.pad[1] + 1
      cend   = min(cstart + layer.kernel - 1, channels)
      cstart = max(1, cstart)

      output_ptr_pc = output_ptr + (pc-1)*spatial_dim

      CuBLAS.copy(backend.cublas_ctx, T, spatial_dim_T, integral_ptr + (cend-1)*spatial_dim, 1,
          output_ptr_pc, 1)
      if cstart > 1
        CuBLAS.axpy(backend.cublas_ctx, spatial_dim_T, neg_one,
            integral_ptr + (cstart-2)*spatial_dim, 1, output_ptr_pc, 1)
      end
      CuBLAS.scal(backend.cublas_ctx, spatial_dim_T, scale, output_ptr_pc, 1)
    end
  end
end

function cuda_mean_channel_pooling_backward{T}(backend::GPUBackend, input::CuTensorBlob{T},
    output::CuTensorBlob{T}, layer, op_dim)

  spatial_dim_T, channels, num = split_dims(input, op_dim)
  pooled_chann = size(output, op_dim)
  scale = 1/convert(T, layer.kernel)

  fill!(input, 0)

  spatial_dim = spatial_dim_T * sizeof(T)
  fea_dim = spatial_dim * channels
  output_fea_dim = spatial_dim * pooled_chann

  for n = 1:num
    input_ptr = convert(Ptr{T}, input.ptr.p) + fea_dim*(n-1)
    output_ptr = convert(Ptr{T}, output.ptr.p) + output_fea_dim*(n-1)

    for pc = 1:pooled_chann
      cstart = (pc-1)*layer.stride - layer.pad[1] + 1
      cend   = min(cstart + layer.kernel - 1, channels)
      cstart = max(1, cstart)
      output_ptr_pc = output_ptr + (pc-1)*spatial_dim

      for c = cstart:cend
        CuBLAS.axpy(backend.cublas_ctx, spatial_dim_T, scale, output_ptr_pc, 1,
            input_ptr + (c-1)*spatial_dim, 1)
      end
    end
  end
end

function cuda_geometry_max_chann_pool(sp_dim::Int, num::Int)
  x_block = round(Int, ceil(convert(Float64, num)/CUDA.THREADS_PER_BLOCK_X));
  y_block = 1;
  z_block = round(Int, ceil(convert(Float64, sp_dim)/CUDA.THREADS_PER_BLOCK_Z));
  return ((x_block,y_block,z_block),
          (CUDA.THREADS_PER_BLOCK_X,1,CUDA.THREADS_PER_BLOCK_Z))

end
function cuda_max_channel_pooling_forward{T}(backend::GPUBackend, input::CuTensorBlob{T},
    output::CuTensorBlob{T}, mask::CuPtr, layer, op_dim)

  sp_dim, channels, num = split_dims(input, op_dim)
  pooled_chann = size(output, op_dim)

  cuda_dim = cuda_geometry_max_chann_pool(sp_dim, num);
  if T == Float32
    kernel = backend.mocha.max_channel_pooling_forward_float
  elseif T == Float64
    kernel = backend.mocha.max_channel_pooling_forward_double
  else
    error("Unsupported data type for channel pooling: $T")
  end

  CUDA.launch(kernel, cuda_dim..., (input.ptr.p, output.ptr.p, mask.p, sp_dim, channels, num,
      pooled_chann, layer.kernel, layer.stride, layer.pad[1]))
end

function cuda_max_channel_pooling_backward{T}(backend::GPUBackend, input::CuTensorBlob{T},
    output::CuTensorBlob{T}, mask::CuPtr, layer, op_dim)

  sp_dim, channels, num = split_dims(input, op_dim)
  pooled_chann = size(output, op_dim)

  cuda_dim = cuda_geometry_max_chann_pool(sp_dim, num);
  if T == Float32
    kernel = backend.mocha.max_channel_pooling_backward_float
  elseif T == Float64
    kernel = backend.mocha.max_channel_pooling_backward_double
  else
    error("Unsupported data type for channel pooling: $T")
  end
  erase!(input)

  CUDA.launch(kernel, cuda_dim..., (input.ptr.p, output.ptr.p, mask.p, sp_dim, channels, num,
      pooled_chann, layer.kernel, layer.stride, layer.pad[1]))
end


function get_corrected_measurement(obj::KalmanEstimation)
    return jcall(obj, "getCorrectedMeasurement", EstimatedMeasurement, ())
end

function get_corrected_spacecraft_states(obj::KalmanEstimation)
    return jcall(obj, "getCorrectedSpacecraftStates", Vector{SpacecraftState}, ())
end

function get_current_date(obj::KalmanEstimation)
    return jcall(obj, "getCurrentDate", AbsoluteDate, ())
end

function get_current_measurement_number(obj::KalmanEstimation)
    return jcall(obj, "getCurrentMeasurementNumber", jint, ())
end

function get_estimated_measurements_parameters(obj::KalmanEstimation)
    return jcall(obj, "getEstimatedMeasurementsParameters", ParameterDriversList, ())
end

function get_estimated_orbital_parameters(obj::KalmanEstimation)
    return jcall(obj, "getEstimatedOrbitalParameters", ParameterDriversList, ())
end

function get_estimated_propagation_parameters(obj::KalmanEstimation)
    return jcall(obj, "getEstimatedPropagationParameters", ParameterDriversList, ())
end

function get_physical_estimated_covariance_matrix(obj::KalmanEstimation)
    return jcall(obj, "getPhysicalEstimatedCovarianceMatrix", RealMatrix, ())
end

function get_physical_estimated_state(obj::KalmanEstimation)
    return jcall(obj, "getPhysicalEstimatedState", RealVector, ())
end

function get_physical_innovation_covariance_matrix(obj::KalmanEstimation)
    return jcall(obj, "getPhysicalInnovationCovarianceMatrix", RealMatrix, ())
end

function get_physical_kalman_gain(obj::KalmanEstimation)
    return jcall(obj, "getPhysicalKalmanGain", RealMatrix, ())
end

function get_physical_measurement_jacobian(obj::KalmanEstimation)
    return jcall(obj, "getPhysicalMeasurementJacobian", RealMatrix, ())
end

function get_physical_state_transition_matrix(obj::KalmanEstimation)
    return jcall(obj, "getPhysicalStateTransitionMatrix", RealMatrix, ())
end

function get_predicted_measurement(obj::KalmanEstimation)
    return jcall(obj, "getPredictedMeasurement", EstimatedMeasurement, ())
end

function get_predicted_spacecraft_states(obj::KalmanEstimation)
    return jcall(obj, "getPredictedSpacecraftStates", Vector{SpacecraftState}, ())
end


module NavAbilitySDK

const NVA = NavAbilitySDK
export NVA

# Global imports
using DocStringExtensions
using LinearAlgebra
using JSON
using UUIDs
using Downloads
using HTTP

# for overloading with visualization helpers
import Base: show

# LinearAlgebra pass through exports
export diagm, norm
# UUIDs pass through exports
export uuid4

# Graphql
include("./navability/graphql/Factor.jl")
include("./navability/graphql/Status.jl")
include("./navability/graphql/Variable.jl")
include("./navability/graphql/DataBlobs.jl")

include("./navability/graphql/QueriesDeprecated.jl")
export QUERY_VARIABLE_LABELS
export QUERY_FILES, QUERY_CALIBRATION
export MUTATION_ADDVARIABLE, MUTATION_ADDFACTOR, MUTATION_ADDSESSIONDATA
export MUTATION_SOLVESESSION, MUTATION_SOLVEFEDERATED
export MUTATION_DEMOCANONICALHEXAGONAL
export MUTATION_CREATE_UPLOAD, MUTATION_ABORT_UPLOAD, MUTATION_COMPLETE_UPLOAD
export MUTATION_PROC_CALIBRATION
export SUBSCRIPTION_UPDATES

# Entities
include("./navability/entities/NavAbilityClient.jl")
include("./navability/entities/Client.jl")
include("./navability/entities/Common.jl")
include("./navability/entities/Distributions.jl")
include("./navability/entities/Variable.jl")
include("./navability/entities/InferenceTypes.jl")
include("./navability/entities/Factor.jl")
export NavAbilityClient, NavAbilityWebsocketClient, NavAbilityHttpsClient, QueryOptions, MutationOptions
export Client, Scope
export QueryDetail, LABEL, SKELETON, SUMMARY, FULL
export Distribution, Normal, Rayleigh, FullNormal, Uniform, Categorical
export ManifoldKernelDensity
export VariableType, Variable
export FactorData, PriorData, PriorPose2Data, PriorPoint2Data, LinearRelativeData, Pose2Pose2Data, Pose2AprilTag4CornersData, Pose2Point2BearingRangeData, Point2Point2RangeData, MixtureData
export ScatterAlignPose2Data
export FactorType, Factor

# Services
include("./navability/services/Variable.jl")
include("./navability/services/Factor.jl")
include("./navability/services/Solve.jl")
include("./navability/services/Status.jl")
include("./navability/services/Utils.jl")
include("./navability/services/StandardAPI.jl")
include("./navability/services/DataBlobs.jl")
export getVariable, getVariables, listVariables, ls
export addVariable, addPackedVariable
export getFactor, getFactors, listFactors, lsf
export addFactor, addPackedFactor
export solveSession, solveFederated
export getStatusMessages, getStatusLatest, getStatusesLatest
export waitForCompletion
export GraphVizApp, MapVizApp

end


module selection
importall common, species
using core, fitness, individual
export Selection, select, select_ids, SelectionDefinition

"""
The base type used by all selection operator definitions.
"""
abstract SelectionDefinition <: OperatorDefinition

"""
The base type used by all selection operators.
"""
abstract Selection <: Operator

"""
Selects a given number of individuals from a collection of candidate individuals
according to a given selection method, returning the selected individuals as a
new individual collection.
"""
select(s::Selection, buffer::IndividualCollection, sp::Species, ic::IndividualCollection, n::Integer) =
  write_to_individual_collection(buffer, ic, select_ids(s, sp.fitness, indexed_fitnesses(ic), n))

"""
Selects a specified number of individuals from a list of candidates, provided
as a vector of tuples, specifying the id and fitness of each candidate.

Returns the internal IDs of the selected individuals.
"""
select_ids{F}(s::Selection, ::Species, ::Vector{Tuple{Int, F}}, ::Integer) =
  error("Unimplemented `select_ids' method for this selection operator: $(typeof(s)).")

"""
Prepares a given selection operator for the breeding process.
"""
prepare(s::Selection, ic::IndividualCollection) = ic

# Load each of the selection operators.
include("selection/random.jl")
include("selection/tournament.jl")
end


for (func, typ) in [(:clblasChemv, CL_float2),
                    (:clblasZhemv, CL_double2)]

    @eval @blasfun $func(order::clblasOrder, uplo::clblasUplo,
                              N::Csize_t,
                              alpha::$typ,
                              A::cl.CL_mem, offa::Csize_t, lda::Csize_t,
                              X::cl.CL_mem, offx::Csize_t, incx::Cint,
                              beta::$typ,
                              Y::cl.CL_mem, offy::Csize_t, incy::Cint,
                              n_queues::cl.CL_uint,
                              queues::Ptr{cl.CL_command_queue},
                              n_events_in_wait_list::cl.CL_uint,
                              event_wait_list::Ptr{cl.CL_event},
                              events::Ptr{cl.CL_event})

    @eval @blasfun2 $func(order::clblasOrder, uplo::clblasUplo,
                               N::Csize_t,
                               alpha::$typ,
                               A::cl.CL_mem, offa::Csize_t, lda::Csize_t,
                               X::cl.CL_mem, offx::Csize_t, incx::Cint,
                               beta::$typ,
                               Y::cl.CL_mem, offy::Csize_t, incy::Cint)

end