module Vis
push!(LOAD_PATH,".")

using ImageView, GtkReactive, Colors

export vis, vis_col, vis2

function vis2(mri_,mask1, mask2, thresh=0.1)
  mri = Gray.(copy(mri_)/ maximum(mri_))
  mriseg1 = RGBA.(mri)
  mriseg1[mask1 .> thresh] += RGBA(0.,0.,1.,0.5)
  mriseg2 = RGBA.(mri)
  mriseg2[mask2 .> thresh] += RGBA(0.,0.,1.,0.5)
  zr, slicedata = roi(mri, (1,2))
  gd = imshow_gui((200, 200), slicedata, (1,2))
  imshow(gd["frame"][1,1], gd["canvas"][1,1], mriseg1, nothing, zr, slicedata)
  imshow(gd["frame"][1,2], gd["canvas"][1,2], mriseg2, nothing, zr, slicedata)
  showall(gd["window"])
end

function vis(mri_, mask, thresh=0.1)
  mri = Gray.(copy(mri_)/ maximum(mri_))
  mriseg = RGBA.(mri)
  mriseg[mask .> thresh] += RGBA(0.,0.,1.,0.5)
  zr, slicedata = roi(mri, (1,2))
  gd = imshow_gui((200, 200), slicedata, (1,2))
  imshow(gd["frame"][1,1], gd["canvas"][1,1], mri, nothing, zr, slicedata)
  imshow(gd["frame"][1,2], gd["canvas"][1,2], mriseg, nothing, zr, slicedata)
  showall(gd["window"])
end

function vis_col(mri_ ,mask,thresh = 0.1)
  mri = Gray.(copy(mri_)/ maximum(mri_))
  mriseg = RGB.(mri)
  cm = colormap("Oranges", 255)
  mriseg[mask .> thresh] = cm[Int.(ceil.(255*mask[mask .> thresh]))]
  zr, slicedata = roi(mri, (1,2))
  gd = imshow_gui((200, 200), slicedata, (1,2))
  imshow(gd["frame"][1,1], gd["canvas"][1,1], mri, nothing, zr, slicedata)
  imshow(gd["frame"][1,2], gd["canvas"][1,2], mriseg, nothing, zr, slicedata)
  showall(gd["window"])
end

end


const DEMCMC = DifferentialEvolutionMCMC

"""
    sample(model::DEModel, de::DE, n_iter::Int; progress=false, kwargs...)

Samples from the posterior distribution.

# Arguments

- `model`: a model containing likelihood function with data and priors
- `de`: differential evolution object
- `n_iter`: number of iterations or samples

# Keywords

- `progress=false`: show progress of sampler
- `kwargs...`: optional keyword arguments
"""
sample(model::DEModel, de::DE, n_iter::Int; progress=false, kwargs...) = _sample(model::DEModel, de::DE, n_iter::Int; progress, stepfun=step!, kwargs...)

function _sample(model::DEModel, de::DE, n_iter::Int; progress=false, stepfun=step!, kwargs...)
    meter = Progress(n_iter)
    # initialize particles based on prior distribution
    groups = sample_init(model, de, n_iter)
    for iter in 1:n_iter
        de.iter = iter + de.n_initial
        # explicitly pass groups so parallel works
        groups = stepfun(model, de, groups)
        progress ? next!(meter) : nothing
    end
    # convert to chain object
    chain = bundle_samples(model, de, groups, n_iter)
    return chain
end

"""
    sample(model::DEModel, de::DE, ::MCMCThreads, n_iter::Int; progress=false, kwargs...)

Samples from the posterior distribution with each group of particles on a seperarate thread for
the mutation and crossover steps.

# Arguments

- `model`: a model containing likelihood function with data and priors
- `de`: differential evolution object
- `MCMCThreads`: pass MCMCThreads() object to run on multiple threads
- `n_iter`: number of iterations or samples

# Keywords

- `progress=false`: show progress of sampler
- `kwargs...`: optional keyword arguments
"""
function sample(model::DEModel, de::DE, ::MCMCThreads, n_iter::Int; progress=false, kwargs...)
    _sample(model::DEModel, de::DE, n_iter::Int; progress, stepfun=pstep!, kwargs...)
end

"""
    step!(model::DEModel, de::DE, groups)

Perform a single step for DE-MCMC.

# Arguments

- `model`: model containing a likelihood function with data and priors
- `de`: DE-MCMC sampler object
- `groups`: Array of vectors of particles
"""
function step!(model::DEModel, de::DE, groups)
    rand() <= de.α ? migration!(de, groups) : nothing
    groups = update!(model, de, groups)
    store_samples!(de, groups)
    return groups
end

"""
    pstep!(model::DEModel, de::DE, groups)

Perform a single step for DE-MCMC with each particle group on a separate thread.

# Arguments

- `model`: model containing a likelihood function with data and priors
- `de`: DE-MCMC sampler object
- `groups`: Array of vectors of particles
"""
function pstep!(model::DEModel, de::DE, groups)
    rand() <= de.α ? migration!(de, groups) : nothing
    groups = p_update!(model, de, groups)
    store_samples!(de, groups)
    return groups
end

"""
    p_update!(model, de, groups)

Multithreaded update of particles. Particles are updated by block or simultaneously.

# Arguments

- `model`: model containing a likelihood function with data and priors
- `de`: differential evolution object
- `group`: a vector of interacting particles (e.g. chains)
"""
function p_update!(model, de, groups)
    seeds = rand(UInt, length(groups))
    if de.blocking_on(de)
        Threads.@threads for g in 1:de.n_groups
            group = block_update!(model, de, groups[g], seeds[g])
        end
        return groups
    else
        Threads.@threads for g in 1:de.n_groups
            group = mutate_or_crossover!(model, de, groups[g], seeds[g])
        end
        return groups
    end
end

"""
    update!(model, de, groups)(model, de, groups)

Particles are updated by block or simultaneously on a single thread.

# Arguments

- `model`: model containing a likelihood function with data and priors
- `de`: differential evolution object
- `group`: a vector of interacting particles (e.g. chains)
"""
function update!(model, de, groups)
    if de.blocking_on(de)
        return map(g -> block_update!(model, de, g), groups)
    else
        return map(g -> mutate_or_crossover!(model, de, g), groups)
    end
end

function block_update!(model, de, group, seed)
    Random.seed!(seed)
    return block_update!(model, de, group)
end

function block_update!(model, de, group)
    for block_idx in de.blocks
        mutate_or_crossover!(model, de, group, block_idx)
    end
    return group
end

"""
    mutate_or_crossover!(model, de, group, seed::Number)

Selects between mutation and crossover step with probability β.

# Arguments

- `model`: model containing a likelihood function with data and priors
- `de`: differential evolution object
- `group`: a vector of interacting particles (e.g. chains)
- `seed::Number`: RNG seed
"""
function mutate_or_crossover!(model, de, group, seed::Number)
    Random.seed!(seed)
    mutate_or_crossover!(model, de, group)
    return group
end

function mutate_or_crossover!(model, de, group)
    rand() <= de.β ? mutation!(model, de, group) : crossover!(model, de, group)
    return group
end

function mutate_or_crossover!(model, de, group, block_idx)
    rand() <= de.β ? mutation!(model, de, group) : crossover!(model, de, group, block_idx)
    return group
end

"""
    bundle_samples(model::DEModel, de::DE, groups, n_iter)

Converts group particles to a chain object capable of generating convergence diagnostics
and posterior summaries.

# Arguments

- `model`: model containing a likelihood function with data and priors
- `de`: differential evolution object
- `groups`: a vector of groups of particles
- `n_iter`: number of iterations
"""
function bundle_samples(model::DEModel, de::DE, groups, n_iter)
    @unpack burnin, discard_burnin = de
    particles = vcat(groups...)
    Np = length(particles)
    Ns = discard_burnin ? n_iter - burnin : n_iter
    offset = discard_burnin ? burnin : 0
    all_names = get_names(model, particles[1])
    n_parms = length(all_names)
    n_names = length(model.names)
    v = fill(0.0, Ns, n_parms, Np)
    for (c,p) in enumerate(particles)
        for s in 1:Ns
            sΔ = s + offset
            temp = Float64[]
            for ni in 1:n_names
                push!(temp, de.samples[sΔ,ni,c]...)
            end
            push!(temp, p.accept[sΔ], p.lp[sΔ])
            v[s,:,c] = temp'
        end
    end
    chains = Chains(v, all_names, (parameters = [model.names...],
        internals = ["acceptance", "lp"]))
    return chains
end

"""
    sample_init(model::DEModel, de::DE, n_iter)

Creates vectors of particles and samples initial parameter values from priors.

# Arguments

- `model`: model containing a likelihood function with data and priors
- `de`: differential evolution object
- `n_iter`: number of iterations
"""
function sample_init(model::DEModel, de::DE, n_iter)
    de.samples = initialize_samples(de, model, n_iter)
    id = 0
    groups = [[init_particle(de, model, id+=1, n_iter) for p in 1:de.Np]
        for c in 1:de.n_groups]
    return groups
end


"""
Contains all constraints and data to solve the global optimization optimization problem. Add linear constraints to a `GlobalModel` through `JuMP.@constraint(gm.model, ...)`, and nonlinear constraints using `add_nonlinear_constraint`.
"""
@with_kw mutable struct GlobalModel
    model::JuMP.Model                                            # Associated JuMP.Model
    name::String = "Model"                                       # Name
    bbls::Array{BlackBoxLearner} = BlackBoxLearner[]             # Constraints to be learned
    vars::Array{JuMP.VariableRef} = JuMP.all_variables(model)    # JuMP variables
    objective = JuMP.objective_function(model)                # Original objective function
    solution_history::DataFrame = DataFrame(string.(vars) .=> [Float64[] for i=1:length(vars)]) # Solution history
    cost::Array = []                                             # List of costs.
    soldict::Dict = Dict()                                       # For solution extraction
    params::Dict = gm_defaults()                                 # GM settings
end

function Base.show(io::IO, gm::GlobalModel)
    println(io, "GlobalModel " * gm.name * " with $(length(gm.vars)) variables: ")
    println(io, "Has $(length(gm.bbls)) BlackBoxLearners.")
    if get_param(gm, :ignore_feasibility)
        if get_param(gm, :ignore_accuracy)
            println(io, "Ignores training accuracy and data_feasibility thresholds.")
        else
            println(io, "Ignores data feasibility thresholds.")
        end
    else
        if get_param(gm, :ignore_accuracy)
            println("Ignores training accuracy thresholds.")
        end
    end
end

function set_param(gm::GlobalModel, key::Symbol, val)
    set_param(gm.params, key, val)
    for bbl in gm.bbls
        set_param(bbl.params, key, val, false)
    end
end

get_param(gm::GlobalModel, key::Symbol) = get_param(gm.params, key)

"""
    (gm::GlobalModel)(name::String)

Finds BlackBoxLearner in GlobalModel by name.
"""
function (gm::GlobalModel)(name::String)
    fn_names = getfield.(gm.bbls, :name)
    fns = gm.bbls[findall(x -> x == name, fn_names)]
    if length(fns) == 1
        return fns[1]
    elseif length(fns) == 0
        @warn("No constraint with name ", name)
        return
    else
        @warn("Multiple constraints with name ", name)
        return fns
    end
end

""" Finds active leaves of all constraints in GlobalModel."""
function active_leaves(gm::GlobalModel)
    for bbl in gm.bbls
        active_leaves(bbl)
        for ll in bbl.lls
            active_leaves(ll)
        end
    end
end

"""
    $(TYPEDSIGNATURES)

Extends JuMP.all_variables to GlobalModels and BlackBoxLearners, but only returns non-auxiliary variables
"""
JuMP.all_variables(bbo::Union{GlobalModel, BlackBoxLearner, LinkedLearner}) = bbo.vars

function JuMP.all_variables(bbls::Union{Array{BlackBoxLearner}, Array{LinkedLearner}})
    return unique(Iterators.flatten(([JuMP.all_variables(bbl) for bbl in bbls])))
end

""" Extends JuMP.set_optimizer to GlobalModels. """
function JuMP.set_optimizer(gm::GlobalModel, optimizer_factory)
    JuMP.set_optimizer(gm.model, optimizer_factory)
end

count_types(gm::GlobalModel) = count_types(gm.model)
count_variables(gm::GlobalModel) = count_variables(gm.model)
count_constraints(gm::GlobalModel) = count_constraints(gm.model)

"""
    $(TYPEDSIGNATURES)

Returns bounds of all variables.
"""
function get_bounds(model::Union{GlobalModel, JuMP.Model, BlackBoxLearner,
                                 Array{BlackBoxLearner}, LinkedLearner,
                                 Array{LinkedLearner}})
    return get_bounds(JuMP.all_variables(model))
end

""" Returns bounds of all variables, flattened in order of .vars attribute. """
function flattened_bounds(bbc::Union{GlobalModel, BlackBoxLearner, LinkedLearner})
    var_bounds = get_bounds(bbc)
    return [var_bounds[var] for var in bbc.vars]
end

"""
    $(TYPEDSIGNATURES)

Returns only unbounded variables.
"""
function get_unbounds(model::Union{GlobalModel, JuMP.Model, BlackBoxLearner,
                      Array{BlackBoxLearner}})
    return get_unbounds(JuMP.all_variables(model))
end

"""
    $(TYPEDSIGNATURES)

Takes on parsing and allocation of variables depending on user input.
"""
function determine_vars(gm::GlobalModel,
                        constraint::Union{JuMP.ConstraintRef, Expr};
                        vars::Union{Nothing, Array{JuMP.VariableRef, 1}} = nothing,
                        expr_vars::Union{Nothing, Array} = nothing)
    if constraint isa JuMP.ConstraintRef
        if isnothing(vars) && isnothing(expr_vars)
            vars = vars_from_constraint(constraint)
            return vars, vars
        elseif !isnothing(vars)
            return vars, vars
        end
    else
        if isnothing(vars) && isnothing(expr_vars)
            expr_vars = vars_from_expr(constraint, gm.model)
            vars = flat(expr_vars)
            return vars, expr_vars
        elseif isnothing(expr_vars)
            expr_vars = vars_from_expr(constraint, gm.model)
            @assert length(flat(expr_vars)) >= length(vars)
            return vars, expr_vars
        elseif isnothing(vars)
            vars = flat(expr_vars)
            return vars, expr_vars
        else
            @assert length(flat(expr_vars)) >= length(vars)
            return vars, expr_vars
        end
    end
end

"""
    add_nonlinear_constraint(gm::GlobalModel,
                     constraint::Union{JuMP.ConstraintRef, Expr};
                     vars::Union{Nothing, Array{JuMP.VariableRef, 1}} = nothing,
                     expr_vars::Union{Nothing, Array} = nothing,
                     dependent_var::Union{Nothing, JuMP.VariableRef} = nothing,
                     name::String = gm.name * " " * string(length(gm.bbls) + 1),
                     equality::Bool = false)

 Adds a new nonlinear constraint to Global Model. Standard method for adding BlackBoxClassifiers and BlackBoxRegressors.

 Note: If adding constraints via an `Expr`, please define the appropriate vector indices. Eg.
    `expr = :((x) -> -x[1:\$(N)]'*\$(Q)*x[1:\$(N)] - \$(c)'*x[1:\$(N)])`,
instead of
    `expr = :((x) -> -x'*\$(Q)*x - \$(c)'*x)`,
to avoid vectorization errors.
"""
function add_nonlinear_constraint(gm::GlobalModel,
                     constraint::Union{JuMP.ConstraintRef, Expr};
                     vars::Union{Nothing, Array{JuMP.VariableRef, 1}} = nothing,
                     expr_vars::Union{Nothing, Array} = nothing,
                     dependent_var::Union{Nothing, JuMP.VariableRef} = nothing,
                     name::String = "bbl" * string(length(gm.bbls) + 1),
                     equality::Bool = false)
    vars, expr_vars = determine_vars(gm, constraint, vars = vars, expr_vars = expr_vars)
    if constraint isa Expr
        if isnothing(dependent_var)
            new_bbl = BlackBoxClassifier(constraint = constraint, vars = vars, expr_vars = expr_vars,
                                         equality = equality, name = name)
            set_param(new_bbl, :n_samples, Int(ceil(get_param(gm, :sample_coeff)*sqrt(length(vars)))))
            push!(gm.bbls, new_bbl)
            return
        else
            new_bbl = BlackBoxRegressor(constraint = constraint, vars = vars, expr_vars = expr_vars,
                                        dependent_var = dependent_var, equality = equality, name = name)
            set_param(new_bbl, :n_samples, Int(ceil(get_param(gm, :sample_coeff)*sqrt(length(vars)))))
            push!(gm.bbls, new_bbl)
            return
        end
    elseif constraint isa JuMP.ConstraintRef
        !isnothing(dependent_var) && throw(OCTHaGOnException("Constraint " * name * " is of type $(string(typeof(constraint))) " *
                                                        "and cannot have a dependent variable " * string(dependent_var) * "."))
        new_bbl = BlackBoxClassifier(constraint = constraint, vars = vars, expr_vars = expr_vars,
                                        equality = equality, name = name)
        set_param(new_bbl, :n_samples, Int(ceil(get_param(gm, :sample_coeff)*sqrt(length(vars)))))
        push!(gm.bbls, new_bbl)
        JuMP.delete(gm.model, constraint)
    end
end

"""
    function add_nonlinear_or_compatible(gm::GlobalModel,
        constraint::Union{JuMP.ConstraintRef, Expr};
        vars::Union{Nothing, Array{JuMP.VariableRef, 1}} = nothing,
        expr_vars::Union{Nothing, Array} = nothing,
        dependent_var::Union{Nothing, JuMP.VariableRef} = nothing,
        name::String = gm.name * "_" * string(length(gm.bbls) + 1),
        equality::Bool = false)

Extends `add_nonlinear_constraint` to recognize JuMP compatible linear constraints and add them as linear JuMP constraints instead.
"""
function add_nonlinear_or_compatible(gm::GlobalModel,
                     constraint::Union{JuMP.ConstraintRef, Expr};
                     vars::Union{Nothing, Array{JuMP.VariableRef, 1}} = nothing,
                     expr_vars::Union{Nothing, Array} = nothing,
                     dependent_var::Union{Nothing, JuMP.VariableRef} = nothing,
                     name::String = gm.name * "_" * string(length(gm.bbls) + 1),
                     equality::Bool = false)
    vars, expr_vars = determine_vars(gm, constraint, vars = vars, expr_vars = expr_vars)
    fn = functionify(constraint)
    if fn isa Function
        try
            constr_expr = Base.invokelatest(fn, expr_vars...)
            if constr_expr isa JuMP.GenericAffExpr || get_param(gm, :convex_constrs)
                if equality && !isnothing(dependent_var)
                    @constraint(gm.model, dependent_var == constr_expr)
                elseif equality
                    @constraint(gm.model, constr_expr == 0)
                elseif !isnothing(dependent_var)
                    @constraint(gm.model, dependent_var >= constr_expr)
                else
                    @constraint(gm.model, constr_expr >= 0)
                end
            else
                throw(ErrorException())
            end
        catch
            add_nonlinear_constraint(gm, constraint, vars = vars, expr_vars = expr_vars,
                                     dependent_var = dependent_var,
                                     name = name, equality = equality)
        end
    else
        add_nonlinear_constraint(gm, constraint, vars = vars, expr_vars = expr_vars,
                                    dependent_var = dependent_var,
                                    name = name, equality = equality)
    end
end

"""
    add_linked_constraint(gm::GlobalModel, bbc::BlackBoxClassifier, linked_vars::Array{JuMP.Variable})
    add_linked_constraint(gm::GlobalModel, bbr::BlackBoxRegressor, linked_vars::Array{JuMP.Variable}, linked_dependent::JuMP.Variable)

Adds a linked constraint of the same structure as the BlackBoxLearner.
When a nonlinear constraint is repeated more than once, this function allows the underlying
approximator to be replicated without retraining trees for each constraint.
Note that the bounds used for sampling are for the original variables of the BlackBoxLearner, so be careful!
"""
function add_linked_constraint(gm::GlobalModel, bbc::BlackBoxClassifier, vars::Array{JuMP.VariableRef})
    length(vars) == length(bbc.vars) || throw(OCTHaGOnException("BBC $(bbc.name) does not" *
    " have the same number of variables as linked variables $(vars)."))
    if !isempty(bbc.mi_constraints)
        clear_tree_constraints!(gm, bbc)
        @info "Cleared constraints from BBC $(bbc.name) since it was relinked."
    end
    lc = LinkedClassifier(vars = vars, equality = bbc.equality)
    push!(bbc.lls, lc)
    return
end

function add_linked_constraint(gm::GlobalModel, bbr::BlackBoxRegressor, vars::Array{JuMP.VariableRef},
                                dependent_var::JuMP.VariableRef)
    length(vars) == length(bbr.vars) || throw(OCTHaGOnException("BBR $(bbr.name) does not" *
    " have the same number of variables as linked variables $(vars)"))
    if !isempty(bbr.mi_constraints)
        clear_tree_constraints!(gm, bbr)
        @info "Cleared constraints from BBR $(bbr.name) since it was relinked."
    end
    lr = LinkedRegressor(vars = vars, dependent_var = dependent_var,
                                   equality = bbr.equality)
    push!(bbr.lls, lr)
    return
end

"""
    nonlinearize!(gm::GlobalModel, bbls::Array{BlackBoxLearner})
    nonlinearize!(gm::GlobalModel)

Turns gm.model into the nonlinear representation of global optimization problem.
NOTE: to get back to MI-compatible forms, must regenerate GlobalModel from scratch.
"""
function nonlinearize!(gm::GlobalModel, bbls::Array{BlackBoxLearner})
    for (i, bbl) in enumerate(bbls)
        if bbl.constraint isa JuMP.ConstraintRef
            JuMP.add_constraint(gm.model, bbl.constraint)
        elseif bbl.constraint isa Expr
            symb = Symbol(bbl.name)
            vars = flat(bbl.expr_vars) # We want flattening of dense vars.
            var_ranges = get_var_ranges(bbl.expr_vars)
            expr = bbl.constraint
            flat_expr = :((x...) -> $(expr)([x[i] for i in $(var_ranges)]...))
            fn = eval(flat_expr)
            JuMP.register(gm.model, symb, length(vars), fn; autodiff = true)
            expr = Expr(:call, symb, vars...)
            if bbl.equality && bbl isa BlackBoxRegressor
                JuMP.add_NL_constraint(gm.model, :($(expr) == $(bbl.dependent_var)))
            elseif bbl.equality
                JuMP.add_NL_constraint(gm.model, :($(expr) == 0))
            elseif bbl isa BlackBoxRegressor
                JuMP.add_NL_constraint(gm.model, :($(expr) <= $(bbl.dependent_var)))
            else
                JuMP.add_NL_constraint(gm.model, :($(expr) >= 0))
            end
        end
    end
    return
end

function nonlinearize!(gm::GlobalModel)
    nonlinearize!(gm, gm.bbls)
end

function bound!(model::GlobalModel, bounds::Union{Pair,Dict})
    bound!(model.model, bounds)
end

""" Separates and returns linear and nonlinear constraints in a model. """
function classify_constraints(model::Union{GlobalModel, JuMP.Model})
    jump_model = model
    if model isa GlobalModel
        jump_model = model.model
    end
    all_types = list_of_constraint_types(jump_model)
    nl_constrs = [];
    l_constrs = [];
    l_vartypes = [JuMP.VariableRef, JuMP.GenericAffExpr{Float64, VariableRef}]
    l_constypes = [MOI.GreaterThan{Float64}, MOI.LessThan{Float64}, MOI.EqualTo{Float64}]
    for (vartype, constype) in all_types
        constrs_of_type = JuMP.all_constraints(jump_model, vartype, constype)
        if any(vartype .== l_vartypes) && any(constype .== l_constypes)
            append!(l_constrs, constrs_of_type)
        else
            append!(nl_constrs, constrs_of_type)
        end
    end
    if !isnothing(jump_model.nlp_data)
        append!(nl_constrs, jump_model.nlp_data.nlconstr)
    end
    return l_constrs, nl_constrs
end

""" Returns the feasibility of data points in a bbl or GM. """
feasibility(bbc::BlackBoxClassifier) = bbc.feas_ratio
feasibility(bbr::BlackBoxRegressor) = size(bbr.X, 1) / (size(bbr.X, 1) + size(bbr.infeas_X, 1))
feasibility(bbls::BlackBoxLearner) = feasibility.(bbls)
feasibility(gm::GlobalModel) = feasibility.(gm.bbls)

""" Returns the accuracy of learners in a bbl or GM. """
function evaluate_accuracy(bbc::BlackBoxClassifier)
    if bbc.convex
        return 1.
    elseif bbc.feas_ratio in [1., 0]
        @warn(string("Accuracy of BlackBoxClassifier ", bbc.name, " is tautological."))
        return 1.
    elseif isempty(bbc.learners)
        throw(OCTHaGOnException(string("BlackBoxClassifier ", bbc.name, " has not been trained yet.")))
    else
        return bbc.accuracies[end]
    end
end

function evaluate_accuracy(bbr::BlackBoxRegressor)
    if bbr.convex && !bbr.equality
        return 1.
    elseif isempty(bbr.learners)
        throw(OCTHaGOnException(string("BlackBoxRegressor ", bbr.name, " has not been trained yet.")))
    else
        return bbr.accuracies[active_lower_tree(bbr)]
    end
end

evaluate_accuracy(gm::GlobalModel) = evaluate_accuracy.(gm.bbls)

"""
    JuMP.optimize!(gm::GlobalModel; kwargs...)

Applies JuMP.optimize! to GlobalModels, and saves solution history.
"""
function JuMP.optimize!(gm::GlobalModel; kwargs...)
    JuMP.optimize!(gm.model, kwargs...) # Solves the optimization problem.
    append!(gm.solution_history, solution(gm), cols=:intersect) # Saves the solution.
    feas_gap(gm) # Computes the feasibility gaps of all constraints.
    push!(gm.cost, JuMP.getobjectivevalue(gm.model)) # Updates the final cost.
    active_leaves(gm) # Updates the active leaves of all approximations.
    gm.soldict = Dict(key => JuMP.getvalue.(gm.model[key]) for (key, value) in gm.model.obj_dict)
    return
end

"""
    solution(gm::GlobalModel)
    solution(m::JuMP.Model)

Returns the optimal solution of the GlobalModel/JuMP.Model in a DataFrame.
"""
function solution(gm::GlobalModel)
    vals = getvalue.(gm.vars)
    return DataFrame(vals', string.(gm.vars))
end

function solution(m::JuMP.Model)
    variables = JuMP.all_variables(m)
    vals = getvalue.(variables)
    return DataFrame(vals', string.(variables))
end

"""
    feas_gap(gm::GlobalModel, soln = solution(gm))

Evaluates relative feasibility gap of each nonlinear constraint at the given solution.
Negative values -> constraint violation for BBCs,
                    regression underestimation for BBRs.
Positive values -> constraint violation for BBC equalities,
                    regression overestimation for BBRs.
"""
function feas_gap(gm::GlobalModel, soln = solution(gm))
    for bbl in gm.bbls
        if bbl isa BlackBoxClassifier && !isnothing(bbl.constraint)
            for ll in bbl.lls # LL feas_gaps evaluated first, for descent function
                eval!(bbl, DataFrame(string.(bbl.vars) .=> values(soln[1, string.(ll.vars)])))
                push!(ll.feas_gap, bbl.Y[end] ./ (bbl.max_Y - bbl.min_Y))
            end
            eval!(bbl, soln)
            push!(bbl.feas_gap, bbl.Y[end] ./ (bbl.max_Y - bbl.min_Y))
        elseif bbl isa BlackBoxRegressor && !isnothing(bbl.constraint)
            for ll in bbl.lls # LL feas_gaps evaluated first, for descent function
                eval!(bbl, DataFrame(string.(bbl.vars) .=> values(soln[1, string.(ll.vars)])))
                optimum = soln[:, string(ll.dependent_var)][1]
                actual = bbl.Y[end]
                push!(ll.optima, optimum)
                push!(ll.actuals, actual)
                push!(ll.feas_gap, (optimum-actual) / ((bbl.max_Y - bbl.min_Y)))
            end
            eval!(bbl, soln)
            optimum = soln[:, string(bbl.dependent_var)][1]
            actual = bbl.Y[end]
            push!(bbl.optima, optimum)
            push!(bbl.actuals, actual)
            push!(bbl.feas_gap, (optimum-actual) / ((bbl.max_Y - bbl.min_Y)))
        elseif bbl isa BlackBoxClassifier && isnothing(bbl.constraint)
            for ll in bbl.lls # LL feas_gaps evaluated first, for descent function
                push!(ll.feas_gap, 0)
            end
            push!(bbl.feas_gap, 0) # data constraints are always feasible
        elseif bbl isa BlackBoxRegressor && isnothing(bbl.constraint)
            for ll in bbl.lls # LL feas_gaps evaluated first, for descent function
                optimum = soln[:, string(ll.dependent_var)][1]
                push!(ll.optima, optimum)
                push!(ll.feas_gap, 0) # data constraints are always feasible
            end
            optimum = soln[:, string(bbl.dependent_var)][1]
            push!(bbl.optima, optimum)
            push!(bbl.feas_gap, 0) # data constraints are always feasible
        end
    end
    return
end

""" Prints the last feasibility gap of each constraint. """
function print_feas_gaps(gm::GlobalModel)
    @info "Feasibility gaps:"
    for bbl in gm.bbls
        println("$(bbl.name): $(bbl.feas_gap[end])")
        for ll in bbl.lls
            println("Linked: $(ll.feas_gap[end])")
        end
    end
    return
end

""" Shows feasibility of last solution w.r.t. each approximated constraint. """
function is_feasible(bbl::Union{BlackBoxLearner, LinkedLearner}, tighttol = 1e-8)
    if bbl.equality
        return abs(bbl.feas_gap[end]) <= tighttol
    else
        return bbl.feas_gap[end] >= -tighttol
    end
end

""" Returns the feasibility of the GlobalModel. """
function is_feasible(gm::GlobalModel)
    tighttol = get_param(gm, :tighttol)
    for bbl in gm.bbls
        is_feasible(bbl, tighttol) || return false
        for ll in bbl.lls
            is_feasible(ll, tighttol) || return false
        end
    end
    return true
end

""" Checks whether a BBL or GM is sampled. """
is_sampled(bbl::BlackBoxLearner) = size(bbl.X,1) != 0

is_sampled(gm::GlobalModel) = all(is_sampled(bbl) for bbl in gm.bbls)

""" Clears all sampling, training and optimization data from GlobalModel. """
function clear_data!(gm::GlobalModel)
    clear_tree_constraints!(gm, gm.bbls)
    clear_data!.(gm.bbls)
    gm.solution_history = DataFrame(string.(gm.vars) .=> [Float64[] for i=1:length(gm.vars)])
end

update_vexity(gm::GlobalModel) = update_vexity.(gm.bbls)

""" Prints relevant details of a GlobalModel. """
function print_details(gm::GlobalModel)
    @info "GlobalModel $(gm.name) has:"
    n_vars = length(gm.vars)
    @info "$(n_vars) variables,"
    all_types = list_of_constraint_types(gm.model)
    l_vartypes = [JuMP.VariableRef, JuMP.GenericAffExpr{Float64, VariableRef}]
    l_constrs = []
    nl_constrs = []
    l_constypes = [MOI.GreaterThan{Float64}, MOI.LessThan{Float64}, MOI.EqualTo{Float64}]
    for (vartype, constype) in all_types
        constrs_of_type = JuMP.all_constraints(gm.model, vartype, constype)
        if any(vartype .== l_vartypes) && any(constype .== l_constypes)
            append!(l_constrs, constrs_of_type)
        else
            append!(nl_constrs, constrs_of_type)
        end
    end
    @info "$(length(l_constrs)) linear constraints."
    @info "$(length(nl_constrs)) nonlinear constraints."

    n_eqs = length(findall(x -> x.equality, gm.bbls))
    n_ineqs = length(gm.bbls) - n_eqs
    obj_bbl = filter(x -> x.dependent_var == gm.objective,
        [bbl for bbl in gm.bbls if bbl isa BlackBoxRegressor])
    @info "$(n_ineqs - length(obj_bbl)) black box inequalities,"
    @info "$(n_eqs) black box equalities,"
    n_lls = sum(length(bbl.lls) for bbl in gm.bbls)
    if n_lls > 0
        @info "$(n_lls) linked constraints,"
    end
    if isempty(obj_bbl)
        @info "And a linear objective."
    else
        @info "And a nonlinear objective."
    end
    return
end

""" Sets whether a BlackBoxLearner should be represented with big-M or big-M free constraints. """
function set_bigM(bbl::BlackBoxLearner, b::Bool)
    bbl.bigM = b
    return
end

set_bigM(gm::GlobalModel, b::Bool) = set_bigM.(gm.bbls, b)

"""
    $(TYPEDSIGNATURES)

Detects BlackBoxLearner with similar structure, for speeding up constraint learning.
"""
function detect_linked_constraints!(gm::GlobalModel)
    bounds = get_bounds(gm.bbls)
    vars = all_variables(gm.bbls)
    isnothing(check_bounds(bounds)) || throw(OCTHaGOnException("Not all variables in nonlinear constraints are bounded in $(gm.name)."))
    bins = Dict()
    queue = copy(gm.bbls)
    while !isempty(queue)
        bbl = popfirst!(queue)
        n_vars = length(bbl.vars)
        reg_bool = bbl isa BlackBoxRegressor
        similars = [q for q in queue if (n_vars == length(q.vars)
                        && reg_bool == (q isa BlackBoxRegressor) && bbl.equality == q.equality)]
        if !isempty(similars)
            samples = lh_sample(bbl, lh_iterations = 5, n_samples = 10)
            vals = evaluate(bbl, samples)
            act_sims = []
            for i = 1:length(similars)
                sim = similars[i]
                rename!(samples, string.(sim.vars))
                qs = evaluate(sim, samples)
                if all(isapprox.(vals, qs, atol = 1e-8))
                    push!(act_sims, sim)
                end
            end
            if !isempty(act_sims)
                queue = filter!(e -> !(e in act_sims), queue)
                bbl_bounds = flattened_bounds(bbl)
                ll_bounds = [flattened_bounds(sim) for sim in act_sims]
                outer_bounds = []
                for i = 1:length(bbl_bounds)
                    var_bound = [
                        minimum([minimum(bbl_bounds[i]),
                                [minimum(ll_bound[i]) for ll_bound in ll_bounds]...])
                        maximum([maximum(bbl_bounds[i]),
                                [maximum(ll_bound[i]) for ll_bound in ll_bounds]...])]
                    push!(outer_bounds, var_bound)
                end
                bins[bbl] = act_sims
            end
        else
            continue
        end
    end
    return bins
end

"""
    $(TYPEDSIGNATURES)

Merges detected linked constraints via LinkedLearners.
"""
function merge_linked_constraints!(gm::GlobalModel, bins::Dict)
    @info "Merging $(length(flat(values(bins)))) LinkedLearners into $(length(keys(bins))) BlackBoxLearners. "
    for (bbl, linked_bbls) in bins
        for linked_bbl in linked_bbls
            if linked_bbl isa BlackBoxClassifier
                add_linked_constraint(gm, bbl, linked_bbl.vars)
            else
                add_linked_constraint(gm, bbl, linked_bbl.vars, linked_bbl.dependent_var)
            end
        end
    end
    gm.bbls = filter!(x -> !(x in flat(values(bins))), gm.bbls)
    return
end


function read_data(name, dims, side)
    s = Tuple(side for _ in 1:dims)
    array = zeros(UInt8, s)

    open(name) do io
        read!(io, array)
    end

    return BitArray(array)
end

function test_rec(array)
    _, imp1 = phaserec(array; maxsteps = 1)
    _, imp2 = phaserec(array; maxsteps = 100)

    @test imp2 < imp1
end

@testset "Reconstruct ceramic" begin
    test_rec(read_data("ceramic200", 3, 200))
end

@testset "Reconstruct soil" begin
    test_rec(read_data("soil160", 3, 160))
end


export DDECPSolver

# Solver object to indicate a solver is not set
type UnsetSolver <: MathProgBase.AbstractMathProgSolver
end

# DD-ECP solver
type DDECPSolver <: MathProgBase.AbstractMathProgSolver
    lp_solver::MathProgBase.AbstractMathProgSolver   # LP solver
end

function DDECPSolver(;
    lp_solver = UnsetSolver(),
    )

    if lp_solver == UnsetSolver()
        error("LP solver not specified; set lp_solver")
    end
end


function TAME_LRT_increasing_cliques(testing=true)

    if testing

        n_sizes = [25]
        orders_to_test = [4]
        samples = 10^3
        trials = 2
        betas=[0.0]
        alphas=[.5]
        results_path = all_results_path*"RG_DupNoise/test/"

    else

        n_sizes = [100]
        orders_to_test = [3,4,5,6,7]
        samples = 10^4
        trials = 25
        betas=[0.0,1.0]
        alphas=[.5,1.0]
        results_path = all_results_path*"RG_DupNoise/"

    end
    step_percentage = [.2]
    edge_inclusion_p = .5

    if !isdir(results_path)
        mkpath(results_path)
    end

    LambdaTAME_results = []
    for order in orders_to_test
        X = distributed_random_trials(trials,DuplicationNoise(),true;postProcessing=noPostProcessing(),
                                        n_sizes=n_sizes,edge_inclusion_p,step_percentage,
                                        method=ΛTAME_MultiMotif_M(),graph=RandomGeometric(),degreedist=LogNormal(log(5),1),
                                        betas,profile=true,orders=[order],samples=samples)
        push!(LambdaTAME_results,(order,X))
    end
    filename=filter(x-> !isspace(x),"LambdaTAME_RandomGeometric_degreedist:LogNormal_alpha:$(alphas)_beta:$(betas)_n:$(n_sizes)_orders:$(orders_to_test)_p:$(edge_inclusion_p)_sample:$(samples)_sp:$(step_percentage)_trials:$(trials)_results.json")
    save_results(results_path*filename,LambdaTAME_results)

    LowRankTAME_results = []
    for order in orders_to_test
        X = distributed_random_trials(trials,DuplicationNoise(),true;postProcessing=noPostProcessing(),
                                      n_sizes=n_sizes,edge_inclusion_p,step_percentage,method=TAME_MultiMotif_M(),
                                      graph=RandomGeometric(),degreedist=LogNormal(log(5),1),betas,profile=true,
                                      orders=[order],samples=samples)
        push!(LowRankTAME_results,(order,X))
    end
    filename=filter(x-> !isspace(x),"LowRankTAME_RandomGeometric_degreedist:LogNormal_alpha:$(alphas)_beta:$(betas)_n:$(n_sizes)_orders:$(orders_to_test)_p:$(edge_inclusion_p)_sample:$(samples)_sp:$(step_percentage)_trials:$(trials)_results.json")
    save_results(results_path*filename,LowRankTAME_results)

    TAME_results = []
    for order in orders_to_test
        X = distributed_random_trials(trials,DuplicationNoise(),true;postProcessing=noPostProcessing(),
                                      n_sizes=n_sizes,edge_inclusion_p,step_percentage,
                                      method=LowRankTAME_MultiMotif_M(),graph=RandomGeometric(),degreedist=LogNormal(log(5),1),
                                      betas,profile=true,orders=[order],samples=samples)
        push!(TAME_results,(order,X))
    end
    filename=filter(x-> !isspace(x),"TAME_RandomGeometric_degreedist:LogNormal_alpha:$(alphas)_beta:$(betas)_n:$(n_sizes)_orders:$(orders_to_test)_p:$(edge_inclusion_p)_sample:$(samples)_sp:$(step_percentage)_trials:$(trials)_results.json")
    save_results(results_path*filename,TAME_results)

end

function increased_clique_size_duplication_noise_experiments(testing=true)

    if testing
        #  --  test params  --  #
        n_sizes = [500]
        step_percentage = [.1]
        edge_inclusion_p = [.5]
        trials = 1
        samples = 1000
        maxiter = 100
        betas = [1.0]
        alphas = [.5]
        orders_to_test = [3]
        k_vals = [5]
        save = true
        result_path = all_results_path*"RG_DupNoise/test/"

    else
        edge_inclusion_p = [.5]
        n_sizes = [500]
        step_percentage = [.25]
        trials = 25
        samples = 1000000
        maxiter = 1000

        orders_to_test = [2,3,4,5,6,7,8,9]
        #orders_to_test = [3,4,5]
        k_vals = [15,30,45,60,75,90]
        betas = [0.0,1.0,10.0,100.0]
        save = true
        result_path = all_results_path*"RG_DupNoise/"

    end

    if !isdir(result_path)
        mkpath(result_path)
    end

    LambdaTAME_Klau_results = []
    for order in orders_to_test
        for k in k_vals

            X = distributed_random_trials(trials,DuplicationNoise(),true;postProcessing=KlauAlgo(k=k,maxiter=maxiter,rtype=1),n_sizes=n_sizes,edge_inclusion_p=edge_inclusion_p,step_percentage=step_percentage,
                                                            method=ΛTAME_MultiMotif_M(),graph=RandomGeometric(),degreedist=LogNormal(log(5),1),motif=Clique(),matchingMethod=ΛTAME_GramMatching(),
                                                            orders=[order],samples=samples,betas,profile=true)
            push!(LambdaTAME_Klau_results,(order,k,X))
        end
    end
    filename = filter(x-> !isspace(x),"RandomGeometric_degreedist:LogNormal_KlauAlgokvals:$(k_vals)_noiseModel:Duplication_n:$(n_sizes)_KAmiter:$(maxiter)_orders:$(orders_to_test)_p:$(edge_inclusion_p)_samples:$(samples)_sp:$(step_percentage)_trials:$(trials)_data.json")
    save_results(result_path*filename,LambdaTAME_Klau_results)

    LambdaTAME_LocalSearch_results = []
    for order in orders_to_test
        for k in k_vals

            X = distributed_random_trials(trials,DuplicationNoise(),true;postProcessing=LocalSearch(k=k),n_sizes=n_sizes,edge_inclusion_p=edge_inclusion_p, step_percentage=step_percentage,
                                                            method=ΛTAME_MultiMotif_M(),graph=RandomGeometric(),degreedist=LogNormal(log(5),1),motif=Clique(),matchingMethod=ΛTAME_GramMatching(),
                                                            orders=[order],samples=samples,betas,profile=true)
            push!(LambdaTAME_LocalSearch_results,(order,k,X))
        end
    end
    filename = filter(x-> !isspace(x),"RandomGeometric_degreedist:LogNormal_LocalSearchkvals:$(k_vals)_noiseModel:Duplication_n:$(n_sizes)_orders:$(orders_to_test)_p:$(edge_inclusion_p)_samples:$(samples)_sp:$(step_percentage)_trials:$(trials)_data.json")
    save_results(result_path*filename,LambdaTAME_LocalSearch_results)

end


module MeasureSystems

#   This file is part of MeasureSystems.jl. It is licensed under the MIT license
#   MeasureSystems Copyright (C) 2021 Michael Reed

import Base: @pure
import UnitSystems
import UnitSystems: UnitSystem, Systems, Constants, Physics, Convert
import UnitSystems: Coupling, measure, unit, universe
export UnitSystems, Measure, measure, cache

# measure

using Measurements
struct Measure{N} end
const measure_cache = Measurement{Float64}[]
@pure measure(::Measure{N}) where N = measure_cache[N]
function cache(M::Measurement{Float64})
    N = findfirst(x->x==M,measure_cache)
    if isnothing(N)
        push!(measure_cache,M)
        N = Base.length(measure_cache)
    end
    return Measure{N}()
end
Base.show(io::IO,M::Measure{N}) where N = show(io,measure(M))

# unit systems

@pure mass(U::UnitSystem,S::UnitSystem) = electronmass(U,S)
@pure electronmass(𝘩,R∞) = αinv^2*R∞*2𝘩/𝘤
@pure electronmass(𝘩,R∞,C::Coupling) = inv(finestructure(C))^2*R∞*2𝘩/𝘤
@pure planckmass(U::UnitSystem,C::Coupling=universe(U)) = electronmass(U,C)/√coupling(C)
@pure planck(U::UnitSystem,C::Coupling=universe(U)) = 2π*planckreduced(U,C)
@pure newton(U::UnitSystem,C::Coupling=universe(U)) = lightspeed(U,C)*planckreduced(U,C)/planckmass(U,C)^2
@pure charge(U::UnitSystem,C::Coupling=universe(U)) = sqrt(2planck(U)/(permeability(U)/finestructure(U))/(lightspeed(U)*rationalization(U)*lorentz(U)^2))

for unit ∈ (:coupling,:finestructure,:electronunit,:protonunit,:protonelectron)
    @eval @pure $unit(C::Coupling) = UnitSystems.$unit(C)
    @eval @pure $unit(U::UnitSystem) = UnitSystems.$unit(universe(U))
end
for unit ∈ (:boltzmann,:planckreduced,:lightspeed,:permeability,:electronmass,:rationalization,:lorentz,:molarmass)
    @eval @pure $unit(U::UnitSystem) = UnitSystems.$unit(U)
end
for unit ∈ (:boltzmann,:planckreduced,:lightspeed,:permeability,:electronmass,:molarmass)
    @eval @pure $unit(U::UnitSystem,C::Coupling) = $unit(U)
end
for unit ∈ Constants
    @eval @pure $unit(U::UnitSystem,S::UnitSystem) = unit($unit(S)/$unit(U))
end
for unit ∈ Convert
    @eval begin
        @pure @inline $unit(v::Real,U::UnitSystem) = $unit(v,U,Metric)
        @pure @inline $unit(v::Real,U::UnitSystem,S::UnitSystem) = (u=$unit(U,S);isone(u) ? v : v/u)
        @pure @inline $unit(v::Real,U::UnitSystem{kB,ħ,𝘤,μ₀,mₑ},S::UnitSystem{kB,ħ,𝘤,μ₀,mₑ}) where {kB,ħ,𝘤,μ₀,mₑ} = v
        @pure @inline $unit(U::UnitSystem{kB,ħ,𝘤,μ₀,mₑ},S::UnitSystem{kB,ħ,𝘤,μ₀,mₑ}) where {kB,ħ,𝘤,μ₀,mₑ} = 1
    end
    if unit ∉ (Constants...,:permittivity,:charge,:magneticflux,:impedance,:conductance)
        @eval @pure @inline $unit(U::UnitSystem) = $unit(U,Metric)
    end
end
for unit ∈ (Systems...,Constants...,Physics...,Convert...)
    unit ∉ (:length,:time) && @eval export $unit
end

# common conversion factors

import UnitSystems: atm,g₀,lbm,slug,ft,ftUS,rankine,kelvin
for CAL ∈ (:cal,:calₜₕ,:cal₄,:cal₁₀,:cal₂₀,:calₘ,:calᵢₜ)
    KCAL = Symbol(:k,CAL)
    @eval import UnitSystems: $CAL, $KCAL
end

# fundamental constants, αinv = (34259-1/4366.8123)/250 # 137.036 exactly?

const mP = measurement("0.00000002176434(24)")
const αinv = measurement("137.035999084(21)")
const μₑᵤ = 1/measurement("1822.888486209(53)")
const μₚᵤ = measurement("1.007276466621(53)")
const R∞ = measurement("10973731.5681601(210)")
const KJ2014 = measurement("483597.8525(30)")*1e9
const RK2014 = measurement("25812.8074555(59)")

import UnitSystems: ΔνCs,Kcd,NA,kB,𝘩,𝘤,𝘦,ħ,Rᵤ,αL,RK1990,KJ1990,ħ1990
const mₑ,μ₀ = electronmass(𝘩,R∞),2𝘩/𝘤/αinv/𝘦^2 # ≈ 4π*(1e-7+5.5e-17), exact charge
const δμ₀,μₚₑ,αG,ħ2014 = μ₀-4π*1e-7,μₚᵤ/μₑᵤ,(mₑ/mP)^2,2/RK2014/KJ2014^2/π
const mₑ1990,mₑ2014 = electronmass(2π*ħ1990,R∞),electronmass(2π*ħ2014,R∞)

# engineering units # Thomson: αL = 1/2

const Universe = Coupling{cache(αG),cache(1/αinv),cache(μₑᵤ),cache(μₚᵤ)}()
const Gauss = UnitSystem{cache(1e10*Rᵤ*mₑ/μₑᵤ),1e7*ħ,100𝘤,1,cache(1000mₑ),4π,0.01/𝘤}()
const LorentzHeaviside = UnitSystem{cache(1e10*Rᵤ*mₑ/μₑᵤ),1e7*ħ,100𝘤,1,cache(1000mₑ),1,0.01/𝘤}()
const Thomson = UnitSystem{cache(1e10*Rᵤ*mₑ/μₑᵤ),1e7*ħ,100𝘤,1,cache(1000mₑ),4π,1/2}()
const Kennelly = UnitSystem{cache(Rᵤ*mₑ/μₑᵤ/0.001),ħ,𝘤,1e-7,cache(mₑ),4π}() # ?
const ESU = UnitSystem{cache(1e10*Rᵤ*mₑ/μₑᵤ),1e7*ħ,100𝘤,(100𝘤)^-2,cache(1000mₑ),4π}()
const ESU2019 = UnitSystem{1e7*kB,1e7*ħ,100𝘤,cache(1e3*μ₀/𝘤^2),cache(1000mₑ)}()
const EMU = UnitSystem{cache(1e10*Rᵤ*mₑ/μₑᵤ),1e7*ħ,100𝘤,1,cache(1000mₑ),4π}()
const EMU2019 = UnitSystem{1e7*kB,1e7*ħ,100𝘤,cache(1e7*μ₀),cache(1000mₑ)}()
const MTS = UnitSystem{cache(1e6*Rᵤ*mₑ/μₑᵤ),1000ħ,𝘤,4π/1e4,cache(mₑ/1000)}()
const Mixed = UnitSystem{cache(Rᵤ*mₑ/μₑᵤ/0.001),ħ,𝘤,cache(μ₀),cache(mₑ)}()
const Metric = UnitSystem{cache(Rᵤ*mₑ/μₑᵤ/0.001),ħ,𝘤,4π*1e-7,cache(mₑ)}()
const SI1976 = UnitSystem{cache(8.31432mₑ/μₑᵤ/0.001),ħ,𝘤,4π*1e-7,cache(mₑ)}()
const SI2019 = UnitSystem{kB,ħ,𝘤,cache(μ₀),cache(mₑ)}()
const CODATA = UnitSystem{cache(Rᵤ*mₑ2014/μₑᵤ/0.001),cache(ħ2014),𝘤,cache(2RK2014/𝘤/αinv),cache(mₑ2014)}()
const Conventional = UnitSystem{cache(Rᵤ*mₑ1990/μₑᵤ/0.001),ħ1990,𝘤,cache(2RK1990/𝘤/αinv),cache(mₑ1990)}()
const English = UnitSystem{kB*rankine/slug/ft^2,ħ/slug/ft^2,𝘤/ft,4π,cache(mₑ/slug)}()
const EnglishUS = UnitSystem{cache(1000Rᵤ*mₑ/μₑᵤ*rankine/slug/ftUS^2),ħ/slug/ftUS^2,𝘤/ftUS,4π,cache(mₑ/slug)}()

# astronomical units

const GMsun = measurement("1.32712440018(9)")*1e20
const GMearth = measurement("3.986004418(8)")*1e14
const GMjupiter = measurement("1.26686534(9)")*1e17

import UnitSystems: au,LD,day,pc,ly
const GG = 𝘤*ħ/mP^2; const mₛ = GMsun/GG; const Jₛ = mₛ*au^2/day^2; export mₛ,Jₛ,au,day
const IAU = UnitSystem{cache(Rᵤ*mₑ/μₑᵤ/0.001/Jₛ),cache(ħ/day/Jₛ),day*𝘤/au,cache(4π*1e-7*day^2/Jₛ),cache(mₑ/mₛ)}()

# aliased & humorous units

const mf = mass(90/lbm,Metric,English); const Jf = mf*(201.168/14day)^2
const FFF = UnitSystem{cache(1000Rᵤ*mₑ/μₑᵤ*rankine/Jf),cache(ħ/14day/Jf),14day*𝘤/201.168,0,cache(mₑ/mf)}()
const units, US, SI, MKS, temp = UnitSystem, UnitSystem, SI2019, Metric, temperature
const CGS, CGS2019, CGSm, CGSe, HLU = Gauss, EMU2019, EMU, ESU, LorentzHeaviside

# natural units

const Planck = UnitSystem{1,1,1,1,cache(√(4π*αG))}()
const PlanckGauss = UnitSystem{1,1,1,4π,cache(√αG)}()
const Stoney = UnitSystem{1,cache(αinv),1,4π,cache(√(αG*αinv))}()
const Hartree = UnitSystem{1,1,cache(αinv),cache(4π/αinv^2),1}()
const Rydberg = UnitSystem{1,1,cache(2αinv),cache(π/αinv^2),1/2}()
const Schrodinger = UnitSystem{1,1,cache(αinv),cache(4π/αinv^2),cache(√(αG*αinv))}()
const Electronic = UnitSystem{1,cache(αinv),1,4π,1}()
const Natural = UnitSystem{1,1,1,1,1}()
const NaturalGauss = UnitSystem{1,1,1,4π,1}()
const QCD = UnitSystem{1,1,1,1,cache(1/μₚₑ)}()
const QCDGauss = UnitSystem{1,1,1,4π,cache(1/μₚₑ)}()
const QCDoriginal = UnitSystem{1,1,1,cache(4π/αinv),cache(1/μₚₑ)}()

# physical constants

@pure electronmass(U::typeof(Planck),C::Coupling) = sqrt(4π*coupling(C))
@pure electronmass(U::typeof(PlanckGauss),C::Coupling) = sqrt(coupling(C))
@pure electronmass(U::UnitSystem{kB,ħ,𝘤,μ₀,cache(√(αG*αinv))},C::Coupling) where {kB,ħ,𝘤,μ₀} = sqrt(coupling(C)/finestructure(C))
@pure electronmass(U::UnitSystem{kB,ħ,𝘤,μ₀,cache(1/μₚₑ)},C::Coupling) where {kB,ħ,𝘤,μ₀} = 1/protonelectron(C)
@pure permeability(U::UnitSystem{kB,ħ,𝘤,cache(4π/αinv^2)},C::Coupling) where {kB,ħ,𝘤} = 4π*finestructure(C)^2
@pure permeability(U::UnitSystem{kB,ħ,𝘤,cache(π/αinv^2)},C::Coupling) where {kB,ħ,𝘤} = π*finestructure(C)^2
@pure lightspeed(U::UnitSystem{kB,ħ,cache(αinv)},C::Coupling) where {kB,ħ} = 1/finestructure(C)
@pure lightspeed(U::UnitSystem{kB,ħ,cache(2αinv)},C::Coupling) where {kB,ħ} = 2/finestructure(C)
@pure planckreduced(U::UnitSystem{kB,cache(αinv)},C::Coupling) where kB = 1/finestructure(C)

@pure electronmass(U::UnitSystem{kB,ħ,𝘤,μ₀,cache(mₑ)},C::Coupling) where {kB,ħ,μ₀} = electronmass(planck(U),R∞,C)
@pure electronmass(U::UnitSystem{kB,ħ,100𝘤,μ₀,cache(1000mₑ)},C::Coupling) where {kB,ħ,μ₀} = 1000electronmass(SI,C)
@pure electronmass(U::UnitSystem{kB,ħ,𝘤,μ₀,cache(mₑ/1000)},C::Coupling) where {kB,ħ,μ₀} = electronmass(SI,C)/1000
@pure electronmass(U::UnitSystem{kB,ħ,𝘤,μ₀,cache(mₑ2014)},C::Coupling) where {kB,ħ,μ₀} = electronmass(planck(U),R∞,C)
@pure electronmass(U::UnitSystem{kB,ħ,𝘤,μ₀,cache(mₑ1990)},C::Coupling) where {kB,ħ,μ₀} = electronmass(planck(U),R∞,C)
@pure electronmass(U::UnitSystem{kB,ħ,𝘤/ftUS,μ₀,cache(mₑ/slug)},C::Coupling) where {kB,ħ,μ₀} = electronmass(SI,C)/slug
@pure permeability(U::UnitSystem{kB,ħ,𝘤,cache(μ₀)},C::Coupling) where {kB,ħ,𝘤} = finestructure(C)*2𝘩/𝘤/𝘦^2
@pure permeability(U::typeof(ESU2019),C::Coupling) = 1e3*permeability(SI,C)/𝘤^2
@pure permeability(U::typeof(EMU2019),C::Coupling) = 1e7*permeability(SI,C)
@pure permeability(U::typeof(CODATA),C::Coupling) = 2RK2014*finestructure(C)/𝘤
@pure permeability(U::typeof(Conventional),C::Coupling) = 2RK1990*finestructure(C)/𝘤

@pure molarmass(U::UnitSystem{kB},C::Coupling=universe(U)) = NA*electronmass(U,C)/electronunit(C)
@pure molarmass(U::UnitSystem{1e7*kB},C::Coupling=universe(U)) = 1000molarmass(SI2019,C)
@pure molarmass(U::UnitSystem{1e3*kB},C::Coupling=universe(U)) = molarmass(SI2019,C)/1000
@pure molarmass(U::UnitSystem{cache(boltzmann(MTS))}) = molarmass(CGS)/1e6
@pure molarmass(U::UnitSystem{cache(boltzmann(CGS))}) = molarmass(Natural)
@pure molarmass(U::UnitSystem{cache(boltzmann(FFF))}) = molarmass(Natural)
@pure molarmass(U::UnitSystem{boltzmann(English)},C::Coupling=universe(U)) = 1000molarmass(SI2019,C)
@pure molarmass(U::UnitSystem{cache(boltzmann(EnglishUS))}) = molarmass(Natural)
@pure molarmass(U::UnitSystem{cache(boltzmann(IAU))}) = 1/1000mₛ

@pure luminousefficacy(U::UnitSystem{1}) = 1
@pure luminousefficacy(U::UnitSystem) = power(Kcd,SI2019,U)

@pure universe(::typeof(FFF)) = Universe
for unit ∈ UnitSystems.Systems
    if unit ∉ (:Natural,:NaturalGauss)
        @eval universe(::typeof($unit)) = Universe
    end
end

const dir = dirname(pathof(UnitSystems))
include("$dir/kinematic.jl")
include("$dir/electromagnetic.jl")
include("$dir/thermodynamic.jl")
include("$dir/physics.jl")
include("$dir/systems.jl")

end # module


#####
##### Tendency specification
#####

#####
##### Sources
#####

eq_tends(pv::AbstractPrognosticVariable, m::AtmosModel, tt::Source) =
    (m.source[pv]..., eq_tends(pv, turbconv_model(m), tt)...)

#####
##### First order fluxes
#####

eq_tends(::Mass, ::Anelastic1D, ::Flux{FirstOrder}) = ()

eq_tends(::Mass, ::Compressible, ::Flux{FirstOrder}) = (Advect(),)

# Mass
eq_tends(pv::Mass, atmos::AtmosModel, tt::Flux{FirstOrder}) =
    (eq_tends(pv, compressibility_model(atmos), tt))

# Momentum
eq_tends(pv::Momentum, ::Compressible, ::Flux{FirstOrder}) =
    (PressureGradient(),)

eq_tends(pv::Momentum, ::Anelastic1D, ::Flux{FirstOrder}) = ()

eq_tends(pv::Momentum, m::AtmosModel, tt::Flux{FirstOrder}) =
    (Advect(), eq_tends(pv, compressibility_model(m), tt)...)

# Energy
eq_tends(::Energy, m::TotalEnergyModel, tt::Flux{FirstOrder}) =
    (Advect(), Pressure())

eq_tends(::ρθ_liq_ice, m::θModel, tt::Flux{FirstOrder}) = (Advect(),)

# TODO: make radiation aware of which energy formulation is used:
# eq_tends(pv::PV, m::AtmosModel, tt::Flux{FirstOrder}) where {PV <: AbstractEnergyVariable} =
#     (eq_tends(pv, energy_model(m), tt)..., eq_tends(pv, energy_model(m), radiation_model(m), tt)...)
eq_tends(pv::AbstractEnergyVariable, m::AtmosModel, tt::Flux{FirstOrder}) = (
    eq_tends(pv, energy_model(m), tt)...,
    eq_tends(pv, radiation_model(m), tt)...,
)

# AbstractMoistureVariable
eq_tends(::AbstractMoistureVariable, ::AtmosModel, ::Flux{FirstOrder}) =
    (Advect(),)

# AbstractPrecipitationVariable
eq_tends(
    pv::AbstractPrecipitationVariable,
    m::AtmosModel,
    tt::Flux{FirstOrder},
) = (eq_tends(pv, precipitation_model(m), tt)...,)

# Tracers
eq_tends(pv::Tracers{N}, ::AtmosModel, ::Flux{FirstOrder}) where {N} =
    (Advect(),)

#####
##### Second order fluxes
#####

eq_tends(
    ::Union{Mass, Momentum, AbstractMoistureVariable},
    ::DryModel,
    ::Flux{SecondOrder},
) = ()
eq_tends(
    ::Union{Mass, Momentum, AbstractMoistureVariable},
    ::AbstractMoistureModel,
    ::Flux{SecondOrder},
) = (MoistureDiffusion(),)

# Mass
eq_tends(pv::Mass, m::AtmosModel, tt::Flux{SecondOrder}) =
    (eq_tends(pv, moisture_model(m), tt)...,)

# Momentum
eq_tends(pv::Momentum, m::AtmosModel, tt::Flux{SecondOrder}) = (
    ViscousStress(),
    eq_tends(pv, moisture_model(m), tt)...,
    eq_tends(pv, turbconv_model(m), tt)...,
    eq_tends(pv, hyperdiffusion_model(m), tt)...,
)

# Energy
eq_tends(::Energy, m::TotalEnergyModel, tt::Flux{SecondOrder}) =
    (ViscousFlux(), DiffEnthalpyFlux())

eq_tends(::ρθ_liq_ice, m::θModel, tt::Flux{SecondOrder}) = (ViscousFlux(),)

eq_tends(pv::AbstractEnergyVariable, m::AtmosModel, tt::Flux{SecondOrder}) = (
    eq_tends(pv, energy_model(m), tt)...,
    eq_tends(pv, turbconv_model(m), tt)...,
    eq_tends(pv, hyperdiffusion_model(m), tt)...,
)

# AbstractMoistureVariable
eq_tends(pv::AbstractMoistureVariable, m::AtmosModel, tt::Flux{SecondOrder}) = (
    eq_tends(pv, moisture_model(m), tt)...,
    eq_tends(pv, turbconv_model(m), tt)...,
    eq_tends(pv, hyperdiffusion_model(m), tt)...,
)

# AbstractPrecipitationVariable
eq_tends(
    pv::AbstractPrecipitationVariable,
    m::AtmosModel,
    tt::Flux{SecondOrder},
) = (eq_tends(pv, precipitation_model(m), tt)...,)

# Tracers
eq_tends(pv::Tracers{N}, m::AtmosModel, tt::Flux{SecondOrder}) where {N} =
    (eq_tends(pv, tracer_model(m), tt)...,)


type Rnd <: Op; rgen; Rnd(;rgen=Uniform(0,1), o...)=new(rgen); end
Kenv.kdef(:rnd,Rnd)
ninputs(::Rnd)=0
canoverwrite(::Rnd)=false
back_reads_x(::Rnd)=false
back_reads_y(::Rnd)=false
infersize(::Rnd,y)=tuple(y)
forw(r::Rnd, y; o...)=rgen!(r.rgen, y)
back(r::Rnd, dy; o...)=error("rnd has no backward pass")


using Flux
using Random
using MinimalRLCore

Base.@kwdef struct EvaluateAgent{A, SB, Φ, P}
    agent::A
    prev_s::Φ
    policy::AP
end

function EvaluateAgent(agent::DQNAgent, eval_policy::AP)

    EvaluateAgent(agent,
                  zeros(eltype(agent.state_buffer),
                        size(agent.state_buffer)[1:end-1]...,
                        hist_length(agent.state_buffer)),
                  eval_policy)
end

function MinimalRLCore.start!(agent::EvaluateAgent{DQNAgent},
                              env_s_tp1,
                              rng::AbstractRNG=Random.GLOBAL_RNG)

    agent.prev_s

    agent.state_postproc(agent.state_preproc(env_s_tp1))

end

function MinimalRLCore.step!(agent::EvaluateAgent{DQNAgent},
                             env_s_tp1,
                             r,
                             terminal,
                             rng::AbstractRNG=Random.GLOBAL_RNG)
end


using Distributions
using IncrementalInference

import Unmarshal: unmarshal
import Base: convert

# Already in IIF DefaultNodeTypes.jl Ln151
# """
# Converter: Prior -> PackedPrior::Dict{String, Any}
# """
# function convert(::Type{Dict{String, Any}}, prior::IncrementalInference.Prior)
#     z = convert(Type{Dict{String, Any}}, prior.Z)
#     return Packed_Factor([z], "Prior")
# end

# """
# Converter: PackedPrior::Dict{String, Any} -> Prior
# """
# function convert(::Type{Prior}, prior::Dict{String, Any})
#     # Genericize to any packed type next.
#     z = prior["measurement"][1]
#     z = convert(_evalType(z["distType"]), z)
#     return Prior(z)
# end

"""
Converter: PackedPrior::Dict{String, Any} -> PriorPose2
"""
function convert(::Type{PriorPose2}, prior::Dict{String, Any})
    # Genericize to any packed type next.
    z = prior["measurement"][1]
    z = convert(_evalType(z["distType"]), z)
    return PriorPose2(z)
end

"""
Converter: PackedPrior::Dict{String, Any} -> PriorPose3
"""
function convert(::Type{PriorPose3}, prior::Dict{String, Any})
    # Genericize to any packed type next.
    z = prior["measurement"][1]
    z = convert(_evalType(z["distType"]), z)
    return PriorPose3(z)
end

"""
Converter: Pose2Pose2::Dict{String, Any} -> Pose2Pose2
"""
function convert(::Type{Pose2Pose2}, prior::Dict{String, Any})
    # Genericize to any packed type next.
    z = prior["measurement"][1]
    z = convert(_evalType(z["distType"]), z)
    return Pose2Pose2(z)
end

"""
Converter: Pose2Point2BearingRange::Dict{String, Any} -> Pose2Point2BearingRange
"""
function convert(::Type{Pose2Point2BearingRange}, fact::Dict{String, Any})
    # Genericize to any packed type next.
    b = fact["measurement"][1]
    r = fact["measurement"][2]
    b = convert(_evalType(b["distType"]), b)
    r = convert(_evalType(r["distType"]), r)
    return Pose2Point2BearingRange(b, r)
end

"""
Converter: Pose3Pose3::Dict{String, Any} -> Pose3Pose3
"""
function convert(::Type{Pose3Pose3}, fact::Dict{String, Any})
    # Genericize to any packed type next.
    z = fact["measurement"][1]
    zdist = convert(_evalType(z["distType"]), z)
    return Pose3Pose3(zdist)
end

# already added in RoME PartialPose3.jl Ln230
# """
# Converter: PartialPose3XYYaw::Dict{String, Any} -> PartialPose3XYYaw
# """
# function convert(::Type{PartialPose3XYYaw}, fact::Dict{String, Any})
#     # Genericize to any packed type next.
#     xy = fact["measurement"][1]
#     yaw = fact["measurement"][2]
#     xydistr = convert(_evalType(xy["distType"]), xy)
#     yawdistr = convert(_evalType(yaw["distType"]), yaw)
#     return PartialPose3XYYaw(xydistr, yawdistr)
# end

# already added in RoME PartialPose3.jl Ln89
# """
# Converter: PartialPriorRollPitchZ::Dict{String, Any} -> PartialPriorRollPitchZ
# """
# function convert(::Type{PartialPriorRollPitchZ}, fact::Dict{String, Any})
#     # Genericize to any packed type next.
#     z = fact["measurement"][1]
#     rp = fact["measurement"][2]
#     rpdistr = convert(_evalType(rp["distType"]), rp)
#     zdistr = convert(_evalType(z["distType"]), z)
#     return PartialPriorRollPitchZ(rpdistr, zdistr)
# end

# TODO: Write the other side of the converters for above before moving.


isdefined(Base, :__precompile__) && __precompile__(false)

module SMC
    using BenchmarkTools, DataFrames, Dates, Distributed, Distributions
    using FileIO, HDF5, JLD2, LinearAlgebra, Random, Test
    using ModelConstructors, SparseArrays

    using Roots: fzero, ConvergenceFailed
    using StatsBase: sample, Weights

    import Base.<, Base.isempty, Base.min, Base.max
    import Calculus, ModelConstructors
    import SparseArrays.SparseMatrixCSC

    export
        compute_parameter_covariance, get_estimation_output_files,
        compute_moments, find_density_bands, mutation, resample, smc,
        mvnormal_mixture_draw, nearest_spd, marginal_data_density,
        initial_draw!, Cloud, get_cloud

    const VERBOSITY   = Dict(:none => 0, :low => 1, :high => 2)
    const DATE_FORMAT = "yymmdd"

    include("particle.jl")
    include("initialization.jl")
    include("helpers.jl")
    include("util.jl")
    include("mutation.jl")
    include("resample.jl")
    include("smc_main.jl")
end


"""
This function reduces the input tasks to a single task by applying `func` to each input task pair.
"""
function reducetasks(func, tasks)
    @assert !isempty(tasks)
    todo = collect(tasks)
    while length(todo) > 1
        a, b = popfirst!(todo), popfirst!(todo)
        newtask = Threads.@spawn func(fetch(a), fetch(b))
        push!(todo, newtask)
    end
    return todo[1]
end

function verify(::Type{SpawnThreads}, sig::Model.Signature, pkhashpairs)
    # kick off workers in parallel
    task_rhside = Threads.@spawn Curve.curve_miller(Config.G2, sig.sig)
    tasks_lhside = [Threads.@spawn(Curve.curve_miller(pair.first.pk, pair.second.hash)) for pair in pkhashpairs]
    task_lhside = reducetasks(*, tasks_lhside)
    res = Curve.field_final_exp(fetch(task_lhside) // fetch(task_rhside))
    return res == one(typeof(res))
end


evaluate(H::AbstractHomotopy, x, t) = error("`evaluate` not implemented for $(typeof(H))<:AbstractHomotopy")
substitute(H::AbstractHomotopy, pair::Pair{Symbol,<:Number}) = error("`substitute` not implemented for $(typeof(H))<:AbstractHomotopy")
startsystem(H::AbstractHomotopy) = error("`startsystem` not implemented for $(typeof(H))<:AbstractHomotopy")
targetsystem(H::AbstractHomotopy) = error("`targetsystem` not implemented for $(typeof(H))<:AbstractHomotopy")
nvariables(H::AbstractHomotopy) = error("`nvars` not implemented for $(typeof(H))<:AbstractHomotopy")
removepoly(H::AbstractHomotopy, i::Int) = error("`removepoly` not implemented for $(typeof(H))<:AbstractHomotopy")
degrees(H::AbstractHomotopy) = error("`degrees` not implemented for $(typeof(H))<:AbstractHomotopy")
differentiate(H::AbstractHomotopy) = error("`differentiate` not implemented for $(typeof(H))<:AbstractHomotopy")
dt(H::AbstractHomotopy) = error("`dt` not implemented for $(typeof(H))<:AbstractHomotopy")

ishomogenous(H::AbstractHomotopy) = error("`ishomogenous` not implemented for $(typeof(H))<:AbstractHomotopy")
homogenize(H::AbstractHomotopy) = error("`homogenize` not implemented for $(typeof(H))<:AbstractHomotopy")
homogenized(H::AbstractHomotopy) = error("`homogenized` not implemented for $(typeof(H))<:AbstractHomotopy")

weylnorm(H::AbstractHomotopy, t::Float64) = error("`weyl_norm` not implemented for $(typeof(H)) for AbstractHomotopy")

@inline ∂t(H::AbstractHomotopy) = dt(H)
eltype(H::AbstractHomotopy{T}) where {T<:Number} = T
coefftype(H::AbstractHomotopy{T}) where {T<:Number} = T


connector_comp_name(i::Int) = Symbol("ConnectorComp$i")

# Return the datatype to use for instance variables/parameters
function _instance_datatype(md::ModelDef, def::AbstractDatumDef)
    dtype = def.datatype == Number ? number_type(md) : def.datatype
    dims = dim_names(def)
    num_dims = dim_count(def)

    ti = get_time_index_position(def)

    if num_dims == 0
        T = ScalarModelParameter{dtype}

    elseif ti === nothing     # there's no time dimension
        T = Array{dtype, num_dims}

    else
        if isuniform(md)
            first, stepsize = first_and_step(md)
            first === nothing && @warn "_instance_datatype: first === nothing"
            T = TimestepArray{FixedTimestep{first, stepsize}, Union{dtype, Missing}, num_dims, ti}
        else
            times = time_labels(md)
            T = TimestepArray{VariableTimestep{(times...,)}, Union{dtype, Missing}, num_dims, ti}
        end
    end

    # @info "_instance_datatype returning $T"
    return T
end

# Create the Ref or Array that will hold the value(s) for a Parameter or Variable
function _instantiate_datum(md::ModelDef, def::AbstractDatumDef)
    dtype = _instance_datatype(md, def)
    dims = dim_names(def)
    num_dims = length(dims)

    # Scalar datum
    if num_dims == 0
        value = dtype(0)

    # Array datum, with :time dimension
    elseif dims[1] == :time

        if num_dims == 1
            value = dtype(dim_count(md, :time))
        else
            counts = dim_counts(md, Vector{Symbol}(dims))
            value = dtype <: AbstractArray ? dtype(undef, counts...) : dtype(counts...)
        end

    # Array datum, without :time dimension
    else
        # TBD: Handle unnamed indices properly
        counts = dim_counts(md, Vector{Symbol}(dims))
        value = dtype <: AbstractArray ? dtype(undef, counts...) : dtype(counts...)
    end

    return value
end

"""
    _instantiate_component_vars(md::ModelDef, comp_def::ComponentDef)

Instantiate a component `comp_def` in the model `md` and its variables (but not its
parameters). Return the resulting ComponentInstanceVariables.
"""
function _instantiate_component_vars(md::ModelDef, comp_def::ComponentDef)
    var_defs = variables(comp_def)

    names  = Symbol[nameof(def) for def in var_defs]
    values = Any[_instantiate_datum(md, def) for def in var_defs]
    types  = DataType[_instance_datatype(md, def) for def in var_defs]
    paths  = repeat(Any[comp_def.comp_path], length(names))

    return ComponentInstanceVariables(names, types, values, paths)
end

# Create ComponentInstanceVariables for a composite component from the list of exported vars
function _combine_exported_vars(comp_def::AbstractCompositeComponentDef, var_dict::Dict{ComponentPath, Any})
    names  = Symbol[]
    values = Any[]

    types = DataType[typeof(val) for val in values]
    paths = repeat(Any[comp_def.comp_path], length(names))
    ci_vars = ComponentInstanceVariables(names, types, values, paths)
    # @info "ci_vars: $ci_vars"]
    return ci_vars
end

function _combine_exported_pars(comp_def::AbstractCompositeComponentDef, par_dict::Dict{Tuple{ComponentPath, Symbol}, Any})
    names  = Symbol[]
    values = Any[]
    paths = repeat(Any[comp_def.comp_path], length(names))
    types = DataType[typeof(val) for val in values]
    return ComponentInstanceParameters(names, types, values, paths)
end

function _instantiate_vars(comp_def::ComponentDef, md::ModelDef, var_dict::Dict{ComponentPath, Any})
    var_dict[comp_def.comp_path] = _instantiate_component_vars(md, comp_def)
end

# Creates the top-level vars for the model
function _instantiate_vars(md::ModelDef, var_dict::Dict{ComponentPath, Any})
    _instantiate_vars(md, md, var_dict)
end

# Recursively instantiate all variables and store refs in the given dict.
function _instantiate_vars(comp_def::AbstractCompositeComponentDef, md::ModelDef, var_dict::Dict{ComponentPath, Any})
    comp_path = comp_def.comp_path
    # @info "_instantiate_vars composite $comp_path"

    for cd in compdefs(comp_def)
        _instantiate_vars(cd, md, var_dict)
    end
    var_dict[comp_path] = _combine_exported_vars(comp_def, var_dict)
end

# Do nothing if called on a leaf component
_collect_params(comp_def::ComponentDef, var_dict, par_dict) = nothing

# Recursively collect all parameters with connections to allocated storage for variables
function _collect_params(comp_def::AbstractCompositeComponentDef,
                         var_dict::Dict{ComponentPath, Any},
                         par_dict::Dict{Tuple{ComponentPath, Symbol}, Any})
    # depth-first search of composites
    for cd in compdefs(comp_def)
        _collect_params(cd, var_dict, par_dict)
    end

    # @info "Collecting params for $(comp_def.comp_id)"

    # Iterate over connections to create parameters, referencing storage in vars
    for ipc in internal_param_conns(comp_def)
        src_vars = var_dict[ipc.src_comp_path]
        var_value_obj = get_property_obj(src_vars, ipc.src_var_name)
        par_dict[(ipc.dst_comp_path, ipc.dst_par_name)] = var_value_obj
        # @info "internal conn: $(ipc.src_comp_path):$(ipc.src_var_name) => $(ipc.dst_comp_path):$(ipc.dst_par_name)"
    end

    for ext in external_param_conns(comp_def)
        param = external_param(comp_def, ext.external_param)
        par_dict[(ext.comp_path, ext.param_name)] = (param isa ScalarModelParameter ? param : value(param))
        # @info "external conn: $(ext.comp_name).$(ext.param_name) => $(param)"
    end

    # Make the external parameter connections for the hidden ConnectorComps.
    # Connect each :input2 to its associated backup value.
    for (i, backup) in enumerate(comp_def.backups)
        conn_comp = compdef(comp_def, connector_comp_name(i))
        conn_path = conn_comp.comp_path

        param = external_param(comp_def, backup)
        par_dict[(conn_path, :input2)] = (param isa ScalarModelParameter ? param : value(param))
    end
end

function _instantiate_params(comp_def::ComponentDef, par_dict::Dict{Tuple{ComponentPath, Symbol}, Any})
    # @info "Instantiating params for $(comp_def.comp_path)"
    comp_path = comp_def.comp_path
    names = parameter_names(comp_def)
    vals  = Any[par_dict[(comp_path, name)] for name in names]
    types = DataType[typeof(val) for val in vals]
    paths = repeat([comp_def.comp_path], length(names))

    return ComponentInstanceParameters(names, types, vals, paths)
end

function _instantiate_params(comp_def::AbstractCompositeComponentDef, par_dict::Dict{Tuple{ComponentPath, Symbol}, Any})
    _combine_exported_pars(comp_def, par_dict)
end

# Return a built leaf or composite LeafComponentInstance
function _build(comp_def::ComponentDef,
                var_dict::Dict{ComponentPath, Any},
                par_dict::Dict{Tuple{ComponentPath, Symbol}, Any},
                time_bounds::Tuple{Int, Int})
    # @info "_build leaf $(comp_def.comp_id)"
    # @info "  var_dict $(var_dict)"
    # @info "  par_dict $(par_dict)"

    pars = _instantiate_params(comp_def, par_dict)
    vars = var_dict[comp_def.comp_path]

    return LeafComponentInstance(comp_def, vars, pars, time_bounds)
end

function _build(comp_def::AbstractCompositeComponentDef,
                var_dict::Dict{ComponentPath, Any},
                par_dict::Dict{Tuple{ComponentPath, Symbol}, Any},
                time_bounds::Tuple{Int, Int})
    # @info "_build composite $(comp_def.comp_id)"
    # @info "  var_dict $(var_dict)"
    # @info "  par_dict $(par_dict)"

    comps = [_build(cd, var_dict, par_dict, time_bounds) for cd in compdefs(comp_def)]
    return CompositeComponentInstance(comps, comp_def, time_bounds)
end

function _build(md::ModelDef)
    # @info "_build(md)"
    add_connector_comps!(md)

    # check if all parameters are set
    not_set = unconnected_params(md)

    if ! isempty(not_set)
        params = join(not_set, "\n  ")
        error("Cannot build model; the following parameters are not set:\n  $params")
    end

    var_dict = Dict{ComponentPath, Any}()                 # collect all var defs and
    par_dict = Dict{Tuple{ComponentPath, Symbol}, Any}()  # store par values as we go

    _instantiate_vars(md, var_dict)
    _collect_params(md, var_dict, par_dict)

    # @info "var_dict: $var_dict"
    # @info "par_dict: $par_dict"

    t = dimension(md, :time)
    time_bounds = (firstindex(t), lastindex(t))

    propagate_time!(md, t)

    ci = _build(md, var_dict, par_dict, time_bounds)
    mi = ModelInstance(ci, md)
    return mi
end

function build(m::Model)
    # fix paths and propagate imports
    fix_comp_paths!(m.md)

    # Reference a copy in the ModelInstance to avoid changes underfoot
    md = deepcopy(m.md)
    m.mi = _build(md)
    m.md.dirty = false
    return nothing
end

"""
    create_marginal_model(base::Model, delta::Float64=1.0)

Create a `MarginalModel` where `base` is the baseline model and `delta` is the
difference used to create the `marginal` model.  Return the resulting `MarginaModel`
which shares the internal `ModelDef` between the `base` and `marginal`.
"""
function create_marginal_model(base::Model, delta::Float64=1.0)
    # Make sure the base has a ModelInstance before we copy since this
    # copies the ModelDef to avoid being affected by later changes.
    if ! is_built(base)
        build(base)
    end

    # Create a marginal model, which shares the internal ModelDef between base and marginal
    mm = MarginalModel(base, delta)
end

function Base.run(mm::MarginalModel; ntimesteps::Int=typemax(Int))
    run(mm.base, ntimesteps=ntimesteps)
    run(mm.marginal, ntimesteps=ntimesteps)
end

function build(mm::MarginalModel)
    build(mm.base)
    build(mm.marginal)
end


module TestLinkTree

using Phylo
using DataFrames
using JuliaDB

using Test
using IterableTables: getiterator

species = ["Dog", "Cat", "Human"]
ntips = 10
df = DataFrame(species = species, count = [10, 20, 3])
observations = ["Dog", "Cat", "Dog", "Dog"]
jdb = table((species = observations, count = 1:4))

@testset "LinkTree()" begin
    lts = RootedTree(species)
    @test nodedatatype(typeof(lts)) ≡ Dict{String, Any}
    @test branchdatatype(typeof(lts)) ≡ Dict{String, Any}
    @test leafinfotype(typeof(lts)) ≡ Dict{String, Any}

    ltdf = Phylo.LT{OneRoot, DataFrame, Float64}(df)
    @test nodedatatype(typeof(ltdf)) ≡ Dict{String, Any}
    @test branchdatatype(typeof(ltdf)) ≡ Dict{String, Any}
    @test leafinfotype(typeof(ltdf)) ≡ DataFrame

    rt = Unrooted
    lb = LinkBranch{rt, String, Nothing}
    ln = LinkNode{rt, String, Vector{Int}, lb}
    ltjdb = LinkTree{rt, String, ln, lb, typeof(jdb)}(jdb)
    @test nodedatatype(typeof(ltjdb)) ≡ Vector{Int}
    @test branchdatatype(typeof(ltjdb)) ≡ Nothing
    @test leafinfotype(typeof(ltjdb)) ≡ typeof(jdb)
end

end


# Another "low-level" example, this one incorporating per-vertex color
# and element arrays
using ModernGL, GeometryTypes, GLAbstraction, GLWindow

# Create the window. This sets all the hints and makes the context current.
window = create_glcontext("Drawing polygons 4", resolution=(800,600))

# Create the Vertex Array Object (VAO) and make it current
# Note that while the tutorial describes this after the attributes (below),
# we need to make vao current before calling glVertexAttribPointer.
# You should also do this before creating any element arrays.
vao = Ref(GLuint(0))
glGenVertexArrays(1, vao)
glBindVertexArray(vao[])

# The vertices of our rectangle, with color
vertices = Point{5,Float32}[(-0.5,  0.5, 1, 0, 0),    # top-left
                            ( 0.5,  0.5, 0, 1, 0),    # top-right
                            ( 0.5, -0.5, 0, 0, 1),    # bottom-right
                            (-0.5, -0.5, 1, 1, 1)]    # bottom-left

elements = Vec{3,GLuint}[(0,1,2),          # the first triangle
                         (2,3,0)]          # the second triangle

# Create the Vertex Buffer Object (VBO)
vbo = Ref(GLuint(0))   # initial value is irrelevant, just allocate space
glGenBuffers(1, vbo)
glBindBuffer(GL_ARRAY_BUFFER, vbo[])
BufferData(GL_ARRAY_BUFFER, sizeof(vertices), vertices, GL_STATIC_DRAW)

# Create the Element Buffer Object (EBO)
ebo = Ref(GLuint(0))
glGenBuffers(1, ebo)
glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, ebo[])
BufferData(GL_ELEMENT_ARRAY_BUFFER, sizeof(elements), elements, GL_STATIC_DRAW)

# The shaders. Here we do everything manually, but life will get
# easier with GLAbstraction. See drawing_polygons5.jl for such an
# implementation.

# The vertex shader
vertex_source = """
#version 150

in vec2 position;
in vec3 color;

out vec3 Color;

void main()
{
    Color = color;
    gl_Position = vec4(position, 0.0, 1.0);
}
"""

# The fragment shader
fragment_source = """
# version 150

in vec3 Color;

out vec4 outColor;

void main()
{
    outColor = vec4(Color, 1.0);
}
"""

# Compile the vertex shader
vertex_shader = glCreateShader(GL_VERTEX_SHADER)
glShaderSource(vertex_shader, vertex_source)  # nicer thanks to GLAbstraction
glCompileShader(vertex_shader)
# Check that it compiled correctly
status = Ref(GLint(0))
glGetShaderiv(vertex_shader, GL_COMPILE_STATUS, status)
if status[] != GL_TRUE
    buffer = Array(UInt8, 512)
    glGetShaderInfoLog(vertex_shader, 512, C_NULL, buffer)
    @error "$(bytestring(buffer))"
end

# Compile the fragment shader
fragment_shader = glCreateShader(GL_FRAGMENT_SHADER)
glShaderSource(fragment_shader, fragment_source)
glCompileShader(fragment_shader)
# Check that it compiled correctly
status = Ref(GLint(0))
glGetShaderiv(fragment_shader, GL_COMPILE_STATUS, status)
if status[] != GL_TRUE
    buffer = Array(UInt8, 512)
    glGetShaderInfoLog(fragment_shader, 512, C_NULL, buffer)
    @error "$(bytestring(buffer))"
end

# Connect the shaders by combining them into a program
shader_program = glCreateProgram()
glAttachShader(shader_program, vertex_shader)
glAttachShader(shader_program, fragment_shader)
glBindFragDataLocation(shader_program, 0, "outColor") # optional

glLinkProgram(shader_program)
glUseProgram(shader_program)

# Link vertex data to attributes
pos_attribute = glGetAttribLocation(shader_program, "position")
glEnableVertexAttribArray(pos_attribute)
glVertexAttribPointer(pos_attribute, 2,
                      GL_FLOAT, GL_FALSE, 5*sizeof(Float32), C_NULL)

col_attribute = glGetAttribLocation(shader_program, "color")
glEnableVertexAttribArray(col_attribute)
glVertexAttribPointer(col_attribute, 3,
                      GL_FLOAT, GL_FALSE, 5*sizeof(Float32), Ptr{Nothing}(2*sizeof(Float32)))

# Draw while waiting for a close event
glClearColor(0,0,0,0)
while !GLFW.WindowShouldClose(window)
    glClear(GL_COLOR_BUFFER_BIT)
    glDrawElements(GL_TRIANGLES, 6, GL_UNSIGNED_INT, C_NULL)
    GLFW.SwapBuffers(window)
    GLFW.PollEvents()
    if GLFW.GetKey(window, GLFW.KEY_ESCAPE) == GLFW.PRESS
        GLFW.SetWindowShouldClose(window, true)
    end
end

GLFW.DestroyWindow(window)  # needed if you're running this from the REPL


module LeafAreaIndex

using Optim: optimize, minimizer
using StatsBase: mean, midpoints
import FileIO, JLD2, Netpbm, Parameters, StatsBase, Optim

# import StatsBase #for `midpoints`
import ColorTypes #in crops.jl for AbstractRGB

# in calibration.jl
using Graphics: Point, norm
using DataFrames

export rawblueread, rawcolourread,
    CameraLens, CameraLensParams, PolarImage,
    Slope, SlopeParams, Mask, MaskParams,
    RedMax,#pixels, gapfraction, contactfreqs,
    threshold, #EdgeDetection, MinimumThreshold, RidlerCalvard,
    inverse, #Zenith57, Miller, Lang, EllipsLUT, EllipsOpt,
    langxiang #chencihlar,
    #calibrate_center, calibrate_projfun

include("Types.jl")
include("thresholding.jl")
include("segments.jl")
include("gapfraction.jl")
include("inversion.jl")
include("clumping.jl")
include("crops.jl")
include("ChenCihlar.jl")
include("calibration.jl")
include("io.jl")

# reoccuring argument checks
function checkθ1θ2(θ1, θ2)
    θ1 < 0   && throw(DomainError())
    θ2 > π/2 && throw(DomainError())
    @assert θ1 < θ2
end

#TODO replace with OnlineStats.jl
"Specialed type for immutable streaming sum, based on PR #18 from StreamStats.jl"
struct StreamMean
    streamsum::Float64
    len::Int
end
StreamMean() = StreamMean(0.0, 0)
update(sm::StreamMean, term) = StreamMean(sm.streamsum + term, sm.len + 1)
StatsBase.mean(sm::StreamMean) = sm.streamsum / sm.len
Base.empty!(sm::StreamMean) = StreamMean()
Base.length(sm::StreamMean) = sm.len

# A fast histogram method derived from julialang PR #8952. Used in thresholding
# and slope adjustment.
function fasthist(img::AbstractVector, edg::AbstractRange)
    n = length(edg) - 1
    histcount = zeros(Int, n)
    step(edg) <= 0 && error("step(edg) must be positive")
    for pixel in img
        f = (pixel - first(edg)) / step(edg)
        if 0 < f <= n
            histcount[ceil(Int, f)] += 1
        end
    end
    histcount
end

end # module


# A symbol is a name and the concolicly computed type
struct Sym
    name::Symbol
    _type::DataType
    Sym(base, _type::DataType) = new(Base.gensym(Symbol(base, '#', _type)), _type)
    Sym(_type::DataType) = new(Base.gensym(Symbol(_type)), _type)
end
Cassette.metadatatype(::Type{<:TraceCtx}, ::DataType) = Sym

record(args, ctx) = map(args) do arg
    if istagged(arg, ctx)
        metadata(arg, ctx)
    else
        arg
    end
end

# A callsite is a function with a concolic argument set
# children are all the callsites within the function
mutable struct Callsite
    f::Any
    args::Any
    retval::Union{Some{Any}, Nothing}
    children::Vector{Callsite}
    Callsite(f, args) = new(f, args, nothing, Callsite[])
end
Base.push!(trace::Callsite, call::Callsite) = push!(trace.children, call)

struct Metadata
    trace::Callsite
    # For current execution
    record::Vector{Sym}
    substitutes::Dict{DataType, Any}
end
Metadata() = Metadata(
    Callsite(:toplevel, ()),
    Vector{Sym}(),
    Dict{DataType, Vector}())

Metadata(m::Metadata, call::Callsite) = Metadata(call, m.record, m.substitutes)
Metadata(subs) = Metadata(Callsite(:toplevel, ()), Vector{Sym}(), subs)

function record!(ctx, type)
    m = ctx.metadata
    sym = Sym(:fval, type)
    push!(m.record, sym)

    if haskey(m.substitutes, type)
        subs = m.substitutes[type]
        if !isempty(subs)
            return pop!(subs), sym
        end
    end
    return nothing, sym
end
Base.push!(m::Metadata, call) = push!(m.trace, call)

function augment(record, subs)
    substitutions = Dict{DataType, Vector}()
    function insert!(type, val)
        if !haskey(substitutions, type)
            substitutions[type] = Vector{type}()
        end
        push!(substitutions[type], val)
    end

    for sym in record
        insert!(sym._type, subs[sym.name])
    end
    return substitutions
end

function Cassette.overdub(ctx::TraceCtx, ::typeof(rand), ::Type{T}) where T<:INTEGERS
    val, sym = record!(ctx, T)
    if val === nothing
        val = rand(T)
    end
    return tag(val::T, ctx, sym)
end

# oftype will drop tag on the x value since we are only looking at it'd type.
function Cassette.overdub(ctx::TraceCtx, ::typeof(oftype), x::Tagged, y)
    x = untag(x, ctx)
    return oftype(x, y)
end

##
# We need to manually override for IntrinsicFunctions (which are the leaf-nodes we are interested in)
# Since in tagged contexts there is an automatic fallback available.
##
function Cassette.overdub(ctx::TraceCtx{Metadata, <:Cassette.Tag}, f::Core.IntrinsicFunction, args...)
    call = Callsite(f, record(args, ctx))
    push!(ctx.metadata, call)
    retval = fallback(ctx, f, args...)
    if any(a -> istagged(a, ctx), args)
        retval = tag(retval, ctx, Sym(typeof(retval)))
    end
    call.retval = Some(istagged(retval, ctx) ? metadata(retval, ctx) : retval)
    return retval
end

##
# Recursivly trace though our program.
#
# Note: This won't trace functions that are defined primitive and there are several
#       fallbacks that Cassette provides for tagged contexts.
#
# Question: Shouldn't canoverdub for primitives be false?
##
function Cassette.overdub(ctx::TraceCtx, f, args...)
    # We need to push the callsite first into metadata
    # otherwise we run into issues with functions that throw errors
    call = Callsite(f, record(args, ctx))
    push!(ctx.metadata, call)

    # Recurse into the next function
    retval = if canrecurse(ctx, f, args...)
        newctx = similarcontext(ctx, metadata = Metadata(ctx.metadata, call))
        Cassette.recurse(newctx, f, args...)
    else
        retval = fallback(ctx, f, args...)
        # If any of our inputs was tagged we want to tag the return value so that
        # the symbolic nature of this operation is captured.
        if any(a -> istagged(a, ctx), args) && !istagged(retval, ctx)
            retval = tag(retval, ctx, Sym(typeof(retval)))
        end
        retval
    end
    call.retval = Some(istagged(retval, ctx) ? metadata(retval, ctx) : retval)
    return retval
end


module Tracks
using OffsetArrays, Images, Unitful
export Track, SIMTrack

using ..Acquisitions

abstract type AbstractTrack end

struct Track <: AbstractTrack # {{{
    x::AbstractVector{<:Real}
    y::AbstractVector{<:Real}
    view::AxisArray{A,1} where {A<:OffsetArray{<:AbstractGray}}
end # }}}

# TODO: Should be generic over the SIMacquisition type <20-05-21, kunzaatko> #
struct SIMTrack <: AbstractTrack # {{{
    x::AbstractVector{<:Real}
    y::AbstractVector{<:Real}
    view::AxisArray{A,1} where {A<:OffsetArray{<:AbstractGray}}
    view_reconstructed::AxisArray{A,1} where {A<:OffsetArray{<:AbstractGray}}
end # }}}

function Track( # {{{
    x, y, acquisition::Acquisition, times::AbstractVector{<:Unitful.Time}
)
    # TODO: Add check for times, x and y having the same lengths <20-05-21, kunzaatko> #
    x_offsets, y_offsets = round.(Integer, x), round.(Integer, y)
    view = AxisArray(
        [
            OffsetArray(
                acquisition.images[Axis{:x}(:), Axis{:y}(:), Axis{:time}(t)],
                OffsetArrays.Origin(-x_offset, -y_offset),
            ) for (t, x_offset, y_offset) in zip(times, x_offsets, y_offsets)
        ],
        Axis{:time}(times),
    )
    return Track(x, y, view)
end # }}}

function Track(x, y, acquisition::Acquisition, frames::AbstractVector{<:Integer})
    return Track(x, y, acquisition, times(acquisition)[frames])
end

function SIMTrack( # {{{
    x, y, acquisition::SIMAcquisition, times::AbstractVector{<:Unitful.Time}
)
    x_offsets, y_offsets = round.(Integer, x), round.(Integer, y)
    view = AxisArray(
        [
            OffsetArray(
                acquisition.unreconstructedimages[Axis{:x}(:), Axis{:y}(:), Axis{:time}(t)],
                OffsetArrays.Origin(-x_offset, -y_offset),
            ) for (t, x_offset, y_offset) in zip(times, x_offsets, y_offsets)
        ],
        Axis{:time}(times),
    )
    # FIX: This should be meaningfull times as they will be defined in the constructor <20-05-21, kunzaatko> #
    view_reconstructed = AxisArray(
        [OffsetArray(Gray.(zeros(3, 3)), OffsetArrays.Origin(-1, -1)) for _ in [1, 2]],
        Axis{:time}([1u"s", 2u"s"]),
    ) # DUMMY instance
    return SIMTrack(x, y, view, view_reconstructed)
end # }}}

function SIMTrack(x, y, acquisition::SIMAcquisition, frames::AbstractVector{<:Integer})
    return SIMTrack(x, y, acquisition, times(acquisition)[frames])
end

end