"""
    bounded_aux(x::Array{JuMP.VariableRef}, binary_var::JuMP.VariableRef)
    bounded_aux(x::Array{JuMP.VariableRef}, y::JuMP.VariableRef, binary_var::JuMP.VariableRef)

Generates binary-bounded auxiliary variables and their bounding constraints of the same size as x + y.
"""
function bounded_aux(x::Array{JuMP.VariableRef}, binary_var::JuMP.VariableRef)
    var_bounds = get_bounds(x)
    aux_vars = @variable(x[1].model, [1:length(x)])
    bound_cons = []
    for i = 1:length(x) # Bounding auxiliary variables
        bds = var_bounds[x[i]]
        push!(bound_cons, @constraint(x[i].model, aux_vars[i] >= minimum(bds)*binary_var))
        push!(bound_cons, @constraint(x[i].model, aux_vars[i] <= maximum(bds)*binary_var))
    end
    aux_vars = (binary_var, aux_vars)
    return aux_vars, bound_cons
end

function bounded_aux(x::Array{JuMP.VariableRef}, y::JuMP.VariableRef, binary_var::JuMP.VariableRef)
    aux_vars, bound_cons = bounded_aux(x, binary_var)
    aux_dep = @variable(y.model)
    push!(bound_cons, @constraint(y.model, aux_dep >= binary_var * JuMP.lower_bound(y)))
    push!(bound_cons, @constraint(y.model, aux_dep <= binary_var * JuMP.upper_bound(y)))
    aux_vars = (aux_vars[1], aux_vars[2], aux_dep)
    return aux_vars, bound_cons
end

"""
    $(TYPEDSIGNATURES)

Generates MI constraints from GlobalModel and its BlackBoxLearners(gm.bbls), and adds them to gm.model.
"""
function add_tree_constraints!(gm::GlobalModel, bbc::BlackBoxClassifier, idx = length(bbc.learners))
    isempty(bbc.mi_constraints) || throw(OCTHaGOnException("BBC $(bbc.name) already has associated MI approximation."))
    isempty(bbc.leaf_variables) || throw(OCTHaGOnException("BBC $(bbc.name) already has associated MI variables."))
    if bbc.feas_ratio == 1.0 # Just a placeholder to show that the tree is "trained".
        z_feas = @variable(gm.model, binary = true)
        bbc.mi_constraints = Dict(1 => [@constraint(gm.model, z_feas == 1)])
        bbc.leaf_variables = Dict(1 => (z_feas, []))
        if !isempty(bbc.lls)
            for ll in bbc.lls
                z_feas = @variable(gm.model, binary = true)
                ll.mi_constraints = Dict(1 => [@constraint(gm.model, z_feas == 1)])
                ll.leaf_variables = Dict(1 => (z_feas, []))
            end
        end
    elseif get_param(bbc, :reloaded)
        bbc.mi_constraints, bbc.leaf_variables = add_feas_constraints!(gm.model, bbc.vars, bbc.learners[idx]; equality = bbc.equality, relax_var = bbc.relax_var, bigM = bbc.bigM)
        if !isempty(bbc.lls)
            for ll in bbc.lls
                ll.mi_constraints, ll.leaf_variables = add_feas_constraints!(gm.model, ll.vars, bbc.learners[idx]; equality = bbc.equality, relax_var = ll.relax_var, bigM = bbc.bigM, symbs = Symbol.(bbc.vars))
            end
        end
    elseif size(bbc.X, 1) == 0
        throw(OCTHaGOnException("Constraint " * string(bbc.name) * " has not been sampled yet, and is thus untrained."))
    elseif isempty(bbc.learners)
        throw(OCTHaGOnException("Constraint " * string(bbc.name) * " must be learned before tree constraints
                            can be generated."))
    elseif bbc.feas_ratio == 0.0
        throw(OCTHaGOnException("Constraint " * string(bbc.name) * " is INFEASIBLE but you tried to include it in
            your global problem. Find at least one feasible solution, train and try again."))
    elseif !get_param(gm, :ignore_accuracy) && !check_accuracy(bbc)
        throw(OCTHaGOnException("Constraint " * string(bbc.name) * " is inaccurately approximated. "))
    else
        bbc.mi_constraints, bbc.leaf_variables = add_feas_constraints!(gm.model, bbc.vars, bbc.learners[idx]; equality = bbc.equality, relax_var = bbc.relax_var, bigM = bbc.bigM)
        if !isempty(bbc.lls)
            for ll in bbc.lls
                ll.mi_constraints, ll.leaf_variables = add_feas_constraints!(gm.model, ll.vars, bbc.learners[idx]; equality = bbc.equality, relax_var = ll.relax_var, bigM = bbc.bigM, symbs = Symbol.(bbc.vars))
            end
        end
    end
    bbc.active_leaves = []
    return
end

function add_tree_constraints!(gm::GlobalModel, bbr::BlackBoxRegressor, idx = length(bbr.learners))
    if size(bbr.X, 1) == 0 && !get_param(bbr, :reloaded)
        throw(OCTHaGOnException("Constraint " * string(bbr.name) * " has not been sampled yet, and is thus untrained."))
    elseif bbr.convex && !bbr.equality
        clear_tree_constraints!(gm, bbr)
        mi_constraints = Dict(1 => [])
        for i = Int64.(ceil.(size(bbr.X,1) .* rand(10)))
            push!(mi_constraints[1], @constraint(gm.model, bbr.dependent_var >= sum(Array(bbr.gradients[i,:]) .* (bbr.vars .- Array(bbr.X[i, :]))) + bbr.Y[i]))
        end
        merge!(append!, bbr.mi_constraints, mi_constraints)
        if !isempty(bbr.lls)
            for lr in bbr.lls
                clear_tree_constraints!(gm, lr)
                mc = Dict(1 => [])
                # Number of initial cuts depends on the dimension of the constraint
                pt_idxs = Int64.(ceil.(size(bbr.X,1) .* rand(maximum([10, Int64(10*ceil(log(length(bbr.vars))))]))))
                update_gradients(bbr, pt_idxs)
                for i in pt_idxs
                    push!(mc[1], @constraint(gm.model, lr.dependent_var >= sum(Array(bbr.gradients[i,:]) .* (lr.vars .- Array(bbr.X[i, :]))) + bbr.Y[i]))
                end
                merge!(append!, lr.mi_constraints, mc)
                lr.active_leaves = [1]
            end
        end
        bbr.active_leaves = [1]
        return
    elseif isempty(bbr.learners)
        throw(OCTHaGOnException("Constraint " * string(bbr.name) * " must be learned before tree constraints
                            can be generated."))
    elseif isempty(bbr.ul_data)
        throw(OCTHaGOnException("Constraint " * string(bbr.name) * " is a Regressor,
        but doesn't have a ORT and/or OCT with upper/lower bounding approximators!"))
    end
    if !isempty(bbr.lls)
        bbr.thresholds[idx].first == "reg" ||
            throw(OCTHaGOnException("BBRs with LinkedRegressors are only allowed ORT approximators."))
        length(bbr.active_trees) == 0 ||
            throw(OCTHaGOnException("BBRs with LinkedRegressors cannot have more than one active tree."))
    elseif bbr.thresholds[idx].first == "reg"
        (isempty(bbr.leaf_variables) || !isnothing(bbr.thresholds[idx].second)) ||
            throw(OCTHaGOnException("Please clear previous tree constraints from $(gm.name) " *
                "before adding an unthresholded regression constraints."))
        if !isnothing(bbr.thresholds[idx].second) && !isempty(bbr.active_trees)
            bbr.thresholds[active_upper_tree(bbr)].second == bbr.thresholds[idx].second ||
                throw(OCTHaGOnException("Upper-thresholded regressors must be preceeded by upper classifiers " *
                                    "of the same threshold for $(bbr.name)."))
        end
    elseif bbr.thresholds[idx].first == "rfreg"
        isempty(bbr.leaf_variables) ||
            throw(OCTHaGOnException("Please clear previous tree constraints from $(gm.name) " *
            "before adding an random forest regressor constraints."))
            !isempty(bbr.lls) && throw(OCTHaGOnException("Random Forests cannot be used to approximate " *
                        "linked BBR $(bbr.name)."))
        isnothing(bbr.thresholds[idx].second) ||
            throw(OCTHaGOnException("RandomForestRegressors are not allowed upper bounds."))
        bbr.equality &&
            throw(OCTHaGOnException("RandomForest cannot approximate BBR $(bbr.name) since it is an equality constraint."))
    elseif bbr.thresholds[idx].first == "upper"
        all(collect(keys(bbr.mi_constraints)) .>= 0)  || throw(OCTHaGOnException("Please clear previous upper tree constraints from $(gm.name) " *
                                                          "before adding new constraints."))
    elseif bbr.thresholds[idx].first == "lower"
        all(collect(keys(bbr.mi_constraints)) .<= 0) || throw(OCTHaGOnException("Please clear previous lower tree constraints from $(gm.name) " *
                                                            "before adding new constraints."))
    end
    if bbr.thresholds[idx].first == "rfreg"
        isempty(bbr.lls) || throw(ErrorException("RandomForestRegressors augment LinkedLearners when MI approximations are made. To avoid bugs, BBRs approximated via RFs cannot have LinkedLearners."))
        trees = get_random_trees(bbr.learners[idx])
        bbr.mi_constraints, bbr.leaf_variables = add_regr_constraints!(gm.model, bbr.vars, bbr.dependent_var,
                    trees[1], bbr.ul_data[idx][1]; equality = bbr.equality, bigM = bbr.bigM, relax_var = bbr.relax_var,)
                    # mic = bbr.mi_constraints, lv = bbr.leaf_variables)
        for i=2:length(trees)
            nll = LinkedRegressor(vars = bbr.vars, dependent_var = bbr.dependent_var)
            nll.mi_constraints, nll.leaf_variables = add_regr_constraints!(gm.model, bbr.vars, bbr.dependent_var,
                                                trees[i], bbr.ul_data[idx][i];
                                                equality = bbr.equality, bigM = bbr.bigM,)
                                                # mic = nll.mi_constraints,
                                                # lv = nll.leaf_variables)
            push!(bbr.lls, nll)
        end
        bbr.active_trees[idx] = bbr.thresholds[idx]
    else
       mi_constraints, leaf_variables = add_regr_constraints!(gm.model, bbr.vars, bbr.dependent_var,
                                                bbr.learners[idx], bbr.ul_data[idx];
                                                equality = bbr.equality, bigM = bbr.bigM, relax_var = bbr.relax_var,)
                                                # mic = bbr.mi_constraints,
                                                # lv = bbr.leaf_variables)

        merge!(append!, bbr.mi_constraints, mi_constraints)
        merge!(bbr.leaf_variables, leaf_variables)
        if !isempty(bbr.lls)
            for lr in bbr.lls
                lr_mi, lr_lv = add_regr_constraints!(gm.model, lr.vars, lr.dependent_var,
                                                bbr.learners[idx], bbr.ul_data[idx];
                                                equality = bbr.equality, bigM = bbr.bigM, relax_var = lr.relax_var,
                                                symbs = Symbol.(bbr.vars),)
                                                # mic = lr.mi_constraints,
                                                # lv = lr.leaf_variables)
                merge!(append!, lr.mi_constraints, lr_mi)
                merge!(lr.leaf_variables, lr_lv)
            end
        end
        if bbr.thresholds[idx].first == "upper"
            push!(bbr.mi_constraints[-1], @constraint(gm.model, bbr.dependent_var <= bbr.thresholds[idx].second))
            [push!(lr.mi_constraints[-1], @constraint(gm.model, lr.dependent_var <= bbr.thresholds[idx])) for lr in bbr.lls]
        elseif bbr.thresholds[idx].first == "lower"
            push!(bbr.mi_constraints[1], @constraint(gm.model, bbr.dependent_var >= bbr.thresholds[idx].second))
            [push!(lr.mi_constraints[1], @constraint(gm.model, lr.dependent_var >= bbr.thresholds[idx])) for lr in bbr.lls]
        end
        bbr.active_trees[idx] = bbr.thresholds[idx]
    end
    bbr.active_leaves = []
    [lr.active_leaves = [] for lr in bbr.lls]
    return
end

function add_tree_constraints!(gm::GlobalModel, bbls::Vector{BlackBoxLearner})
    for bbl in bbls
        add_tree_constraints!(gm, bbl)
    end
    return
end

add_tree_constraints!(gm::GlobalModel) = add_tree_constraints!(gm, gm.bbls)

""" Computes the smallest possible big-M for a α'x ≥ threshold split in a tree. """
function compute_hyperplane_bigM_lower(threshold::Real, α::Vector, var_bounds::Vector)
    return threshold - sum(minimum(α[i] .* var_bounds[i]) for i=1:length(α))
end

""" Computes the smallest possible big-M for a α'x ≤ threshold split in a tree. """
function compute_hyperplane_bigM_upper(threshold::Real, α::Vector, var_bounds::Vector)
    return -threshold + sum(maximum(α[i] .* var_bounds[i]) for i=1:length(α))
end

""" Computes the smallest possible big-M for lower bounding regression in a tree leaf. """
function compute_regression_bigM_lower(β0::Real, β::Vector, var_bounds::Vector, y_bounds::Vector)
    return β0 + sum(maximum(β[i] .* var_bounds[i]) for i=1:length(β)) - minimum(y_bounds)
end

""" Computes the smallest possible big-M for upper bounding regression in a tree leaf. """
function compute_regression_bigM_upper(β0::Real, β::Vector, var_bounds::Vector, y_bounds::Vector)
    return -β0 - sum(minimum(β[i] .* var_bounds[i]) for i=1:length(β)) + maximum(y_bounds)
end

"""
Creates a set of binary feasibility constraints from a binary classification tree.
Arguments:
- m: JuMP Model
- x: independent JuMP.Variables (features in lnr)
- lnr: A fitted OptimalTreeClassifier
- equality: whether the constraint is an equality.
NOTE: mic and lv are only nonempty if we are adding an OCT approximation on top of an existing set of MI constraints and variables respectively. Leave defaults empty for basic usage.
symbs helps add constraints to linked learners, since the varkeys in lnr are different to the variables in the linked learner.
"""
function add_feas_constraints!(m::JuMP.Model, x::Array{JuMP.VariableRef}, lnr::IAI.OptimalTreeLearner;
                               equality::Bool = false, bigM::Bool = false,
                               relax_var::Union{Real, JuMP.VariableRef} = 0,
                               mic::Dict = Dict(), lv::Dict = Dict(),
                               symbs = Symbol.(x))
    check_if_trained(lnr);
    all_leaves = find_leaves(lnr)
    # Add a binary variable for each leaf
    feas_leaves = []
    if lnr isa IAI.OptimalTreeClassifier
        feas_leaves =
        [i for i in all_leaves if Bool(IAI.get_classification_label(lnr, i))];
    else
        feas_leaves = all_leaves
    end
    mi_constraints = Dict(leaf => [] for leaf in feas_leaves)
    mi_constraints[1] = []
    leaf_variables = Dict{Int64, Tuple{JuMP.VariableRef, Array}}()
    var_bounds = get_bounds(x)
    var_bounds = [var_bounds[var] for var in x] # ordered vector form
    if isempty(lv) # if the variables haven't already been created and passed!
        if bigM
            # big-M formulation variable generation and coupling
            for leaf in feas_leaves
                leaf_variables[leaf] = (@variable(m, binary = true), [])
            end
            push!(mi_constraints[1], @constraint(m, sum(leaf_variables[leaf][1] for leaf in feas_leaves) == 1))
        else
            # big-M free formulation variable generation and coupling
            for leaf in feas_leaves
                leaf_variables[leaf], bd_mi = bounded_aux(x, @variable(m, binary=true))
                append!(mi_constraints[leaf], bd_mi)
            end
            push!(mi_constraints[1], @constraint(m, sum(leaf_variables[leaf][1] for leaf in feas_leaves) == 1))
            [push!(mi_constraints[1], constr)
                for constr in @constraint(m, sum(leaf_variables[leaf][2] for leaf in feas_leaves) .== x)]
        end
    else
        leaf_variables = lv
    end
    # Getting lnr data
    upperDict, lowerDict = trust_region_data(lnr, symbs)
    for leaf in feas_leaves
        # ADDING TRUST REGIONS
        for (threshold, α) in upperDict[leaf] # upperDict is lower-bounded.
            if bigM
                M = compute_hyperplane_bigM_lower(threshold, α, var_bounds)
                push!(mi_constraints[leaf], @constraint(m, threshold ≤ sum(α .* x) + M * (1 - leaf_variables[leaf][1] + relax_var)))
            else
                push!(mi_constraints[leaf], @constraint(m, threshold * leaf_variables[leaf][1] <= sum(α .* leaf_variables[leaf][2]) + relax_var))
            end
        end
        for (threshold, α) in lowerDict[leaf] # lowerDict is upper-bounded.
            if bigM
                M = compute_hyperplane_bigM_upper(threshold, α, var_bounds)
                push!(mi_constraints[leaf], @constraint(m, M * (1 - leaf_variables[leaf][1] + relax_var) + threshold ≥ sum(α .* x)))
            else
                push!(mi_constraints[leaf], @constraint(m, threshold * leaf_variables[leaf][1] + relax_var >= sum(α .* leaf_variables[leaf][2])))
            end
        end
    end
    if equality
        infeas_leaves = [i for i in all_leaves if !Bool(IAI.get_classification_label(lnr, i))];
        [mi_constraints[leaf] = [] for leaf in infeas_leaves]
        if isempty(lv)
            if bigM
                # big-M formulation variable generation and coupling
                for leaf in infeas_leaves
                    leaf_variables[leaf] = (@variable(m, binary = true), [])
                end
                push!(mi_constraints[1], @constraint(m, sum(leaf_variables[leaf][1] for leaf in infeas_leaves) == 1))
            else
                for leaf in infeas_leaves
                    leaf_variables[leaf], bd_mi = bounded_aux(x, @variable(m, binary=true))
                    append!(mi_constraints[leaf], bd_mi)
                end
                push!(mi_constraints[1], @constraint(m, sum(leaf_variables[leaf][1] for leaf in infeas_leaves) == 1))
                [push!(mi_constraints[1], constr)
                    for constr in @constraint(m, sum(leaf_variables[leaf][2] for leaf in infeas_leaves) .== x)]
            end
        end
        for leaf in infeas_leaves
            # ADDING TRUST REGIONS
            for (threshold, α) in upperDict[leaf] # upperDict is lower-bounded.
                if bigM
                    M = compute_hyperplane_bigM_lower(threshold, α, var_bounds)
                    push!(mi_constraints[leaf], @constraint(m, threshold ≤ sum(α .* x) + M * (1 - leaf_variables[leaf][1] + relax_var)))
                else
                    push!(mi_constraints[leaf], @constraint(m, threshold * leaf_variables[leaf][1] <= sum(α .* leaf_variables[leaf][2]) + relax_var))
                end
            end
            for (threshold, α) in lowerDict[leaf] # lowerDict is upper-bounded.
                if bigM
                    M = compute_hyperplane_bigM_upper(threshold, α, var_bounds)
                    push!(mi_constraints[leaf], @constraint(m, threshold + M * (1 - leaf_variables[leaf][1] + relax_var) >= sum(α .* x)))
                else
                    push!(mi_constraints[leaf], @constraint(m, threshold * leaf_variables[leaf][1] + relax_var >= sum(α .* leaf_variables[leaf][2])))
                end
            end
        end
    end
    merge!(append!, mi_constraints, mic)
    return mi_constraints, leaf_variables
end

"""
    $(TYPEDSIGNATURES)

Creates a set of MIO constraints from a OptimalTreeClassifier that thresholds a BlackBoxRegressor.

Arguments:
- m: JuMP.Model
- x:: independent JuMP.Variables (features of learner)
- y:: dependent JuMP.Variables
- lnr:: A fitted OptimalTreeLearner
- ul_data:: Upper and lower bounding hyperplanes for data in leaves of lnr (empty by default)
- symbs:: The varkeys corresponding to the lnr.
"""
function add_regr_constraints!(m::JuMP.Model, x::Array{JuMP.VariableRef}, y::JuMP.VariableRef,
                               lnr::IAI.OptimalTreeLearner,
                               ul_data::Dict; equality::Bool = false, bigM::Bool = false,
                               relax_var::Union{Real, JuMP.VariableRef} = 0, symbs = Symbol.(x))
    # TODO: Determine whether I should relax approximator or trust regions or both.
    if lnr isa OptimalTreeRegressor
        check_if_trained(lnr)
        all_leaves = find_leaves(lnr)
        # Add a binary variable for each leaf
        mi_constraints = Dict(leaf => [] for leaf in all_leaves)
        mi_constraints[1] = []
        leaf_variables = Dict{Int64, Tuple{JuMP.VariableRef, Array, JuMP.VariableRef}}()
        var_bounds = get_bounds(x)
        var_bounds = [var_bounds[var] for var in x]
        y_bounds = get_bound(y)
        if bigM
            for leaf in all_leaves
                leaf_variables[leaf] = (@variable(m, binary = true), [], @variable(m)) # third item in tuple is placeholder. TODO: find a more elegant bookkeeping method.
            end
            push!(mi_constraints[1], @constraint(m, sum(leaf_variables[leaf][1] for leaf in all_leaves) == 1))
        else
            for leaf in all_leaves
                leaf_variables[leaf], bd_mi = bounded_aux(x, y, @variable(m, binary=true))
                append!(mi_constraints[leaf], bd_mi)
            end
            push!(mi_constraints[1], @constraint(m, sum(leaf_variables[leaf][1] for leaf in all_leaves) == 1))
            [push!(mi_constraints[1], constr)
                for constr in @constraint(m, sum(leaf_variables[leaf][2] for leaf in all_leaves) .== x)]
            push!(mi_constraints[1], @constraint(m, sum(leaf_variables[leaf][3] for leaf in all_leaves) == y))
        end
        # Getting lnr data
        pwlDict = pwl_constraint_data(lnr, symbs)
        upperDict, lowerDict = trust_region_data(lnr, symbs)
        for leaf in all_leaves
            # ADDING PWL APPROXIMATIONS
            β0, β = pwlDict[leaf]
            if bigM
                if equality
                    M_upper = compute_regression_bigM_upper(β0, β, var_bounds, y_bounds)
                    M_lower = compute_regression_bigM_lower(β0, β, var_bounds, y_bounds)
                    # Use the ridge regressor!
                    push!(mi_constraints[leaf], @constraint(m, sum(β .* x) + β0 + M_upper * (1 - leaf_variables[leaf][1] + relax_var) >= y))
                    push!(mi_constraints[leaf], @constraint(m, sum(β .* x) + β0 <= y + M_lower * (1 - leaf_variables[leaf][1] + relax_var)))
                elseif !isempty(ul_data)
                    # Use only the lower approximator
                    β0, β = ul_data[leaf] # update the lower approximator
                    M_lower = compute_regression_bigM_lower(β0, β, var_bounds, y_bounds)
                    push!(mi_constraints[leaf], @constraint(m, sum(β .* x) + β0 <= y + M_lower * (1 - leaf_variables[leaf][1] + relax_var)))
                else
                    # Use only the ridge regressor as the lower approximator.
                    M_lower = compute_regression_bigM_lower(β0, β, var_bounds, y_bounds)
                    push!(mi_constraints[leaf], @constraint(m, sum(β .* x) + β0 <= y + M_lower * (1 - leaf_variables[leaf][1] + relax_var)))
                end
            else
                if equality
                    # Use the ridge regressor!
                    push!(mi_constraints[leaf], @constraint(m, sum(β .* leaf_variables[leaf][2]) + β0 * leaf_variables[leaf][1] + relax_var >= leaf_variables[leaf][3]))
                    push!(mi_constraints[leaf], @constraint(m, sum(β .* leaf_variables[leaf][2]) + β0 * leaf_variables[leaf][1] <= leaf_variables[leaf][3] + relax_var))
                elseif !isempty(ul_data)
                    # Use only the lower approximator
                    β0, β = ul_data[leaf] # update the lower approximator
                    push!(mi_constraints[leaf], @constraint(m, sum(β .* leaf_variables[leaf][2]) + β0 * leaf_variables[leaf][1] <= leaf_variables[leaf][3] + relax_var))
                else
                    # Use only the ridge regressor as the lower approximator.
                    push!(mi_constraints[leaf], @constraint(m, sum(β .* leaf_variables[leaf][2]) + β0 * leaf_variables[leaf][1] <= leaf_variables[leaf][3] + relax_var))
                end
            end
            # ADDING TRUST REGIONS
            for (threshold, α) in upperDict[leaf] # upperDict is lower-bounded.
                if bigM
                    M = compute_hyperplane_bigM_lower(threshold, α, var_bounds)
                    push!(mi_constraints[leaf], @constraint(m, threshold ≤ sum(α .* x) + M * (1 - leaf_variables[leaf][1] + relax_var)))
                else
                    push!(mi_constraints[leaf], @constraint(m, threshold * leaf_variables[leaf][1] <= sum(α .* leaf_variables[leaf][2]) + relax_var))
                end
            end
            for (threshold, α) in lowerDict[leaf] # lowerDict is upper-bounded.
                if bigM
                    M = compute_hyperplane_bigM_upper(threshold, α, var_bounds)
                    push!(mi_constraints[leaf], @constraint(m, M * (1 - leaf_variables[leaf][1] + relax_var) + threshold ≥ sum(α .* x)))
                else
                    push!(mi_constraints[leaf], @constraint(m, threshold * leaf_variables[leaf][1] + relax_var >=
                    sum(α .* leaf_variables[leaf][2])))
                end
            end
        end
        return mi_constraints, leaf_variables
    elseif lnr isa OptimalTreeClassifier
        # Add a binary variable for each leaf
        all_leaves = find_leaves(lnr)
        # Add a binary variable for each leaf
        feas_leaves = [i for i in all_leaves if Bool(IAI.get_classification_label(lnr, i))]
        mi_constraints = Dict(leaf => [] for leaf in feas_leaves)
        mi_constraints[1] = []
        leaf_variables = Dict{Int64, Tuple{JuMP.VariableRef, Array, JuMP.VariableRef}}()
        for leaf in feas_leaves
            leaf_variables[leaf], bd_mi = bounded_aux(x, y, @variable(m, binary=true))
            append!(mi_constraints[leaf], bd_mi)
        end
        push!(mi_constraints[1], @constraint(m, sum(leaf_variables[leaf][1] for leaf in feas_leaves) == 1))
        [push!(mi_constraints[1], constr)
            for constr in @constraint(m, sum(leaf_variables[leaf][2] for leaf in feas_leaves) .== x)]
        push!(mi_constraints[1], @constraint(m, sum(leaf_variables[leaf][3] for leaf in feas_leaves) == y))
        mi_constraints, leaf_variables = add_feas_constraints!(m, x, lnr, equality = equality, bigM = bigM, mic = mi_constraints, lv = leaf_variables)
        if !isempty(ul_data)
            if all(keys(ul_data) .<= 0) || !all(keys(ul_data) .>= 0) # means an upper or upperlower bounding tree
                mi_constraints = Dict(-key => value for (key, value) in mi_constraints) # hacky sign flipping for upkeep.
                leaf_variables = Dict(-key => value for (key, value) in leaf_variables)
            end
            for leaf in collect(keys(ul_data))
                γ0, γ = ul_data[leaf]
                if !haskey(mi_constraints, leaf) # occurs with ORTS with bounding hyperplanes
                    mi_constraints[leaf] = [@constraint(m, leaf_variables[-leaf][3] <= γ0 * leaf_variables[-leaf][1] +
                        sum(γ .* leaf_variables[-leaf][2]) + relax_var)]
                elseif leaf <= 0
                    push!(mi_constraints[leaf], @constraint(m, leaf_variables[leaf][3] <= γ0 * leaf_variables[leaf][1] +
                        sum(γ .* leaf_variables[leaf][2]) + relax_var))
                else
                    push!(mi_constraints[leaf], @constraint(m, leaf_variables[leaf][3] + relax_var >=
                    γ0 * leaf_variables[leaf][1] + sum(γ .* leaf_variables[leaf][2])))
                end
            end
        end
        return mi_constraints, leaf_variables
    end
end

""" Clears upper-bounding constraints from a BBR and its associated GM. """
function clear_upper_constraints!(gm, bbr::Union{BlackBoxRegressor, LinkedRegressor})
    for (leaf_key, leaf_constrs) in bbr.mi_constraints
        if leaf_key <= 0
            while !isempty(leaf_constrs)
                constr = pop!(leaf_constrs)
                if is_valid(gm.model, constr)
                    delete(gm.model, constr)
                else
                    push!(leaf_constrs, constr) # make sure to put the constraint back.
                    throw(OCTHaGOnException("Bug: Constraints could not be removed."))
                end
            end
            delete!(bbr.mi_constraints, leaf_key)
        end
    end
    for (leaf_key, (bin_var, leaf_vars, aux_dep)) in bbr.leaf_variables
        if leaf_key <= 0
            for leaf_var in leaf_vars
                if is_valid(gm.model, leaf_var)
                    delete(gm.model, leaf_var)
                else
                    throw(OCTHaGOnException("Bug: Variables could not be removed."))
                end
            end
            if is_valid(gm.model, bin_var)
                delete(gm.model, bin_var)
            else
                throw(OCTHaGOnException("Bug: Binary variable could not be removed."))
            end
            if is_valid(gm.model, aux_dep)
                delete(gm.model, aux_dep)
            else
                throw(OCTHaGOnException("Bug: Aux dependent variable could not be removed."))
            end
            delete!(bbr.leaf_variables, leaf_key)
        end
    end
    if bbr isa BlackBoxRegressor
        idx = active_upper_tree(bbr)
        delete!(bbr.active_trees, idx)
        for lr in bbr.lls
            clear_upper_constraints!(gm, lr)
        end
    end
    return
end

""" Clears lower-bounding constraints from a BBR and its associated GM. """
function clear_lower_constraints!(gm, bbr::Union{BlackBoxRegressor, LinkedRegressor})
    for (leaf_key, leaf_constrs) in bbr.mi_constraints
        if leaf_key >= 0
            while !isempty(leaf_constrs)
                constr = pop!(leaf_constrs)
                if is_valid(gm.model, constr)
                    delete(gm.model, constr)
                else
                    push!(leaf_constrs, constr) # make sure to put the constraint back.
                    throw(OCTHaGOnException("Bug: Constraints could not be removed."))
                end
            end
            delete!(bbr.mi_constraints, leaf_key)
        end
    end
    for (leaf_key, (bin_var, leaf_vars, aux_dep)) in bbr.leaf_variables
        if leaf_key >= 0
            for leaf_var in leaf_vars
                if is_valid(gm.model, leaf_var)
                    delete(gm.model, leaf_var)
                else
                    throw(OCTHaGOnException("Bug: Variables could not be removed."))
                end
            end
            if is_valid(gm.model, bin_var)
                delete(gm.model, bin_var)
            else
                throw(OCTHaGOnException("Bug: Binary variable could not be removed."))
            end
            if is_valid(gm.model, aux_dep)
                delete(gm.model, aux_dep)
            else
                throw(OCTHaGOnException("Bug: Aux dependent variable could not be removed."))
            end
            delete!(bbr.leaf_variables, leaf_key)
        end
    end
    bbr.active_leaves = []
    if bbr isa BlackBoxRegressor
        idx = active_lower_tree(bbr)
        delete!(bbr.active_trees, idx)
        for lr in bbr.lls
            clear_lower_constraints!(gm, lr)
        end
    end
    return
end

"""
    $(TYPEDSIGNATURES)

Clears the MI-approximating constraints and variables in GlobalModel and its sub-structs.
"""
function clear_tree_constraints!(gm::GlobalModel, bbc::Union{BlackBoxClassifier, LinkedClassifier})
    for (leaf_key, leaf_constrs) in bbc.mi_constraints
        while !isempty(leaf_constrs)
            constr = pop!(leaf_constrs)
            if is_valid(gm.model, constr)
                delete(gm.model, constr)
            else
                push!(leaf_constrs, constr) # make sure to put the constraint back.
                throw(OCTHaGOnException("Bug: Constraints could not be removed."))
            end
        end
        delete!(bbc.mi_constraints, leaf_key)
    end
    for (leaf_key, (bin_var, leaf_vars)) in bbc.leaf_variables
        for leaf_var in leaf_vars
            if is_valid(gm.model, leaf_var)
                delete(gm.model, leaf_var)
            else
                throw(OCTHaGOnException("Bug: Variables could not be removed."))
            end
        end
        if is_valid(gm.model, bin_var)
            delete(gm.model, bin_var)
        else
            throw(OCTHaGOnException("Bug: Binary variable could not be removed."))
        end
        delete!(bbc.leaf_variables, leaf_key)
    end
    bbc.active_leaves = []
    if bbc isa BlackBoxClassifier
        for lc in bbc.lls
            clear_tree_constraints!(gm, lc)
        end
    end
    return
end

function clear_tree_constraints!(gm::GlobalModel, bbr::Union{BlackBoxRegressor, LinkedRegressor})
    clear_lower_constraints!(gm, bbr)
    clear_upper_constraints!(gm, bbr)
    return
end

function clear_tree_constraints!(gm::GlobalModel, bbls::Array{BlackBoxLearner})
    for bbl in bbls
        clear_tree_constraints!(gm, bbl)
    end
    return
end

clear_tree_constraints!(gm::GlobalModel) = clear_tree_constraints!(gm, gm.bbls)

"""
    update_tree_constraints!(gm::GlobalModel, bbr::BlackBoxRegressor, idx = length(bbr.learners))
    update_tree_constraints!(gm::GlobalModel, bbc::BlackBoxClassifier, idx = length(bbc.learners))

Updates the MI constraints associated with a BBL.
For BBRs, makes sure to replace the appropriate lower/upper/regressor approximations.
"""
function update_tree_constraints!(gm::GlobalModel, bbr::BlackBoxRegressor, idx = length(bbr.learners))
    if isempty(bbr.active_trees) # no mi constraints yet.
        add_tree_constraints!(gm, bbr, idx)
    elseif length(bbr.active_trees) == 1 # one set of mi_constraints
        new_threshold = bbr.thresholds[idx]
        active_idx = collect(keys(bbr.active_trees))[1]
        last_threshold = bbr.active_trees[active_idx]
        if new_threshold.first == last_threshold.first
            clear_tree_constraints!(gm, bbr)
            add_tree_constraints!(gm, bbr, idx)
            return
        elseif new_threshold == Pair("reg", nothing) || last_threshold == Pair("reg", nothing) ||
            last_threshold.first == "upperlower" || new_threshold.first == "upperlower"
            clear_tree_constraints!(gm, bbr)
            add_tree_constraints!(gm, bbr, idx)
            return
        elseif new_threshold.first == "reg" && last_threshold.first == "upper" ||
                last_threshold.first == "reg" && new_threshold.first == "upper"
            if new_threshold.second == last_threshold.second # if corresponding upper thresholds.
                add_tree_constraints!(gm, bbr, idx)
                return
            else
                clear_tree_constraints!(gm, bbr)
                add_tree_constraints!(gm, bbr, idx)
                return
            end
        elseif last_threshold.first in valid_uppers && new_threshold.first in valid_lowers ||
            last_threshold.first in valid_lowers && new_threshold.first in valid_uppers
            add_tree_constraints!(gm, bbr, idx)
            return
        else # updating a lower or upper threshold
            clear_tree_constraints!(gm, bbr)
            add_tree_constraints!(gm, bbr, idx)
            return
        end
    elseif length(bbr.active_trees) == 2 # two sets of mi_constraints
        if bbr.thresholds[idx] == Pair("reg", nothing) || bbr.thresholds[idx].first == "upperlower"
            clear_tree_constraints!(gm, bbr)
            add_tree_constraints!(gm, bbr, idx)
            return
        end
        hypertype = bbr.thresholds[idx].first # otherwise, check approximation type
        if hypertype in valid_lowers
            clear_lower_constraints!(gm, bbr)
        elseif hypertype in valid_uppers
            clear_upper_constraints!(gm, bbr)
        else
            throw(OCTHaGOnException("Hypertype $(hypertype) not recognized."))
        end
        add_tree_constraints!(gm, bbr, idx)
        return
    else
        throw(OCTHaGOnException("$(gm.name) has too many active lower/upper bounding regressors."))
    end
end

function update_tree_constraints!(gm::GlobalModel, bbc::BlackBoxClassifier, idx = length(bbc.learners))
    clear_tree_constraints!(gm, bbc)
    add_tree_constraints!(gm, bbc, idx)
    return
end


export PrioritizedDQNLearner

"""
    PrioritizedDQNLearner(;kwargs...)

See paper: [Prioritized Experience Replay](https://arxiv.org/abs/1511.05952)
And also https://danieltakeshi.github.io/2019/07/14/per/

# Keywords

- `approximator`::[`AbstractApproximator`](@ref): used to get Q-values of a state.
- `target_approximator`::[`AbstractApproximator`](@ref): similar to `approximator`, but used to estimate the target (the next state).
- `loss_func`: the loss function.
- `γ::Float32=0.99f0`: discount rate.
- `batch_size::Int=32`
- `update_horizon::Int=1`: length of update ('n' in n-step update).
- `min_replay_history::Int=32`: number of transitions that should be experienced before updating the `approximator`.
- `update_freq::Int=4`: the frequency of updating the `approximator`.
- `target_update_freq::Int=100`: the frequency of syncing `target_approximator`.
- `stack_size::Union{Int, Nothing}=4`: use the recent `stack_size` frames to form a stacked state.
- `default_priority::Float64=100.`: the default priority for newly added transitions.
- `rng = Random.GLOBAL_RNG`
"""
mutable struct PrioritizedDQNLearner{
    Tq<:AbstractApproximator,
    Tt<:AbstractApproximator,
    Tf,
    R<:AbstractRNG,
} <: AbstractLearner
    approximator::Tq
    target_approximator::Tt
    loss_func::Tf
    sampler::NStepBatchSampler
    min_replay_history::Int
    update_freq::Int
    target_update_freq::Int
    update_step::Int
    default_priority::Float32
    β_priority::Float32
    rng::R
    # for logging
    loss::Float32
end

function PrioritizedDQNLearner(;
    approximator::Tq,
    target_approximator::Tt,
    loss_func::Tf,
    stack_size::Union{Int,Nothing} = nothing,
    γ::Float32 = 0.99f0,
    batch_size::Int = 32,
    update_horizon::Int = 1,
    min_replay_history::Int = 32,
    update_freq::Int = 1,
    target_update_freq::Int = 100,
    update_step::Int = 0,
    default_priority::Float32 = 100.0f0,
    β_priority::Float32 = 0.5f0,
    traces = SARTS,
    rng = Random.GLOBAL_RNG,
) where {Tq,Tt,Tf}
    copyto!(approximator, target_approximator)
    sampler = NStepBatchSampler{traces}(;γ=γ, n=update_horizon,stack_size=stack_size,batch_size=batch_size)
    PrioritizedDQNLearner(
        approximator,
        target_approximator,
        loss_func,
        sampler,
        min_replay_history,
        update_freq,
        target_update_freq,
        update_step,
        default_priority,
        β_priority,
        rng,
        0.0f0,
    )
end

Flux.functor(x::PrioritizedDQNLearner) = (Q = x.approximator, Qₜ = x.target_approximator),
y -> begin
    x = @set x.approximator = y.Q
    x = @set x.target_approximator = y.Qₜ
    x
end

"""

!!! note
    The state of the observation is assumed to have been stacked,
    if `!isnothing(stack_size)`.
"""
function (learner::PrioritizedDQNLearner)(env)
    env |>
    get_state |>
    x -> Flux.unsqueeze(x, ndims(x) + 1) |>
    x -> send_to_device(device(learner), x) |>
    learner.approximator |>
    vec |>
    send_to_host
end

function RLBase.update!(learner::PrioritizedDQNLearner, batch::NamedTuple)
    Q = learner.approximator
    Qₜ = learner.target_approximator
    γ = learner.sampler.γ
    β = learner.β_priority
    loss_func = learner.loss_func
    n = learner.sampler.n
    batch_size = learner.sampler.batch_size

    D = device(Q)
    s, a, r, t, s′ = (send_to_device(D,batch[x]) for x in SARTS)
    a = CartesianIndex.(a, 1:batch_size)

    updated_priorities = Vector{Float32}(undef, batch_size)
    w = 1.0f0 ./ ((batch.priority .+ 1f-10) .^ β)
    w ./= maximum(w)
    w = send_to_device(D, w)

    target_q = Qₜ(s′)
    if haskey(batch, :next_legal_actions_mask)
        l′ = send_to_device(D, batch[:next_legal_actions_mask])
        target_q .+= ifelse.(l′, 0.f0, typemin(Float32))
    end

    q′ = dropdims(maximum(target_q; dims = 1), dims = 1)
    G = r .+ γ^n .* (1 .- t) .* q′

    gs = gradient(params(Q)) do
        q = Q(s)[a]
        batch_losses = loss_func(G, q)
        loss = dot(vec(w), vec(batch_losses)) * 1 // batch_size
        ignore() do
            updated_priorities .= send_to_host(vec((batch_losses .+ 1f-10) .^ β))
            learner.loss = loss
        end
        loss
    end

    update!(Q, gs)
    updated_priorities
end


using Statistics: mean

using Parameters: @with_kw, @unpack
using RecipesBase
using Setfield: @set, print_application, Lens, PropertyLens

using .ArrayUtils: nan_

using .Continuations: as, ContinuationSweep, ContinuationSolution,
    ContinuationSolver, sweeps_as_vectors
using .BifurcationsBase: BifurcationSweep, BifurcationSolution,
    BifurcationSolver, special_intervals, measure, problem_of,
    get_param_axis1, get_param_axis2
using .Codim1: Codim1Sweep, Codim1Solution, Codim1Solver,
    AbstractCodim1SpecialPoint, stabilities, curves_by_stability,
    SpecialPoint, SpecialPointInterval, resolved_points, special_points!
using .Codim2: Codim2Sweep, Codim2Solution, Codim2Solver,
    AbstractCodim2SpecialPoint
using .Codim1LimitCycle: Codim1LCSweep, Codim1LCSolution, Codim1LCSolver,
    CWStateMeasurement, limitcycles
using .Codim2LimitCycle: Codim2LCSweep, Codim2LCSolution, Codim2LCSolver

const AbstractSolver = Union{ContinuationSolver, BifurcationSolver}
const AbstractSolution = Union{ContinuationSolution, BifurcationSolution}
const AbstractSweep = Union{ContinuationSweep, BifurcationSweep}
const Codim1Ctx = Union{Codim1Sweep, Codim1Solution, Codim1Solver,
                        AbstractCodim1SpecialPoint}
const Codim2Ctx = Union{Codim2Sweep, Codim2Solution, Codim2Solver,
                        AbstractCodim2SpecialPoint}
const Codim1LCCtx = Union{Codim1LCSweep, Codim1LCSolution, Codim1LCSolver}
const Codim2LCCtx = Union{Codim2LCSweep, Codim2LCSolution, Codim2LCSolver}
const LCCtx = Union{Codim1LCCtx, Codim2LCCtx}

@with_kw struct StateSpacePlotter{TC, TM}
    ctx::TC
    mapping::TM
end

StateSpacePlotter(ctx) =
    StateSpacePlotter(ctx, (:x => 1, :y => 2))
StateSpacePlotter(ctx::LCCtx) =
    StateSpacePlotter(ctx, (:x => 1, :y => 2, :color => :period))

plot_state_space(args...; kwargs...) =
    RecipesBase.plot(StateSpacePlotter(args...); kwargs...)

plot_state_space!(args...; kwargs...) =
    RecipesBase.plot!(StateSpacePlotter(args...); kwargs...)

plottables(ctx, key) = plottables(measure(ctx, key))
plottables(ctx, keys::Tuple) = plottables.((ctx,), keys)
plottables(::Any, ::Nothing) = nothing
# plottables(ctx, keys) = [plottables(c, keys) for c in curves(ctx)]

# post-processing:
plottables(x::Number) = x
plottables(xs::AbstractArray) = xs
plottables(xs::Vector{<: Tuple}) = transpose(hcat(collect.(xs)...))

process_key(::Any, key::Integer) = key
mapping_to_tuple(::Any; x, y) = (x, y)

"""
    get_keys(ctx, mapping; kwargs...) :: Tuple

Transform `mapping` to measurement keys.
"""
get_keys(ctx, mapping; kwargs...) =
    process_key.((ctx,), mapping_to_tuple(ctx; mapping...);
                 kwargs...)
# TODO: Maybe get rid of get_keys?  There are too much indirections here.

map_to_plottables(ctx, mapping; kwargs...) =
    plottables(ctx, get_keys(ctx, mapping; kwargs...))

dim_domain(ctx) = length(domain_prototype(ctx))
domain_prototype(solver::AbstractSolver) = domain_prototype(solver.sol)
domain_prototype(sol::AbstractSolution) = domain_prototype(sol.sweeps[1])
domain_prototype(sweep::AbstractSweep) = as(sweep, ContinuationSweep).u[1]
domain_prototype(point::SpecialPoint) = point.u
domain_prototype(point::SpecialPointInterval) = point.u0

# TODO: don't define process_key(::Codim1Ctx, ...); use measure() interface.
process_key(::Codim1Ctx, key::Integer) = key
function process_key(ctx::Codim1Ctx, key::Symbol)
    if key in (:p1, :parameter)
        return length(domain_prototype(ctx))
    else
        error("Unsupported mapping key: ", key)
    end
end
# - [[./fixedpoint.jl::^function.*get_u0]]

# TODO: don't define process_key(::Codim2Ctx, ...); use measure() interface.
process_key(::Codim2Ctx, key::Integer) = key
function process_key(ctx::Codim2Ctx, key::Symbol)
    if key == :p1
        return length(domain_prototype(ctx)) - 1
    elseif key == :p2
        return length(domain_prototype(ctx))
    elseif key in (:ω, :w)
        if ! (contkind(problem_of(ctx)) isa HopfCont)
            error("Key $key is invalid for $(contkind(prob))")
        end
        return length(domain_prototype(ctx)) - 2
    else
        error("Unsupported mapping key: ", key)
    end
end
# - [[./codim2/diffeq.jl::^function.*get_u0]]
# - [[./codim2/diffeq.jl::get(param_axis]]

assert_non_empty(solver::AbstractSolver) = assert_non_empty(solver.sol)
function assert_non_empty(sol::AbstractSolution)
    if isempty(sol.sweeps)
        error("Solution does not have a sweep:\n",
              sprint(showcompact, sol))
    end
    for sweep in sol.sweeps
        assert_non_empty(sweep)
    end
end
function assert_non_empty(sweep::AbstractSweep)
    if isempty(as(sweep, ContinuationSweep).u)
        error("Sweep does not have a point:\n",
              sprint(showcompact, sweep))
    end
end

default_mapping(::Any) = (:x => 1, :y => 2)
default_mapping(::Union{Codim1Ctx, Codim1LCCtx}) = (:x => :p1, :y => 1)
default_mapping(::Union{Codim2Ctx, Codim2LCCtx}) = (:x => :p1, :y => :p2)

# TODO: handle "color = :stability" etc. everywhere:
#=
default_mapping(::Union{Codim1Ctx, Codim1LCCtx}) =
    (x = :p1, y = 1,
     color = :stability, linestyle = :stability,
     marker = :special_point)
=#
# Supporting "marker" this way is ugly since using marker for the
# curves (non-special points) also makes sense.  (Maybe don't use
# named tuple then?)

lens_name(l::Lens) = sprint(print_application, l)
lens_name(l::PropertyLens) = lstrip(invoke(lens_name, Tuple{Lens}, l), '.')

plottable_name(::Any, key::Integer) = "u$key"

function plottable_name(ctx, key::Symbol)
    if key in (:parameter, :p1, :p2)
        lens = if key in (:parameter, :p1)
            get_param_axis1(ctx)
        else
            get_param_axis2(ctx)
        end
        return lens_name(lens)
    else
        return key
    end
end

plottable_name(ctx, key::CWStateMeasurement) = plottable_name(ctx, key.index)

function var_name(ctx::Codim1Ctx, i::Integer)
    prob = problem_of(ctx)
    if i == dim_domain(ctx)
        return lens_name(prob.p.param_axis)
    else
        return "u$i"
    end
end
# - [[./fixedpoint.jl::^function.*get_u0]]

function var_name(ctx::Codim2Ctx, i::Integer)
    prob = problem_of(ctx)
    if i == dim_domain(ctx)
        return lens_name(prob.param_axis2)
    elseif i == dim_domain(ctx) - 1
        return lens_name(prob.param_axis1)
    elseif contkind(prob) isa HopfCont && i == dim_domain(ctx) - 2
        return "ω"
    else
        return "u$i"
    end
end
# - [[./codim2/diffeq.jl::^function.*get_u0]]
# - [[./codim2/diffeq.jl::get(param_axis]]

const STYLE = Dict(
    :stability => Dict(
        # is_stable => plotattributes
        true => Dict(:linestyle => :solid, :linecolor => 1),
        false => Dict(:linestyle => :dot, :linecolor => 2),
    ),
    :point => Dict(
        :default => Dict(
            :markersize => 5,
        ),
        Codim1.PointTypes.saddle_node => Dict(
            :markercolor => 3,
            :markershape => :cross,
        ),
        Codim1.PointTypes.hopf => Dict(
            :markercolor => 4,
            :markershape => :circle,
        ),
        Codim1.PointTypes.period_doubling => Dict(
            :markercolor => 5,
            :markershape => :triangle,
        ),
        Codim2.PointTypes.cusp => Dict(
            :markercolor => 6,
            :markershape => :diamond,
        ),
        Codim2.PointTypes.bautin => Dict(
            :markercolor => 7,
            :markershape => :circle,
        ),
        Codim2.PointTypes.bogdanov_takens => Dict(
            :markercolor => 8,
            :markershape => :hexagon,
        ),
        Codim2.PointTypes.fold_hopf => Dict(
            :markercolor => 9,
            :markershape => :star8,
        ),
        Codim2.PointTypes.hopf_hopf => Dict(
            :markercolor => 10,
            :markershape => :star5,
        ),
    ),
)

struct PlottableData{T, K}
    data::T
    keys::K
end

@recipe function plot(
        pltbl::PlottableData,
        )
    color_key = pltbl.keys[3]
    if color_key !== nothing
        clims = nan_(extrema, (
            zs
            for data in pltbl.data
            for zs in plottables(data, color_key)
        ))
        clims --> clims
        colorbar_title --> string(color_key)
    end
    for data in pltbl.data
        xs, ys, zs = plottables(data, pltbl.keys)
        @series begin
            label --> ""
            if zs !== nothing
                line_z := zs
                marker_z := zs
            else
                color --> 1
            end
            xlabel --> plottable_name(data, pltbl.keys[1])
            ylabel --> plottable_name(data, pltbl.keys[2])
            (xs, ys)
        end
    end
end

@recipe function plot(
        wrapper::AbstractSweep,
        mapping = default_mapping(wrapper);
        # To be compatible with Codim1Sweep plotter. They are ignored:
        bif_style = nothing,
        resolve_points = nothing,
        include_points = nothing,
        )
    sweep = as(wrapper, ContinuationSweep)
    ix, iy = get_keys(sweep, mapping) :: Tuple{<: Integer, <: Integer}
    xs = [u[ix] for u in sweep.u]
    ys = [u[iy] for u in sweep.u]
    (xs, ys)
end

@recipe function plot(sol::ContinuationSolution,
                      mapping)
    ix, iy = get_keys(sol, mapping) :: Tuple{<: Integer, <: Integer}
    xs = sweeps_as_vectors(sol, ix)
    ys = sweeps_as_vectors(sol, iy)

    (xs, ys)
end

plottables(point::SpecialPoint, key::Integer) = [point.u[key]]
plottables(point::SpecialPointInterval, key::Integer) =
    [point.u0[key], point.u1[key]]
# TODO: plot SpecialPointInterval as an interval

@recipe function plot(
        point::Union{SpecialPoint,
                     SpecialPointInterval},
        mapping = default_mapping(point);
        bif_style = STYLE)

    if haskey(bif_style, :point)
        point_style = bif_style[:point]
        merge!(plotattributes,
               get(point_style, :default, Dict()),
               get(point_style, point.point_type, Dict()))
    end

    label --> ""
    map_to_plottables(point, mapping)
end

function maybe_get_points(sweep, include_points, resolve_points)
    if include_points
        if resolve_points
            return special_points!(sweep)
        else
            return special_intervals(sweep)
        end
    end
    return []
end

@recipe function plot(
        sweep::Codim1Sweep,
        mapping = default_mapping(sweep);
        bif_style = STYLE,
        resolve_points = false,
        include_points = true)

    ix, iy = get_keys(sweep, mapping) :: Tuple{<: Integer, <: Integer}
    (info, (xs, ys)) = curves_by_stability(sweep, (ix, iy))

    # is_stable => Dict(:linestyle => ?, :linecolor => ?)
    style = deepcopy(STYLE[:stability])
    # Empty `style` won't work so copying the default...  Not ideal.
    # TODO: support arbitrary keys, not just linestyle and linecolor
    if haskey(bif_style, :stability)
        stability_style = bif_style[:stability]
        merge!(style[true], get(stability_style, true, Dict()))
        merge!(style[false], get(stability_style, false, Dict()))
    end

    linestyle = []
    linecolor = []
    for stable in info
        push!(linestyle, style[stable][:linestyle])
        push!(linecolor, style[stable][:linecolor])
    end

    for point in maybe_get_points(sweep, include_points, resolve_points)
        @series begin
            point
        end
    end

    xname = var_name(sweep, ix)
    yname = var_name(sweep, iy)
    for (x, y, lc, ls) in zip(xs, ys, linecolor, linestyle)
        @series begin
            label --> ""
            linecolor --> lc
            linestyle --> ls
            xlabel --> xname
            ylabel --> yname
            (x, y)
        end
    end
end

mapping_to_tuple(
    ::LCCtx;
    x,
    y,
    color = nothing,
    # color,  # TODO: use it
) = (x, y, color)

process_key(::LCCtx, key::Any; _...) = key
process_key(::LCCtx, key::Integer; cw_agg=nothing) =
    if cw_agg !== nothing
        CWStateMeasurement(key, cw_agg)
    else
        key
    end

@recipe function plot(
        sweep::Union{Codim1LCSweep, Codim2LCSweep},
        mapping = default_mapping(sweep);
        cw_agg = extrema,  # coordinate-wise aggregator
        resolve_points = false,
        include_points = false,
        bif_style = STYLE)

    # TODO: support special points
    #=
    for point in maybe_get_points(sweep, include_points, resolve_points)
        @series begin
            (point, mapping)
        end
    end
    =#

    keys = get_keys(sweep, mapping; cw_agg = cw_agg)
    curve_list = [sweep]  # TODO: chunk by stability
    # curve_list = curves(sweep, :stability)

    PlottableData(curve_list, keys)
end

@recipe function plot(
        ssp::StateSpacePlotter{<: LCCtx};
        bif_style = STYLE)
    @unpack ctx, mapping = ssp

    # TODO: support special "points"

    keys = get_keys(ctx, mapping)

    PlottableData(limitcycles(ctx), keys)
end

@recipe function plot(
        sweep::Codim2Sweep,
        mapping = default_mapping(sweep);
        resolve_points = false,
        include_points = true)

    for point in maybe_get_points(sweep, include_points, resolve_points)
        @series begin
            (point, mapping)
        end
    end

    ix, iy = get_keys(sweep, mapping) :: Tuple{<: Integer, <: Integer}
    xname = var_name(sweep, ix)
    yname = var_name(sweep, iy)

    label --> ""
    linecolor --> 1
    xlabel --> xname
    ylabel --> yname
    as(sweep, ContinuationSweep), (x=ix, y=iy)  # plot(::AbstractSweep)
end

@recipe function plot(
        sol::BifurcationSolution,
        arguments__...)  # TODO: use `mapping`
    for sweep in sol.sweeps
        @series begin
            # `args` instead of `arguments__` does not work
            (sweep, arguments__...)
        end
    end
end

@recipe function plot(
        solver::BifurcationSolver,
        arguments__...;  # TODO: use `mapping`
        resolve_points = solver isa Union{Codim1Solver, Codim2Solver},
        include_points = true)

    for point in maybe_get_points(solver, include_points, resolve_points)
        @series begin
            point
        end
    end

    for sweep in solver.sol.sweeps
        @series begin
            include_points := false  # already included
            (sweep, arguments__...)
        end
    end
end

#=
@recipe function plot(
        ssp::StateSpacePlotter{<: BifurcationSolution},
        )
    for sweep in ssp.ctx.sweeps
        @series begin
            @set ssp.ctx = sweep
        end
    end
end

@recipe function plot(
        ssp::StateSpacePlotter{<: BifurcationSolver},
        )
    @set ssp.ctx = ssp.ctx.sol
end
=#


@inline function handle_tstop!(integrator)
    tstops = integrator.opts.tstops
    if !isempty(tstops)
      tdir_t = integrator.tdir * integrator.t
      tdir_ts_top = first(tstops)
      if tdir_t == tdir_ts_top
        pop!(tstops)
        integrator.just_hit_tstop = true
      elseif tdir_t > tdir_ts_top
        if !integrator.dtchangeable
          change_t_via_interpolation!(integrator, integrator.tdir * pop!(tstops), Val{true})
          integrator.just_hit_tstop = true
        else
          error("Something went wrong. Integrator stepped past tstops but the algorithm was dtchangeable. Please report this error.")
        end
      end
      isempty(integrator.opts.d_discontinuities) || handle_discontinuities!(integrator)
    end
  end
#=
Dealing with discontinuities

If we hit a discontinuity (this is checked in `apply_step!`), then we remove the
discontinuity, additional discontinuities at the current time point (if present), and
maybe also discontinuities and time stops coming shortly after the current time point
in `handle_discontinuities!`. The order of the discontinuity at the current time point is
defined as the lowest order of all these discontinuities.

If the problem is not neutral, we will only add additional discontinuities if
this order is less or equal to the order of the algorithm in
`add_next_discontinuities!`. If we add discontinuities, we add discontinuities
of the next order caused by constant lags (these we can calculate explicitly and
just add them to `d_discontinuities` and `tstops`) and we add the current
discontinuity to `tracked_discontinuities` which is the array of old
discontinuities that are checked by a `DiscontinuityCallback` (if existent).
=#

# handle discontinuities at the current time point of the `integrator`
function handle_discontinuities!(integrator::SDDEIntegrator)
    # remove all discontinuities at current time point and calculate minimal order
    # of these discontinuities
    d = pop!(integrator.opts.d_discontinuities)
    order = d.order
    while !isempty(integrator.opts.d_discontinuities) &&
        first(integrator.opts.d_discontinuities) == integrator.tdir * integrator.t

        d2 = pop!(integrator.opts.d_discontinuities)
        order = min(order, d2.order)
    end

    # remove all discontinuities close to the current time point as well and
    # calculate minimal order of these discontinuities
    # integrator.EEst has unitless type of integrator.t
    if typeof(integrator.EEst) <: AbstractFloat
        maxΔt = 10eps(integrator.t)

        while !isempty(integrator.opts.d_discontinuities) &&
            abs(first(integrator.opts.d_discontinuities).t - integrator.tdir * integrator.t) < maxΔt

            d2 = pop!(integrator.opts.d_discontinuities)
            order = min(order, d2.order)
        end

        # also remove all corresponding time stops
        while !isempty(integrator.opts.tstops) &&
            abs(first(integrator.opts.tstops) - integrator.tdir * integrator.t) < maxΔt

            pop!(integrator.opts.tstops)
        end
    end

    # add discontinuities of next order to integrator
    add_next_discontinuities!(integrator, order)
    nothing
end

"""
    add_next_discontinuities!(integrator::DDEIntegrator, order[, t=integrator.t])

Add discontinuities of next order that are propagated from discontinuity of
order `order` at time `t` in `integrator`, but only if `order` is less or equal
than the order of the applied method or the problem is neutral.

Discontinuities caused by constant delays are immediately calculated, and
discontinuities caused by dependent delays are tracked by a callback.
"""
function add_next_discontinuities!(integrator, _order, t = integrator.t)
    neutral = integrator.sol.prob.neutral
    order = Rational{Int}(_order)
    next_order = neutral ? order : order + 1//2

  # only track discontinuities up to order of the applied method
    alg_order = StochasticDiffEq.alg_order(getalg(integrator.alg))
    next_order <= alg_order + 1//2 || return

  # discontinuities caused by constant lags
    if has_constant_lags(integrator)
        constant_lags = integrator.sol.prob.constant_lags
        maxlag = integrator.tdir * (integrator.sol.prob.tspan[end] - t)

        for lag in constant_lags
            if integrator.tdir * lag < maxlag
        # calculate discontinuity and add it to heap of discontinuities and time stops
                d = Discontinuity(integrator.tdir * (t + lag), next_order)
                push!(integrator.opts.d_discontinuities, d)
                push!(integrator.opts.tstops, d.t)
            end
        end
    end

  # track propagated discontinuities with callback
    push!(integrator.tracked_discontinuities, Discontinuity(integrator.tdir * t, order))
    nothing
end


function model_selection(input::ModelSelectionInput)

	# For logging
	log_prefix = "GpABC model selection simulation "
	if isa(input, EmulatedModelSelectionInput)
		log_prefix = "GpABC model selection emulation "
	end

	@info "$(log_prefix) Population 1 - ABC Rejection ϵ = $(input.threshold_schedule[1])"

	tracker = initialise_modelselection(input)

	if all_models_dead(tracker)
		@warn "No particles were accepted in population 1 with threshold $(input.threshold_schedule[1])- terminating model selection algorithm"
		return build_modelselection_output(tracker, false)
	end

	if all_but_one_models_dead(tracker)
		@warn "All but one model is dead after population 1 - terminating model selection algorithm"
		return build_modelselection_output(tracker, true)
	end

	for i in 2:length(input.threshold_schedule)
		@info "$(log_prefix) Population $i - ABC SMC ϵ = $(input.threshold_schedule[i])"

		iterate_modelselection!(tracker, input.threshold_schedule[i])

		# Avoid infinite loop if no particles are accepted
		if all_models_dead(tracker)
			@warn "No particles were accepted in population $i with threshold $(input.threshold_schedule[i])- terminating model selection algorithm"
			return build_modelselection_output(tracker, false)
		end

		if all_but_one_models_dead(tracker)
			@warn "All but one model is dead after population $i - terminating model selection algorithm"
			return build_modelselection_output(tracker, true)
		end
	end

	return build_modelselection_output(tracker, true)
end

# Initialises the model selection run and runs the first (rejection) population
function initialise_modelselection(input::SimulatedModelSelectionInput)

	#
	# Check input sizes
	#
	if span(input.model_prior) != input.M
		throw(ArgumentError("There are $(input.M) models but the span of the model prior support is $(span(input.model_prior))"))
	end

	if length(input.distance_simulation_input) != input.M
		throw(ArgumentError("There are $(input.M) models but $(length(input.simulator_functions)) simulator functions"))
	end

	if length(input.parameter_priors) != input.M
		throw(ArgumentError("There are $(input.M) models but $(length(input.parameter_priors)) sets of parameter priors"))
	end

	total_n_accepted = 0
	n_iterations = 1

	#
	# Initialise arrays that will track rejection ABC run for each model - these
	# will be used to create SimulatedCandidateModelTrackers after the rejection ABC run
	#
	cm_trackers = [CandidateModelTracker(length(input.parameter_priors[m]))
								for m in 1:input.M]

    # Initialise variables to hold ABC results
    parameters = zeros(0)
    distance = 0.0
    weight_value = 0.0

	#
	# Compute first population using rejection-ABC
	#
	while total_n_accepted < input.n_particles && n_iterations <= input.max_iter

		m = rand(input.model_prior)
		parameters, weight_value = generate_parameters(input.parameter_priors[m])
		try
			distance = simulate_distance(parameters, input.distance_simulation_input[m])
        catch e
            if isa(e, DimensionMismatch)
                # This prevents the whole code from failing if there is a problem
                # solving the differential equation(s). The exception is thrown by the
                # distance function
                @warn "The summarised simulated data does not have the same size as the summarised reference data. If this is not happening at every iteration it may be due to the behaviour of OrdinaryDiffEq::solve - please check for related warnings. Continuing to the next iteration."
                n_iterations += 1
                continue
            else
                throw(e)
            end
        end

        cm_trackers[m].n_tries += 1

		# If particle accepted
		if distance[1] < input.threshold_schedule[1]
			total_n_accepted += 1
			update_candidatemodeltracker!(cm_trackers[m], parameters, distance[1], weight_value)
		end

		n_iterations += 1
	end

	if n_iterations == input.max_iter+1
		@warn "Simulated model selection reached maximum number of iterations ($(input.max_iter)) on the first population - consider trying more iterations."
	end

	@info join(["GpABC model selection simulation",
		"Completed $(n_iterations-1) iterations, accepting $total_n_accepted particles in total.",
		"Number of accepted parameters by model:",
		join([string("Model ", m, ": ",  cm_trackers[m].n_accepted) for m in 1:input.M], "\t")
		], "\n")

	smc_trackers = [SimulatedABCSMCTracker(
		length(input.parameter_priors[m]),
		[cm_trackers[m].n_accepted],
		[cm_trackers[m].n_tries],
		[input.threshold_schedule[1]],
        cm_trackers[m].n_accepted > 0,
		[cm_trackers[m].population],
		[cm_trackers[m].distances],
		[StatsBase.Weights(cm_trackers[m].weight_values ./ sum(cm_trackers[m].weight_values), 1.0)], 	# Normalise weights for each model
		input.parameter_priors[m],
		input.distance_simulation_input[m],
		1) for m in 1:input.M]

	return SimulatedModelSelectionTracker(input.M,
		input.n_particles,
		[input.threshold_schedule[1]],
		input.model_prior,
		smc_trackers,
		input.max_iter)
end

# Initialises the model selection run and runs the first (rejection) population
function initialise_modelselection(input::EmulatedModelSelectionInput{AF, CUD, ER, EPS, ET}) where {
	AF<:AbstractFloat, CUD<:ContinuousUnivariateDistribution,
	ER<:AbstractEmulatorRetraining, EPS<:AbstractEmulatedParticleSelection, ET<:AbstractEmulatorTraining}
	#
	# Check input sizes
	#
	if span(input.model_prior) != input.M
		throw(ArgumentError("There are $(input.M) models but the span of the model prior support is $(span(input.model_prior))"))
	end

	if length(input.parameter_priors) != input.M
		throw(ArgumentError("There are $(input.M) models but $(length(input.parameter_priors)) sets of parameter priors"))
	end

	#
	# Initialise arrays that will track rejection ABC run for each model - these
	# will be used to create CandidateModelTrackers after the rejection ABC run
	#
	cm_trackers = [CandidateModelTracker(length(input.parameter_priors[m]))
								for m in 1:input.M]

	#
	# Train the emulators
	#
	emulators = []
	for m in 1:input.M
		emulator = abc_train_emulator(input.parameter_priors[m], input.emulator_training_input[m])
		push!(emulators, emulator)
	end

    # Initialise variables to hold ABC results
    total_n_accepted = 0
	n_iterations = 1

	#
	# Emulate first population using rejection-ABC
	#
	while total_n_accepted < input.n_particles && n_iterations <= input.max_iter
		sampled_models = rand(input.model_prior, min(input.max_batch_size, input.n_particles-total_n_accepted))
		tries_this_it = [size(sampled_models[sampled_models.==m],1) for m=1:input.M]

		for m=1:input.M
			if tries_this_it[m] > 0

				parameters, weights = generate_parameters(input.parameter_priors[m], tries_this_it[m])
				distances, accepted_batch_idxs = abc_select_emulated_particles(emulators[m], parameters,
				 	input.threshold_schedule[1], input.emulated_particle_selection)
				cm_trackers[m].n_tries += tries_this_it[m]

				# If at least one particle is accepted
				if length(accepted_batch_idxs) > 0
					total_n_accepted += length(accepted_batch_idxs)
					update_candidatemodeltracker!(cm_trackers[m],
												  accepted_batch_idxs,
												  parameters,
												  distances,
												  weights)

				elseif length(accepted_batch_idxs) > tries_this_it[m]
					error("$(length(accepted_batch_idxs))) particles were accepted when model $m was only sampled $(tries_this_it[m]) times!")
				end
			end
		end

		n_iterations += 1
	end

	if n_iterations == input.max_iter+1
		@warn "Emulated model selection reached maximum number of iterations ($(input.max_iter)) without accepting $(input.n_particles) particles on the first population - consider trying more iterations."
	end

	@info join(["GpABC model selection emulation",
			"Completed $(n_iterations-1) iterations, accepting $total_n_accepted particles in total.",
			"Number of accepted parameters by model: ",
			join([string("Model ", m, ": ",  cm_trackers[m].n_accepted) for m in 1:input.M], "\t")
		], "\n")

	smc_trackers = [EmulatedABCSMCTracker{CUD, typeof(emulators[m]), ER, EPS}(
		length(input.parameter_priors[m]),
		[cm_trackers[m].n_accepted],
		[cm_trackers[m].n_tries],
		[input.threshold_schedule[1]],
		cm_trackers[m].n_accepted > 0,
		[cm_trackers[m].population],
		[cm_trackers[m].distances],
		[StatsBase.Weights(cm_trackers[m].weight_values ./ sum(cm_trackers[m].weight_values), 1.0)],
		input.parameter_priors[m],
		input.emulator_training_input[m],
		input.emulator_retraining,
		input.emulated_particle_selection,
		input.max_batch_size,
		input.max_iter,
		[emulators[m]]) for m = 1:input.M]

	return EmulatedModelSelectionTracker(
				input.M,
				input.n_particles,
				[input.threshold_schedule[1]],
				input.model_prior,
				smc_trackers,
				input.max_batch_size,
				input.max_iter)
end

# Perform a subsequent model selection iteration (based on ABC-SMC)
function iterate_modelselection!(tracker::SimulatedModelSelectionTracker,
	threshold::AbstractFloat)

	# Generate kernels for alive models only
	kernels = Dict{Int,Matrix{ContinuousUnivariateDistribution}}()
	for m in 1:tracker.M
		if tracker.smc_trackers[m].n_accepted[end] > 0
			kernels[m] = generate_kernels(tracker.smc_trackers[m].population[end],
										  tracker.smc_trackers[m].priors)
		end
	end

	# Initialise
	cm_trackers = [CandidateModelTracker(tracker.smc_trackers[m].n_params) for m in 1:tracker.M]
	total_n_accepted = 0
	n_iterations = 1
	parameters = zeros(0)
	distance = 0.0
	weight_value = 0.0

	while total_n_accepted < tracker.n_particles && n_iterations <= tracker.max_iter
		# Sample model from prior
		m = rand(tracker.model_prior)
		smc_tracker = tracker.smc_trackers[m]

		if smc_tracker.n_accepted[end] == 0
			continue
		end

		# Do ABC SMC for a single particle
		parameters, weight_value = generate_parameters(smc_tracker.priors, smc_tracker.weights[end], kernels[m])
		try
            distance = simulate_distance(parameters, smc_tracker.distance_simulation_input)
        catch e
            if isa(e, DimensionMismatch)
                # This prevents the whole code from failing if there is a problem
                # solving the differential equation(s). The exception is thrown by the
                # distance function
                @warn "The summarised simulated data does not have the same size as the summarised reference data. If this is not happening at every iteration it may be due to the behaviour of OrdinaryDiffEq::solve - please check for related warnings. Continuing to the next iteration."
                n_iterations += 1
                continue
            else
                throw(e)
            end
        end

        cm_trackers[m].n_tries += 1

		if distance[1] < threshold
			total_n_accepted += 1
			update_candidatemodeltracker!(cm_trackers[m], parameters, distance[1], weight_value)
		end

		n_iterations += 1

	end

	if n_iterations == tracker.max_iter+1
		@warn "Simulated model selection reached maximum number of iterations ($(tracker.max_iter)) on an SMC run - consider trying more iterations."
	end

	@info join(["GpABC model selection simulation",
			"Completed $(n_iterations-1) iterations, accepting $total_n_accepted particles in total.",
			"Number of accepted parameters by model: ",
			join([string("Model ", m, ": ",  cm_trackers[m].n_accepted) for m in 1:tracker.M], "\t")
		], "\n")

	update_modelselection_tracker!(tracker, cm_trackers, threshold)
end

function iterate_modelselection!(
	tracker::EmulatedModelSelectionTracker{CUD, ET, ER, EPS},
	threshold::AbstractFloat) where {CUD<:ContinuousUnivariateDistribution, ET,
	        ER<:AbstractEmulatorRetraining, EPS<:AbstractEmulatedParticleSelection}

	# Initialise
	cm_trackers = [CandidateModelTracker(length(smc_tracker.priors)) for smc_tracker in tracker.smc_trackers]

	# Generate kernels for alive models only
	kernels = Dict{Int,Matrix{ContinuousUnivariateDistribution}}()
	for m in 1:tracker.M
		if tracker.smc_trackers[m].n_accepted[end] > 0
			kernels[m] = generate_kernels(tracker.smc_trackers[m].population[end],
										  tracker.smc_trackers[m].priors)
		end
	end

	# Retrain emulators for alive models only
	retrained_emulators = Array{ET,1}(undef,tracker.M)
	for m in 1:tracker.M
		smc_tracker = tracker.smc_trackers[m]
		if smc_tracker.n_accepted[end] > 0
			particle_sampling_function(batch_size) = generate_parameters_no_weights(batch_size,
		        smc_tracker.population[end],
		        smc_tracker.weights[end],
		        kernels[m])
			retrained_emulators[m] = abc_retrain_emulator(smc_tracker.emulators[end],
				particle_sampling_function,
				threshold,
				smc_tracker.emulator_training_input,
				smc_tracker.emulator_retraining)
		else
			retrained_emulators[m] = smc_tracker.emulators[end]
		end
	end

    # Initialise
    total_n_accepted = 0
	n_iterations = 1

    # emulate
	while total_n_accepted < tracker.n_particles && n_iterations <= tracker.max_iter

		sampled_models = rand(tracker.model_prior, min(tracker.max_batch_size, tracker.n_particles-total_n_accepted))
		tries_this_it = [size(sampled_models[sampled_models.==m],1) for m=1:tracker.M]

		for m = 1:tracker.M
			# Skip a model that wasn't sampled in this iteration or had no accepted particles in the previous population
			if tries_this_it[m] == 0 || tracker.smc_trackers[m].n_accepted[end] == 0
				continue
			end
			smc_tracker = tracker.smc_trackers[m]
			parameters, weights = generate_parameters(tries_this_it[m],
	                                                 smc_tracker.priors,
	                                                 smc_tracker.weights[end],
	                                                 kernels[m])
			cm_trackers[m].n_tries += tries_this_it[m]

			distances, accepted_indices = abc_select_emulated_particles(retrained_emulators[m], parameters, threshold, smc_tracker.selection)

			# If at least one particle is accepted
			if length(accepted_indices) > 0
				total_n_accepted += length(accepted_indices)
				update_candidatemodeltracker!(cm_trackers[m],
												  accepted_indices,
												  parameters,
												  distances,
												  weights)

			elseif length(accepted_indices) > tries_this_it[m]
				error("$(length(accepted_indices)) particles were accepted when model $m was only sampled $(tries_this_it[m]) times!")
			end
		end

		n_iterations += 1
	end

	if n_iterations == tracker.max_iter+1
		@warn "Emulated model selection reached maximum number of iterations ($(tracker.max_iter)) on an SMC run - consider trying more iterations."
	end

	@info join(["GpABC model selection emulation",
			"Completed $(n_iterations-1) iterations, accepting $total_n_accepted particles in total.",
			"Number of accepted parameters by model: ",
			join([string("Model ", m, ": ",  cm_trackers[m].n_accepted) for m in 1:tracker.M], "\t")],
		"\n")

	update_modelselection_tracker!(tracker, cm_trackers, retrained_emulators, threshold)
end

# not exported
function build_modelselection_output(tracker::ModelSelectionTracker, successful_completion::Bool)
	return ModelSelectionOutput(
		tracker.M,
		[[tracker.smc_trackers[m].n_accepted[i] for m = 1:tracker.M] for i in 1:length(tracker.threshold_schedule)],
		tracker.threshold_schedule,
		[buildAbcSmcOutput(tracker.smc_trackers[m]) for m in 1:tracker.M],
		successful_completion
		)
end

# not exported
function all_models_dead(tracker::ModelSelectionTracker)
	return sum([tracker.smc_trackers[m].n_accepted[end] for m = 1:tracker.M]) == 0
end

# not exported
function all_but_one_models_dead(tracker::ModelSelectionTracker)
	return sum([tracker.smc_trackers[m].n_accepted[end] == 0 for m = 1:tracker.M]) == tracker.M-1
end

# not exported
function update_candidatemodeltracker!(
	cm_tracker::CandidateModelTracker,
	parameters::AbstractArray{Float64,2},
	distance::AbstractFloat,
	weight_value::AbstractFloat)

	cm_tracker.n_accepted += 1
	cm_tracker.population = vcat(cm_tracker.population, parameters)
	push!(cm_tracker.distances, distance)
	push!(cm_tracker.weight_values, weight_value)
end

function update_candidatemodeltracker!(
	cm_tracker::CandidateModelTracker,
	accepted_batch_idxs::AbstractArray{Int64,1},
	parameter_batch::AbstractArray{Float64,2},
	distance_batch::AbstractArray{Float64,1},
	weight_values_batch::AbstractArray{Float64,1})

	cm_tracker.n_accepted += length(accepted_batch_idxs)
	cm_tracker.population = vcat(cm_tracker.population, parameter_batch[accepted_batch_idxs,:])
	cm_tracker.distances = vcat(cm_tracker.distances, distance_batch)
	cm_tracker.weight_values = vcat(cm_tracker.weight_values, weight_values_batch[accepted_batch_idxs])
end

# not exported
function update_modelselection_tracker!(
	tracker::SimulatedModelSelectionTracker,
	cm_trackers::AbstractArray{CandidateModelTracker{AF},1},
	threshold::AF) where {AF<:AbstractFloat}
	push!(tracker.threshold_schedule, threshold)
	for m in 1:tracker.M
		push!(tracker.smc_trackers[m].n_accepted, cm_trackers[m].n_accepted)
		push!(tracker.smc_trackers[m].n_tries, cm_trackers[m].n_tries)
		push!(tracker.smc_trackers[m].threshold_schedule, threshold)
		push!(tracker.smc_trackers[m].population, cm_trackers[m].population)
		push!(tracker.smc_trackers[m].distances, cm_trackers[m].distances)
		push!(tracker.smc_trackers[m].weights, normalise(cm_trackers[m].weight_values))
	end
end

# not exported
function update_modelselection_tracker!(
	tracker::EmulatedModelSelectionTracker,
	cm_trackers::AbstractArray{CandidateModelTracker{AF},1},
	retrained_emulators::AbstractArray{GPModel,1},
	threshold::AF) where {AF<:AbstractFloat}

	push!(tracker.threshold_schedule, threshold)
	for m in 1:tracker.M
		push!(tracker.smc_trackers[m].n_accepted, cm_trackers[m].n_accepted)
		push!(tracker.smc_trackers[m].n_tries, cm_trackers[m].n_tries)
		push!(tracker.smc_trackers[m].threshold_schedule, threshold)
		push!(tracker.smc_trackers[m].population, cm_trackers[m].population)
		push!(tracker.smc_trackers[m].distances, cm_trackers[m].distances)
		push!(tracker.smc_trackers[m].weights, normalise(cm_trackers[m].weight_values))
		push!(tracker.smc_trackers[m].emulators, retrained_emulators[m])
	end
end


"""
    plotseis(S[, fmt=FMT, use_name=false, n=N])

Renormalized, time-aligned trace plot of data in S.x using timestamps in S.t.

Keywords:
* fmt=FMT formats x-axis labels using C-language `strftime` format string `FMT`.
If unspecified, the format is determined by when data in `S` start and end.
* use_name=true uses `S.name`, rather than `S.id`, for labels.
* n=N sets the number of X-axis ticks.

"""
function plotseis(S::SeisData; fmt::String="auto", use_name::Bool=false, nxt::Int64=5)
  xmi = typemax(Int64)
  xma = typemin(Int64)

  fig = PyPlot.figure(figsize=[8.0, 6.0], dpi=150)
  ax = PyPlot.axes([0.20, 0.10, 0.72, 0.85])
  for i = 1:S.n
    x = rescaled(S.x[i].-mean(S.x[i]),i)
    if S.fs[i] > 0
      t = t_expand(S.t[i], S.fs[i])
      plot(t, x, linewidth=1.0)
    else
      t = view(S.t[i], :, 2)
      plot(t, x, "o", linewidth=1, markeredgecolor=[0,0,0])
    end
    xmi = min(xmi, t[1])
    xma = max(xma, t[end])
  end

  xfmt(xmi, xma, fmt, nxt)
  PyPlot.setp(gca().get_yticklabels(), fontsize=8.0, color="black", fontweight="bold", family="serif")
  PyPlot.setp(gca().get_xticklabels(), fontsize=8.0, color="black", fontweight="bold", family="serif")
  PyPlot.yticks(1:S.n, map((i) -> replace(i, " " => ""), use_name ? S.name : S.id))
  PyPlot.ylim(0.5, S.n+0.5)
  return fig
end


module ImPlot

using CEnum

include("libcimplot/libcimplot.jl")

# Types
using .LibCImPlot
import CImGui: ImVec2, ImVec4
import .LibCImPlot: ImPlotStyleVar, ImPlotPoint, ImPlotColormap, ImPlotCol
import .LibCImPlot: ImPlotFlags, ImPlotAxisFlags

# Export plot flags
for i in [ImPlotFlags, ImPlotAxisFlags]
    for j in instances(i)
        @eval export $(Symbol(j))
    end
end

function BeginPlot(title_id::String, x_label, y_label, size::ImVec2;
                    flags::ImPlotFlags = ImPlotFlags_None,
                    x_flags::ImPlotAxisFlags = ImPlotAxisFlags_None,
                    y_flags::ImPlotAxisFlags = ImPlotAxisFlags_None,
                    y2_flags::ImPlotAxisFlags = ImPlotAxisFlags_None,
                    y3_flags::ImPlotAxisFlags = ImPlotAxisFlags_None)::Bool

    LibCImPlot.BeginPlot(title_id, x_label, y_label, size,
                         flags, x_flags, y_flags, y2_flags, y3_flags)
end

import .LibCImPlot: EndPlot

include("lines.jl")
include("shaded.jl")
include("scatter.jl")
include("heatmap.jl")
include("digital.jl")
include("barchart.jl")
include("piechart.jl")
include("errorbars.jl")
include("other.jl")
include("util.jl")
include("color.jl")
include("styling.jl")

end # module


############################
###### Boundary Types ######
############################

abstract type AbstractProblem  end
abstract type BoundaryID       end
abstract type BoundaryFunction end
abstract type BoundaryArray end

struct DirichletID <: BoundaryID id::Vector{Symbol} end
struct NeumannID   <: BoundaryID id::Vector{Symbol} end
struct RobinID     <: BoundaryID id::Vector{Symbol} end

struct DirichletFunction <: BoundaryFunction f::Function end
struct NeumannFunction   <: BoundaryFunction f::Function end
struct RobinFunction     <: BoundaryFunction f::Function end
struct ForcingFunction   <: BoundaryFunction f::Function end

struct DirichletArray <: BoundaryArray v::Vector end
struct NeumannArray   <: BoundaryArray v::Vector end
struct RobinArray     <: BoundaryArray v::Vector end
struct ForcingArray   <: BoundaryArray v::Vector end

#### constructors ####

function Dirichlet(s::Symbol ...)
  return DirichletID([i for i in s])
end

function Neumann(s::Symbol ...)
  return NeumannID([i for i in s])
end

function Robin(s::Symbol ...)
  return RobinID([i for i in s])
end

function Dirichlet(f::Function)
  return DirichletFunction(f)
end
function Dirichlet(v::Vector)
  return DirichletArray(v)
end

function Neumann(f::Function)
  return NeumannFunction(f)
end
function Neumann(v::Vector)
  return NeumannArray(v)
end

function Robin(f::Function)
  return RobinFunction(f)
end
function Robin(v::Vector)
  return RobinArray(v)
end

function Forcing(f::Function)
  return ForcingFunction(f)
end
function Forcing(v::Vector)
  return ForcingArray(v)
end

#### Methods #####

isDirichletID(x)       = (typeof(x)==DirichletID ? true : false)
isNeumannID(x)         = (typeof(x)==NeumannID   ? true : false)
isRobinID(x)           = (typeof(x)==RobinID     ? true : false)

isDirichletFunction(x) = (typeof(x)==DirichletFunction ? true : false)
isNeumannFunction(x)   = (typeof(x)==NeumannFunction   ? true : false)
isRobinFunction(x)     = (typeof(x)==RobinFunction     ? true : false)
isForcingFunction(x)   = (typeof(x)==ForcingFunction   ? true : false)

isDirichletArray(x) = (typeof(x)==DirichletArray ? true : false)
isNeumannArray(x)   = (typeof(x)==NeumannArray   ? true : false)
isRobinArray(x)     = (typeof(x)==RobinArray     ? true : false)
isForcingArray(x)   = (typeof(x)==ForcingArray   ? true : false)

############################
###### Problem Types #######
############################

mutable struct Problem <: AbstractProblem
  OperatorType::Symbol
  bcid::DictList{Int}
  bcval::DictList{Float64}
end

######  Constructors for Problem  ##########

# Assumes that only 1 dNodes and nNodes and rNodes (maximally ...)
function Problem(mesh::ScalarMesh,Nodes,bcfun,OpType)
  bcIDNodes::Vector{Vector{Int}} = []
  bcIDNamesSym::Vector{Symbol} = []
  bcValNodes::Vector{Vector{Float64}} = []
  bcValNamesSym::Vector{Symbol} = []

  # derive boundary node ID information
  dNodesSym::Vector{DirichletID} = filter(isDirichletID,Nodes)
  if length(dNodesSym)>0
    dNodes::Vector{Int} = vcat([mesh.bdry[s] for s in dNodesSym[1].id]...)
    push!(bcIDNodes,dNodes)
    push!(bcIDNamesSym,:dNodes)
  end

  nNodesSym::Vector{NeumannID} = filter(isNeumannID,Nodes)
  if length(nNodesSym)>0
    nNodes::Vector{Int} = vcat([mesh.bdry[s] for s in nNodesSym[1].id] ...)
    push!(bcIDNodes,nNodes)
    push!(bcIDNamesSym,:nNodes)
  end

  rNodesSym = filter(isRobinID,Nodes)
  if length(rNodesSym)>0
    rNodes::Vector{Int} = vcat([mesh.bdry[s] for s in rNodesSym[1].id] ...)
    push!(bcIDNodes,rNodes)
    push!(bcIDNamesSym,:rNodes)
  end

  # derive boundary value information
  # dirichlet
  dNodesFun::Vector{DirichletFunction} = filter(isDirichletFunction,bcfun)
  if length(dNodesFun)>0
    dBCarr = [dNodesFun[1].f(mesh.xy[i].x,mesh.xy[i].y)
                               for i in dNodes]
    push!(bcValNodes,dBCarr)
    push!(bcValNamesSym,:dBC)
  end
  dNodesArr::Vector{DirichletArray} = filter(isDirichletArray,bcfun)
  if length(dNodesArr)>0
    dBCarr = [dNodesArr[1].v[i] for i in dNodes]
    push!(bcValNodes,dBCarr)
    push!(bcValNamesSym,:dBC)
  end

  # neumann
  nNodesFun::Vector{NeumannFunction} = filter(isNeumannFunction,bcfun)
  nNodesArr::Vector{NeumannArray} = filter(isNeumannArray,bcfun)
  if (length(nNodesFun)>0) || (length(nNodesArr)>0)
    if length(nNodesFun)>0
      nBCarr = [nNodesFun[1].f(mesh.xy[i].x,mesh.xy[i].y)
                                for i=1:length(nNodes)]
      push!(bcValNodes,nBCarr)
      push!(bcValNamesSym,:nBC)
    elseif length(nNodesArr)>0
      nBCarr = [nNodesArr[1].v[i] for i=1:length(nNodes)]
      push!(bcValNodes,nBCarr)
      push!(bcValNamesSym,:nBC)
    end
  else
    if :nNodes in bcIDNamesSym
      nBCarr = zeros(Float64,length(nNodes))
      push!(bcValNodes,nBCarr)
      push!(bcValNamesSym,:nBC)
    end
  end

  # robin
  rNodesFun::Vector{RobinFunction} = filter(isRobinFunction,bcfun)
  if length(rNodesFun)>0
    rBCarr = [rNodesFun[1].f(mesh.xy[i].x,mesh.xy[i].y)
                               for i in rNodes]
    push!(bcValNodes,rBCarr)
    push!(bcValNamesSym,:rBC)
  end
  rNodesArr::Vector{DirichletArray} = filter(isRobinArray,bcfun)
  if length(rNodesArr)>0
    rBCarr = [rNodesArr[1].v[i] for i in rNodes]
    push!(bcValNodes,rBCarr)
    push!(bcValNamesSym,:rBC)
  end

  # forcing
  fNodesFun::Vector{ForcingFunction} = filter(isForcingFunction,bcfun)
  fNodesArr::Vector{ForcingArray} = filter(isForcingArray,bcfun)
  if (length(fNodesFun)>0) || (length(fNodesArr)>0)
    if length(fNodesFun)>0
      farr = [fNodesFun[1].f(mesh.xy[i].x,mesh.xy[i].y)
                                for i=1:length(mesh.xy)]
      push!(bcValNodes,farr)
      push!(bcValNamesSym,:forcing)
    elseif length(fNodesArr)>0
      farr = [fNodesArr[1].v[i] for i=1:length(mesh.xy)]
      push!(bcValNodes,farr)
      push!(bcValNamesSym,:forcing)
    end
  else
    farr = zeros(Float64,length(mesh.xy))
    push!(bcValNodes,farr)
    push!(bcValNamesSym,:forcing)
  end

  # convert all these into appropriate types for problem object
  bcIDNames = Index(bcIDNamesSym)
  bcid  = DictList{Int}(bcIDNodes,bcIDNames)

  bcValNames = Index(bcValNamesSym)
  bcval  = DictList{Float64}(bcValNodes,bcValNames)

  # return FluidProblem object
  return Problem(OpType,bcid,bcval)
end

# assumes that dUNodes = dVNodes and anything that isn't dirichlet is neumann
function Problem(mesh::FluidMesh,Nodes,bcfun,OperatorType)
  bcIDNodes::Vector{Vector{Int}} = []
  bcIDNamesSym::Vector{Symbol} = []
  bcValNodes::Vector{Vector{Float64}} = []
  bcValNamesSym::Vector{Symbol} = []

  # derive boundary node ID information
  dNodesSym::Vector{DirichletID} = filter(isDirichletID,Nodes)
  if length(dNodesSym)>0
    dUNodes::Vector{Int} = unique(vcat([mesh.bdry[s]
                                  for s in dNodesSym[1].id]...))
    push!(bcIDNodes,dUNodes)
    push!(bcIDNamesSym,:dUNodes)

    if length(dNodesSym)==2
      dVNodes::Vector{Int} = unique(vcat([mesh.bdry[s]
                                    for s in dNodesSym[2].id]...))
    else
      dVNodes = copy(dUNodes)
    end

    push!(bcIDNodes,dVNodes)
    push!(bcIDNamesSym,:dVNodes)
  end

  nNodesSym::Vector{NeumannID} = filter(isNeumannID,Nodes)
  if length(nNodesSym)==1
    nNodes::Vector{Int} = unique(vcat([mesh.bdry[s]
                                 for s in nNodesSym[1].id] ...))
    push!(bcIDNodes,nNodes)
    push!(bcIDNamesSym,:nNodes)
  end

  # derive boundary value information
  dNodesFun::Vector{DirichletFunction} = filter(isDirichletFunction,bcfun)
  if length(dNodesFun)==2
    dUBCarr = [dNodesFun[1].f(mesh.xy[i].x,mesh.xy[i].y)
                               for i in dUNodes]
    push!(bcValNodes,dUBCarr)
    push!(bcValNamesSym,:dUBC)

    dVBCarr = [dNodesFun[2].f(mesh.xy[i].x,mesh.xy[i].y)
                               for i in dVNodes]
    push!(bcValNodes,dVBCarr)
    push!(bcValNamesSym,:dVBC)
  end
  dNodesArr::Vector{DirichletFunction} = filter(isDirichletArray,bcfun)
  if length(dNodesArr)==2
    dUBCarr = [dNodesArr[1].v[i] for i in dUNodes]
    push!(bcValNodes,dUBCarr)
    push!(bcValNamesSym,:dUBC)

    dVBCarr = [dNodesArr[2].v[i] for i in dVNodes]
    push!(bcValNodes,dVBCarr)
    push!(bcValNamesSym,:dVBC)
  end

  nNodesFun::Vector{NeumannFunction} = filter(isNeumannFunction,bcfun)
  if length(nNodesFun)==1
    nBCarr::Vector{Float64} = [nNodesFun[1].f(mesh.xy[i].x,mesh.xy[i].y)
                               for i in nNodes]
    push!(bcValNodes,nBCarr)
    push!(bcValNamesSym,:nBC)
  else
    if :nNodes in bcIDNamesSym
      nBCarr = zeros(Float64,length(nNodes))
      push!(bcValNodes,nBCarr)
      push!(bcValNamesSym,:nBC)
    end
  end

  fNodesFun::Vector{ForcingFunction} = filter(isForcingFunction,bcfun)
  fNodesArr::Vector{ForcingArray} = filter(isForcingArray,bcfun)
  if length(fNodesFun)==2
    farrX = [fNodesFun[1].f(mesh.xy[i].x,mesh.xy[i].y)
                               for i=1:length(mesh.xy)]
    farrY = [fNodesFun[2].f(node.x,node.y)
                               for node in mesh.xy]
    push!(bcValNodes,farrX)
    push!(bcValNamesSym,:forcingX)
    push!(bcValNodes,farrY)
    push!(bcValNamesSym,:forcingY)
  elseif length(fNodesArr)==2
    farrX = [fNodesArr[1].v[i] for i=1:length(mesh.xy)]
    farrY = [fNodesArr[2].v[i] for i=1:length(mesh.xy)]
    push!(bcValNodes,farrX)
    push!(bcValNamesSym,:forcingX)
    push!(bcValNodes,farrY)
    push!(bcValNamesSym,:forcingY)
  else
    farrX = zeros(Float64,length(mesh.xy))
    farrY = zeros(Float64,length(mesh.xy))
    push!(bcValNodes,farrX)
    push!(bcValNamesSym,:forcingX)
    push!(bcValNodes,farrY)
    push!(bcValNamesSym,:forcingY)
  end

  # convert all these into appropriate types for problem object
  bcIDNames = Index(bcIDNamesSym)
  bcid  = DictList{Int}(bcIDNodes,bcIDNames)

  bcValNames = Index(bcValNamesSym)
  bcval  = DictList{Float64}(bcValNodes,bcValNames)

  # return Problem object
  return Problem(OperatorType,bcid,bcval)
end

####### Problem Methods #######

bc(prob::Problem) = indexlist(prob.bcid)


module PHonon

using REPL.Terminals: TTYTerminal
using REPL.TerminalMenus: RadioMenu, request

using Crayons.Box: GREEN_FG
using QuantumESPRESSOBase.Inputs.PHonon:
    PhNamelist, Q2rNamelist, MatdynNamelist, DynmatNamelist, Q2rInput, DynmatInput

using ..Inputs: namelist_builder, setfield_helper

import QuantumESPRESSOHelpers.Inputs

function Inputs.namelist_builder(terminal::TTYTerminal, ::Type{T}) where {T<:PhNamelist}
    print(
        terminal,
        GREEN_FG("Please input the atomic mass [amu] of each atomic type `amass` (separated by spaces): ") |>
        string,
    )
    amass = map(x -> parse(Float64, x), split(readline(terminal), " ", keepempty = false))
    epsil_pool = pairs((false, true))
    epsil = epsil_pool[request(
        terminal,
        GREEN_FG("Please select the `epsil`: ") |> string,
        RadioMenu([false, true]),
    )]
    q_in_band_form_pool = pairs((false, true))
    q_in_band_form = q_in_band_form_pool[request(
        terminal,
        GREEN_FG("Please select the `q_in_band_form`: ") |> string,
        RadioMenu([false, true]),
    )]
    print(
        terminal,
        GREEN_FG("Please input parameters of the Monkhorst-Pack grid `nq` 1-3 (separated by spaces): ") |>
        string,
    )
    nq1, nq2, nq3 =
        map(x -> parse(Float64, x), split(readline(terminal), " ", keepempty = false))
    print(
        terminal,
        GREEN_FG("Please input parameters of the Monkhorst-Pack grid `nk` 1-3 (separated by spaces): ") |>
        string,
    )
    nk1, nk2, nk3 =
        map(x -> parse(Float64, x), split(readline(terminal), " ", keepempty = false))
    print(
        terminal,
        GREEN_FG("Please input offset `k` 1-3 (separated by spaces): ") |> string,
    )
    k1, k2, k3 =
        map(x -> parse(Float64, x), split(readline(terminal), " ", keepempty = false))
    ph = T(
        amass = amass,
        epsil = epsil,
        q_in_band_form = q_in_band_form,
        nq1 = nq1,
        nq2 = nq2,
        nq3 = nq3,
        nk1 = nk1,
        nk2 = nk2,
        nk3 = nk3,
        k1 = k1,
        k2 = k2,
        k3 = k3,
    )
    return setfield_helper(terminal, ph)
end # function namelist_builder
function Inputs.namelist_builder(terminal::TTYTerminal, ::Type{T}) where {T<:Q2rNamelist}
    print(terminal, GREEN_FG("name of input dynamical matrices `fildyn`: ") |> string)
    fildyn = strip(readline(terminal))
    print(terminal, GREEN_FG("name of output force constants `flfrc`: ") |> string)
    flfrc = strip(readline(terminal))
    zasr_pool = pairs(("no", "simple", "crystal", "one-dim", "zero-dim"))
    zasr = zasr_pool[request(
        terminal,
        GREEN_FG("Please input the type of acoustic sum rules used for the Born effective charges `zasr`: ") |>
        string,
        RadioMenu(collect(values(zasr_pool))),
    )]
    q2r = T(fildyn = fildyn, flfrc = flfrc, zasr = zasr)
    return setfield_helper(terminal, q2r)
end # function namelist_builder
function Inputs.namelist_builder(
    terminal::TTYTerminal,
    ::Type{T},
) where {T<:MatdynNamelist}
    dos_pool = pairs((false, true))
    dos = dos_pool[request(
        terminal,
        GREEN_FG("Please select if calculate phonon density of states `dos`: ") |> string,
        RadioMenu([false, true]),
    )]
    print(
        terminal,
        GREEN_FG("Please input the energy step, in cm^(-1) `deltaE`: ") |> string,
    )
    deltaE = parse(Float64, readline(terminal))
    print(
        terminal,
        GREEN_FG("Please input uniform q-point grid for DOS calculation `nk` 1-3 (separated by spaces): ") |>
        string,
    )
    nk1, nk2, nk3 =
        map(x -> parse(Float64, x), split(readline(terminal), " ", keepempty = false))
    asr_pool = pairs(("no", "simple", "crystal", "one-dim", "zero-dim"))
    asr = asr_pool[request(
        terminal,
        GREEN_FG("Please input the type of acoustic sum rule `asr`: ") |> string,
        RadioMenu(collect(values(asr_pool))),
    )]
    print(terminal, GREEN_FG("name of output force constants `flfrc`: ") |> string)
    flfrc = strip(readline(terminal))
    print(terminal, GREEN_FG("name of input dynamical matrices `fildyn`: ") |> string)
    fildyn = strip(readline(terminal))
    print(
        terminal,
        GREEN_FG("Please input the masses of atoms in the supercell (a.m.u.) `amass` (separated by spaces): ") |>
        string,
    )
    amass = map(x -> parse(Float64, x), split(readline(terminal), " ", keepempty = false))
    print(
        terminal,
        GREEN_FG("Please input the number of atom types in the supercell `ntyp`: ") |>
        string,
    )
    ntyp = parse(Int, readline(terminal))
    q_in_band_form_pool = pairs((false, true))
    q_in_band_form = q_in_band_form_pool[request(
        terminal,
        GREEN_FG("Please select the `q_in_band_form`: ") |> string,
        RadioMenu([false, true]),
    )]
    q_in_cryst_coord_pool = pairs((false, true))
    q_in_cryst_coord = q_in_cryst_coord_pool[request(
        terminal,
        GREEN_FG("Please select the `q_in_cryst_coord`: ") |> string,
        RadioMenu([false, true]),
    )]
    nosym_pool = pairs((false, true))
    nosym = nosym_pool[request(
        terminal,
        GREEN_FG("Please select if impose symmetry and time reversal `nosym`: ") |> string,
        RadioMenu([false, true]),
    )]
    matdyn = T(
        dos = dos,
        deltaE = deltaE,
        nk1 = nk1,
        nk2 = nk2,
        nk3 = nk3,
        asr = asr,
        flfrc = flfrc,
        fildyn = fildyn,
        amass = amass,
        ntyp = ntyp,
        q_in_band_form = q_in_band_form,
        q_in_cryst_coord = q_in_cryst_coord,
        nosym = nosym,
    )
    return setfield_helper(terminal, matdyn)
end # function namelist_builder
function Inputs.namelist_builder(
    terminal::TTYTerminal,
    ::Type{T},
) where {T<:DynmatNamelist}
    asr_pool = pairs(("no", "simple", "crystal", "one-dim", "zero-dim"))
    asr = asr_pool[request(
        terminal,
        GREEN_FG("Please select the type of acoustic sum rule `asr`: ") |> string,
        RadioMenu(collect(values(asr_pool))),
    )]
    print(
        terminal,
        GREEN_FG("Please input mass for each atom type `amass` (separated by spaces): ") |>
        string,
    )
    amass = map(x -> parse(Float64, x), split(readline(terminal), " ", keepempty = false))
    dynmat = T(asr = asr, amass = amass)
    return setfield_helper(terminal, dynmat)
end # function namelist_builder

# function Inputs.input_builder(terminal::TTYTerminal, ::Type{T}) where {T<:PhInput}
#     ph = namelist_builder(terminal, PhNamelist)
# end
function Inputs.input_builder(terminal::TTYTerminal, ::Type{T}) where {T<:Q2rInput}
    return Q2rInput(namelist_builder(terminal, Q2rNamelist))
end
# function Inputs.input_builder(terminal::TTYTerminal, ::Type{T}) where {T<:MatdynInput}
#     ph = namelist_builder(terminal, MatdynNamelist)
# end
function Inputs.input_builder(terminal::TTYTerminal, ::Type{T}) where {T<:DynmatInput}
    return DynmatInput(namelist_builder(terminal, DynmatNamelist))
end

end # module PHonon


using POMDPModels
using Test

let
    problem = BabyPOMDP()
    solver = RandomSolver(rng=Random.MersenneTwister(1))
    policy = solve(solver, problem)
    sim = RolloutSimulator(max_steps=10, rng=Random.MersenneTwister(1))
    r1 = @inferred simulate(sim, problem, policy, updater(policy), initialstate_distribution(problem))

    sim = RolloutSimulator(max_steps=10, rng=Random.MersenneTwister(1))
    dummy = @inferred simulate(sim, problem, policy, updater(policy), nothing, true)

    problem = LegacyGridWorld()
    solver = RandomSolver(rng=Random.MersenneTwister(1))
    policy = solve(solver, problem)
    sim = RolloutSimulator(max_steps=10, rng=Random.MersenneTwister(1))
    r2 = @inferred simulate(sim, problem, policy, initialstate(problem, sim.rng))

    problem = LegacyGridWorld()
    solver = RandomSolver(rng=Random.MersenneTwister(1))
    policy = solve(solver, problem)
    sim = RolloutSimulator(Random.MersenneTwister(1), 10) # new constructor
    r2 = @inferred simulate(sim, problem, policy, initialstate(problem, sim.rng))

    @test isapprox(r1, -27.27829, atol=1e-3)
    @test isapprox(r2, 0.0, atol=1e-3)
end


@deprecate predict(mm::LinearModel, newx::AbstractMatrix, interval::Symbol, level::Real = 0.95) predict(mm, newx; interval=interval, level=level)

@deprecate confint(obj::LinearModel, level::Real) confint(obj, level=level)
@deprecate confint(obj::AbstractGLM, level::Real) confint(obj, level=level)