# Core structure and functions for decision diagrams
# This structure stores both node and arc layers in arrays as fields of DD
# It is fast to compute longest path, but (possibly?!) slow in removing nodes and arcs!
# Commented out sections are not adjusted...

# All higher level functions should use these functions instead of accessing
# the internal structure of a decision diagram directly

# Functions in this file must be careful to keep graph and layers synchronized

using LightGraphs
#using ResumableFunctions
#import LightGraphs.SimpleGraphs.SimpleEdge
import LightGraphs.inneighbors

"""Node structure"""
mutable struct Node{S<:AbstractFloat}
	layer::Int				# node layer
	state::S				# node state
#	out_neighbor::Array{Int,1}		# outgoing arcs of the node, stores arc id in the arc layer
#	in_neighbor::Array{Int,1}		# outgoing arcs of the node, stores arc id in the arc layer
	out_neighbor::Dict{Int, Array{Int, 1}}		# key: the id of node neighbor in the next layer, value: set of arcs (id) connected to that node
	in_neighbor::Dict{Int, Array{Int, 1}}		# key: the id of node neighbor in the previous layer, value: set of arcs (id) connected to that node
	aux_val::Float64		# auxiliary value at the node e.g., the longest path value
end

"""Decision diagram structure"""
struct DecisionDiagram{S<:AbstractFloat}
	graph::SimpleDiGraph{Int}            # represents the directed graph of the DD
    node_layers::Array{Array{Node{S},1},1}        # stores instances of the Node class in each node layer
	arc_layers::Array{Array{Tuple{Int, Int, Float64},1},1}		# stores arc info as tuples (tail, head, label) in arc layers
end

"""
A constrcter for type `DecisionDiagram`

# Input
- `t::Type`: Data type of the state values of nodes
- `nvars::Int`: Number of variables (arc layers)
"""

DecisionDiagram(nvars::Int; t::DataType = Float64)::DecisionDiagram{t} = DecisionDiagram{t}(SimpleDiGraph{Int}(), [Array{Node{t},1}(  0) for i=1:nvars+1], [Array{Tuple{Int, Int, Float64},1}(  0) for i=1:nvars])

"""Return the number of variables (or number of arc layers) in the decision diagram"""
nvars(dd::DecisionDiagram)::Int = length(dd.node_layers) - 1

#nnodes(dd::DecisionDiagram) = nv(dd.graph)
nnodes(dd::DecisionDiagram)::Int = sum(length(dd.node_layers[i]) for i = 1:length(dd.node_layers))

#narcs(dd::DecisionDiagram) = ne(dd.graph)
narcs(dd::DecisionDiagram)::Int = sum(length(dd.arc_layers[i]) for i = 1:length(dd.arc_layers))

"""Set state of an existing node"""
function set_node_state!(dd::DecisionDiagram, layer::Int, id::Int, state::S) where S<:AbstractFloat
	dd.node_layers[layer][id].state = state
end

"""Get state of a node"""
get_node_state(dd::DecisionDiagram, layer::Int, id::Int)::S where S<:AbstractFloat = dd.node_layers[layer][id].state

"""Return arc label"""
get_arc_label(dd::DecisionDiagram, layer::Int, id::Int)::Float64 = dd.arc_layers[layer][id][3]
#get_arc_labels(dd::DecisionDiagram, arc::LightGraphs.SimpleGraphs.SimpleEdge{Int64}) = collect(dd.nodes[src(e)].out_neighbor[dst(e)])

"""Return arc tail and head"""
get_arc_tail(dd::DecisionDiagram, layer::Int, id::Int)::Int = dd.arc_layers[layer][id][1]
get_arc_head(dd::DecisionDiagram, layer::Int, id::Int)::Int = dd.arc_layers[layer][id][2]

get_node_layer_size(dd::DecisionDiagram, layer::Int)::Int = length(dd.node_layers[layer])
get_arc_layer_size(dd::DecisionDiagram, layer::Int)::Int = length(dd.arc_layers[layer])

get_width(dd::DecisionDiagram)::Int = maximum(length(dd.node_layers[i]) for i=1:length(dd.node_layers))

#=
function get_arc_labels(dd::DecisionDiagram, node1::Int, node2::Int)
	i = findfirst(dd.nodes[node1].out_neighbor_nodes, node2)
	#if i == 0
	#	error("No arc exists between node ", node1, " and node ", node2)
	#else
		return dd.nodes[node1].out_neighbor_nodes[i]
	#end
end
=#

"""Return the array of incoming arcs of a node"""
inneighbors(dd::DecisionDiagram, layer::Int, id::Int)::Dict{Int, Array{Int, 1}} = dd.node_layers[layer][id].in_neighbor
#inneighbors(dd::DecisionDiagram, layer::Int, id::Int)::Array{Int, 1} = dd.node_layers[layer][id].in_neighbor

"""Return the array of outgoing nodes of a node"""
outneighbors(dd::DecisionDiagram, layer::Int, id::Int)::Dict{Int, Array{Int, 1}} = dd.node_layers[layer][id].out_neighbor
#outneighbors(dd::DecisionDiagram, layer::Int, id::Int)::Array{Int, 1} = dd.node_layers[layer][id].out_neighbor

"""Add a new node to the decision diagram with a given state and return the node's id at its layer"""
function add_node!(dd::DecisionDiagram, layer::Int, state::S)::Int where S<:AbstractFloat
	#g = dd.graph
	#add_vertex!(g)
	#aux_node::Node{S} = Node(layer, state, Array{Int, 1}(0), Array{Int,1}(0), 0.0)		#This is for the case where neighbor field stores array of arcs only
	aux_node::Node{S} = Node(layer, state, Dict{Int, Array{Int, 1}}(), Dict{Int, Array{Int, 1}}(), 0.0)
	push!(dd.node_layers[layer], aux_node)
	id::Int = length(dd.node_layers[layer])
	return id
end

#=
# TODO: rem_node! is untested
"""Remove a node from the decision diagram"""
function rem_node!(dd::DecisionDiagram, node::Int)

	# ensure the node belongs to the vertex set
	# node in vertices(g) || return false
	node <= length(dd.nodes) || return false

	# Remove incoming and outgoing arcs of the node
	for outnode in keys(dd.nodes[node].out_neighbor)
		delete!(dd.nodes[outnode].in_neighbor, node)
	end
	for innode in keys(dd.nodes[node].in_neighbor)
		delete!(dd.nodes[innode].out_neighbor, node)
	end

	# Move end node in the layer to node's position within the layer
	layer = dd.nodes[node].layer
	layerid = dd.nodes[node].layerid
	dd.layers[layer][layerid] = pop!(dd.layers[layer])
	dd.nodes[dd.layers[layer][layerid]].layerid = layerid

	# Renumber the last node of the graph to the removing node according to LightGraphs implementation
	lastid = length(dd.nodes)
	dd.nodes[node] = pop!(dd.nodes)			# swap the last node with the removed node
	dd.layers[dd.nodes[node].layer][dd.nodes[node].layerid] = node			# update the last node number in dd.layers
	# Update the head/tail number in the incoming/outgoing arcs of the last node
	for outnode in keys(dd.nodes[node].out_neighbor)
		dd.nodes[outnode].in_neighbor[node] = dd.nodes[outnode].in_neighbor[lastid]
		delete!(dd.nodes[outnode].in_neighbor, lastid)
	end
	for innode in keys(dd.nodes[node].in_neighbor)
		dd.nodes[innode].out_neighbor[node] = dd.nodes[innode].out_neighbor[lastid]
		delete!(dd.nodes[innode].out_neighbor, lastid)
	end

	# Remove vertex from graph structure
	#removed = rem_vertex!(dd.graph, node)
	#@assert(removed)
end
=#

#=
"""Remove a node from the decision diagram given its layer and layerid"""
function rem_node!(dd::DecisionDiagram, layer::Int, layerid::Int)
	rem_node!(dd, dd.layers[layer][layerid])
end
=#

"""Add a new arc to the decision diagram between two nodes and retunrs arc id in its layer"""
function add_arc!(dd::DecisionDiagram, layer1::Int, node1_id::Int, layer2::Int, node2_id::Int, label::Float64)::Int
	#@assert((dd.nodes[node1].layer < dd.nodes[node2].layer) && (node1 <= length(dd.nodes)) && (node2 <= length(dd.nodes)))
	#g = dd.graph
	#e = Edge(node1, node2)
	@assert(layer1 == layer2 - 1)
#	if isa(label, Int)					# converts Int label type to floating point
#		lbl = float(label)
#	else
#		lbl = label
#	end
	push!(dd.arc_layers[layer1], (node1_id, node2_id, label))
	id::Int = length(dd.arc_layers[layer1])
#	push!(dd.node_layers[layer1][node1_id].out_neighbor, id)	#for case where neighbors are array of arcs only
#	push!(dd.node_layers[layer2][node2_id].in_neighbor, id)	#for case where neighbors are array of arcs only
	if haskey(dd.node_layers[layer1][node1_id].out_neighbor, node2_id)
		push!(dd.node_layers[layer1][node1_id].out_neighbor[node2_id], id)
		push!(dd.node_layers[layer2][node2_id].in_neighbor[node1_id], id)
	else
		dd.node_layers[layer1][node1_id].out_neighbor[node2_id] = [id]
		dd.node_layers[layer2][node2_id].in_neighbor[node1_id] = [id]
	end
	return id
end

"""
	Add a new arc to the decision diagram between two nodes, where the arc_merging is on. Returns the arc id in the layer.
	If the arc is parallel, it does not add the arc if its label is not lower or upper bound among other arcs.
"""
#NOTE: this code considers the case where inneighbor/outneighbor field is a Dict (NOT array)
function add_reduced_arc!(dd::DecisionDiagram, layer1::Int, node1_id::Int, layer2::Int, node2_id::Int, label::Float64)::Int
	#g = dd.graph
	#e = Edge(node1, node2)
	id::Int = 0			# if the arc is not added, the returning id will be zero
	# check if an arc between two nodes already exists
	if !haskey(dd.node_layers[layer1][node1_id].out_neighbor, node2_id)
		#add_edge!(g, e)
		push!(dd.arc_layers[layer1], (node1_id, node2_id, label))
		id = length(dd.arc_layers[layer1])
		dd.node_layers[layer1][node1_id].out_neighbor[node2_id] = [id]
		dd.node_layers[layer2][node2_id].in_neighbor[node1_id] = [id]
	# if there is only a single arc, the arc is added to in/outneighbor in an increasing order wrt label value
	elseif length(dd.node_layers[layer1][node1_id].out_neighbor[node2_id]) == 1
		push!(dd.arc_layers[layer1], (node1_id, node2_id, label))
		id = length(dd.arc_layers[layer1])
		if dd.arc_layers[layer1][dd.node_layers[layer1][node1_id].out_neighbor[node2_id][1]][3] <= label
			 push!(dd.node_layers[layer1][node1_id].out_neighbor[node2_id], id)	# adds the new label to the parallel arcs
			 push!(dd.node_layers[layer2][node2_id].in_neighbor[node1_id], id)
		 else
			 pushfirst!(dd.node_layers[layer1][node1_id].out_neighbor[node2_id], id)	# adds the new label to the parallel arcs
 			 pushfirst!(dd.node_layers[layer2][node2_id].in_neighbor[node1_id], id)
 		 end
	# if there are two arcs between the nodes, the new arc is added only if it does not belong to the label interval of the current arcs
	# To add the new arc, we only need to change the label of the arc that is being replaced
	else
		@assert(length(dd.node_layers[layer1][node1_id].out_neighbor[node2_id]) == 2)
		id1 = dd.node_layers[layer1][node1_id].out_neighbor[node2_id][1]
		id2 = dd.node_layers[layer1][node1_id].out_neighbor[node2_id][2]
		if label < dd.arc_layers[layer1][id1][3]
			id = id1
			dd.arc_layers[layer1][id] = (node1_id, node2_id, label)
		elseif label > dd.arc_layers[layer1][id2][3]
			id = id2
			dd.arc_layers[layer1][id] = (node1_id, node2_id, label)
		end
	end
	return id
end

#=
# TODO rem_arc! is untested
"""Remove a labeled arc from a decision diagram"""
function rem_arc!(dd::DecisionDiagram, node1::Int, node2::Int, label::Int)
	# checks if the arc exists
	if haskey(dd.nodes[node1].out_neighbor, node2)
		setdiff!(dd.nodes[node1].out_neighbor[node2], label)
		setdiff!(dd.nodes[node2].in_neighbor[node1], label)
		if isempty(dd.nodes[node1].out_neighbor[node2])
			#e = Edge(node1, node2)
			#rem_edge!(dd.graph, e)
			delete!(dd.nodes[node1].out_neighbor, node2)
			delete!(dd.nodes[node2].in_neighbor, node1)
		end
	end
end
=#

#=
# TODO rem_arcs! is untested
"""Remove all (parallel) arcs between two nodes from a decision diagram"""
function rem_arcs!(dd::DecisionDiagram, node1::Int, node2::Int)
	# checks if the arc exists
	if haskey(dd.nodes[node1].out_neighbor, node2)
		#e = Edge(node1, node2)
		#rem_edge!(dd.graph, e)
		delete!(dd.nodes[node1].out_neighbor, node2)
		delete!(dd.nodes[node2].in_neighbor, node1)
	end
end
=#


println("starting!!!")
using DelimitedFiles
using TicketToRide
using Flux
using Flux.Optimise
using TicketToRide: train!

Ls = [4, 6, 8]
Nlayers = collect(1:50)
repeats = [i for i in 1:5]

arg = 1
 #parse(Int, ARGS[1])

list = Iterators.product(Ls, Nlayers, repeats) |> collect
println("length of array jobs is:\t", length(list))

n, nlayers, r = list[arg]

for L in Ls
    for nlayers in Nlayers
        for repeats in 1:10
            @info "\nL is $n"
            @info "nlayers is $nlayers"
            @info "r is $r"
            # make directory if none exists
            dir = "/scratch/mbeach/new_new_tickettoride/L-$n/"
            mkpath(dir)
            file = dir * "layers-$nlayers-r-$r.txt"
            @info file

            # pre-compile once
            opt = Optimise.ADAMW()
            circuit = variational_circuit(n, nlayers);
            ham = heisenberg1D(n)
            dispatch!(circuit, :random);

            ind = findall(x -> x == n, Ls)
            println(ind)

            opt = Optimise.ADAM()
            # @time history = train!(opt, 1, ham, circuit; verbose=true)
            history = train!(opt, 2000, ham, circuit; verbose=false)
            writedlm(file, history)

            opt = Optimise.Descent(1e-3)
            history = train!(opt, 2000, ham, circuit; verbose=false)
            writedlm(file, history)

            opt = Optimise.Descent(1e-4)
            history = train!(opt, 20000, ham, circuit; verbose=false)
            writedlm(file, history)

        end
    end
end


export OneBody

# the type definition of OneBody is in JuLIP to enable its use from
# IPFitting without having to load NBodyIPs. Here we just extend
# the functionality so that it can work effectively with NBodyIPs

using JuLIP: Atoms
import JuLIP.Potentials: OneBody
import JuLIP: energy, forces, virial
import Base: hash
import NBodyIPs: fast, combinebasis

combinebasis(Vs::AbstractVector{<: OneBody}, Cs::AbstractVector{<: Real}) =
             OneBody( sum( V.E0 * c for (V, c) in zip(Vs, Cs) ) )

fast(V::OneBody) = V

hash(::BASIS, ::OneBody) = hash(OneBody)

bodyorder(V::OneBody) = 1

energy(B::AbstractVector{TB}, at::Atoms{T}) where {TB <: OneBody, T} =
   [ energy(b, at) for b in B ]

forces(B::AbstractVector{TB}, at::Atoms{T}) where {TB <: OneBody, T} =
   [ forces(b, at) for b in B ]

virial(B::AbstractVector{TB}, at::Atoms{T}) where {TB <: OneBody, T} =
   [ virial(b, at) for b in B ]

# TODO: maybe the above should be moved to TB <: AbstractCalculator?


module MathOptInterfaceToMathProgBase

using MathOptInterface
const MOI = MathOptInterface
const MOIU = MOI.Utilities

import MathProgBase
const MPB = MathProgBase.SolverInterface

MOIU.@model(InnerModel,
    (MOI.ZeroOne, MOI.Integer),
    (MOI.EqualTo, MOI.GreaterThan, MOI.LessThan, MOI.Interval),
    (),
    (),
    (MOI.SingleVariable,),
    (),
    (),
    ()
)

# Make `Model` a constant for use in the rest of the file.
const Model = MOIU.UniversalFallback{InnerModel{Float64}}

"Attribute for the MathProgBase solver."
struct MPBSolver <: MOI.AbstractOptimizerAttribute end

"Attribute for the MathProgBase status."
struct MPBSolutionAttribute <: MOI.AbstractModelAttribute end

"Struct to contain the MPB solution."
struct MPBSolution
    status::Symbol
    is_minimization::Bool
    objective_value::Float64
    primal_solution::Dict{MOI.VariableIndex, Float64}
end

"""
    Optimizer(; solver)

# Example

    model = Optimizer(solver = AmplNLWriter("/usr/bin/bonmin"))
"""
function Optimizer(; solver)
    model = MOIU.UniversalFallback(InnerModel{Float64}())
    MOI.set(model, MPBSolver(), solver)
    return model
end

Base.show(io::IO, ::Model) = println(io, "A MathProgBase model")

# We re-define is_empty and empty! to prevent the universal fallback from
# deleting the solver that we are caching in it.
function MOI.is_empty(model::Model)
    return MOI.is_empty(model.model) &&
        isempty(model.constraints) &&
        isempty(model.modattr) &&
        isempty(model.varattr) &&
        isempty(model.conattr) &&
        length(model.optattr) == 1 &&
        haskey(model.optattr, MPBSolver())
end

function MOI.empty!(model::Model)
    MOI.empty!(model.model)
    empty!(model.constraints)
    model.nextconstraintid = 0
    empty!(model.con_to_name)
    model.name_to_con = nothing
    empty!(model.modattr)
    empty!(model.varattr)
    empty!(model.conattr)
    mpb_solver = model.optattr[MPBSolver()]
    empty!(model.optattr)
    model.optattr[MPBSolver()] = mpb_solver
    return
end

set_to_bounds(set::MOI.LessThan) = (-Inf, set.upper)
set_to_bounds(set::MOI.GreaterThan) = (set.lower, Inf)
set_to_bounds(set::MOI.EqualTo) = (set.value, set.value)
set_to_bounds(set::MOI.Interval) = (set.lower, set.upper)
set_to_cat(set::MOI.ZeroOne) = :Bin
set_to_cat(set::MOI.Integer) = :Int

struct NLPEvaluator{T <: MOI.AbstractNLPEvaluator} <: MPB.AbstractNLPEvaluator
    inner::T
    variable_map::Dict{MOI.VariableIndex, Int}
end

"""
MathProgBase expects expressions with variables denoted by `x[i]` for contiguous
`i`. However, JuMP 0.19 creates expressions with `x[MOI.VariableIndex(i)]`. So
we have to recursively walk the expression replacing instances of
MOI.VariableIndex by a corresponding integer.
"""
function replace_variableindex_by_int(variable_map, expr::Expr)
    for (i, arg) in enumerate(expr.args)
        expr.args[i] = replace_variableindex_by_int(variable_map, arg)
    end
    return expr
end
function replace_variableindex_by_int(variable_map, expr::MOI.VariableIndex)
    return variable_map[expr]
end
replace_variableindex_by_int(variable_map, expr) = expr

function MPB.initialize(d::NLPEvaluator, requested_features::Vector{Symbol})
    MOI.initialize(d.inner, requested_features)
    return
end

function MPB.features_available(d::NLPEvaluator)
    return MOI.features_available(d.inner)
end

function MPB.eval_f(d::NLPEvaluator, x)
    return MOI.eval_objective(d.inner, x)
end

function MPB.eval_g(d::NLPEvaluator, g, x)
    return MOI.eval_constraint(d.inner, g, x)
end

function MPB.eval_grad_f(d::NLPEvaluator, g, x)
    return MOI.eval_objective_gradient(d.inner, g, x)
end

function MPB.jac_structure(d::NLPEvaluator)
    return MOI.jacobian_structure(d.inner)
end

function MPB.obj_expr(d::NLPEvaluator)
    expr = MOI.objective_expr(d.inner)
    return replace_variableindex_by_int(d.variable_map, expr)
end

function MPB.constr_expr(d::NLPEvaluator, i)
    expr = MOI.constraint_expr(d.inner, i)
    return replace_variableindex_by_int(d.variable_map, expr)
end

function MOI.optimize!(model::Model)
    mpb_solver = MOI.get(model, MPBSolver())

    # Get the variable data.
    variables = MOI.get(model, MOI.ListOfVariableIndices())
    num_var = length(variables)
    x_l = fill(-Inf, num_var)
    x_u = fill(Inf, num_var)
    x_cat = fill(:Cont, num_var)

    variable_map = Dict{MOI.VariableIndex, Int}()
    for (i, variable) in enumerate(variables)
        variable_map[variable] = i
    end

    for set_type in (MOI.LessThan{Float64}, MOI.GreaterThan{Float64})
        for c_ref in MOI.get(model,
            MOI.ListOfConstraintIndices{MOI.SingleVariable, set_type}())
            c_func = MOI.get(model, MOI.ConstraintFunction(), c_ref)
            c_set = MOI.get(model, MOI.ConstraintSet(), c_ref)
            v_index = variable_map[c_func.variable]
            lower, upper = set_to_bounds(c_set)
            x_l[v_index] = lower
            x_u[v_index] = upper
        end
    end

    for set_type in (MOI.ZeroOne, MOI.Integer)
        for c_ref in MOI.get(model,
            MOI.ListOfConstraintIndices{MOI.SingleVariable, set_type}())
            c_func = MOI.get(model, MOI.ConstraintFunction(), c_ref)
            c_set = MOI.get(model, MOI.ConstraintSet(), c_ref)
            v_index = variable_map[c_func.variable]
            x_cat[v_index] = set_to_cat(c_set)
        end
    end

    # Get the optimzation sense.
    opt_sense = MOI.get(model, MOI.ObjectiveSense())
    sense = opt_sense == MOI.MAX_SENSE ? :Max : :Min

    nlp_block = try
        MOI.get(model, MOI.NLPBlock())
    catch ex
        error("Expected a NLPBLock.")
    end

    # Extract constraint bounds.
    num_con = length(nlp_block.constraint_bounds)
    g_l = fill(-Inf, num_con)
    g_u = fill(Inf, num_con)
    for (i, bound) in enumerate(nlp_block.constraint_bounds)
        g_l[i] = bound.lower
        g_u[i] = bound.upper
    end

    mpb_model = MPB.NonlinearModel(mpb_solver)
    MPB.loadproblem!(
        mpb_model, num_var, num_con, x_l, x_u, g_l, g_u, sense,
        NLPEvaluator(nlp_block.evaluator, variable_map)
    )

    MPB.setvartype!(mpb_model, x_cat)

    MPB.optimize!(mpb_model)

    # MathProbBase solution
    primal_solution = Dict{MOI.VariableIndex, Float64}()
    for (variable, sol) in zip(variables, MPB.getsolution(mpb_model))
        primal_solution[variable] = sol
    end
    MOI.set(model, MPBSolutionAttribute(), MPBSolution(
        MPB.status(mpb_model),
        sense == :Min,
        MPB.getobjval(mpb_model),
        primal_solution
    ))
    return
end

function MOI.get(model::Model, ::MOI.VariablePrimal, var::MOI.VariableIndex)
    mpb_solution = MOI.get(model, MPBSolutionAttribute())
    if mpb_solution === nothing
        return nothing
    end
    return mpb_solution.primal_solution[var]
end

function MOI.get(model::Model, ::MOI.ObjectiveValue)
    mpb_solution = MOI.get(model, MPBSolutionAttribute())
    if mpb_solution === nothing
        return nothing
    end
    if mpb_solution.is_minimization
        return -mpb_solution.objective_value
    else
        return mpb_solution.objective_value
    end
end

function MOI.get(model::Model, ::MOI.TerminationStatus)
    mpb_solution = MOI.get(model, MPBSolutionAttribute())
    if mpb_solution === nothing
        return MOI.OPTIMIZE_NOT_CALLED
    end
    status = mpb_solution.status
    if status == :Optimal
        return MOI.OPTIMAL
    elseif status == :Infeasible
        return MOI.INFEASIBLE
    elseif status == :Unbounded
        return MOI.DUAL_INFEASIBLE
    elseif status == :UserLimit
        return MOI.OTHER_LIMIT
    elseif status == :Error
        return MOI.OTHER_ERROR
    end
    return MOI.OTHER_ERROR
end

function MOI.get(model::Model, ::MOI.PrimalStatus)
    mpb_solution = MOI.get(model, MPBSolutionAttribute())
    if mpb_solution === nothing
        return MOI.NO_SOLUTION
    end
    status = mpb_solution.status
    if status == :Optimal
        return MOI.FEASIBLE_POINT
    end
    return MOI.NO_SOLUTION
end

function MOI.get(model::Model, ::MOI.DualStatus)
    return MOI.NO_SOLUTION
end

end


using LinearAlgebra

@testset "Test QXZoo Circ" begin

    @testset "Test circuit creation" begin
        @testset "Test empty circuit" begin
            @test begin
                cct = QXZoo.Circuit.Circ()
                cct.num_qubits == 0
                cct.circ_ops.len == 0
                length(cct.gate_set) == 0
            end
        end

        @testset "Test gate n-qubit circuit" begin
            cct = QXZoo.Circuit.Circ(5)
            cct << QXZoo.DefaultGates.x(1)
            @test length(cct.gate_set) == 1
            @test cct.circ_ops.len == 1
            @test cct.num_qubits == 5
            @test_throws ErrorException cct << QXZoo.DefaultGates.x(6)
            @test_throws ErrorException cct << QXZoo.DefaultGates.c_x(1,6)
            @test_throws ErrorException cct << QXZoo.DefaultGates.c_x(6,1)
        end

        @testset "Test circuit iteration" begin
            cct1 = QXZoo.Circuit.Circ(5)

            for i in 1:5
                cct1 << QXZoo.DefaultGates.x(i)
            end

            for i in cct1.circ_ops
                @test i.gate_symbol === QXZoo.DefaultGates.GateSymbols.x
            end
        end

        @testset "Test circuit IO" begin
            cct = QXZoo.Circuit.Circ(5)

            cct << QXZoo.DefaultGates.x(1)
            cct << QXZoo.DefaultGates.x(2)
            cct << QXZoo.DefaultGates.h(3)
            cct << QXZoo.DefaultGates.c_z(3,1)

            QXZoo.Circuit.export_circuit(cct, "test_circ")
            @test isfile("test_circ.jld2") === true

            data = QXZoo.Circuit.load_circuit("test_circ")
            cct2 = data["circuit"]

            @test cct.circ_ops.len == cct2.circ_ops.len
            @test cct.gate_set == cct2.gate_set

            h1 = cct.circ_ops.head
            h2 = cct2.circ_ops.head

            for i in 1:cct.circ_ops.len
                @test h1.data === h2.data
                h1 = h1.next
                h2 = h2.next
            end
        end

        @testset "Test defined gateset" begin
            gates = Set([QXZoo.DefaultGates.GateSymbols.x, QXZoo.DefaultGates.GateSymbols.h, QXZoo.DefaultGates.GateSymbols.r_phase(pi/5)])
            cct = QXZoo.Circuit.Circ(5, gates)
            @test cct.num_qubits == 5
            @test cct.gate_set == gates
        end

        @testset "Test circuit combine" begin
            cct1 = QXZoo.Circuit.Circ(5)
            cct2 = QXZoo.Circuit.Circ(5)
            cct3 = QXZoo.Circuit.Circ(5)
            cct4 = QXZoo.Circuit.Circ(6)

            for i in 1:5
                cct1 << QXZoo.DefaultGates.x(i)
                cct3 << QXZoo.DefaultGates.x(i)
                cct4 << QXZoo.DefaultGates.x(i)
            end
            for i in 1:4
                cct2 << QXZoo.DefaultGates.c_x(i, i+1)
                cct3 << QXZoo.DefaultGates.c_x(i, i+1)
                cct4 << QXZoo.DefaultGates.c_x(i, i+1)
            end

            @test cct3.circ_ops.len == cct1.circ_ops.len + cct2.circ_ops.len
            @test cct3.gate_set == union(cct1.gate_set, cct2.gate_set)

            cct1 << cct2

            @test cct3.circ_ops.len == cct1.circ_ops.len
            @test cct2.circ_ops.len == 0
            @test cct3.gate_set == cct1.gate_set

            h1 = cct1.circ_ops.head
            h3 = cct3.circ_ops.head

            for i in 1:cct1.circ_ops.len
                @test h1.data === h3.data
                h1 = h1.next
                h3 = h3.next
            end

            @test_throws ErrorException cct3 << cct4

        end
    end
end


"Create an empty blueprint for `n` copies of a chromosome of `len` base pairs."
plan_chrom(len::Int, n::Int) = ChromosomeBlueprint(len, n)

function plan_chrom(; len::Int = 500, ploidy::Int = 1, ngenome::Int = 1, het = 0.0, div::Float64 = 0.0)
    bp = ChromosomeBlueprint(len, ploidy * ngenome)
    @info "- ✔ Created blank blueprint"
    if div > 0.0
        if ngenome > 1
            nsites = Int(ceil(len * div))
            bp = plan_het(bp, nsites, repeat(1:ngenome, inner = ploidy))
            @info "- ✔ Planned genome diverging polymorphisms"
        else
            @error "- ✘ Could not plan genome diverging polymorphisms: only one genome is planned"
        end
    end

    bp = _try_plan_het(bp, ploidy, ngenome, het)

    return bp
end

function _try_plan_het(bp::ChromosomeBlueprint, ploidy::Int, ngenome::Int, het::Float64)
    @info "A single value $het has been provided to plan heterozygosity"
    if het > 0.0 && ngenome === 1 && ploidy === 2
        nsites = Int(floor(length(bp) * het))
        bp = plan_het(bp, nsites, 2)
        @info "- ✔ Planned heterozygosity"
    else
        @error """- ✘ Could not plan heterozygosity:
                      Don't know how to plan heterozygosity when:
                      Ploidy is $ploidy
                      The number of genomes is $ngenome
                      het has been set as $het"""
    end
    return bp
end


"""
    module Hamiltonians

This module defines Hamiltonian types and functions for working with
Hamiltonians.

## Exported concrete Hamiltonian types include
- [`MatrixHamiltonian`](@ref)
- [`HubbardReal1D`](@ref)
- [`ExtendedHubbardReal1D`](@ref)
- [`HubbardMom1D`](@ref)
- [`BoseHubbardMom1D2C`](@ref)
- [`BoseHubbardReal1D2C`](@ref)

## Wrappers
- [`GutzwillerSampling`](@ref)
- [`GuidingVectorSampling`](@ref)

## Other
- [`G2Correlator`](@ref)

## Interface for working with Hamiltonians
- [`AbstractHamiltonian`](@ref): defined in the module [`Interfaces`](@ref)
"""
module Hamiltonians

using Parameters, StaticArrays, LinearAlgebra, SparseArrays
using FFTW
using Setfield

using ..StochasticStyles
using ..DictVectors
using ..BitStringAddresses
using ..ConsistentRNG

using ..Interfaces
import ..Interfaces: diagonal_element, num_offdiagonals, get_offdiagonal, starting_address,
    offdiagonals, random_offdiagonal, LOStructure

export AbstractHamiltonian, TwoComponentBosonicHamiltonian
export dimension, rayleigh_quotient, momentum
export BasisSetRep

export MatrixHamiltonian
export HubbardReal1D, HubbardMom1D, ExtendedHubbardReal1D, HubbardRealSpace
export HubbardReal1DEP, shift_lattice, shift_lattice_inv
export HubbardMom1DEP
export BoseHubbardMom1D2C, BoseHubbardReal1D2C
export GutzwillerSampling, GuidingVectorSampling
export Transcorrelated1D
export hubbard_dispersion, continuum_dispersion

export G2Correlator

export LatticeGeometry, PeriodicBoundaries, HardwallBoundaries, LadderBoundaries
export num_neighbours, neighbour_site

export sparse # from SparseArrays

include("abstract.jl")
include("offdiagonals.jl")
include("operations.jl")
include("geometry.jl")
include("excitations.jl")

include("MatrixHamiltonian.jl")

include("HubbardReal1D.jl")
include("HubbardReal1DEP.jl")
include("HubbardMom1D.jl")
include("HubbardMom1DEP.jl")
include("HubbardRealSpace.jl")
include("ExtendedHubbardReal1D.jl")

include("BoseHubbardReal1D2C.jl")
include("BoseHubbardMom1D2C.jl")

include("GutzwillerSampling.jl")
include("GuidingVectorSampling.jl")

include("Transcorrelated1D.jl")

include("TwoBodyCorrelation.jl")

end


"""
    SeisWindowPatch(in,out;<keyword arguments>)
Window a seis file using header words.

# Arguments
- `in::String`: filename of input
- `out::String`: filename of output
- `key=[]`: note that windowing along the time axis is achieved by using the key "t".
- `minval=[]`
- `maxval=[]`
- `it_nt=9e9`

*Credits: AS, FC, 2017*
"""
function SeisWindowPatch(in::String,out::String;key=[],minval=[],maxval=[],it_nt=9e9)
  println("processing patch ",out)
  DATAPATH = get(ENV,"DATAPATH",join([pwd(),"/"]))
  extent = ReadTextHeader(in)
  tmin = extent.o1
  tmax = extent.o1 + extent.d1*(extent.n1-1)
  for ikey=1:length(key)
	   if key[ikey] == "t"
	      tmin = minval[ikey]
	      tmax = maxval[ikey]
	   end
  end

  #itmin is the sample at which tmin is saved
  itmin = convert(Int32,round((tmin - extent.o1)/extent.d1) + 1)

  if (itmin < 1)
	   itmin = 1
  end
  itmax = convert(Int32,round((tmax - extent.o1)/extent.d1) + 1)
  if (itmax) > extent.n1
	   itmax = extent.n1
  end

  SeisWindowHeadersPatch(in, out; key=key, minval=minval, maxval=maxval, tmin=tmin,
                      tmax=tmax,it_nt=it_nt)
  FetchTracesPatch(in,out;itmin=itmin,itmax=itmax ,key=key, minval=minval, maxval=maxval)
  tmp = join(["tmp_SeisWindow_",string(round(Int,rand()*100000))])

 #>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
 #@compat SeisProcessHeaders(out, tmp, [UpdateHeader],[Dict(:itmin=>itmin,:itmax=>itmax)])

  itrace_in = 1
  itrace_out = 1
  ntrace=1000000
  NX = GetNumTraces(out)
  while itrace_in <= NX
    nx = NX - itrace_in + 1
	  ntrace = nx > ntrace ? ntrace : nx
	  h1 = SeisReadHeaders(out, group="some", itrace=itrace_in,
                                 ntrace=ntrace)
    ext=ReadTextHeader(out)

	  num_traces_in = length(h1)
	  h2=UpdateHeaderPatch(h1;itrace=itrace_in,itmin=itmin,itmax=itmax)
	  h1 = copy(h2)
	  num_traces_out = length(h1)
	  SeisWriteHeaders(tmp, h1, itrace=itrace_out)
    itrace_in += num_traces_in
	  itrace_out += num_traces_out
  end

  filename_h_tmp = join([DATAPATH tmp "@headers@"])
  filename_h_out = join([DATAPATH out "@headers@"])
  cp(filename_h_tmp,filename_h_out,force=true);
  rm(filename_h_tmp);
    #rm(tmp);
end

########################################################################################
#####################################################################################
function FetchTracesPatch(in::String, out::String; ntrace=500, itmin=round(Int,1),
    itmax=round(Int,9999999999),key=[], minval=[], maxval=[])

    NX = GetNumTraces(out)

    itrace = round(Int,1)
    filename_data_in = ParseDataName(in)

    DATAPATH = get(ENV,"DATAPATH",join([pwd(),"/"]))
    filename_data_out = join([DATAPATH out "@data@"])

    stream_in  = open(filename_data_in)

    stream_out  = open(filename_data_out,"w")
    extent = ReadTextHeader(in)
    nt = extent.n1

    if itmin < 1
      itmin = round(Int,1)
    end
    if itmax > nt
      itmax = round(Int,nt)
    end

    nhead = length(fieldnames(Header))

    while itrace <= NX
      if (itrace > 1)
	       stream_out = open(filename_data_out,"a")
	    end
      nx = NX - itrace + 1

      ntrace = nx > ntrace ? ntrace : nx
      h  = SeisReadHeaders(out,group="some",itrace=itrace,ntrace=ntrace)
      d = zeros(Float32,itmax-itmin+1,ntrace)

      SeekTracesPatch!(d,stream_in,h,itmin,itmax,nt,round(Int,ntrace))

    write(stream_out,d)

    close(stream_out)
    itrace += ntrace
  end
  close(stream_in)
end

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

function SeekTracesPatch!(d::AbstractArray{T,2}, stream_in::IOStream,
                        h::Array{Header,1},itmin,itmax,nt,ntrace) where T

    d1 = zeros(Float32,nt)
    for ix = 1 : ntrace
	     position = 4*nt*(h[ix].tracenum-1)
	     seek(stream_in,position)
         d1 = read!(stream_in,Array{Float32}(undef, nt))
         for it = itmin : itmax
	        d[it - itmin + 1,ix] = d1[it]
	     end
    end
    nothing
end

##############################################################################
#function SeekTracesPatch{T}(d::AbstractArray{T,2}, h::Array{Header,1},out::String, stream_in::IOStream,stream_out::IOStream, itmin,itmax,nt,ntrace;key=[],minval=[],maxval=[])
# println(itmin," ",itmax," ",nt)
#    d1 = zeros(Float32,nt)
#extent_out=ReadTextHeader(out)
#nx2=extent_out.n3
#nx3=extent_out.n4
#nx4=extent_out.n5
#    for ix = 1 : ntrace
#	position = 4*nt*(h[ix].tracenum-1)
#	seek(stream_in,Int(position))

#	d1 = read(stream_in,Float32,nt)
#	for it = itmin : itmax
#	    d[it - itmin + 1,ix] = d1[it]

#	end

#        itrace_out=(h[ix].imx-minval[2])*nx2*nx3*nx4+(h[ix].imy-minval[3])*nx3*nx4+(h[ix].ihx-minval[4])*nx4+h[ix].ihy-minval[5]+1

#        position_out=4*nt*(itrace_out-1)
#println(out," ",h[ix].imx," ",minval[2]," ",itrace_out)
#        seek(stream_out,Int(position_out))
#        write(stream_out,convert(Array{Float32,1},d[1:itmax-itmin+1,ix]))
#    end
#    nothing
#end

##############################################################################
##############################################################################
function UpdateHeaderPatch(h;itrace=1,itmin=1,itmax=9e9)
    ot = (itmin-1)*h[1].d1 + h[1].o1

    nt = itmax - itmin + 1
    for j = 1 : length(h)
	h[j].o1 = ot
	h[j].n1 = nt
	h[j].tracenum = j + itrace - 1
    end
    return h
end


export hamming, profile, consensus, numseq

"""
	countgaps(s::BioSequence)
"""
function countgaps(s::BioSequence)
	return BioSequences.count(x->isgap(x)||isambiguous(x), s)
end
countgaps(s::AbstractStrain) = countgaps(s.seq)

function hamming_ignore_ambiguous(a::BioSequence, b::BioSequence)
	h = 0
	for (x,y) in zip(a,b)
		if x != y && !isambiguous(x) && !isambiguous(y)
			h += 1
		end
	end
	return h
end
function hamming_ambiguous_mismatch(a::BioSequence, b::BioSequence)
	h = 0
	for (x,y) in zip(a,b)
		if x != y
			h += 1
		end
	end
	return h
end
"""
	hamming(a::BioSequence, b::BioSequence, [normalization=false]; ambiguous=:ignore)
	hamming(a::AbstractStrain, b::AbstractStrain [, normalization=false]; ambiguous=:ignore)

`ambiguous` keyword:
- `:ignore` : ambiguous symbols are considered equal to anything .
- `:mismatch` : ambiguous symbols behave like regular symbols.
"""
function hamming(a::BioSequence, b::BioSequence,
	normalization=false;
	ambiguous=:ignore)
	if ambiguous == :ignore && normalization
		return hamming_ignore_ambiguous(a,b)/length(a)
	elseif ambiguous == :ignore && !normalization
		return Float64(hamming_ignore_ambiguous(a,b))
	elseif ambiguous == :mismatch && normalization
		return hamming_ambiguous_mismatch(a,b)/length(a)
	elseif ambiguous == :mismatch && !normalization
		return Float64(hamming_ambiguous_mismatch(a,b))
	else
		@error "Incorrect arguments."
	end
end
hamming(a::AbstractStrain, b::AbstractStrain, normalization=false; ambiguous=:ignore) = hamming(a.seq, b.seq, normalization, ambiguous=ambiguous)
"""
	hamming(P::Profile{A}, Q::Profile{A}, normalization=false;
				ambiguous=:ignore) where A <: BioSymbol
"""
function hamming(P::Profile{A}, Q::Profile{A}, normalization=false;
				ambiguous=:ignore) where A <: BioSymbol
	length(P) != length(Q) && @error "Profiles of different lengths"
	out = 0.
	for (i,(p,q)) in enumerate(zip(P,Q))
		val = 1.
		for a in BioSequences.alphabet(A)
			if ambiguous == :ignore && isambiguous(a)
				val -= get(p,a,0.) + get(q,a,0.) - get(p,a,0.)*get(q,a,0.)
			else
				val -= get(p,a,0.)*get(q,a,0.)
			end
		end
		out += val
	end
	if normalization
		out /= length(P)
	end
	return out
end
function hamming(P::Profile, s::BioSequence, normalization = false; ambiguous = :ignore)
	length(P) != length(s) && @error "Sequence and profile of different lengths"
	out = 0.
	for (i,(a,p)) in enumerate(zip(s,P))
		out += 1.
		for b in keys(p)
			if b == a || (ambiguous==:ignore && (isambiguous(a)||isambiguous(b)))
				out -= p[b]
			end
		end
	end
	if normalization
		out /= length(P)
	end
	return out
end
hamming(s::BioSequence, P::Profile) = hamming(P,s)
hamming(s::Strain, P::Profile) = hamming(s.seq, P)
hamming(P::Profile, s::Strain) = hamming(s.seq, P)

"""
	consensus(P::Profile{A}) where A<:BioSymbol
	consensus(X::Array{<:BioSequence,1})

Return the consensus *sequence*, of type `LongSequence{alphabet}` where `eltype(alphabet)==A`. Does **NOT** handle ties. In case of a tie, the chosen symbol is undetermined.

	consensus(X::Array{Strain} [, data = Dict(:strain=>"consensus")])

Return consensus of `X`, with field `data`. Does **NOT** handle ties. In case of a tie, the chosen symbol is undetermined.
"""
function consensus(P::Profile{A}) where A<:BioSymbol
	seq = LongSequence{symbol_to_alphabet(A)}(length(P))
	for (i,p) in enumerate(P)
		if !isempty(p)
			seq[i] = findmax(p)[2]
		else
			seq[i] = ambiguous(A)
		end
	end
	return seq
end
consensus(X::Array{<:BioSequence,1}) = consensus(Profile(X))
function consensus(X::Array{<:Strain,1}, data = Dict(:strain=>"consensus"))
	p = Profile(X)
	seq = consensus(p)
	return Strain(seq, data)
end

"""
	numseq(s::LongSequence{AminoAcidAlphabet})
	numseq(S::Array{LongSequence{AminoAcidAlphabet}})
"""
function numseq(s::LongSequence{AminoAcidAlphabet})
	σ = Array{Int64,1}(undef, length(s))
	for (i,a) in enumerate(s)
		σ[i] = haskey(aa2int, a) ? aa2int[a] : 1
	end
	return σ
end
function numseq(S::Array{LongSequence{AminoAcidAlphabet}})
	if isempty(S)
		out = Array{Int64,2}(undef, 0, 0)
	else
		out = Array{Int64,2}(undef, length(S), length(S[1]))
		for (i,s) in enumerate(S)
			out[i,:] .= numseq(s)
		end
	end
	return out
end


import .Interpolations: AbstractInterpolation, AbstractInterpolationWrapper, BSplineInterpolation,
                        LanczosInterpolation, ScaledInterpolation, GriddedInterpolation, MonotonicInterpolation,
                        Extrapolation, itptype, coefficients,  itpinfo, value_weights, WeightedAdjIndex,
                        allbetween, interp_getindex, WeightedIndex, weightedindex_parts, maybe_weightedindex,
                        BoundsCheckStyle, _checkbounds
import Base: @propagate_inbounds
export UnSafeInterp, preweight

# a wrapper to turn off the boundary check when we use broadcast
struct UnSafeInterp{T,N,ITPT,IT} <: AbstractInterpolationWrapper{T,N,ITPT,IT}
    itp::ITPT
    UnSafeInterp(itp::AbstractInterpolation{T,N,IT}) where {T,N,IT} = new{T,N,typeof(itp),IT}(itp)
end
Base.parent(usitp::UnSafeInterp) = usitp.itp
@inline (usitp::UnSafeInterp{T,N})(args::Vararg{Any,N}) where {T,N} = @inbounds usitp.itp(args...)

struct WeightedInterp{A}
    coefs::A
end
Base.parent(witp::WeightedInterp) = witp.coefs
@inline (witp::WeightedInterp{<:AbstractArray{T,N}})(args::Vararg{Any,N}) where {T,N} = @inbounds witp.coefs[args...]

# pre_weight
unsaled(r::AbstractUnitRange, x) = @lzb x .- first(r) .+ oneunit(eltype(r))
unsaled(r::AbstractRange, x) = @lzb (x .- first(r)) .* inv(step(r)) .+ oneunit(eltype(r))

@inline getcoefs(x::AbstractInterpolationWrapper) = parent(x) |> getcoefs
@inline getcoefs(x::AbstractInterpolation) = coefficients(x)
@inline getcoefs(x::WeightedInterp) = parent(x)

# preweight
@inline preweight(itp::Extrapolation{T,N}, args::Vararg{Any,N}) where {T,N} =
    throw(ArgumentError("Extrapolation is not supported"))
@inline preweight(sitp::ScaledInterpolation{T,N}, args::Vararg{Any,N}) where {T,N} = begin
    @boundscheck _checkbounds(BoundsCheckStyle(sitp), sitp, args) ||
        Base.throw_boundserror(sitp, args)
    @inbounds preweight(sitp.itp, unsaled.(sitp.ranges, args)...)
end
@inline preweight(itp::BSplineInterpolation{T,N}, args::Vararg{Any,N}) where {T,N} = begin
    @boundscheck _checkbounds(BoundsCheckStyle(itp), itp, args) ||
        Base.throw_boundserror(itp, args)
    function weight_ind(itpflag, knotvec, x)
        makewi(y, ::Any) = begin
            pos, (coefs,) = weightedindex_parts((value_weights,), itpflag, knotvec, y)
            maybe_weightedindex(pos, coefs)
        end
        makewi.(x, getcoefs(itp) |> Ref)
    end
    (WeightedInterp(itp.coefs), weight_ind.(itpinfo(itp)..., args)...)
end

@require CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" begin

    import Adapt: adapt_structure, adapt
    adapt_structure(to, itp::BSplineInterpolation{T,N}) where {T,N} = begin
        coefs, parentaxes, it = itp.coefs, itp.parentaxes, itp.it
        coefs′ = adapt(to, coefs)
        Para = map(typeof, (coefs′,it,parentaxes))
        BSplineInterpolation{T,N,Para...}(coefs′, parentaxes, it)
    end

    adapt_structure(to, itp::LanczosInterpolation{T,N}) where {T,N} = begin
        coefs, parentaxes, it = itp.coefs, itp.parentaxes, itp.it
        coefs′ = adapt(to, coefs)
        parentaxes′ = adapts(to, parentaxes...)
        Para = map(typeof, (it,coefs′,parentaxes′))
        LanczosInterpolation{T,N,Para...}(coefs′, parentaxes′, it)
    end

    adapt_structure(to, itp::ScaledInterpolation{T,N}) where {T,N} = begin
        ranges = itp.ranges
        itp′ = adapt(to, itp.itp)
        IT = itptype(itp)
        ScaledInterpolation{T,N,typeof(itp′),IT,typeof(ranges)}(itp′, ranges)
    end

    adapt_structure(to, itp::Extrapolation{T,N}) where {T,N} = begin
        et = itp.et
        itp′ = adapt(to,  itp.itp)
        IT = itptype(itp)
        Extrapolation{T,N,typeof(itp′),IT,typeof(et)}(itp′, et)
    end

    adapt_structure(to, itp::UnSafeInterp{T,N}) where {T,N} =
        adapt(to, itp.itp) |> UnSafeInterp

    adapt_structure(to, itp::WeightedInterp) =
        adapt(to, itp.coefs) |> WeightedInterp

    # With Adapt v"3.3.1", there's no need to use Ref to force adapt
    # but I think the style hack is still needed.
    broadcasted(itp::Union{WeightedInterp, AbstractInterpolation}, args...) = begin
        args′ = broadcastable.(args)
        style = combine_styles(Ref(getcoefs(itp)), args′...)
        broadcasted(style, itp, args′...)
    end

    #make checkbounds work on gpu
    allbetween(l::Real, xs::AbstractArray, u::Real) = begin
        device(xs) isa GPU && return all(x -> l <= x <= u, xs)
        invoke(allbetween, Tuple{typeof(l),Any,typeof(u)}, l, xs, u)
    end

    # Type interfer and display fix
    import GPUArrays
    import GPUArrays: AbstractGPUArray
    GPUArrays._getindex(A::AbstractGPUArray{T,N}, I::Vararg{Union{Int,WeightedIndex},N}) where {T,N} =
        interp_getindex(A, I, ntuple(Returns(0), Val(N))...)

end


# edges = [[2000, [35, 40], [355,500,600], 1000],
#         [2001, [35, 50], [100,500,600], 2000],
#         [2002, [37, 40], [100,400,600], 3000]];
# edges = hcat(edges...)';

# create_graph_dicts(edges)

# # pre_to_post, post_to_pre = create_post_pre_dicts(segs)
# @test length(pre_to_post[35]) == 2
# @test pre_to_post[35][1] == 40
# @test pre_to_post[35][2] == 50
# @test pre_to_post[37][1] == 40
# @test post_to_pre[40][1] == 35
# @test post_to_pre[40][2] == 37
# @test post_to_pre[50][1] == 35


"""
    AbstractSolver{T} <: MathOptInterface.AbstractNLPEvaluator

Abstract solver for trajectory optimization problems

# Interface
Any type that inherits from `AbstractSolver` must define the following methods:
```julia
model = get_model(::AbstractSolver)::AbstractModel
obj   = get_objective(::AbstractSolver)::AbstractObjective
E     = get_cost_expansion(::AbstractSolver)::QuadraticExpansion  # quadratic error state expansion
Z     = get_trajectory(::AbstractSolver)::Traj
n,m,N = Base.size(::AbstractSolver)
x0    = get_initial_state(::AbstractSolver)::StaticVector
solve!(::AbstractSolver)
```

Optional methods for line search and merit function interface. Note that these do not
    have to return `Traj`
```julia
Z     = get_solution(::AbstractSolver)  # current solution (defaults to get_trajectory)
Z     = get_primals(::AbstractSolver)   # current primals estimate used in the line search
dZ    = get_step(::AbstractSolver)      # current step in the primal variables
```

Optional methods
```julia
opts  = options(::AbstractSolver)       # options struct for the solver. Defaults to `solver.opts`
st    = solver_stats(::AbstractSolver)  # dictionary of solver statistics. Defaults to `solver.stats`
iters = iterations(::AbstractSolver)    #
```
"""
abstract type AbstractSolver{T} end

"$(TYPEDEF) Unconstrained optimization solver. Will ignore
any constraints in the problem"
abstract type UnconstrainedSolver{T} <: AbstractSolver{T} end

"""$(TYPEDEF)
Abstract solver for constrained trajectory optimization problems

In addition to the methods required for `AbstractSolver`, all `ConstrainedSolver`s
    must define the following method
```julia
get_constraints(::ConstrainedSolver)::ConstrainSet
```
"""
abstract type ConstrainedSolver{T} <: AbstractSolver{T} end

# methods
@inline iterations(solver::AbstractSolver) = solver.stats.iterations
@inline status(solver::AbstractSolver) = solver.stats.status
@inline is_constrained(::Type{<:AbstractSolver})::Bool = true
@inline is_constrained(::Type{<:ConstrainedSolver})::Bool = true
@inline is_constrained(::Type{<:UnconstrainedSolver})::Bool = false
@inline is_constrained(solver::AbstractSolver) = is_constrained(typeof(solver)) && !isempty(get_constraints(solver))

"""
    terminate!(solver::AbstractSolver)

Perform any necessary actions after finishing the solve.
"""
function terminate!(solver::AbstractSolver)
    # Time solve
    stat = solver.stats
    stat.tsolve = (time_ns() - stat.tstart)*1e-6  # in ms

    # Delete extra stats entries, only if terminal solver
    trim!(solver.stats, solver.solver_name)

    # Print solve summary
    if solver.opts.show_summary && is_parentsolver(solver)
        print_summary(solver)
    end
end

"""
    TerminationStatus

* `UNSOLVED`: Initial value. Solve either hasn't been attempted or is in process..
* `SOLVE_SUCCEEDED`: Solve met all the required convergence criteria.
* `MAX_ITERATIONS`: Solve was unable to meet the required convergence criteria within the maximum number of iterations.
* `MAX_ITERATIONS_OUTER`: Solve was unable to meet the required constraint satisfaction the maximum number of outer loop iterations.
* `MAXIMUM_COST`: Cost exceeded maximum allowable cost.
* `STATE_LIMIT`: State values exceeded the imposed numerical limits.
* `CONTROL_LIMIT`: Control values exceeded the imposed numerical limits.
* `NO_PROGRESS`: iLQR was unable to make any progress for `dJ_counter_limit` consecutive iterations.
* `COST_INCREASE`: The cost increased during the iLQR forward pass.
"""
@enum TerminationStatus begin
    UNSOLVED = 1
    MAX_ITERATIONS_OUTER = 2
    MAX_ITERATIONS = 3
    SOLVE_SUCCEEDED = 4
    MAXIMUM_COST = 5
    STATE_LIMIT = 6
    CONTROL_LIMIT = 7
    NO_PROGRESS = 8
    COST_INCREASE = 9
end

function print_summary(solver::S) where S <: AbstractSolver
    stat = stats(solver)
    col_h1 = crayon"bold green"
    col_h2 = crayon"bold blue"
    col0 = Crayon(reset=true)
    get_color(v::Bool) = v ? crayon"green" : crayon"red"

    # Info header
    println(col_h1, "\nSOLVE COMPLETED")
    print(col0," solved using the ")
    print(col0, crayon"bold cyan", solvername(solver))
    print(col0, " Solver,\n part of the Altro.jl package developed by the REx Lab at Stanford and Carnegie Mellon Universities\n")

    # Stats
    println(col_h2, "\n  Solve Statistics")
    println(col0, "    Total Iterations: ", iterations(solver))
    println(col0, "    Solve Time: ", stat.tsolve, " (ms)")

    # Convergence
    println(col_h2, "\n  Covergence")
    if iterations(solver) == 0
        println(crayon"red", "    Solver failed to make it through the first iteration.")
    else
        println(col0, "    Terminal Cost: ", stat.cost[end])
        println(col0, "    Terminal dJ: ", get_color(stat.dJ[end] < solver.opts.cost_tolerance), stat.dJ[end])
        println(col0, "    Terminal gradient: ", get_color(stat.gradient[end] < solver.opts.gradient_tolerance), stat.gradient[end])
        if is_constrained(solver)
            println(col0, "    Terminal constraint violation: ", get_color(stat.c_max[end] < solver.opts.constraint_tolerance), stat.c_max[end])
        end
    end
    println(col0, "    Solve Status: ", crayon"bold", get_color(status(solver) == SOLVE_SUCCEEDED), status(solver))
    print(Crayon(reset=true))  # reset output color
end

# logging

log_level(solver::AbstractSolver) = OuterLoop

is_verbose(solver::AbstractSolver) =
    log_level(solver) >= LogLevel(-100*solver.opts.verbose)

function set_verbosity!(solver::AbstractSolver)
    llevel = log_level(solver)
    if is_verbose(solver)
        set_logger()
        Logging.disable_logging(LogLevel(llevel.level-1))
    else
        Logging.disable_logging(llevel)
    end
end

function clear_cache!(solver::AbstractSolver)
    llevel = log_level(solver)
    if is_verbose(solver)
        SolverLogging.clear_cache!(global_logger().leveldata[llevel])
    end
end


function evaluation_performed(obj::KalmanObserver, arg0::KalmanEstimation)
    return jcall(obj, "evaluationPerformed", void, (KalmanEstimation,), arg0)
end


module BenchDictMigration

using BenchmarkTools
using ConcurrentCollections
using ConcurrentCollections.Implementations:
    LINEAR_PROBING_DICT_EXPAND_BASESIZE, migrate_serial!

pad16(x) = string(x; pad = 16)

function generate(f = pad16; datasize = LINEAR_PROBING_DICT_EXPAND_BASESIZE[])
    vs = UInt64.(1:datasize)
    ks = f.(vs)
    dict = ConcurrentDict{eltype(ks),eltype(vs)}(zip(ks, vs))
    return dict
end

const CACHE = Ref{Any}()

function setup(; generate_options...)
    CACHE[] = Dict(
        "String-UInt64" => generate(; generate_options...),
        "Uint32-UInt64" => generate(UInt32; generate_options...),
        "UInt64-UInt64" => generate(UInt64; generate_options...),
        # "Uint32-Nothing" => ???,
    )

    suite = BenchmarkGroup()
    for key in keys(CACHE[])
        CacheType = typeof(CACHE[][key])
        suite[key] = @benchmarkable(
            migrate_serial!(newslots, slots),
            setup = begin
                dict = CACHE[][$key]::$CacheType
                slots = copy(dict.slots)
                newslots = similar(slots, length(slots) * 2)
            end,
            evals = 1,
        )
    end
    return suite
end

function clear()
    CACHE[] = nothing
end

end  # module


# -------------------- RANDOM EFFECTS ------------------------------

function rkey(key::AbstractString, ilevel::AbstractString="") ilevel != "" ? key*" ("*ilevel*")" : key  end
function rkey(row::DataFrameRow) return rkey(row[:class], row[:level]) end

abstract RanEffect <: ModelEffect

function getRandomFormula(rlist::Vector{Symbol}, excludeCol::Symbol=:NA, df_name::AbstractString="df_data")
    fmula=""
    for s in setdiff(rlist, [excludeCol])
        fmula=fmula*"+"*df_name*"[:"*string(s)*"]"
    end
    return fmula[2:end]
end

function genRandomTotals(reff::RanEffect, df_data::DataFrame)
    model=reff.v_model
    idf=reff.src
    rcols=Symbol[]
    for ranfx in unique(idf[:class])    # -- Loop through each random effect :: create random df_data columns - and add list to reff
        df_data[Symbol(ranfx)] = map(x->string(x), df_data[Symbol(ranfx)])    # Convert all ranfx to string ... since they are categorical anyway
        dest_colname=Symbol("reff_"*lowercase(ranfx)*"_"*model)
        push!(rcols,dest_colname)
        df_data[dest_colname]=0.0
        for row in eachrow(idf[idf[:class].==ranfx,:])   # Loop through each level for outter random effect
            level=string(row[:level])
            println("populating : ",dest_colname," : for :: (",level,")")
            df_data[ (df_data[Symbol(ranfx)].==level) & (df_data[:group].==0),dest_colname]=row[:B0]
            df_data[ (df_data[Symbol(ranfx)].==level) & (df_data[:group].==1),dest_colname]=row[:B1]
        end
    end
    reff.rcols=rcols
    reff.fmula=getRandomFormula(rcols)

    function calc_reff(row::DataFrameRow)
        tot=0
        for x in 1:length(row)
            tot=tot+row[x]
        end
        return tot
    end
    df_data[Symbol("reff_"*model)] = map(x->calc_reff(x),   eachrow(df_data[rcols]) )
    df_data[Symbol("reff_"*model)] = map(Float64,df_data[Symbol("reff_"*model)])      #convert to float

    # --- Random Effect Mean Scor cols ---
    for row = eachrow(reff.src)
        k = row[:key]
        #needs review
        dest_col0=Symbol(k*"_"*model*"_0")
        dest_col1=Symbol(k*"_"*model*"_1")
        rcols = setdiff(reff.rcols, [Symbol("reff_"*lowercase(row[:class])*"_"*model)])
        println("generating : ",k*"_"*model*"_X := sum(",rcols,", :pre_"*model*"_scoreX,+ BX) ")
        function tot_reff(row::DataFrameRow, B::Float64)
            tot=0
            for x in 1:length(row)
                tot=tot+row[x]
            end
            tot = exp(tot+B)
            if model== "pen"
                return tot = tot / (1+tot)  #exp(eval(parse(fmula0))) ./ (1+exp(eval(parse(fmula0))))
            end
            return tot
        end
        #model
        df_data[dest_col0] = map(x->tot_reff(x, Float64(row[:B0]) ),   eachrow(df_data[vcat(rcols,Symbol("pre_"*model*"_score0"))]) )
        df_data[dest_col1] = map(x->tot_reff(x, row[:B1] ),   eachrow(df_data[vcat(rcols,Symbol("pre_"*model*"_score1"))]) )
        df_data[dest_col0] = map(Float64,df_data[dest_col0])      #convert to float
        df_data[dest_col1] = map(Float64,df_data[dest_col1])      #convert to float
    end

end

type ROcc <: RanEffect
    sdf::DataFrame
    src::DataFrame
    rcols::Vector{Symbol}
    fmula::AbstractString
    v_model::AbstractString
    function ROcc(df_data::DataFrame, idf::DataFrame)  this=new(); this.v_model="occ"; this.sdf=deepcopy(SDF);
         this.src = idf
         for row in eachrow(idf)
             pushSDFrow!(this.sdf,OrderedDict(:key => row[:key], :model => this.v_model))
         end
        delete!(this.sdf, :B)
        delete!(this.sdf, :SE)
        delete!(this.sdf, :P)
        if length(this.sdf[1]) > 0 genRandomTotals(this,df_data)  end
        return this
    end
end

type RDolOcc <: RanEffect
    sdf::DataFrame
    src::DataFrame
    rcols::Vector{Symbol}
    fmula::AbstractString
    v_model::AbstractString
    function RDolOcc(df_data::DataFrame, idf::DataFrame)  this=new(); this.v_model="dolocc"; this.sdf=deepcopy(SDF);
         this.src = idf
         for row in eachrow(idf)
             pushSDFrow!(this.sdf,OrderedDict(:key => row[:key], :model => this.v_model))
         end
        delete!(this.sdf, :B)
        delete!(this.sdf, :SE)
        delete!(this.sdf, :P)
        if length(this.sdf[1]) > 0  genRandomTotals(this,df_data) end
        return this
    end
end

type RPen <: RanEffect
    sdf::DataFrame
    src::DataFrame
    rcols::Vector{Symbol}
    fmula::AbstractString
    v_model::AbstractString
    function RPen(df_data::DataFrame, idf::DataFrame)  this=new(); this.v_model="pen"; this.sdf=deepcopy(SDF);
         this.src = idf
         for row in eachrow(idf)
             pushSDFrow!(this.sdf,OrderedDict(:key => row[:key], :model => this.v_model))
         end
        delete!(this.sdf, :B)
        delete!(this.sdf, :SE)
        delete!(this.sdf, :P)
        if length(this.sdf[1]) > 0 genRandomTotals(this,df_data)  end
        return this
    end
end

function Base.show(io::IO, raneffect::RanEffect)
    show(io,raneffect.sdf)
end


using ImageView, ImageMagick
using Images, ImageDraw, TestImages
using LCMCore, CaesarLCMTypes
using JSON
using SynchronySDK

function lcm_ShowData(channel, msg::brookstone_supertype_t, images::Vector{BigDataElementRequest})
    img = msg.img
    dataElementRequest = BigDataElementRequest("CamImage", "Mongo", "Brookstone camera data", base64encode(img.data), "image/jpeg")
    push!(images, dataElementRequest)
end

# Get the imagery from the LCM log
# TODO: Get it from the DB when we have access
cd(Pkg.dir("SynchronySDK"))
lcm = LCMLog("examples/data/brookstone_3.lcmlog")
images = Vector{BigDataElementRequest}()
subscribe(lcm, "BROOKSTONE_ROVER", (c,m)->lcm_ShowData(c,m, images), brookstone_supertype_t)
while handle(lcm)
end

# Let's run an AprilTags detector on the sequence.
using AprilTags
# Simple method to show the image with the tags
function showImage(image, tags, canvas)
    # Convert image to RGB
    imageCol = RGB.(image)
    #draw color box on tag corners
    if tags != nothing && length(tags) != 0
        foreach(tag->drawTagBox!(imageCol, tag, 4), tags)
    end
    imshow(canvas["gui"]["canvas"], imageCol)
    sleep(tags != nothing  && length(tags) != 0 ? 1 : 0.03)
    return nothing
end

canvas = imshow(load(Pkg.dir("SynchronySDK") * "/examples/pexels_small.png"))
sleep(2)
detector = AprilTagDetector()
for imageBlob in images
    imData = base64decode(imageBlob.data)
    image = readblob(imData)
    tags = detector(image)
    @show tags
    showImage(image, tags, canvas)
end
# Close the images
ImageView.closeall()
# Free the detector
freeDetector!(detector)

# TODO - add in the bearing+range data.
# https://github.com/JuliaRobotics/AprilTags.jl/blob/master/test/runtests.jl#L73
# https://github.com/JuliaRobotics/Caesar.jl/blob/example/apriltagvisualizer/examples/apriltagserver/tag_exporter.jl
# https://github.com/JuliaRobotics/Caesar.jl/blob/example/apriltagvisualizer/examples/apriltagserver/tag2slam.jl#L29


using Distributions

mutable struct Simulation{N}
    heatmap::Array{Float64, N}
    stats::NamedTuple{(:mean, :std)}
end

function simulate(distribution, dims, n)
    tp = ntuple(i -> n, dims)
    heatmap = rand(distribution, tp...)
    return Simulation{dims}(heatmap, (mean = mean(heatmap), std = std(heatmap)))
end

sim = simulate(Normal(), 2, 1000);

#=
julia> sim = simulate(Normal(), 2, 1000);

julia> sim.heatmap
1000×1000 Array{Float64,2}:
  1.22768    -0.316595    1.53702     …   0.944001    1.14345
 -1.09027    -0.553394    1.23286     …   0.779135   -0.966926
  1.52175     0.297313    1.02013         0.6078     -0.299004
  ⋮                                   ⋱
  0.912278   -0.219911    0.513086       -1.37209     0.126663
  1.05601     0.650547   -0.490138       -0.321995    0.455638
 -0.764425    0.175759   -0.245518        0.724038   -0.637692

julia> sim.stats
(mean = 0.0003489176881672829, std = 1.0002976677059898)
=#

# accessors - getters
# they are just forwarding to the properties

get_heatmap(s::Simulation) = s.heatmap
get_stats(s::Simulation) = s.stats

# alternative naming convention
# note the extension of Base.size generic function
heatmap(s::Simulation) = s.heatmap
stats(s::Simulation) = s.stats

# accessors - setters
# we only want a setter for `points` but not `size`
# we can also add validation.

function heatmap!(
        s::Simulation{N},
        new_heatmap::AbstractArray{Float64, N}) where {N}
    s.heatmap = new_heatmap
    s.stats = (mean = mean(new_heatmap), std = std(new_heatmap))
    return nothing
end

# add validation
function heatmap!(
            s::Simulation{N},
            new_heatmap::AbstractArray{Float64, N}) where {N}
    if length(unique(size(new_heatmap))) != 1
        error("All dimensions must have same size")
    end
    s.heatmap = new_heatmap
    s.stats = (mean = mean(new_heatmap), std = std(new_heatmap))
    return nothing
end

#=
julia> heatmap!(sim, rand(10, 3))
ERROR: All dimensions must have same size
=#

# ------------------------------------------------------------------------
# marking fields private

mutable struct Simulation{N}
    _heatmap::Array{Float64, N}
    _stats::NamedTuple{(:mean, :std)}
end


using Flux
using Random

function get_batch(dataset, max_sequence::Integer, batch_size::Integer; set::Symbol = :train, max_regions=nothing)
    data = shuffle(dataset[set])

    # Positional data
    word_size = size(data[1])[1]
    region_length = max_sequence
    num_ins = size(data[1])[3]

    # Chop songs up along sequence length
    for j in 1:length(data)
        num_regions = Integer(ceil(size(data[j])[2] / max_sequence))
        if max_regions != nothing
            num_regions = min(num_regions, max_regions)
        end

        data[j] = [data[j][:,max(1, min(i, size(data[j])[2] - max_sequence)):i+min(max_sequence-1, size(data[j])[2]-i-1),:]
                    for i in 1:max_sequence:num_regions*max_sequence]
    end

    # Batch songs
    num_batches = Integer(ceil(length(data) / batch_size))
    data = [data[i:i+min(batch_size-1, length(data)-i)]
                for i in 1:batch_size:batch_size*num_batches]

    return data
end


# -----------------------------------------------------------------------------
# IMPORTS
# -----------------------------------------------------------------------------

# External libraries
using Logging                       # Logging is used for operation diagnostics
using Parameters                    # MetaICVIOpts are Parameters structs
using StatsBase: corspearman        # Rank correlation for cvi criterion values
using PyCall                        # PyCall object definition
using JLD                           # JLD is currently recommended for saving/loading ScikitLearn objects
using PyCallJLD                     # PyCall definition for serialization with JLD
using ScikitLearn                   # Classifiers are scikit-learn pyobjects
using ProgressMeter: @showprogress  # Data loading progress for training
using DelimitedFiles: readdlm       # Loading cvi data
# using BSON
# using ScikitLearn.Skcore: FitBit

# Custom libraries
using ClusterValidityIndices        # All Julia-implemented CVI definitions

# Local libraries
# Get the rocket kernel definitions
include("rocket.jl")
using .Rocket

# -----------------------------------------------------------------------------
# CONSTANTS
# -----------------------------------------------------------------------------

# const MetaICVIClassifier = ScikitLearn.Skcore.FitBit
const MetaICVIClassifier = PyCall.PyObject

# Top of the module for default paths
const module_dir(paths...) = joinpath(dirname(pathof(@__MODULE__)), "..", paths...)

# -----------------------------------------------------------------------------
# STRUCTURES
# -----------------------------------------------------------------------------

"""
    MetaICVIOpts()

Meta-ICVI module options.

# Examples
```julia-repl
julia> MetaICVIOpts()
```
"""
@with_kw mutable struct MetaICVIOpts @deftype Int
    # Which scikitlearn classifier to load
    classifier_selection::Symbol = :SGDClassifier
    # Scikitlearn classifier keyword arguments
    classifier_opts::NamedTuple = (loss="log", max_iter=30)
    # Size of ICVI window: [1, infty)
    icvi_window = 5; @assert icvi_window >= 1
    # Size of correlation window: [1, infty)
    correlation_window = 5; @assert correlation_window >= 1
    # Number of rocket kernels: [1, infty)
    n_rocket = 5; @assert n_rocket >= 1
    # Rocket file location
    rocket_file::String = module_dir("data", "models", "rocket.jld2")
    # Classifier file location
    classifier_file::String = module_dir("data", "models", "classifier.jld")
    # Display flag
    display::Bool = true
    # Flag to fail if any file is missing (rather than creating new objects)
    fail_on_missing::Bool = false
end # MetaICVIOpts

"""
    MetaICVIModule

Stateful information for a single MetaICVI module.

# Fields
- `opts::MetaICVIOpts`: options for construction.
- `cvis::Vector{AbstractCVI}`: list of cvis used for computing the CVIs.
- `criterion_values::RealVector`: list of outputs of the cvis used for computing correlations.
- `correlations::RealVector`: list of outputs of the rank correlations.
- `features::RealVector`: list of outputs of the rocket feature kernels.
- `rocket::RocketModule`: time-series random feature kernels module.
- `classifier::MetaICVIClassifier`: ScikitLearn classifier.
- `performance::RealFP`: final output of the most recent the Meta-ICVI step.
- `is_pretrained::Bool`: internal flag for if the classifier is trained and ready for inference.
"""
mutable struct MetaICVIModule
    opts::MetaICVIOpts
    cvis::Vector{AbstractCVI}
    criterion_values::Vector{RealVector}
    correlations::RealVector
    features::RealVector
    rocket::RocketModule
    classifier::MetaICVIClassifier
    performance::RealFP
    probabilities::RealVector
    is_pretrained::Bool
end # MetaICVIModule

"""
    MetaICVIModule(opts::MetaICVIOpts)

Instantiate a MetaICVIModule with given options.

# Arguments
- `opts::MetaICVIOpts`: options struct for the MetaICVI object.
"""
function MetaICVIModule(opts::MetaICVIOpts)
    # Load the correct classifier from the opts
    @eval @sk_import linear_model: $(opts.classifier_selection)

    # Construct the CVIs
    cvis = [
        PS(),
        GD43()
    ]

    # Initialize the empty vectors for each criterion value
    cvi_values = [Array{RealFP}(undef, 0) for _ = 1:length(cvis)]

    # Construct the rocket kernels
    if isfile(opts.rocket_file)
        # If we have a file, load the module
        @info "Loading rocket kernels: $(opts.rocket_file)"
        rocket_module = load_rocket(opts.rocket_file)
        # Correct the expected number of kernels if necessary
        if opts.n_rocket != length(rocket_module.kernels)
            @warn "Provided incorrect number of kernels with loaded file, correcting MetaICVIModule options."
            opts.n_rocket = length(rocket_module.kernels)
        end
    else
        # Otherwise, construct a module
        @warn "Missing/incorrect path for the rocket module file."
        @info "Constructing a new rocket module"
        rocket_module = RocketModule(opts.correlation_window, opts.n_rocket)
    end

    # Load the classifier
    if isfile(opts.classifier_file)
        @info "Loading classifier: $(opts.classifier_file)"
        # classifier = JLD.load(opts.classifier_file, "learner")
        classifier = load_classifier(opts.classifier_file)
        is_pretrained = true
    # If we didn't get a valid file, check for errors or construct a new object
    else
        error_msg = "Missing/incorrect path for pretrained classifier file."
        # If we want to fail on a missing file, throw an error
        if opts.fail_on_missing
            error(error_msg)
        # Otherwise, put up a warning and create a new classifier.
        else
            # If a file is even provided, construct a classifier
            if !isempty(opts.classifier_file)
                @warn error_msg
                classifier = construct_classifier(opts)
                is_pretrained = false
            # Otherwise, throw an error that no filename is provided
            else
                error("No filename provided for classifier saving/loading.")
            end
        end
    end

    # Construct and return the module
    return MetaICVIModule(
        opts,                       # opts
        cvis,                       # cvis
        cvi_values,                 # criterion_values
        Array{RealFP}(undef, 0),    # correlations
        Array{RealFP}(undef, 0),    # features
        rocket_module,              # rocket
        classifier,                 # classifier
        0.0,                        # performance
        zeros(3),                   # probabilities
        is_pretrained               # is_pretrained
    )
end # MetaICVIModule(opts::MetaICVIOpts)

"""
    MetaICVIModule()

Default constructor for the MetaICVIModule.
"""
function MetaICVIModule()
    # Create the default options
    opts = MetaICVIOpts()
    # Return the Meta-ICVI module constructed with the default options
    return MetaICVIModule(opts)
end # MetaICVIModule()

# -----------------------------------------------------------------------------
# METHODS
# -----------------------------------------------------------------------------

"""
    Base.show(io::IO, metaicvi::MetaICVIModule)

Display a metaicvi module to the command line.

# Arguments
- `io::IO`: default io stream.
- `metaicvi::MetaICVIModule`: metaicvi object about which to display info.
"""
function Base.show(io::IO, metaicvi::MetaICVIModule)
    compact = get(io, :compact, false)
    if compact
        # print(metaicvi.opts)
        print("MetaICVIModule")
    else
        # show(io, metaicvi.opts)
        pretrain_display_flag = MetaICVI.is_pretrained(metaicvi)
        opts = metaicvi.opts
        print(
            """
            MetaICVIModule:
            \tperformance: $(metaicvi.performance)
            \tpretrained: $pretrain_display_flag
            Options:
            \tclassifier: $(opts.classifier_selection)
            \tclassifier_opts: $(opts.classifier_opts)
            \ticvi_window: $(opts.icvi_window)
            \tcorrelation_window: $(opts.correlation_window)
            \tn_rocket: $(opts.n_rocket)
            \trocket_file: $(opts.rocket_file)
            \tclassifier_file: $(opts.classifier_file)
            Flags:
            \tdisplay: $(opts.display)
            \tfail_on_missing: $(opts.display)
            """
        )
    end
end # Base.show(io::IO, metaicvi::MetaICVIModule)

"""
    construct_classifier(opts::MetaICVIOpts)

Construct a new classifier for the MetaICVI module with metaprogramming.

# Arguments
- `opts::MetaICVIOpts`: options containing the classifier type and options for instantiation.
"""
function construct_classifier(opts::MetaICVIOpts)
    @info "Constructing a new classifier"
    @eval classifier = $(opts.classifier_selection)(;$(opts.classifier_opts)...)
    return MetaICVIClassifier(classifier)
end # construct_classifier(opts::MetaICVIOpts)

"""
    safe_save_classifier(metaicvi::MetaICVIModule)

Error handle saving of the metaicvi classifier.

# Arguments
- `metaicvi::MetaICVIModule`: metaicvi module containing the classifier and path for saving.
"""
function safe_save_classifier(metaicvi::MetaICVIModule)
    # If we specified a file but none was there, then save to that file
    if !isempty(metaicvi.opts.classifier_file)
        @info "Saving trained classifier"
        save_classifier(metaicvi.classifier, metaicvi.opts.classifier_file)
    else
        error("No filename provided for classifier file saving/loading.")
    end
end # safe_save_classifier(metaicvi::MetaICVIModule)

"""
    safe_save_rocket(metaicvi::MetaICVIModule)

Error handle saving of the metaicvi rocket kernels.

# Arguments
- `metaicvi::MetaICVIModule`: metaicvi module containing the classifier and path for saving.
"""
function safe_save_rocket(metaicvi::MetaICVIModule)
    # If we specified a file but none was there, then save to that file
    if !isempty(metaicvi.opts.rocket_file)
        @info "Saving current rocket kernels"
        save_rocket(metaicvi.rocket, metaicvi.opts.rocket_file)
    else
        error("No filename provided for rocket file saving/loading.")
    end
end # safe_save_rocket(metaicvi::MetaICVIModule)

"""
    load_classifier(filepath::String)

Load the classifier at the filepath.

# Arguments
- `filepath::String`: location of the classifier .jld file.
"""
function load_classifier(filepath::String)
    return MetaICVIClassifier(JLD.load(filepath, "classifier"))
    # return BSON.load(filepath, @__MODULE__)["classifier"]
end # load_classifier(filepath::String)

"""
    save_classifier(classifier::MetaICVIClassifier, filepath::String)

Save the classifier at the filepath.

# Arguments
- `classifier::MetaICVIClassifier`: classifier object to save.
- `filepath::String`: name/path to save the classifier .jld file.
"""
function save_classifier(classifier::MetaICVIClassifier, filepath::String)
    JLD.save(filepath, "classifier", classifier)
    # bson(filepath, Dict("classifier" => classifier))
end # save_classifier(classifier::MetaICVIClassifier, filepath::String)

"""
    train_and_save(metaicvi::MetaICVIModule, x::RealMatrix, y::IntegerVector)

Train the classifier on x/y and save the kernels and classifier.

# Arguments
- `metaicvi::MetaICVIModule`: metaicvi module to save with.
- `x::RealMatrix`: features to train on.
- `y::IntegerVector`: correct/over/under partition targets.
"""
function train_and_save(metaicvi::MetaICVIModule, x::RealMatrix, y::IntegerVector)
    # Create a new classifier
    # TODO: option to train/retrain loaded/serialized models
    classifier = construct_classifier(metaicvi.opts)
    metaicvi.classifier = classifier

    # Train the classifier
    @info "Training classifier"
    fit!(metaicvi.classifier, x, y)
    metaicvi.is_pretrained = true

    # Save the rocket kernels used
    safe_save_rocket(metaicvi)
    # Save the classifier used
    safe_save_classifier(metaicvi)
end # train_and_save(metaicvi::MetaICVIModule, x::RealMatrix, y::IntegerVector)

"""
    get_icvis(metaicvi::MetaICVIModule, sample::RealVector, label::Integer)

Compute and store the icvi criterion values.

# Arguments
- `metaicvi::MetaICVIModule`: the Meta-ICVI module.
- `sample::RealVector`: the sample used for clustering.
- `label::Integer`: the label prescribed to the sample by the clustering algorithm.
"""
function get_icvis(metaicvi::MetaICVIModule, sample::RealVector, label::Integer)
    # Update all of the cvis incrementally
    for ix = 1:length(metaicvi.cvis)
        # Compute and push the criterion value
        value = get_icvi!(metaicvi.cvis[ix], sample, label)
        push!(metaicvi.criterion_values[ix], value)
        # FIFO the list to size
        while length(metaicvi.criterion_values[ix]) > metaicvi.opts.icvi_window
            popfirst!(metaicvi.criterion_values[ix])
        end
    end
end # get_icvis(metaicvi::MetaICVIModule, sample::RealVector, label::Integer)

"""
    get_correlations(metaicvi::MetaICVIModule)

Compute and store the rank correlations from the cvi values.

# Arguments
- `metaicvi::MetaICVIModule`: the Meta-ICVI module.
"""
function get_correlations(metaicvi::MetaICVIModule)
    # If the cvi window is big enough, compute the correlations
    if length(metaicvi.criterion_values[1]) >= metaicvi.opts.icvi_window
        # Get the spearman correlation
        correlation = corspearman(metaicvi.criterion_values[1], metaicvi.criterion_values[2])
        # Sanitize a potential NaN response
        # metaicvi.performance = isequal(performance, NaN) ? 0 : performance
        push!(metaicvi.correlations, correlation)
        # FIFO the list to size
        while length(metaicvi.correlations) > metaicvi.opts.correlation_window
            popfirst!(metaicvi.correlations)
        end
    end
end # get_correlations(metaicvi::MetaICVIModule)

"""
    get_rocket_features(metaicvi::MetaICVIModule)

Compute and store the rocket features.

# Arguments
- `metaicvi::MetaICVIModule`: the Meta-ICVI module.
"""
function get_rocket_features(metaicvi::MetaICVIModule)
    # If there are enough correlations, compute compute the meta-icvi value
    if length(metaicvi.correlations) >= metaicvi.opts.correlation_window
        metaicvi.features = apply_kernels(metaicvi.rocket, metaicvi.correlations)[:, 1]
    else
        metaicvi.features = zeros(metaicvi.opts.n_rocket)
    end
end # get_rocket_features(metaicvi::MetaICVIModule)

"""
    is_pretrained(metaicvi::MetaICVIModule)

Checks if the classifier is pretrained to permit inference.

# Arguments
- `metaicvi::MetaICVIModule`: metaicvi module containing the classifier to check.
"""
function is_pretrained(metaicvi::MetaICVIModule)
    # Check if the model is pretrained with the isfit function
    # return ScikitLearn.Utils.isfit(metaicvi.classifier)
    return metaicvi.is_pretrained
end # is_pretrained(metaicvi::MetaICVIModule)

"""
    get_probability(metaicvi::MetaICVIModule)

Compute and store the metaicvi value from the classifier.

# Arguments
- `metaicvi::MetaICVIModule`: the Meta-ICVI module.
"""
function get_probability(metaicvi::MetaICVIModule)
    # If we have previously computed features
    if !isempty(metaicvi.features) && is_pretrained(metaicvi)
        # Compute the class 'probabilities'
        probs = predict_proba(metaicvi.classifier, transpose([metaicvi.features]))
        metaicvi.probabilities = probs[:]

        # Store only the probability of correct partitioning
        metaicvi.performance = probs[1]
    else
        metaicvi.performance = 0.0
    end
end # get_probability(metaicvi::MetaICVIModule)

"""
    get_features(metaicvi::MetaICVIModule, sample::RealVector, label::Integer)

Compute only the features on the sample and label without classifier inference.

# Arguments
- `metaicvi::MetaICVIModule`: the Meta-ICVI module.
- `sample::RealVector`: the sample used for clustering.
- `label::Integer`: the label prescribed to the sample by the clustering algorithm.
"""
function get_features(metaicvi::MetaICVIModule, sample::RealVector, label::Integer)
    # Compute the icvi values
    get_icvis(metaicvi, sample, label)

    # Compute the rank correlations
    get_correlations(metaicvi)

    # Compute the rocket features
    get_rocket_features(metaicvi)
end # get_features(metaicvi::MetaICVIModule, sample::RealVector, label::Integer)

"""
    get_metaicvi(metaicvi::MetaICVIModule, sample::RealVector, label::Integer)

Compute and return the meta-icvi value.

# Arguments
- `metaicvi::MetaICVIModule`: the Meta-ICVI module.
- `sample::RealVector`: the sample used for clustering.
- `label::Integer`: the label prescribed to the sample by the clustering algorithm.
"""
function get_metaicvi(metaicvi::MetaICVIModule, sample::RealVector, label::Integer)
    # If the sample was not misclassified
    if label > 0
        # Compute the rocket features
        get_features(metaicvi, sample, label)

        # Compute and store the
        get_probability(metaicvi)
    else
        # Default to 0
        metaicvi.performance = 0.0
    end

    return metaicvi.performance
end # get_metaicvi(metaicvi::MetaICVIModule, sample::RealVector, label::Integer)

"""
    get_cvi_data(data_file::String)

Get the cvi data specified by the data_file path.

# Arguments
- `data_file::String`: file containing clustered data for cvi processing.
"""
function get_cvi_data(data_file::String)
    # Parse the data
    data = readdlm(data_file, ',')
    data = permutedims(data)
    train_x = convert(Matrix{Float64}, data[1:2, :])
    train_y = convert(Vector{Integer}, data[3, :])

    return train_x, train_y
end # get_cvi_data(data_file::String)

"""
    get_training_features(metaicvi::MetaICVIModule, data_path::String)

# Arguments

"""
function get_training_features(metaicvi::MetaICVIModule, data_path::String)
    # Point to data
    data_dir(args...) = joinpath(data_path, args...)

    # Load the data
    correct_x, correct_y = get_cvi_data(data_dir("correct_partition.csv"))
    under_x, under_y = get_cvi_data(data_dir("under_partition.csv"))
    over_x, over_y = get_cvi_data(data_dir("over_partition.csv"))

    # Package the data conveniently
    data = Dict(
        "correct" => Dict(
            "x" => correct_x,
            "y" => correct_y
        ),
        "under" => Dict(
            "x" => under_x,
            "y" => under_y
        ),
        "over" => Dict(
            "x" => over_x,
            "y" => over_y
        )
    )

    # Create the target containers
    data_lengths = [length(correct_y), length(under_y), length(over_y)]
    offset_lengths = [0, length(correct_y), length(under_y)]
    data_length = sum(data_lengths)
    features_data = zeros(metaicvi.opts.n_rocket, data_length)
    features_targets = zeros(Int, data_length)

    # Mapping of datatype to numeral target for classification
    type_to_num = Dict(
        "correct" => 1,
        "under" => 2,
        "over" => 3
    )

    # Itereate over all data to get features
    for (type, subdata) in data
        # Get the offset directly from the type mapping and explicit definition
        data_offset = offset_lengths[type_to_num[type]]

        # Iterate over all entries of the type of partitioning (i.e., correct, under, and over)
        @showprogress for i = 1:length(subdata["y"])
            # Extract the sample and label
            sample = subdata["x"][:, i]
            label = subdata["y"][i]

            # Compute and retrieve the features
            get_features(metaicvi, sample, label)
            features = metaicvi.features

            # Save the results
            features_data[:, i + data_offset] = features
            features_targets[i + data_offset] = type_to_num[type]
        end
    end

    return features_data, features_targets
end # get_training_features(metaicvi::MetaICVIModule, data_path::String)


using Statistics
using DataAPI

abstract type FitResult end
"""
A measured set of correlated k-ratios.
"""
struct BasicFitResult <: FitResult
    label::Label
    kratios::UncertainValues
end

function findlabel(ffr::FitResult, cxr::CharXRay)
    lbls = labels(ffr.kratios)
    fa = findall(lbl -> (lbl isa CharXRayLabel) && (cxr in lbl.xrays), lbls)
    @assert length(fa) ≠ 0 "No label found for $cxr."
    @assert length(fa) == 1 "Multiple labels found for $cxr."
    return lbls[fa[1]]
end

Base.show(io::IO, ffr::FitResult) = print(io, "$(ffr.label)")
NeXLUncertainties.NeXLUncertainties.value(ffr::FitResult, label::ReferenceLabel) =
    NeXLUncertainties.value(ffr.kratios, label)
NeXLUncertainties.σ(ffr::FitResult, label::ReferenceLabel) = σ(ffr.kratios, label)
NeXLUncertainties.uncertainty(ffr::FitResult, label::ReferenceLabel) =
    uncertainty(ffr.kratios, label)
NeXLUncertainties.getindex(ffr::FitResult, label::ReferenceLabel) =
    getindex(ffr.kratios, label)
Base.keys(ffr::FitResult) = keys(ffr.kratios)
NeXLUncertainties.labels(ffr::FitResult) = labels(ffr.kratios)
function kratio(cxr::CharXRay, ffr::FitResult)
    lbl = findlabel(ffr, cxr)
    return uv(NeXLUncertainties.value(ffr, lbl), σ(ffr, lbl))
end
unknown(ffr::FitResult)::Label = ffr.label
Base.values(ffrs::Vector{<:FitResult}, lbl::ReferenceLabel) =
    [NeXLUncertainties.value(ffr.kratios, lbl, 0.0) for ffr in ffrs]
σs(ffrs::Vector{<:FitResult}, lbl::ReferenceLabel) =
    [σ(ffr.kratios, lbl, 0.0) for ffr in ffrs]

"""
    kratios(ffr::FitResult)::Vector{KRatio}

The k-ratios associated with each `CharXRayLabel` as a vector 'KRatio' objects.
"""
function kratios(ffr::FitResult)::Vector{KRatio}
    lbls = filter(labels(ffr.kratios)) do l
        (l isa CharXRayLabel) && haskey(properties(l)[:Composition], element(l))
    end
    return map(lbls) do lbl
        props = properties(lbl)
        KRatio(lbl.xrays, properties(unknown(ffr)), props, props[:Composition], ffr.kratios[lbl])
    end
end

"""
    kratio(ffr::FitResult, cxr::CharXRay)::KRatio

Extract the k-ratio associated with the specified CharXRay (zero if not measured)
"""
function kratio(ffr::FitResult, cxr::CharXRay)
    lbls = labels(ffr.kratios)
    i = findfirst(lbl->(lbl isa CharXRayLabel) && (cxr in lbl.xrays), lbls)
    @assert !isnothing(i) "There is no k-ratio associated with $cxr."
    lbl, props = lbls[i], properties(lbls[i])
    KRatio(lbl.xrays, unknown(ffr).spectrum.properties, props, props[:Composition], ffr.kratios[lbl])
  end

"""
    extractStandards(ffr::FitResult, elm::Element, mat::Material)::Vector{KRatio}

Extract a `Vector{KRatio}` for `elm::Element` from a `ffr::FilterFitResult` measured from `mat::Material`.
"""
function extractStandards(ffr::FitResult, elm::Element, mat::Material)::Vector{KRatio}
    @assert haskey(mat, elm) "The element $elm is not contained within $mat."
    lbls = filter(labels(ffr.kratios)) do l
        (l isa CharXRayLabel) && (elm == element(l))
    end
    krs = map(lbls) do lbl
        props = properties(lbl)
        kr = KRatio(lbl.xrays, properties(unknown(ffr)), props, props[:Composition], ffr.kratios[lbl])
        kr.unkProps[:Composition] = mat
        kr
    end
    length(krs) == 0 && @info  "There are no k-ratios associated with the $elm in $ffr."
    return krs
end

"""
    extractStandards(ffr::FitResult, cxrs::AbstractVector{CharXRay}, mat::Material)::Vector{KRatio}

Extract a `KRatio` for the `CharXRay` from a `FilterFitResult` associated with
the `Material`.
"""
function extractStandards(ffr::FitResult, cxrs::AbstractVector{CharXRay}, mat::Material)::Vector{KRatio}
    lbls = filter(labels(ffr.kratios)) do lbl
        (lbl isa CharXRayLabel) && any(cxr->cxr in xrays(lbl), cxrs)
    end
    krs = map(lbls) do lbl
        props = properties(lbl)
        kr = KRatio(lbl.xrays, properties(unknown(ffr)), props, props[:Composition], ffr.kratios[lbl])
        kr.unkProps[:Composition] = mat
        kr
    end
    length(krs) == 0 && @info "There are no k-ratios associated with any of $cxrs in $ffr."
    return krs
end

"""
    heterogeneity(ffrs::Vector{FilterFitResult}, lbl::ReferenceLabel)

Computes the ratio of the standard deviation of the measured values over the mean calculated uncertainty
from the fit.  A value near 1 means the sample appears homogeneous and a value greater than 1 means the sample
appears heterogeneous.
"""
heterogeneity(ffrs::Vector{<:FitResult}, lbl::ReferenceLabel) =
    std(values(ffrs, lbl)) / mean(σs(ffrs, lbl))

"""
    NeXLUncertainties.asa(::Type{DataFrame}, ffrs::AbstractVector{<:FitResult}; charOnly = true, withUnc = false, format = :normal # :pivot or :long)

Return generic `FitResult` as a DataFrame.

Format:
  * :normal - One row per spectrum, one column per k-ratio
  * :pivot - One row per ROI, one column per spectrum (optional: column for 1σ uncertainty on k-ratio)
  * :long - One row per spectrum, feature, measured k-ratio

"""
function NeXLUncertainties.asa(
    ::Type{DataFrame},
    ffrs::AbstractVector{<:FitResult};
    charOnly = true,
    withUnc = false,
    format = :normal # :pivot or :long
)::DataFrame
    lbls = sort(unique(mapreduce(labels, append!, ffrs)))
    lbls = charOnly ? filter(lbl -> lbl isa CharXRayLabel, lbls) : lbls
    if format==:pivot
        res = DataFrame(ROI = lbls)
        for ffr in ffrs
            vals = [ value(ffr.kratios, lbl, missing) for lbl in lbls ]
            insertcols!(res, ncol(res) + 1, Symbol(repr(ffr.label)) => vals)
            if withUnc
                unc = [σ(ffr.kratios, lbl, missing) for lbl in lbls]
                insertcols!(res, ncol(res) + 1, Symbol('Δ' * repr(ffr.label)) => unc)
            end
        end
    elseif format==:long
        if withUnc
            res = DataFrame(Spectrum=String[], ROI = String[], k = Float64[], dk = Float64[])
            for ffr in ffrs, lbl in lbls
                push!(res, ( repr(ffr.label), repr(lbl), value(ffr.kratios, lbl, missing), σ(ffr.kratios, lbl, missing) ))
            end
        else
            res = DataFrame(Spectrum=String[], ROI = String[], k = Float64[])
            for ffr in ffrs, lbl in lbls
                push!(res, ( repr(ffr.label), repr(lbl), value(ffr.kratios, lbl, missing)))
            end
        end
    else
        rowLbls = [repr(ffr.label) for ffr in ffrs]
        res = DataFrame(Symbol("Spectra") => rowLbls)
        for lbl in lbls
            vals = [value(ffr.kratios, lbl, missing) for ffr in ffrs]
            insertcols!(res, ncol(res) + 1, Symbol(repr(lbl)) => vals)
            if withUnc
                unc = [σ(ffr.kratios, lbl, missing) for ffr in ffrs]
                insertcols!(res, ncol(res) + 1, Symbol('Δ' * repr(lbl)) => unc)
            end
        end
    end
    return res
end

"""
    NeXLUncertainties.asa(::Type{DataFrame}, fr::FitResult; withUnc = false)

Summarize the ROI and k-ratio data within a `FitResult` structure as a `DataFrame`.
"""
function NeXLUncertainties.asa(::Type{DataFrame}, fr::FitResult; withUnc = false)
    rois = labels(fr)
    res =  DataFrame(
        ROI = rois,
        K = map(l -> value(fr.kratios, l), rois),
    )
    withUnc && insertcols!(res, :dK => map(l -> σ(fr.kratios, l), rois) )
    return res
end

function DataAPI.describe(ffrs::Vector{<:FitResult})::DataFrame
    df = asa(DataFrame, ffrs)[:, 2:end]
    desc = describe(df, :mean, :std, :min, :q25, :median, :q75, :max)
    lbls = filter(
        lbl -> lbl isa CharXRayLabel,
        sort(collect(Set(mapreduce(labels, append!, ffrs)))),
    )
    insertcols!(desc, 4, :hetero => heterogeneity.(Ref(ffrs), lbls))
    return desc
end

"""
    FilterFitResult

Represents the result of fitting a FilteredUnknownW to a FilteredUnknown.

Struct elements

    label::UnknownLabel  # Identifies the unknown
    kratios::UncertainValues # Labeled with ReferenceLabel objects
    roi::UnitRange{Int} # Range of channels fit
    raw::Vector{Float64} # Raw spectrum data
    residual::Vector{Float64} # Residual spectrum
    peakback::Dict{ReferenceLabel,NTuple{3,Float64}} # peak counts, background counts and counts/(nAs)

Use asa(DataFrame, ffr::FilterFitResult) to summarize in tabular form.
"""
struct FilterFitResult{T <: AbstractFloat} <: FitResult
    label::UnknownLabel
    kratios::UncertainValues
    roi::UnitRange{Int}
    raw::Vector{T}
    residual::Vector{T}
    peakback::Dict{ReferenceLabel,NTuple{3,T}}
end

"""
    residual(ffr::FilterFitResult)::Spectrum

A Spectrum containing the histogram representing the unknown spectrum
minus the fitted characteristic peaks shapes times the best fit coefficient.
"""
function residual(ffr::FilterFitResult)
    props = copy(properties(ffr.label.spectrum))
    props[:Name] = "Residual[$(props[:Name])]"
    return Spectrum(ffr.label.spectrum.energy, ffr.residual, props)
end

"""
    spectrum(ffr::FilterFitResult)::Spectrum

Returns the original unknown spectrum.
"""
spectrum(ffr::FilterFitResult) = ffr.label.spectrum

"""
    characteristiccounts(ffr::FiterFitResult, strip)

Number of spectrum counts that were accounted for by the fitted elements with the `strip` Element(s) removed.
"""
characteristiccounts(ffr::FilterFitResult, strip) =
    sum(element(ref) in strip ? 0.0 : v[1] - v[2] for (ref, v) in ffr.peakback) # sum(ffr.raw[ffr.roi]-ffr.residual[ffr.roi])

"""
    peaktobackground(ffr::FilterFitResult, backwidth::Float64=10.0)::Float64

The peak-to-background ratio as determined from the raw and residual spectra integrated over
the fit region-of-interest and scaled to `backwidth` eV of continuum (nominally 10 eV).
"""
function peaktobackground(
    ffr::FilterFitResult{T},
    klabel::ReferenceLabel,
    backwidth::T = convert(T, 10.0),
) where { T <: AbstractFloat }
    unk = spectrum(unknown(ffr))
    peak, back, _ = ffr.peakback[klabel]
    return (peak * (energy(klabel.roi.stop, unk) - energy(klabel.roi.start, unk))) /
           (backwidth * back)
end

"""
    NeXLUncertainties.asa(
        ::Type{DataFrame},
        ffr::FilterFitResult;
        charOnly::Bool=true,
        material=nothing,
        mc = XPP, fc=ReedFluorescence,
        columns = () # Selected from ( :roi, :peakback, :counts, :dose )
    )::DataFrame

Tabulate details about each region-of-interest in the 'FilterFitResult' in a 'DataFrame'.
  * If charOnly then only display characteristic X-ray data (not escapes etc.)
  * If `material` is a Material then the computed k-ratio (KCalc) will also be tabulated along with kmeas/kcalc (KoKCalc).
  * columns - Select optional column outputs (see below)

Columns:

  * Spectrum : UnknownLabel - Identifies the fit spectrum
  * Feature  : Label - Identifies the fit feature (Vector{CharXRay} or other)
  * Reference: String - Name of reference against which :Spectrum was fit over :Feature
  * K : Float64 - The multiplicative fit constant
  * dK : Float64 - The 1σ uncertainty in :K
  * :Start : Int - Start index for fit channels (:roi ∈ columns)
  * :Stop : Int - Stop index for fit channels  (:roi ∈ columns)
  * :Peak : Float64 - Total counts in characteristic peak (:peakback ∈ columns)
  * :Back : Float64 - Total counts in background under the characteristic peak (:peakback ∈ columns)
  * :PtoB : Float64 - Peak-to-Background assuming 10 eV/channel (:peakback ∈ columns)
  * :KCalc : Float64 - Computed k-ratio assuming a composition. (Requires `material` argument to be specified.)
  * :KoKcalc : Float64 - Ratio of measured/computed k-ratio.  (Requires `material` argument to be specified.)
  * :LiveTime : Float64 - Acquisiton live time (s) (:dose ∈ columns)
  * :ProbeCurrent : Float64 - Probe current (nA) (:dose ∈ columns)
  * :DeadPct : Float64 - Dead time in ProbeCurrent (:dose ∈ columns)
  * :Counts : Float64 - Total counts in characteristic peak (:counts ∈ columns)
  * :RefCountsPernAs : Float64 - Estimated counts in :Reference in :Feature per unit dose.  (:counts ∈ columns)
  * :CountsPernAs : Float64 - Estimated counts in :Spectrum in :Feature per unit dose.  (:counts ∈ columns)
"""
function NeXLUncertainties.asa(
    ::Type{DataFrame},
    ffr::FilterFitResult;
    charOnly::Bool = true,
    material::Union{Material,Nothing} = nothing,
    columns::Tuple = (), # ( :roi, :peakback, :counts, :dose)
    mc = XPP, fc=ReedFluorescence,
)::DataFrame
    sl = charOnly ? #
        filter(lbl->lbl isa CharXRayLabel, NeXLUncertainties.sortedlabels(ffr.kratios)) : #
        NeXLUncertainties.sortedlabels(ffr.kratios)
    unkspec = ffr.label.spectrum
    # Build the base result DataFrame
    res = DataFrame(
        Spectrum = [ ffr.label for _ in sl ],
        Feature = sl,
        Reference = [properties(lbl)[:Name] for lbl in sl],
        K = [ NeXLUncertainties.value(ffr.kratios, lbl) for lbl in sl],
        dK = [ NeXLUncertainties.σ(ffr.kratios, lbl) for lbl in sl]
    )
    # Add optional columns
    if :roi ∈ columns
        insertcols!(res, 4, :Start => [ lbl.roi.start for lbl in sl])
        insertcols!(res, 5, :Stop => [ lbl.roi.stop for lbl in sl])
    end
    if :peakback ∈ columns
        insertcols!(res, :Peak => [ (ffr.peakback[lbl][1]-ffr.peakback[lbl][2]) for lbl in sl] )
        insertcols!(res, :Back => [ ffr.peakback[lbl][2] for lbl in sl] )
        insertcols!(res, :PtoB => [ peaktobackground(ffr, lbl) for lbl in sl] )
    end
    if material isa Material
        function kcalc(lbl)
            stdspec = lbl.spectrum
            zc = zafcorrection(
                mc,
                fc,
                NullCoating,
                material,
                stdspec[:Composition],
                lbl.xrays,
                unkspec[:BeamEnergy],
            )
            k(zc..., unkspec[:TakeOffAngle], stdspec[:TakeOffAngle])
        end
        insertcols!(res, :KCalc => kcalc.(sl))
        insertcols!(res, :KoKcalc => [ r[:K]/r[:KCalc] for r in eachrow(res) ])
    end
    if :dose ∈ columns
        dt(spec) = 100.0*(spec[:RealTime]-spec[:LiveTime])/spec[:RealTime]
        insertcols!(res, 2, :LiveTime => [ get(unkspec, :LiveTime, missing) for _ in sl ])
        insertcols!(res, 3, :ProbeCurrent => [ get(unkspec, :ProbeCurrent, missing) for _ in sl ])
        insertcols!(res, 4, :DeadPct => [ dt(unkspec) for lbl in sl ])

    end
    if :counts ∈ columns
        insertcols!(res, :Counts => [ ffr.peakback[lbl][1]-ffr.peakback[lbl][2] for lbl in sl] )
        insertcols!(res, :RefCountsPernAs => [ ffr.peakback[lbl][3] for lbl in sl] )
        insertcols!(res, :CountsPernAs => [ (ffr.peakback[lbl][1]-ffr.peakback[lbl][2]) / dose(unkspec) for lbl in sl ])
    end
    return res
end

"""
    filteredresidual(fit::FilterFitResult, unk::FilteredUnknown, ffs::AbstractVector{FilteredReference})::Vector{Float64}

Computes the difference between the best fit and the unknown filtered spectral data.
"""
function filteredresidual(
    fit::FilterFitResult{T},
    unk::FilteredUnknown,
    ffs::AbstractVector{FilteredReference{T}},
) where { T <: AbstractFloat }
    return unk.filtered - mapreduce(+, ffs) do ff
        (value(fit, ff.label) * (ff.scale / unk.scale)) * extract(ff, unk.ffroi)
    end
end

"""
    NeXLMatrixCorrection.estimatecoating(fr::FitResult, substrate::Material, coating::Material, cxr::CharXRay, mc::Type{<:MatrixCorrection}=XPP)::Film

Estimate the mass-thickness of `coating` on bulk `substrate` using the X-ray `cxr` for which there is KRatio in `fr`.  The result is
assigned to the properties of `fr` so that it can be account for in the matrix correction process.

This method is intended for use on standards where the substrate composition is known a priori.
"""
function NeXLMatrixCorrection.estimatecoating(fr::FitResult, substrate::Material, coating::Material, cxr::CharXRay, mc::Type{<:MatrixCorrection}=XPP)::Film
    properties(fr.label.spectrum)[:Coating] = estimatecoating(substrate, coating, kratio(fr, cxr), mc)
end