# MIT License
#
# Copyright (c) 2018 Martin Biel
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

@with_kw mutable struct ParametricGomoryData{T <: AbstractFloat}
    relaxation::Union{Nothing, SecondStageLP{T}} = nothing
    cglp::Union{Nothing, JuMP.Model} = nothing
    unrelax::Function = () -> nothing
end

@with_kw mutable struct ParametricGomoryParameters
    maximum_iterations::Int = 1
    optimizer = nothing
    strategy::ParametricGomoryStrategy = Gomory()
end
"""
    ParametricGomoryMaster

Master functor object for using weak optimality cuts in an integer L-shaped algorithm. Requires all first-stage decisions to be binary. Create by supplying a [`UseParametricGomory`](@ref) object through `integer_strategy` in `LShaped.Optimizer` or set the [`IntegerStrategy`](@ref) attribute.

"""
struct ParametricGomoryMaster <: AbstractIntegerAlgorithm
    integer_variables::Vector{MOI.VariableIndex}
end

function initialize_integer_algorithm!(::ParametricGomoryMaster, first_stage::JuMP.Model)
    # Sanity check
    if all(is_binary, all_decision_variables(first_stage, 1))
       @warn "All first-stage decisions are binary. Convexification or CombinatorialCuts are more approriate integer strategies"
    end
    # Gather integer variables
    append!(master.integer_variables, gather_integer_variables(subproblem))
    # Sanity check
    if isempty(master.integer_variables)
        @warn "No integer variables in subproblem $(subproblem.id). Integer strategy is superfluous."
    end
    return nothing
end

function handle_integrality!(lshaped::AbstractLShaped, ::ParametricGomoryMaster)
    @unpack τ = lshaped.parameters
    if gap(lshaped) <= τ
        set_metadata!(lshaped.execution.metadata, :converged, true)
    end
    return nothing
end
"""
    ParametricGomoryWorker

Worker functor object for using weak optimality cuts in an integer L-shaped algorithm. Create by supplying a [`UseParametricGomory`](@ref) object through `integer_strategy` in `LShaped.Optimizer` or set the [`IntegerStrategy`](@ref) attribute.

"""
struct ParametricGomoryWorker{T <: AbstractFloat} <: AbstractIntegerAlgorithm
    data::ParametricGomoryData{T}
    parameters::ParametricGomoryParameters
    integer_variables::Vector{MOI.VariableIndex}
    cutting_planes::Vector{MOI.ConstraintIndex}

    function ParametricGomoryWorker(::Type{T}; kw...) where T <: AbstractFloat
        T_ = promote_type(T, Float32)
        worker = new{T_}(ParametricGomoryData{T_}(),
                         ParametricGomoryParameters(; kw...),
                         Vector{MOI.VariableIndex}(),
                         Vector{MOI.ConstraintIndex}())
        return worker
    end
end

function initialize_integer_algorithm!(worker::ParametricGomoryWorker, subproblem::SubProblem)
    # Sanity check
    if any(is_integer, all_decision_variables(subproblem.model, StochasticPrograms.stage(subproblem.model))) ||
       any(is_integer, all_auxiliary_variables(subproblem.model))
        if !(worker.parameters.strategy isa Gomory)
            if subproblem.id == 1
                @warn "Switching convexification strategy to Gomory to handle integer variables."
            end
            worker.parameters.strategy = Gomory()
        end
    end
    # Gather integer variables
    append!(worker.integer_variables, gather_integer_variables(subproblem))
    # Sanity check
    if isempty(worker.integer_variables)
        @warn "No integer variables in subproblem $(subproblem.id). Integer strategy is superfluous."
    end
    # Relax integer restrictions
    worker.data.unrelax = relax_decision_integrality(subproblem.model)
    # Extract model information
    standard_form = worker.parameters.strategy isa Gomory
    worker.data.relaxation = SecondStageLP(subproblem; standard_form = standard_form)
    return nothing
end

function restore!(model::MOI.ModelLike, worker::ParametricGomoryWorker)
    # Delete any added cutting planes
    if !isempty(worker.cutting_planes)
        MOI.delete(model, worker.cutting_planes)
    end
    # Restore integral restrictions
    worker.data.unrelax()
    return nothing
end

function integer_variables(worker::ParametricGomoryWorker)
    return worker.integer_variables
end

function solve_subproblem(subproblem::SubProblem{T},
                          metadata,
                          ::NoFeasibilityAlgorithm,
                          worker::ParametricGomoryWorker,
                          x::AbstractVector) where T <: AbstractFloat
    # Initial solve
    cut = solve_subproblem(subproblem, x)
    # Check if integer restrictions are satisfied
    # and collect any fractional variables
    fractional = MOI.VariableIndex[]
    set_metadata!(metadata,
                  subproblem.id,
                  :integral_solution,
                  check_integrality_restrictions(subproblem, fractional))
    # Check if solution satisfies integer requirements
    if isempty(fractional)
        # No cutting planes needed, just return cut
        return cut
    end
    # Check if L-shaped has converged
    if has_metadata(metadata, :converged)
        # Solve subproblem with integer restrictions
        return solve_unrelaxed(subproblem, metadata, worker, x)
    end
    # Otherwise, add cutting planes until solution is integer
    # or user-limit is reached
    num_iterations = 1
    while true
        inds = MOI.get(subproblem.optimizer, MOI.ListOfVariableIndices())
        ȳ = MOI.get.(subproblem.optimizer, MOI.VariablePrimal(), inds[length(x)+1:end])
        try
            # Generate and add cutting plane
            πx, πy, g = cutting_plane(worker,
                                      worker.parameters.strategy,
                                      subproblem,
                                      x,
                                      ȳ,
                                      fractional)
            f, set, coeffs = moi_constraint(subproblem, πx, πy, g)
            # Add cutting plane to subproblem
            ci = MOI.add_constraint(subproblem.optimizer, f, set)
            push!(worker.cutting_planes, ci)
            # Update subproblem
            if !isempty(coeffs)
                push!(subproblem.linking_constraints, ci)
                push!(subproblem.masterterms, coeffs)
            end
        catch
            @warn "Cutting plane could not be generated in subproblem $(subproblem.id)"
            # If cutting plane could not be generated,
            # return latest working cut and try again next iteration
            return cut
        end
        # Re-solve subproblem
        cut = solve_subproblem(subproblem, x)
        fractional = MOI.VariableIndex[]
        set_metadata!(metadata,
                      subproblem.id,
                      :integral_solution,
                      check_integrality_restrictions(subproblem, fractional))
        if isempty(fractional) || num_iterations >= worker.parameters.maximum_iterations
            # Return cut
            return cut
        end
        # Increment iteration count
        num_iterations += 1
    end
end

function solve_subproblem(subproblem::SubProblem,
                          metadata,
                          feasibility_algorithm::FeasibilityCutsWorker,
                          worker::ParametricGomoryWorker,
                          x::AbstractVector)
    # Prepare auxiliary problem
    model = subproblem.optimizer
    if !prepared(feasibility_algorithm)
        prepare!(model, feasibility_algorithm)
    end
    # Relax integrality
    unrelax = relax_decision_integrality(subproblem.model)
    # Optimize auxiliary problem
    MOI.optimize!(model)
    # Sanity check that aux problem could be solved
    status = MOI.get(subproblem.optimizer, MOI.TerminationStatus())
    if !(status ∈ AcceptableTermination)
        error("Subproblem $(subproblem.id) was not solved properly during feasibility check, returned status code: $status")
    end
    # check objective
    sense = MOI.get(subproblem.optimizer, MOI.ObjectiveSense())
    correction = (sense == MOI.MIN_SENSE || sense == MOI.FEASIBILITY_SENSE) ? 1.0 : -1.0
    w = correction * MOI.get(model, MOI.ObjectiveValue())
    # Ensure correction is available in master
    set_metadata!(metadata, subproblem.id, :correction, correction)
    # Check feasibility
    if w > sqrt(eps())
        # Subproblem is infeasible, create feasibility cut
        cut = FeasibilityCut(subproblem, x)
        # Unrelax integrality restrictions again
        unrelax()
        return cut
    end
    # Unrelax integrality restrictions again
    unrelax()
    # Restore subproblem and solve as usual
    restore_subproblem!(subproblem)
    return solve_subproblem(subproblem, metadata, NoFeasibilityAlgorithm(), worker, x)
end

function solve_unrelaxed(subproblem::SubProblem, metadata, worker::ParametricGomoryWorker, x::AbstractVector)
    # Delete any added cutting planes
    if !isempty(worker.cutting_planes)
        MOI.delete(subproblem.optimizer, worker.cutting_planes)
    end
    empty!(worker.cutting_planes)
    # Restore integral restrictions
    worker.data.unrelax()
    worker.data.unrelax = () -> nothing
    # Solve subproblem with integer restrictions
    MOI.optimize!(subproblem.optimizer)
    status = MOI.get(subproblem.optimizer, MOI.TerminationStatus())
    if status ∈ AcceptableTermination
        # Integer restrictions are satisfied if optimal
        set_metadata!(metadata,
                      subproblem.id,
                      :integral_solution,
                      true)
        # Get sense
        sense = MOI.get(subproblem.optimizer, MOI.ObjectiveSense())
        correction = (sense == MOI.MIN_SENSE || sense == MOI.FEASIBILITY_SENSE) ? 1.0 : -1.0
        # Create sense-corrected optimality cut
        Q = correction * MOI.get(subproblem.optimizer, MOI.ObjectiveValue())
        return OptimalityCut(spzeros(length(x)), Q, subproblem.id)
    elseif status == MOI.INFEASIBLE
        return Infeasible(subproblem)
    elseif status == MOI.DUAL_INFEASIBLE
        return Unbounded(subproblem)
    else
        error("Subproblem $(subproblem.id) was not solved properly, returned status code: $status")
    end
end

function moi_constraint(subproblem::SubProblem{T}, πx::AbstractVector, πy::AbstractVector, g::AbstractFloat) where T <: AbstractFloat
    # Create cutting plane
    f = zero(AffineDecisionFunction{T})
    coeffs = Vector{Tuple{Int, Int, T}}()
    N = StochasticPrograms.stage(subproblem.model)
    all_known = mapreduce(vcat, 1:N-1) do s
        index.(all_known_decision_variables(subproblem.model, s))
    end
    all_decisions = index.(all_decision_variables(subproblem.model, N))
    ix = 0
    iy = 0
    for vi in MOI.get(subproblem.optimizer, MOI.ListOfVariableIndices())
        if vi in all_known
            ix += 1
            coeff = T(πx[ix])
            if abs(coeff) >= sqrt(eps())
                push!(f.known_part.terms, MOI.ScalarAffineTerm(coeff, vi))
                push!(coeffs, (1, vi.value, coeff))
            else
                πx[ix] = zero(T)
            end
        elseif vi in all_decisions
            iy += 1
            coeff = T(πy[iy])
            if abs(coeff) >= sqrt(eps())
                push!(f.decision_part.terms, MOI.ScalarAffineTerm(coeff, vi))
            else
                πy[iy] = zero(T)
            end
        else
            iy += 1
            coeff = T(πy[iy])
            if abs(coeff) >= sqrt(eps())
                push!(f.variable_part.terms, MOI.ScalarAffineTerm(coeff, vi))
            else
                πy[iy] = zero(T)
            end
        end
    end
    set = MOI.GreaterThan{Float64}(g)
    return f, set, coeffs
end

function cutting_plane(worker::ParametricGomoryWorker,
                       ::Gomory,
                       subproblem::SubProblem,
                       x̄::AbstractVector,
                       ȳ::AbstractVector,
                       fractional::Vector{MOI.VariableIndex})
    # Pick smallest fractional index
    fractional_index = fractional[1].value - length(x̄)
    # Get matrices
    T = worker.data.relaxation.T
    W = worker.data.relaxation.W
    h = worker.data.relaxation.h
    ny = length(ȳ)
    # Get current optimal basis
    basis = optimal_basis(subproblem)
    nonbasis = setdiff(1:size(W,2),basis)
    # Compute row of basis inverse corresponding
    # to fractional index
    i = something(findfirst(j -> j == fractional_index, basis), 0)
    if i == 0
        error("Fractional variable not in basis.")
    end
    ei = zeros(size(W, 1))
    ei[i] = 1.0
    Bᵢ⁻¹ = W[:,basis]' \ ei
    # Get source rows
    nonbasic = setdiff(1:ny, basis)
    nonbasic_slack = setdiff(ny+1:size(W,2), basis)
    w = (Bᵢ⁻¹'*W[:,nonbasis])
    γ = (Bᵢ⁻¹'*T)[:]
    ρ = Bᵢ⁻¹'*h
    # Translate x to the origin
    for (j,x) in enumerate(x̄)
        if x == 1.
            ρ -= γ[j]
            γ[j] *= -1
        end
    end
    # Calculate Gomory coefficients
    ϕ(x) = x - floor(x)
    #πx = ϕ.(γ)
    #w = ϕ.(w)
    #ρ = ϕ(ρ)
    ϕγ = ϕ.(γ)
    ϕw = ϕ.(w)
    ρ  = ϕ(ρ)
    πx = min.(ϕγ, (ρ .* (1 .- ϕγ) ./ (1 .- ρ)))
    w  = min.(ϕw, (ρ .* (1 .- ϕw) ./ (1 .- ρ)))
    # Translate back
    for (j,x) in enumerate(x̄)
        if x == 1.
            ρ -= πx[j]
            πx[j] *= -1
        end
    end
    πy = zero(ȳ)
    πy[nonbasic] .= w[nonbasic]
    # Map slacks to rows
    for (coeff,slack) in zip(w[length(nonbasic)+1:end], nonbasic_slack)
        row = slack - ny
        πx .-= W[row,slack] * coeff * T[row,:]
        πy .-= W[row,slack] * coeff * W[row,1:ny]
        ρ -= W[row,slack] * coeff * h[row]
    end
    # Extend standard form LP with new cuts
    add_row!(worker.data.relaxation, πx, πy, ρ; standard_form = true)
    # Return cutting plane
    return πx, πy, ρ
end

function cutting_plane(worker::ParametricGomoryWorker,
                       ::LiftAndProject,
                       subproblem::SubProblem{T},
                       x̄::AbstractVector,
                       ȳ::AbstractVector,
                       fractional::Vector{MOI.VariableIndex}) where T <: AbstractFloat
    # Check optimizer
    if worker.parameters.optimizer === nothing
        error("Cannot generate CGLP without optimizer.")
    end
    # Largest fractional index
    fractional_index = fractional[end].value - length(x̄)
    model = subproblem.model
    nx = length(x̄)
    ny = length(ȳ)
    nw = nx + ny - 1
    # Prepare data
    Gᵢ = Vector{Int}()
    Gⱼ = Vector{Int}()
    Gᵥ = Vector{T}()
    Hᵢ = Vector{Int}()
    Hⱼ = Vector{Int}()
    Hᵥ = Vector{T}()
    Fᵢ = Vector{Int}()
    Fⱼ = Vector{Int}()
    Fᵥ = Vector{T}()
    lhs = Vector{T}()
    rhs = Vector{T}()
    i = 0
    # Loop over constraints
    for (F,S) in list_of_constraint_types(model)
        for cref in all_constraints(model, F, S)
            if S <: StochasticPrograms.SingleDecisionSet ||
                S <: StochasticPrograms.MultipleDecisionSet ||
                S <: StochasticPrograms.SingleKnownSet ||
                S <: StochasticPrograms.MultipleKnownSet
                continue
            elseif F <: VariableRef || F <: DecisionRef
                continue
            elseif F <: AffExpr
                i += 1
                f = MOI.get(model, MOI.ConstraintFunction(), cref)
                for term in f.terms
                    idx = term.variable_index.value - nx
                    if idx == fractional_index
                        # yⱼ
                        push!(Hᵢ, i)
                        push!(Hᵥ, term.coefficient)
                        push!(Hⱼ, idx)
                    else
                        # yⱼ
                        push!(Fᵢ, i)
                        push!(Fᵥ, term.coefficient)
                        push!(Fⱼ, nx + ny + nx + idx)
                        # 1-yⱼ
                        push!(Hᵢ, i+1)
                        push!(Hᵥ, term.coefficient)
                        push!(Hⱼ, idx)
                        push!(Fᵢ, i+1)
                        push!(Fᵥ, -term.coefficient)
                        push!(Fⱼ, nx + ny + nx + idx)
                    end
                end
                set = MOI.get(model, MOI.ConstraintSet(), cref)
                if set isa MOI.LessThan
                    # yⱼ
                    push!(lhs, -Inf)
                    push!(rhs, zero(T))
                    push!(Hᵢ, i)
                    push!(Hᵥ, -set.upper)
                    push!(Hⱼ, fractional_index)
                    # 1-yⱼ
                    push!(lhs, -Inf)
                    push!(rhs, set.upper)
                    push!(Hᵢ, i+1)
                    push!(Hᵥ, set.upper)
                    push!(Hⱼ, fractional_index)
                elseif set isa MOI.GreaterThan
                    # yⱼ
                    push!(lhs, zero(T))
                    push!(rhs, Inf)
                    push!(Hᵢ, i)
                    push!(Hᵥ, -set.lower)
                    push!(Hⱼ, fractional_index)
                    # 1-yⱼ
                    push!(lhs, set.lower)
                    push!(rhs, Inf)
                    push!(Hᵢ, i+1)
                    push!(Hᵥ, set.lower)
                    push!(Hⱼ, fractional_index)
                elseif set isa MOI.EqualTo
                    # yⱼ
                    push!(lhs, zero(T))
                    push!(rhs, zero(T))
                    push!(Hᵢ, i)
                    push!(Hᵥ, -set.value)
                    push!(Hⱼ, fractional_index)
                    # 1-yⱼ
                    push!(lhs, set.value)
                    push!(rhs, set.value)
                    push!(Hᵢ, i+1)
                    push!(Hᵥ, set.value)
                    push!(Hⱼ, fractional_index)
                end
                i += 1
            elseif F <: DecisionAffExpr
                i += 1
                f = MOI.get(model, MOI.ConstraintFunction(), cref)
                for term in f.variable_part.terms
                    idx = term.variable_index.value - nx
                    if idx == fractional_index
                        # yⱼ
                        push!(Hᵢ, i)
                        push!(Hᵥ, term.coefficient)
                        push!(Hⱼ, idx)
                    else
                        # yⱼ
                        push!(Fᵢ, i)
                        push!(Fᵥ, term.coefficient)
                        push!(Fⱼ, nx + idx - 1)
                        # 1-yⱼ
                        push!(Hᵢ, i+1)
                        push!(Hᵥ, term.coefficient)
                        push!(Hⱼ, idx)
                        push!(Fᵢ, i+1)
                        push!(Fᵥ, -term.coefficient)
                        push!(Fⱼ, nx + idx - 1)
                    end
                end
                for term in f.decision_part.terms
                    idx = term.variable_index.value - nx
                    if idx == fractional_index
                        # yⱼ
                        push!(Hᵢ, i)
                        push!(Hᵥ, term.coefficient)
                        push!(Hⱼ, idx)
                    else
                        # yⱼ
                        push!(Fᵢ, i)
                        push!(Fᵥ, term.coefficient)
                        push!(Fⱼ, nx + idx - 1)
                        # 1-yⱼ
                        push!(Hᵢ, i+1)
                        push!(Hᵥ, term.coefficient)
                        push!(Hⱼ, idx)
                        push!(Fᵢ, i+1)
                        push!(Fᵥ, -term.coefficient)
                        push!(Fⱼ, nx + idx - 1)
                    end
                end
                for term in f.known_part.terms
                    idx = term.variable_index.value
                    # yⱼ
                    push!(Fᵢ, i)
                    push!(Fᵥ, term.coefficient)
                    push!(Fⱼ, idx)
                    # 1-yⱼ
                    push!(Gᵢ, i+1)
                    push!(Gᵥ, term.coefficient)
                    push!(Gⱼ, idx)
                    push!(Fᵢ, i+1)
                    push!(Fᵥ, -term.coefficient)
                    push!(Fⱼ, idx)
                end
                set = MOI.get(model, MOI.ConstraintSet(), cref)
                if set isa MOI.LessThan
                    # yⱼ
                    push!(lhs, -Inf)
                    push!(rhs, zero(T))
                    push!(Hᵢ, i)
                    push!(Hᵥ, -set.upper)
                    push!(Hⱼ, fractional_index)
                    # 1-yⱼ
                    push!(lhs, -Inf)
                    push!(rhs, set.upper)
                    push!(Hᵢ, i+1)
                    push!(Hᵥ, set.upper)
                    push!(Hⱼ, fractional_index)
                elseif set isa MOI.GreaterThan
                    # yⱼ
                    push!(lhs, zero(T))
                    push!(rhs, Inf)
                    push!(Hᵢ, i)
                    push!(Hᵥ, -set.lower)
                    push!(Hⱼ, fractional_index)
                    # 1-yⱼ
                    push!(lhs, set.lower)
                    push!(rhs, Inf)
                    push!(Hᵢ, i+1)
                    push!(Hᵥ, set.lower)
                    push!(Hⱼ, fractional_index)
                elseif set isa MOI.EqualTo
                    # yⱼ
                    push!(lhs, zero(T))
                    push!(rhs, zero(T))
                    push!(Hᵢ, i)
                    push!(Hᵥ, -set.value)
                    push!(Hⱼ, fractional_index)
                    # 1-yⱼ
                    push!(lhs, set.value)
                    push!(rhs, set.value)
                    push!(Hᵢ, i+1)
                    push!(Hᵥ, set.value)
                    push!(Hⱼ, fractional_index)
                end
                i += 1
            else
                error("Cannot extract linear representation from $F-$S constraint.")
            end
        end
    end
    # x ≥ 0
    for j in 1:length(x̄)
        i += 1
        # yⱼ
        push!(Fᵢ, i)
        push!(Fᵥ, 1.)
        push!(Fⱼ, j)
        push!(lhs, zero(T))
        push!(rhs, Inf)
        # 1-yⱼ
        push!(Gᵢ, i+1)
        push!(Gᵥ, 1.)
        push!(Gⱼ, j)
        push!(Fᵢ, i+1)
        push!(Fᵥ, -1.)
        push!(Fⱼ, j)
        push!(lhs, zero(T))
        push!(rhs, Inf)
        i += 1
    end
    # x ≤ 1
    for j in 1:length(x̄)
        i += 1
        # yⱼ
        push!(Fᵢ, i)
        push!(Fᵥ, 1.)
        push!(Fⱼ, j)
        push!(Hᵢ, i)
        push!(Hᵥ, -1.)
        push!(Hⱼ, fractional_index)
        push!(lhs, -Inf)
        push!(rhs, zero(T))
        # 1-yⱼ
        push!(Gᵢ, i+1)
        push!(Gᵥ, 1.)
        push!(Gⱼ, j)
        push!(Fᵢ, i+1)
        push!(Fᵥ, -1.)
        push!(Fⱼ, j)
        push!(Hᵢ, i+1)
        push!(Hᵥ, 1.)
        push!(Hⱼ, fractional_index)
        push!(lhs, -Inf)
        push!(rhs, 1.)
        i += 1
    end
    # y ≥ 0
    for j in 1:length(ȳ)
        i += 1
        if j == fractional_index
            # yⱼ
            push!(Hᵢ, i)
            push!(Hᵥ, 1.)
            push!(Hⱼ, j)
            push!(lhs, zero(T))
            push!(rhs, Inf)
        else
            # yⱼ
            push!(Fᵢ, i)
            push!(Fᵥ, 1.)
            push!(Fⱼ, nx + j - 1)
            push!(lhs, zero(T))
            push!(rhs, Inf)
            # 1-yⱼ
            push!(Hᵢ, i+1)
            push!(Hᵥ, 1.)
            push!(Hⱼ, j)
            push!(Fᵢ, i+1)
            push!(Fᵥ, -1.)
            push!(Fⱼ, nx + j - 1)
            push!(lhs, zero(T))
            push!(rhs, Inf)
            i += 1
        end
    end
    # y ≤ 1
    for j in 1:length(ȳ)
        i += 1
        if j == fractional_index
            # 1-yⱼ
            push!(Hᵢ, i)
            push!(Hᵥ, 1.)
            push!(Hⱼ, j)
            push!(lhs, -Inf)
            push!(rhs, 1.)
        else
            # yⱼ
            push!(Fᵢ, i)
            push!(Fᵥ, 1.)
            push!(Fⱼ, nx + j - 1)
            push!(Hᵢ, i)
            push!(Hᵥ, -1.)
            push!(Hⱼ, fractional_index)
            push!(lhs, -Inf)
            push!(rhs, zero(T))
            # 1-yⱼ
            push!(Hᵢ, i+1)
            push!(Hᵥ, 1.)
            push!(Hⱼ, j)
            push!(Fᵢ, i+1)
            push!(Fᵥ, -1.)
            push!(Fⱼ, nx + j - 1)
            push!(Hᵢ, i+1)
            push!(Hᵥ, 1.)
            push!(Hⱼ, fractional_index)
            push!(lhs, -Inf)
            push!(rhs, 1.)
            i += 1
        end
    end
    G̃ = sparse(Gᵢ, Gⱼ, Gᵥ, i, nx)
    H̃ = sparse(Hᵢ, Hⱼ, Hᵥ, i, ny)
    F̃ = sparse(Fᵢ, Fⱼ, Fᵥ, i, nw)
    G, f = StochasticPrograms.canonical(G̃, lhs, rhs; direction = :geq)
    H, f = StochasticPrograms.canonical(H̃, lhs, rhs; direction = :geq)
    F, f = StochasticPrograms.canonical(F̃, lhs, rhs; direction = :geq)
    # Generate CGLP
    CGLP = Model(worker.parameters.optimizer)
    n = size(F, 1)
    @variable(CGLP, π[1:n] >= 0)
    @objective(CGLP, Min, π⋅(H*ȳ) - π⋅(f - G*x̄))
    @constraint(CGLP, F'*π .== 0)
    @constraint(CGLP, sum(π) == 1)
    # Solve CGLP
    optimize!(CGLP)
    if termination_status(CGLP) == MOI.OPTIMAL
        π = value.(π)
        πx = G'*π
        πy = H'*π
        g = π⋅f
        # Extend standard form LP with new cuts
        add_row!(worker.data.relaxation, πx, πy, g)
        # Return cutting plane
        return πx, πy, g
    else
        error("CGLP could not be solved.")
    end
end

function cutting_plane(worker::ParametricGomoryWorker,
                       ::CuttingPlaneTree,
                       subproblem::SubProblem,
                       x̄::AbstractVector,
                       ȳ::AbstractVector,
                       fractional::Vector{MOI.VariableIndex})
    # Check optimizer
    if worker.parameters.optimizer === nothing
        error("Cannot generate CGLP without optimizer.")
    end
    # Largest fractional index
    fractional_index = fractional[end].value - length(x̄)
    # Dimensions
    lp = worker.data.relaxation
    m  = size(lp.T, 1)
    nx = size(lp.T, 2)
    ny = size(lp.W, 2)
    length(x̄) == nx || error("Given first-stage decision $x̄ does not match second-stage problem dimensions.")
    length(ȳ) == ny || error("Given second-stage decision $ȳ does not match second-stage problem dimensions.")
    ȳ[abs.(ȳ) .<= sqrt(eps())] .= 0.0
    x_lb = fill(0., nx)
    x_ub = fill(1., nx)
    lb_finite = isfinite.(lp.lb)
    ub_finite = isfinite.(lp.ub)
    # Generate CGLP
    CGLP = Model(worker.parameters.optimizer)
    @variable(CGLP, πx[1:nx])
    @variable(CGLP, πy[1:ny])
    @variable(CGLP, tx >= 0)
    @variable(CGLP, ty >= 0)
    @variable(CGLP, λ₂₁[1:m] >= 0)
    @variable(CGLP, λ₂₂[1:m] >= 0)
    @variable(CGLP, ν₁₁[1:nx] >= 0)
    @variable(CGLP, ν₁₂[1:nx] >= 0)
    @variable(CGLP, ν₂₁[1:ny] >= 0)
    @variable(CGLP, ν₂₂[1:ny] >= 0)
    @variable(CGLP, μ₁₁[1:nx] >= 0)
    @variable(CGLP, μ₁₂[1:nx] >= 0)
    @variable(CGLP, μ₂₁[1:ny] >= 0)
    @variable(CGLP, μ₂₂[1:ny] >= 0)
    @objective(CGLP, Min, tx + ty)
    @constraint(CGLP, vcat(tx, πx) in MOI.NormOneCone(nx+1))
    @constraint(CGLP, vcat(ty, πy) in MOI.NormOneCone(ny+1))
    @constraint(CGLP, πx .== lp.T'*λ₂₁ + μ₁₁ - ν₁₁)
    @constraint(CGLP, πx .== lp.T'*λ₂₂ + μ₁₂ - ν₁₂)
    @constraint(CGLP, πy .== lp.W'*λ₂₁ + μ₂₁ - ν₂₁)
    @constraint(CGLP, πy .== lp.W'*λ₂₂ + μ₂₂ - ν₂₂)
    lp.ub[fractional_index] = 0.
    @constraint(CGLP, lp.h⋅λ₂₁ + x_lb⋅μ₁₁ + lp.lb[lb_finite]⋅μ₂₁[lb_finite] - x_ub⋅ν₁₁ - lp.ub[ub_finite]⋅ν₂₁[ub_finite] >= πx⋅x̄ + πy⋅ȳ + 1)
    lp.ub[fractional_index] = 1.
    lp.lb[fractional_index] = 1.
    @constraint(CGLP, lp.h⋅λ₂₂ + x_lb⋅μ₁₂ + lp.lb[lb_finite]⋅μ₂₂[lb_finite] - x_ub⋅ν₁₂ - lp.ub[ub_finite]⋅ν₂₂[ub_finite] >= πx⋅x̄ + πy⋅ȳ + 1)
    lp.lb[fractional_index] = 0.
    # Solve CGLP
    optimize!(CGLP)
    if termination_status(CGLP) == MOI.OPTIMAL
        # Extract cutting plane parameters
        πx = value.(CGLP[:πx])
        πy = value.(CGLP[:πy])
        g  = πx⋅x̄ + πy⋅ȳ + 1.
        # Extend standard form LP with new cuts
        add_row!(worker.data.relaxation, πx, πy, g)
        # Return cutting plane
        return πx, πy, g
    else
        error("CGLP could not be solved to optimality.")
    end
end

# API
# ------------------------------------------------------------
"""
    ParametricGomory

Factory object for using convexification to handle integer recourse. Pass to `integer_strategy` in `LShaped.Optimizer` or set the [`IntegerStrategy`](@ref) attribute.

"""
struct ParametricGomory <: AbstractIntegerStrategy
    parameters::ParametricGomoryParameters
end
ParametricGomory(; kw...) = ParametricGomory(ParametricGomoryParameters(; kw...))

function master(wc::ParametricGomory, ::Type{T}) where T <: AbstractFloat
    return ParametricGomoryMaster()
end

function worker(wc::ParametricGomory, ::Type{T}) where T <: AbstractFloat
    return ParametricGomoryWorker(T; type2dict(wc.parameters)...)
end
function worker_type(::ParametricGomory)
    return ParametricGomoryWorker{Float64}
end


#Define SunType and object
struct SunType <: abstractStar end
Sun = Star(SunType)

#Data, as previously used in AstroDynBase
naif_id(::Type{SunType}) = 10
mu(::Type{SunType}) = 1.32712440041e11
mean_radius(::Type{SunType}) = 696000.0
polar_radius(::Type{SunType}) = 696000.0
equatorial_radius(::Type{SunType}) = 696000.0
parent(::Type{SunType}) = SSB

#Extra data
alpha(::Type{SunType}) = deg2rad.([286.13, 0., 0.])
delta(::Type{SunType}) = deg2rad.([63.87, 0., 0.])
omega(::Type{SunType}) = deg2rad.([84.176, 14.18440, 0.])

alpha_nut_prec(::Type{SunType}) = 0.
delta_nut_prec(::Type{SunType}) = 0.
omega_nut_prec(::Type{SunType}) = 0.
theta0_nut_prec(::Type{SunType}) = 0.
theta1_nut_prec(::Type{SunType}) = 0.


typealias Scene Frame{Vehicle}
Scene(n::Int=100) = Frame(Vehicle, n)
Scene(arr::Vector{Vehicle}) = Frame{Vehicle}(arr, length(arr))

Base.show(io, scene::Scene) = print(io, "Scene(with $(length(scene)) cars)")

"""
    get_neighbor_index_fore(scene::Scene, vehicle_index::Int, roadway::Roadway)
Return the index of the vehicle that is in the same lane as scene[vehicle_index] and
in front of it with the smallest distance along the lane

    The method will search on the current lane first, and if no vehicle is found it
    will continue to travel along the lane following next_lane(lane, roadway).
    If no vehicle is found within `max_distance_fore,` a value of 0 is returned instead.
"""
immutable NeighborLongitudinalResult
    ind::Int # index in scene of the neighbor
    Δs::Float64 # positive distance along lane between vehicles' positions
end

export
    VehicleTargetPoint,
    VehicleTargetPointFront,
    VehicleTargetPointCenter,
    VehicleTargetPointRear,
    get_targetpoint_delta

abstract VehicleTargetPoint
immutable VehicleTargetPointFront <: VehicleTargetPoint end
get_targetpoint_delta(::VehicleTargetPointFront, veh::Vehicle) = veh.def.length/2*cos(veh.state.posF.ϕ)
immutable VehicleTargetPointCenter <: VehicleTargetPoint end
get_targetpoint_delta(::VehicleTargetPointCenter, veh::Vehicle) = 0.0
immutable VehicleTargetPointRear <: VehicleTargetPoint end
get_targetpoint_delta(::VehicleTargetPointRear, veh::Vehicle) = -veh.def.length/2*cos(veh.state.posF.ϕ)

const VEHICLE_TARGET_POINT_CENTER = VehicleTargetPointCenter()

function get_neighbor_fore_along_lane(
    scene::Scene,
    roadway::Roadway,
    tag_start::LaneTag,
    s_base::Float64,
    targetpoint_primary::VehicleTargetPoint, # the reference point whose distance we want to minimize
    targetpoint_valid::VehicleTargetPoint; # the reference point, which if distance to is positive, we include the vehicle
    max_distance_fore::Float64 = 250.0, # max distance to search forward [m]
    index_to_ignore::Int=-1,
    )

    best_ind = 0
    best_dist = max_distance_fore
    tag_target = tag_start

    dist_searched = 0.0
    while dist_searched < max_distance_fore

        lane = roadway[tag_target]

        for (i,veh) in enumerate(scene)
            if i != index_to_ignore

                s_adjust = NaN
                if veh.state.posF.roadind.tag == tag_target
                    s_adjust = 0.0

                elseif is_between_segments_hi(veh.state.posF.roadind.ind, lane.curve) &&
                       is_in_entrances(roadway[tag_target], veh.state.posF.roadind.tag)

                    distance_between_lanes = abs(roadway[tag_target].curve[1].pos - roadway[veh.state.posF.roadind.tag].curve[end].pos)
                    s_adjust = -(roadway[veh.state.posF.roadind.tag].curve[end].s + distance_between_lanes)

                elseif is_between_segments_lo(veh.state.posF.roadind.ind) &&
                       is_in_exits(roadway[tag_target], veh.state.posF.roadind.tag)

                    distance_between_lanes = abs(roadway[tag_target].curve[end].pos - roadway[veh.state.posF.roadind.tag].curve[1].pos)
                    s_adjust = roadway[tag_target].curve[end].s + distance_between_lanes
                end

                if !isnan(s_adjust)
                    s_valid = veh.state.posF.s + get_targetpoint_delta(targetpoint_valid, veh) + s_adjust
                    dist_valid = s_valid - s_base + dist_searched
                    if dist_valid ≥ 0.0
                        s_primary = veh.state.posF.s + get_targetpoint_delta(targetpoint_primary, veh) + s_adjust
                        dist = s_primary - s_base + dist_searched
                        if dist < best_dist
                            best_dist = dist
                            best_ind = i
                        end
                    end
                end
            end
        end

        if best_ind != 0
            break
        end

        if !has_next(lane) ||
           (tag_target == tag_start && dist_searched != 0.0) # exit after visiting this lane a 2nd time
            break
        end

        dist_searched += (lane.curve[end].s - s_base)
        s_base = -abs(lane.curve[end].pos - next_lane_point(lane, roadway).pos) # negative distance between lanes
        tag_target = next_lane(lane, roadway).tag
    end

    NeighborLongitudinalResult(best_ind, best_dist)
end
function get_neighbor_fore_along_lane(
    scene::Scene,
    vehicle_index::Int,
    roadway::Roadway,
    targetpoint_ego::VehicleTargetPoint,
    targetpoint_primary::VehicleTargetPoint, # the reference point whose distance we want to minimize
    targetpoint_valid::VehicleTargetPoint; # the reference point, which if distance to is positive, we include the vehicle
    max_distance_fore::Float64 = 250.0 # max distance to search forward [m]
    )

    veh_ego = scene[vehicle_index]
    tag_start = veh_ego.state.posF.roadind.tag
    s_base = veh_ego.state.posF.s + get_targetpoint_delta(targetpoint_ego, veh_ego)

    get_neighbor_fore_along_lane(scene, roadway, tag_start, s_base,
        targetpoint_primary, targetpoint_valid,
        max_distance_fore=max_distance_fore, index_to_ignore=vehicle_index)
end
function get_neighbor_fore_along_left_lane(
    scene::Scene,
    vehicle_index::Int,
    roadway::Roadway,
    targetpoint_ego::VehicleTargetPoint,
    targetpoint_primary::VehicleTargetPoint, # the reference point whose distance we want to minimize
    targetpoint_valid::VehicleTargetPoint; # the reference point, which if distance to is positive, we include the vehicle
    max_distance_fore::Float64 = 250.0 # max distance to search forward [m]
    )

    retval = NeighborLongitudinalResult(0, max_distance_fore)

    veh_ego = scene[vehicle_index]
    lane = roadway[veh_ego.state.posF.roadind.tag]
    if n_lanes_left(lane, roadway) > 0
        lane_left = roadway[LaneTag(lane.tag.segment, lane.tag.lane + 1)]
        roadproj = proj(veh_ego.state.posG, lane_left, roadway)
        tag_start = roadproj.tag
        s_base = lane_left[roadproj.curveproj.ind, roadway].s + get_targetpoint_delta(targetpoint_ego, veh_ego)

        retval = get_neighbor_fore_along_lane(scene, roadway, tag_start, s_base,
                                              targetpoint_primary, targetpoint_valid,
                                              index_to_ignore=vehicle_index,
                                              max_distance_fore=max_distance_fore)
    end

    retval
end
function get_neighbor_fore_along_right_lane(
    scene::Scene,
    vehicle_index::Int,
    roadway::Roadway,
    targetpoint_ego::VehicleTargetPoint,
    targetpoint_primary::VehicleTargetPoint, # the reference point whose distance we want to minimize
    targetpoint_valid::VehicleTargetPoint; # the reference point, which if distance to is positive, we include the vehicle
    max_distance_fore::Float64 = 250.0 # max distance to search forward [m]
    )

    retval = NeighborLongitudinalResult(0, max_distance_fore)

    veh_ego = scene[vehicle_index]
    lane = roadway[veh_ego.state.posF.roadind.tag]
    if n_lanes_right(lane, roadway) > 0
        lane_right = roadway[LaneTag(lane.tag.segment, lane.tag.lane - 1)]
        roadproj = proj(veh_ego.state.posG, lane_right, roadway)
        tag_start = roadproj.tag
        s_base = lane_right[roadproj.curveproj.ind, roadway].s + get_targetpoint_delta(targetpoint_ego, veh_ego)

        retval = get_neighbor_fore_along_lane(scene, roadway, tag_start, s_base,
                                              targetpoint_primary, targetpoint_valid,
                                              index_to_ignore=vehicle_index,
                                              max_distance_fore=max_distance_fore)
    end

    retval
end

function get_neighbor_fore_along_lane(
    scene::Scene,
    roadway::Roadway,
    tag_start::LaneTag,
    s_base::Float64;
    max_distance_fore::Float64 = 250.0, # max distance to search forward [m]
    index_to_ignore::Int=-1,
    )

    get_neighbor_fore_along_lane(scene, roadway, tag_start, s_base,
        VEHICLE_TARGET_POINT_CENTER, VEHICLE_TARGET_POINT_CENTER,
        max_distance_fore=max_distance_fore,
        index_to_ignore=index_to_ignore)
end
function get_neighbor_fore_along_lane(scene::Scene, vehicle_index::Int, roadway::Roadway;
    max_distance_fore::Float64 = 250.0 # max distance to search forward [m]
    )

    get_neighbor_fore_along_lane(scene, vehicle_index, roadway,
        VEHICLE_TARGET_POINT_CENTER, VEHICLE_TARGET_POINT_CENTER,
        VEHICLE_TARGET_POINT_CENTER, max_distance_fore=max_distance_fore)
end
function get_neighbor_fore_along_left_lane(scene::Scene, vehicle_index::Int, roadway::Roadway;
    max_distance_fore::Float64 = 250.0 # max distance to search forward [m]
    )

    get_neighbor_fore_along_left_lane(scene, vehicle_index, roadway,
        VEHICLE_TARGET_POINT_CENTER, VEHICLE_TARGET_POINT_CENTER,
        VEHICLE_TARGET_POINT_CENTER, max_distance_fore=max_distance_fore)
end
function get_neighbor_fore_along_right_lane(scene::Scene, vehicle_index::Int, roadway::Roadway;
    max_distance_fore::Float64 = 250.0 # max distance to search forward [m]
    )

    get_neighbor_fore_along_right_lane(scene, vehicle_index, roadway,
        VEHICLE_TARGET_POINT_CENTER, VEHICLE_TARGET_POINT_CENTER,
        VEHICLE_TARGET_POINT_CENTER, max_distance_fore=max_distance_fore)
end

function get_neighbor_rear_along_lane(
    scene::Scene,
    roadway::Roadway,
    tag_start::LaneTag,
    s_base::Float64,
    targetpoint_primary::VehicleTargetPoint, # the reference point whose distance we want to minimize
    targetpoint_valid::VehicleTargetPoint; # the reference point, which if distance to is positive, we include the vehicle
    max_distance_rear::Float64 = 250.0, # max distance to search rearward [m]
    index_to_ignore::Int=-1,
    )

    best_ind = 0
    best_dist = max_distance_rear
    tag_target = tag_start

    ignore = Set{Int}()

    dist_searched = 0.0
    while dist_searched < max_distance_rear

        lane = roadway[tag_target]

        for (i,veh) in enumerate(scene)
            if i != index_to_ignore && !in(veh.id, ignore)

                s_adjust = NaN

                if veh.state.posF.roadind.tag == tag_target
                    s_adjust = 0.0

                elseif is_between_segments_hi(veh.state.posF.roadind.ind, lane.curve) &&
                       is_in_entrances(roadway[tag_target], veh.state.posF.roadind.tag)

                    distance_between_lanes = abs(roadway[tag_target].curve[1].pos - roadway[veh.state.posF.roadind.tag].curve[end].pos)
                    s_adjust = -(roadway[veh.state.posF.roadind.tag].curve[end].s + distance_between_lanes)

                elseif is_between_segments_lo(veh.state.posF.roadind.ind) &&
                       is_in_exits(roadway[tag_target], veh.state.posF.roadind.tag)

                    distance_between_lanes = abs(roadway[tag_target].curve[end].pos - roadway[veh.state.posF.roadind.tag].curve[1].pos)
                    s_adjust = roadway[tag_target].curve[end].s + distance_between_lanes
                end

                if !isnan(s_adjust)
                    s_valid = veh.state.posF.s + get_targetpoint_delta(targetpoint_valid, veh) + s_adjust
                    dist_valid = s_base - s_valid + dist_searched
                    if dist_valid ≥ 0.0
                        s_primary = veh.state.posF.s + get_targetpoint_delta(targetpoint_primary, veh) + s_adjust
                        dist = s_base - s_primary + dist_searched
                        if dist < best_dist
                            best_dist = dist
                            best_ind = i
                        end
                    else
                        push!(ignore, veh.id)
                    end
                end
            end
        end

        if best_ind != 0
            break
        end

        if !has_prev(lane) ||
           (tag_target == tag_start && dist_searched != 0.0) # exit after visiting this lane a 2nd time
            break
        end

        dist_searched += s_base
        s_base = lane.curve[end].s + abs(lane.curve[end].pos - prev_lane_point(lane, roadway).pos) # end of prev lane plus crossover
        tag_target = prev_lane(lane, roadway).tag
    end

    NeighborLongitudinalResult(best_ind, best_dist)
end
function get_neighbor_rear_along_lane(
    scene::Scene,
    vehicle_index::Int,
    roadway::Roadway,
    targetpoint_ego::VehicleTargetPoint,
    targetpoint_primary::VehicleTargetPoint, # the reference point whose distance we want to minimize
    targetpoint_valid::VehicleTargetPoint; # the reference point, which if distance to is positive, we include the vehicle
    max_distance_rear::Float64 = 250.0 # max distance to search rearward [m]
    )

    veh_ego = scene[vehicle_index]
    tag_start = veh_ego.state.posF.roadind.tag
    s_base = veh_ego.state.posF.s + get_targetpoint_delta(targetpoint_ego, veh_ego)

    get_neighbor_rear_along_lane(scene, roadway, tag_start, s_base,
        targetpoint_primary, targetpoint_valid,
        max_distance_rear=max_distance_rear, index_to_ignore=vehicle_index)
end
function get_neighbor_rear_along_left_lane(
    scene::Scene,
    vehicle_index::Int,
    roadway::Roadway,
    targetpoint_ego::VehicleTargetPoint,
    targetpoint_primary::VehicleTargetPoint, # the reference point whose distance we want to minimize
    targetpoint_valid::VehicleTargetPoint; # the reference point, which if distance to is positive, we include the vehicle
    max_distance_rear::Float64 = 250.0 # max distance to search forward [m]
    )

    retval = NeighborLongitudinalResult(0, max_distance_rear)

    veh_ego = scene[vehicle_index]
    lane = roadway[veh_ego.state.posF.roadind.tag]
    if n_lanes_left(lane, roadway) > 0
        lane_left = roadway[LaneTag(lane.tag.segment, lane.tag.lane + 1)]
        roadproj = proj(veh_ego.state.posG, lane_left, roadway)
        tag_start = roadproj.tag
        s_base = lane_left[roadproj.curveproj.ind, roadway].s + get_targetpoint_delta(targetpoint_ego, veh_ego)

        retval = get_neighbor_rear_along_lane(scene, roadway, tag_start, s_base,
                                              targetpoint_primary, targetpoint_valid,
                                              index_to_ignore=vehicle_index,
                                              max_distance_rear=max_distance_rear)
    end

    retval
end
function get_neighbor_rear_along_right_lane(
    scene::Scene,
    vehicle_index::Int,
    roadway::Roadway,
    targetpoint_ego::VehicleTargetPoint,
    targetpoint_primary::VehicleTargetPoint, # the reference point whose distance we want to minimize
    targetpoint_valid::VehicleTargetPoint; # the reference point, which if distance to is positive, we include the vehicle
    max_distance_rear::Float64 = 250.0 # max distance to search forward [m]
    )

    retval = NeighborLongitudinalResult(0, max_distance_rear)

    veh_ego = scene[vehicle_index]
    lane = roadway[veh_ego.state.posF.roadind.tag]
    if n_lanes_right(lane, roadway) > 0
        lane_right = roadway[LaneTag(lane.tag.segment, lane.tag.lane - 1)]
        roadproj = proj(veh_ego.state.posG, lane_right, roadway)
        tag_start = roadproj.tag
        s_base = lane_right[roadproj.curveproj.ind, roadway].s + get_targetpoint_delta(targetpoint_ego, veh_ego)

        retval = get_neighbor_rear_along_lane(scene, roadway, tag_start, s_base,
                                              targetpoint_primary, targetpoint_valid,
                                              index_to_ignore=vehicle_index,
                                              max_distance_rear=max_distance_rear)
    end

    retval
end

function get_neighbor_rear_along_lane(
    scene::Scene,
    roadway::Roadway,
    tag_start::LaneTag,
    s_base::Float64;
    max_distance_rear::Float64 = 250.0, # max distance to search rearward [m]
    index_to_ignore::Int=-1,
    )

    get_neighbor_rear_along_lane(scene, roadway, tag_start, s_base,
        VEHICLE_TARGET_POINT_CENTER, VEHICLE_TARGET_POINT_CENTER,
        max_distance_rear=max_distance_rear,
        index_to_ignore=index_to_ignore)
end
function get_neighbor_rear_along_lane(scene::Scene, vehicle_index::Int, roadway::Roadway;
    max_distance_rear::Float64 = 250.0 # max distance to search forward [m]
    )

    get_neighbor_rear_along_lane(scene, vehicle_index, roadway,
        VEHICLE_TARGET_POINT_CENTER, VEHICLE_TARGET_POINT_CENTER,
        VEHICLE_TARGET_POINT_CENTER, max_distance_rear=max_distance_rear)
end
function get_neighbor_rear_along_left_lane(
    scene::Scene,
    vehicle_index::Int,
    roadway::Roadway;
    max_distance_rear::Float64 = 250.0 # max distance to search forward [m]
    )

    get_neighbor_rear_along_left_lane(scene, vehicle_index, roadway,
        VEHICLE_TARGET_POINT_CENTER, VEHICLE_TARGET_POINT_CENTER,
        VEHICLE_TARGET_POINT_CENTER, max_distance_rear=max_distance_rear)
end
function get_neighbor_rear_along_right_lane(
    scene::Scene,
    vehicle_index::Int,
    roadway::Roadway;
    max_distance_rear::Float64 = 250.0 # max distance to search forward [m]
    )

    get_neighbor_rear_along_right_lane(scene, vehicle_index, roadway,
        VEHICLE_TARGET_POINT_CENTER, VEHICLE_TARGET_POINT_CENTER,
        VEHICLE_TARGET_POINT_CENTER, max_distance_rear=max_distance_rear)
end

"""
    Project the given point to the same lane as the given RoadIndex.
This will return the projection of the point, along with the Δs along the lane from the RoadIndex.
"""
immutable FrenetRelativePosition
    origin::RoadIndex
    target::RoadIndex
    Δs::Float64
    t::Float64
    ϕ::Float64
end
function get_frenet_relative_position(posG::VecSE2, roadind::RoadIndex, roadway::Roadway;
    max_distance_fore::Float64 = 250.0, # max distance to search forward [m]
    max_distance_rear::Float64 = 250.0, # max distance to search backward [m]
    improvement_threshold::Float64 = 1e-4,
    )

    # project to current lane first
    tag_start = roadind.tag
    lane_start = roadway[tag_start]
    curveproj_start = proj(posG, lane_start, roadway, move_along_curves=false).curveproj
    curvept_start = lane_start[curveproj_start.ind, roadway]
    s_base = lane_start[roadind.ind, roadway].s
    s_proj = curvept_start.s
    Δs = s_proj - s_base
    sq_dist_to_curve = abs2(posG - curvept_start.pos)
    retval = FrenetRelativePosition(roadind,
                RoadIndex(curveproj_start.ind, tag_start), Δs, curveproj_start.t, curveproj_start.ϕ)

    # search downstream
    if has_next(lane_start)
        dist_searched = lane_start.curve[end].s - s_base
        s_base = -abs(lane_start.curve[end].pos - next_lane_point(lane_start, roadway).pos) # negative distance between lanes
        tag_target = next_lane(lane_start, roadway).tag
        while dist_searched < max_distance_fore

            lane = roadway[tag_target]
            curveproj = proj(posG, lane, roadway, move_along_curves=false).curveproj
            sq_dist_to_curve2 = abs2(posG - lane[curveproj.ind, roadway].pos)

            if sq_dist_to_curve2 < sq_dist_to_curve - improvement_threshold

                sq_dist_to_curve = sq_dist_to_curve2
                s_proj = lane[curveproj.ind, roadway].s
                Δs = s_proj - s_base + dist_searched
                retval = FrenetRelativePosition(
                            roadind, RoadIndex(curveproj.ind, tag_target),
                            Δs, curveproj.t, curveproj.ϕ)
            end

            if !has_next(lane)
                break
            end

            dist_searched += (lane.curve[end].s - s_base)
            s_base = -abs(lane.curve[end].pos - next_lane_point(lane, roadway).pos) # negative distance between lanes
            tag_target = next_lane(lane, roadway).tag

            if tag_target == tag_start
                break
            end
        end
    end

    # search upstream
    if has_prev(lane_start)
        dist_searched = s_base
        tag_target = prev_lane(lane_start, roadway).tag
        s_base = roadway[tag_target].curve[end].s + abs(lane_start.curve[1].pos - prev_lane_point(lane_start, roadway).pos) # end of the lane

        while dist_searched < max_distance_fore

            lane = roadway[tag_target]
            curveproj = proj(posG, lane, roadway, move_along_curves=false).curveproj
            sq_dist_to_curve2 = abs2(posG - lane[curveproj.ind, roadway].pos)

            if sq_dist_to_curve2 < sq_dist_to_curve - improvement_threshold

                sq_dist_to_curve = sq_dist_to_curve2
                s_proj = lane[curveproj.ind, roadway].s
                dist = s_base - s_proj  + dist_searched
                Δs = -dist
                retval = FrenetRelativePosition(
                            roadind, RoadIndex(curveproj.ind, tag_target),
                            Δs, curveproj.t, curveproj.ϕ)
            end

            if !has_prev(lane)
                break
            end

            dist_searched += s_base
            s_base = lane.curve[end].s + abs(lane.curve[1].pos - prev_lane_point(lane, roadway).pos) # length of this lane plus crossover
            tag_target = prev_lane(lane, roadway).tag

            if tag_target == tag_start
                break
            end
        end
    end

    retval
end
get_frenet_relative_position(veh_fore::Vehicle, veh_rear::Vehicle, roadway::Roadway) = get_frenet_relative_position(veh_fore.state.posG, veh_rear.state.posF.roadind, roadway)


cpu(x) = x
gpu(x) = x

cpu(x::CuArray) = adapt(Array, x)
gpu(x::Array) = CuArrays.cu(x)
cpu(x::StructArray) = replace_storage(Array, x)
gpu(x::StructArray) = replace_storage(x) do v
    typeof(v) <: StructArray ? gpu(v) :
    isbitstype(eltype(v)) ? gpu(v) : v
end

cpu(x::Neuron) = Neuron(cpu(x.synapses), x.soma)
gpu(x::Neuron) = Neuron(gpu(x.synapses), x.soma)

function cpu(x::STDP{T}) where T<:Real
    clastpre = cpu(x.lastpre)
    clastpost = cpu(x.lastpost)

    return STDP{T, typeof(clastpre)}(x.A₊, x.A₋, x.τ₊, x.τ₋, clastpre, clastpost)
end
function gpu(x::STDP{T}) where T<:Real
    glastpre = gpu(x.lastpre)
    glastpost = gpu(x.lastpost)

    return STDP{T, typeof(glastpre)}(x.A₊, x.A₋, x.τ₊, x.τ₋, glastpre, glastpost)
end

function cpu(x::Population)
    csynapses = cpu(x.synapses)
    csomas = cpu(x.somas)
    cweights = cpu(x.weights)
    clearner = cpu(x.learner)

    return Population(csomas, cweights, csynapses, clearner)
end

function gpu(x::Population)
    gsynapses = gpu(x.synapses)
    gsomas = gpu(x.somas)
    gweights = gpu(x.weights)
    glearner = gpu(x.learner)

    return Population(gsomas, gweights, gsynapses, glearner)
end

function cpu(x::InputPopulation)
    cinputs = cpu(x.inputs)

    return InputPopulation{typeof(cinputs)}(cinputs)
end

function gpu(x::InputPopulation)
    ginputs = gpu(x.inputs)

    return InputPopulation{typeof(ginputs)}(ginputs)
end

function cpu(x::Network)
    y = deepcopy(x)

    @inbounds for (k, v) in y.pops
        y.pops[k] = cpu(v)
    end

    @inbounds for (k, v) in y.connections
        y.pops[k] = NetworkEdge(cpu(v.weights), cpu(v.synapses), cpu(v.learner))
    end

    return y
end

function gpu(x::Network)
    y = deepcopy(x)

    @inbounds for (k, v) in y.pops
        y.pops[k] = gpu(v)
    end

    @inbounds for (k, v) in y.connections
        y.connections[k] = NetworkEdge(gpu(v.weights), gpu(v.synapses), gpu(v.learner))
    end

    return y
end


# module NAssets
using Agents: vertices
using MetaGraphs: add_vertex!, has_edge
using Agents, AgentsPlots, Plots, LightGraphs, MetaGraphs, GraphPlot, GraphRecipes, NetworkLayout
using Tables, DataFrames
using CSV, JSON, Serialization, DelimitedFiles
using BritishNationalGrid
using ZipFile, Shapefile
using Random
using Match
using LinearAlgebra
using StatsBase
using Distributions
using StatsPlots
using SparseArrays
using Laplacians
using DataStructures
using RollingFunctions
using BenchmarkTools
using Statistics
using Logging,LoggingExtras, LoggingFacilities
using Dates,TimeZones

#Core structures
include("core/core_structs.jl")
#graph-related
include("core/graph_functions.jl")
#events
include("eve/artificial_events.jl")
#various util functions
include("utils/util_functions.jl")
#logging
include("utils/logging_functions.jl")
#Plotting functions
include("utils/plotting_functions.jl")
# time-to-event functions
include("utils/tte_functions.jl")
#running functions
include("utils/running_functions.jl")

#maintenance model
include("phy/maintenance_model.jl")

include("ntw/of_switch.jl")
include("ntw/of_control.jl")
include("ctl/agent_control.jl")

include("phy/physical_model.jl")

include("ntw/network_model.jl")
include("phy/geo_model.jl")

#Main Functions
include("model/netManFunctions.jl")

#Agents.jl function implementation for this model
include("model/netManModel.jl")

#Basic queries
#include("model/queries_basic.jl")
#Multiple queries
include("ctl/queries_multiple.jl")

export load_run_configs, single_run_with_logging

configs = []
exp_params = deserialize("data/services.bin")

for exp_param in exp_params
    for cg=6:6
        push!(configs, ( seed = first(exp_param)
                    ,ctl_model= GraphModel(cg) #NAssets.GraphModel(6)
                    ,ntw_topo = GraphModel(6) #NAssets.GraphModel(6)
                    ,size=100
                    ,n_steps=200
                    ,deterioration = 0. # physical deterioration parameter. 0: switched off.
                    # random asset drops
                    ,drop_proportion=0.2 #proportion of assets to randomly fail. 0: switched off.
                    ,drop_stabilisation=50 # time (from simulation start) to allow assets to stabilise before trigger random drops.
                    ,prob_random_walks = 1. #BA(6): 0.5#star(5): 1. #complete(3): 0.3 #ring(2): 1. # grid(4): 0.8 #0.5
                    ,benchmark = false
                    ,animation = false
                    ,k=2
                    ,B = 0.5
                    ,ctl_k=2
                    ,ctl_B = 0.5
                    ,custom_topo = nothing
                    ,ctl_custom_topo = nothing
                    ,mnt_policy = 1
                    ,mnt_wc_duration = 20
                    ,mnt_bc_duration = 10
                    ,mnt_wc_cost = 150.
                    ,mnt_bc_cost = 100.
                    ,traffic_dist_params = [1,0.05]
                    ,traffic_proportion = 0.01
                    ,clear_cache_graph_freq = 50#50,#25, # How often the ntw graph is cleared to initial state, 0: no cache. A value of 10, is not enough in a 16 mesh network to find paths when queries are not repeated, prob_eq_query.
                    ,interval_tpt = 10 #interval used to calculate tpt
                    ,pkt_size = 1#0.065, # (in MB) pkt size  IP between 21 is 65536 bytes Ref: Internet Core Protocols: The Definitive Guide by Eric Hall
                    ,pkt_per_tick = 2000#0, # How many packets are processsed per tick. #TODO which number is reasonable?
                    # I am setting this to 2000 as the expectations is nes
                    # are able to process. Check references e.g. Nokia SR 7750.
                    ,max_queue_ne = 300#700 #This indicates how many pkts/msgs can be stored in tick to be processed the next tick
                    ,ofmsg_reattempt=10#4,# greater number to avoid duplicated install flows
                    ,max_cache_paths = 2
                    ,data_dir = "data/exp1/"
                    ,plots_dir = "plots/exp1/"
                    # ,ntw_services = [(1,7),(4,1),(5,14),(12,8)]
                    ,ntw_services = last(exp_param)
                ))
    end
end

BenchmarkTools.DEFAULT_PARAMETERS.samples = 100

Threads.@threads for config in configs
    single_run_with_file_logging(config)
end
# end # module


module Resonance

export airtable_metadata,
       brain_ingest,
       findprevstool,
       count_set,
       upset_dots!,
       load_metabolites,
       pull_row,
       loadings,
       varexplained,
       mdsaxis,
       permanovas,
       plot_permanovas,
       mantel,
       commonname,
       countmap,
       codebreastfeeding!,
       stp_overlap,
       startup

using Reexport
using Airtable
using CairoMakie
using AlgebraOfGraphics
using MultivariateStats
using Dictionaries
using CodecZlib
using Statistics
using Arrow
using Tables
using SparseArrays
using CodecZlib
using FilePaths
using ProgressLogging
using CategoricalArrays
using PERMANOVA
using LinearAlgebra
using Random

@reexport using BiobakeryUtils
@reexport using DataFrames
@reexport using CSV
@reexport using XLSX
@reexport using DataFramesMeta

const transform = DataFrames.transform

include("airtable.jl")
include("wrangle.jl")
include("metabolites.jl")
include("genefamilies.jl")
include("plotting.jl")
include("brain.jl")
include("startup.jl")
include("kneaddata.jl")
include("mantel.jl")

end


include("./params.jl")

function fitgamma(isi)

    ################# WRITE SOMETHING #################

    return lambda, k

end


function get_devices(sys::PowerSystems.PowerSystem,category::Type{PowerSystems.ThermalGen})
    return sys.generators.thermal
end
function get_devices(sys::PowerSystems.PowerSystem,category::Type{PowerSystems.RenewableGen})
    return sys.generators.renewable
end
function get_devices(sys::PowerSystems.PowerSystem,category::Type{PowerSystems.HydroGen})
    return sys.generators.hydro
end
function get_devices(sys::PowerSystems.PowerSystem,category::Type{PowerSystems.PowerSystems.ElectricLoad})
    return sys.loads
end

function constructservice!(m::JuMP.AbstractModel, service::PowerSystems.StaticReserve, category_formulation::Type{PS.RampLimitedReserve},devices::Array{NamedTuple{(:device, :formulation), Tuple{DataType,DataType}}}, sys::PowerSystems.PowerSystem; args...)

    dev_set = Array{NamedTuple{(:device,:formulation),Tuple{PowerSystems.PowerSystemDevice,DataType}}}([])

    for category in devices
        if category != nothing
            D = get_devices(sys,category.device)
            for d in D
                if d in service.contributingdevices
                    push!(dev_set,(device=d,formulation=category.formulation))
                end
            end
        end
    end

    if !isempty(dev_set)

        p_rsv =  reservevariables(m, dev_set, sys.time_periods)
        m =  PS.reserves(m, dev_set, service, sys.time_periods)

    end

    return m

end


struct FullVar end
struct BranchVar end

Base.@kwdef mutable struct GlobalOptimizer{N,T<:Real,S<:ExtensionType} <: MOI.AbstractOptimizer

        _ext_type::S = S()
        _stack::BinaryMinMaxHeap{NodeBB{N,T}} = BinaryMinMaxHeap{NodeBB{N,T}}()
        _current_node::NodeBB{N,T} = NodeBB{N,T}()

        _input_to_local_map::MOIU.IndexMap              = MOIU.IndexMap()
        _constraint_primal::Dict{CI,Union{T,Vector{T}}} = Dict{CI,T}()
        _constraint_offset::Vector{Int}                 = Int[]
        _constraint_index_num::Int = 0
        _constraint_row_num::Int = 0

        _working_problem::ParsedProblem{N,T} = ParsedProblem{N,T}()
        _input_problem::InputModel{T} = InputModel{T}()

        _first_relax_point_set::Bool = false
        _current_xref::Vector{T} = T[]
        _candidate_xref::Vector{T} = T[]

        _use_prior_objective_xref::Bool = false
        _current_objective_xref::Vector{T} = T[]
        _prior_objective_xref::Vector{T} = T[]

        _user_branch_variables::Bool = false
        _fixed_variable::Vector{Bool} = Bool[]
        _branch_variable_num::Int = 0
        _branch_to_variable_map::Vector{Int} = Int[]
        _variable_to_branch_map::Vector{Int} = Int[]

        _lower_variable_bound::Vector{T} = T[]
        _upper_variable_bound::Vector{T} = T[]

        # all subproblem immutable subproblem status are set in global_solve in corresponding routines
        # in optimize_nonconvex.jl
        _preprocess_feasibility::Bool = true
        _preprocess_result_status::MOI.ResultStatusCode = MOI.OTHER_RESULT_STATUS
        _preprocess_termination_status::MOI.TerminationStatusCode = MOI.OPTIMIZE_NOT_CALLED

        _lower_result_status::MOI.ResultStatusCode = MOI.OTHER_RESULT_STATUS
        _lower_termination_status::MOI.TerminationStatusCode = MOI.OPTIMIZE_NOT_CALLED
        _lower_feasibility::Bool = true
        _lower_objective_value::T = -Inf

        # set in TODO
        _lower_solution::Vector{T} = T[]
        _lower_lvd::Vector{T}      = T[]
        _lower_uvd::Vector{T}      = T[]

        _cut_result_status::MOI.ResultStatusCode = MOI.OTHER_RESULT_STATUS
        _cut_termination_status::MOI.TerminationStatusCode = MOI.OPTIMIZE_NOT_CALLED
        _cut_objective_value::T = -Inf
        _cut_feasibility::Bool = true

        # set in TODO
        _cut_solution::Vector{T} = T[]

        _upper_result_status::MOI.ResultStatusCode = MOI.OTHER_RESULT_STATUS
        _upper_termination_status::MOI.TerminationStatusCode = MOI.OPTIMIZE_NOT_CALLED
        _upper_feasibility::Bool = true
        _upper_objective_value::T = Inf

        # array is initialized to correct size in TODO, reset in single_nlp_solve! in optimize_convex.jl
        _upper_variables::Vector{VI} =  VI[]

        # set in TODO
        _upper_solution::Vector{T} = T[]

        _postprocess_feasibility::Bool = true

        _time_left::T = 1000*one(T)
        _start_time::T = zero(T)
        _run_time::T = zero(T)
        _parse_time::T = zero(T)
        _presolve_time::T = zero(T)
        _last_preprocess_time::T = zero(T)
        _last_lower_problem_time::T = zero(T)
        _last_upper_problem_time::T = zero(T)
        _last_postprocessing_time::T = zero(T)

        # reset in initial_parse! in parse.jl
        _min_converged_value::T = typemax(T)
        _global_lower_bound::T = typemin(T)
        _global_upper_bound::T = typemax(T)
        _maximum_node_id::Int = 0
        _iteration_count::Int = 0
        _node_count::Int = 0

        # Storage for output, reset in initial_parse! in parse.jl
        _solution_value::T = zero(T)
        _feasible_solution_found::Bool = false
        _first_solution_node::Int = -1
        _best_upper_value::T = Inf

        # Optimality-Based Bound Tightening (OBBT) Options
        # set in TODO
        _obbt_working_lower_index::Vector{Bool} = Bool[]
        _obbt_working_upper_index::Vector{Bool} = Bool[]
        _lower_indx_diff::Vector{Bool} = Bool[]
        _upper_indx_diff::Vector{Bool} = Bool[]
        _old_low_index::Vector{Bool} = Bool[]
        _old_upp_index::Vector{Bool} = Bool[]
        _new_low_index::Vector{Bool} = Bool[]
        _new_upp_index::Vector{Bool} = Bool[]
        _obbt_variables::Vector{VI} = VI[]
        _obbt_variable_num::Int = 0
        _obbt_performed_flag::Bool = false

        # Feasibility-Based Bound Tightening Options
        # set in set_constraint_propagation_fbbt in domain_reduction.jl
        _cp_improvement::T = zero(T)
        _cp_evaluation_reverse::Bool = false

        _cut_iterations::Int = 0
        _cut_add_flag::Bool = false

        # Options for Repetition (If DBBT Performed Well)
        # set in within preprocess in optimize_nonconvex.jl
        _node_repetitions::Int = 0
        _initial_volume::T = zero(T)
        _final_volume::T = zero(T)

        # Log
        _log::Log = Log()

        # set in TODO
        _buffered_quadratic_ineq_ci::Vector{CI{SAF,LT}} = CI{SAF,LT}[]
        _buffered_quadratic_eq_ci::Vector{CI{SAF,LT}} = CI{SAF,LT}[]
        _buffered_nonlinear_ci::Vector{CI{SAF,LT}} = CI{SAF,LT}[]

        # need to retreive primal _relaxed_variable_index
        # set in TODO
        #"Number of variables actively branched on in B&B routine (excludes linear and fixed)"
        _relaxed_variable_number::Int = 0
        _relaxed_variable_index::Vector{VI} = VI[]
        _relaxed_variable_eq::Vector{Tuple{CI{SV, ET}, Int}} = Tuple{CI{SV, ET}, Int}[]
        _relaxed_variable_lt::Vector{Tuple{CI{SV, LT}, Int}} = Tuple{CI{SV, LT}, Int}[]
        _relaxed_variable_gt::Vector{Tuple{CI{SV, GT}, Int}} = Tuple{CI{SV, GT}, Int}[]

        # set as user-input
        _branch_variables::Vector{Bool} = Bool[]

        _new_eval_constraint::Bool = false
        _new_eval_objective::Bool = false

        _node_to_sv_leq_ci::Vector{CI{SV,LT}} = CI{SV,LT}[]
        _node_to_sv_geq_ci::Vector{CI{SV,GT}} = CI{SV,GT}[]

        #"Set to true if a nonlinear evaluator was created (NLconstraint or NLobjective specified)"
        _nonlinear_evaluator_created::Bool = false

        #_relaxed_evaluator::Evaluator = Evaluator{1,NS}()
        #_relaxed_constraint_bounds::Vector{MOI.NLPBoundsPair} = Vector{MOI.NLPBoundsPair}[]
end
GlobalOptimizer(::Type{T}) where T<:AbstractFloat = GlobalOptimizer{1,T,DefaultExt}()

function MOI.is_empty(m::GlobalOptimizer{N,T,S}) where {N,T<:AbstractFloat,S}
    empty_opt = GlobalOptimizer{N,T,S}()
    is_empty_flag = true
    for f in fieldnames(GlobalOptimizer)
        if f == :_stack
            is_empty_flag &= isempty(getfield(m, f))
        elseif f == :_current_node
            is_empty_flag &= getfield(m, f) == NodeBB{N,T}()
        elseif f == :_log
                # TODO
        elseif f == :_working_problem
                # TODO
        elseif f == :_input_problem
        else
            if getfield(empty_opt, f) != getfield(m, f)
                is_empty_flag = false
                break
            end
        end
    end
    return is_empty_flag
end

#=
TODO: StructArray access is faster here for most...
=#

Base.@propagate_inbounds function _is_unfixed_integer(::FullVar, m::GlobalOptimizer, i)
end
Base.@propagate_inbounds function _is_unfixed_integer(::BranchVar, m::GlobalOptimizer, i)
end

for (f,arr) in ((:_lower_bound, :_lower_variable_bound),
                (:_upper_bound, :_upper_variable_bound),
                (:_lower_solution, :_lower_solution),)
    @eval function ($f)(::FullVar, m::GlobalOptimizer{N,T,S}) where {N,T<:Real,S<:ExtensionType}
        m.$arr
    end
    @eval Base.@propagate_inbounds function ($f)(::FullVar, m::GlobalOptimizer{N,T,S}, i) where {N,T<:Real,S<:ExtensionType}
        m.$arr[i]
    end
    @eval Base.@propagate_inbounds function ($f)(::BranchVar, m::GlobalOptimizer{N,T,S}) where {N,T<:Real,S<:ExtensionType}
        m.$arr[m._branch_variable]
    end
    @eval Base.@propagate_inbounds function ($f)(::BranchVar, m::GlobalOptimizer{N,T,S}, i) where {N,T<:Real,S<:ExtensionType}
        m.$arr[m._branch_variable[i]]
    end
end

Base.@propagate_inbounds function _set_lower_bound!(::BranchVar, m::GlobalOptimizer{N,T,S}, i, v) where {N,T<:Real,S<:ExtensionType}
    m._lower_variable_bound[m._branch_variable[i]] = v
end
Base.@propagate_inbounds function _set_lower_bound!(::FullVar, m::GlobalOptimizer{N,T,S}, i, v) where {N,T<:Real,S<:ExtensionType}
    m._lower_variable_bound[i] = v
end
Base.@propagate_inbounds function _set_upper_bound!(::BranchVar, m::GlobalOptimizer{N,T,S}, i, v) where {N,T<:Real,S<:ExtensionType}
    m._upper_variable_bound[m._branch_variable[i]] = v
end
Base.@propagate_inbounds function _set_upper_bound!(::FullVar, m::GlobalOptimizer{N,T,S}, i, v) where {N,T<:Real,S<:ExtensionType}
    m._upper_variable_bound[i] = v
end

function _mid(::FullVar, m::GlobalOptimizer{N,T,S}) where {N,T<:Real,S<:ExtensionType}
    0.5*(_upper_bound(FullVar, m) - _lower_bound(FullVar, m))
end
Base.@propagate_inbounds function _mid(::FullVar, m::GlobalOptimizer{N,T,S}, i) where {N,T<:Real,S<:ExtensionType}
    0.5*(_upper_bound(FullVar, m, i) - _lower_bound(FullVar, m, i))
end
function _mid(::BranchVar, m::GlobalOptimizer{N,T,S}) where {N,T<:Real,S<:ExtensionType}
end
Base.@propagate_inbounds function _mid(::BranchVar, m::GlobalOptimizer{N,T,S}, i) where {N,T<:Real,S<:ExtensionType}
end

_input_problem(m::GlobalOptimizer) = m._input_problem
_working_problem(m::GlobalOptimizer) = m._working_problem


module ValueInferenceTools

export TaskStats, InferenceAccumulator, ValueAccumulator, Simulator
export accumulatefromdata, inferencesimulator, valuesimulator

using StatsFuns: logistic

include("probabilities.jl")
include("accumulate.jl")
include("agent.jl")

end # module


# optimized parsing functions, copied from julia/base/parse.jl

import Base: tryparse, tryparse_internal

function tryparse(::Type{Float64}, s::DenseStringViewAndSub)
    hasvalue, val = ccall(:jl_try_substrtod, Tuple{Bool, Float64},
                          (Ptr{UInt8},Csize_t,Csize_t), s, 0, sizeof(s))
    hasvalue ? val : nothing
end
function tryparse_internal(::Type{Float64}, s::DenseStringViewAndSub, startpos::Int, endpos::Int)
    hasvalue, val = ccall(:jl_try_substrtod, Tuple{Bool, Float64},
                          (Ptr{UInt8},Csize_t,Csize_t), s, startpos-1, endpos-startpos+1)
    hasvalue ? val : nothing
end
function tryparse(::Type{Float32}, s::DenseStringViewAndSub)
    hasvalue, val = ccall(:jl_try_substrtof, Tuple{Bool, Float32},
                          (Ptr{UInt8},Csize_t,Csize_t), s, 0, sizeof(s))
    hasvalue ? val : nothing
end
function tryparse_internal(::Type{Float32}, s::DenseStringViewAndSub, startpos::Int, endpos::Int)
    hasvalue, val = ccall(:jl_try_substrtof, Tuple{Bool, Float32},
                          (Ptr{UInt8},Csize_t,Csize_t), s, startpos-1, endpos-startpos+1)
    hasvalue ? val : nothing
end

function tryparse_internal(::Type{Complex{T}}, s::DenseStringViewAndSub, i::Int, e::Int, raise::Bool) where {T<:Real}
    # skip initial whitespace
    while i ≤ e && isspace(s[i])
        i = nextind(s, i)
    end
    if i > e
        raise && throw(ArgumentError("input string is empty or only contains whitespace"))
        return nothing
    end

    # find index of ± separating real/imaginary parts (if any)
    i₊ = something(findnext(in(('+','-')), s, i), 0)
    if i₊ == i # leading ± sign
        i₊ = something(findnext(in(('+','-')), s, i₊+1), 0)
    end
    if i₊ != 0 && s[i₊-1] in ('e','E') # exponent sign
        i₊ = something(findnext(in(('+','-')), s, i₊+1), 0)
    end

    # find trailing im/i/j
    iᵢ = something(findprev(in(('m','i','j')), s, e), 0)
    if iᵢ > 0 && s[iᵢ] == 'm' # im
        iᵢ -= 1
        if s[iᵢ] != 'i'
            raise && throw(ArgumentError("expected trailing \"im\", found only \"m\""))
            return nothing
        end
    end

    if i₊ == 0 # purely real or imaginary value
        if iᵢ > i && !(iᵢ == i+1 && s[i] in ('+','-')) # purely imaginary (not "±inf")
            x = tryparse_internal(T, s, i, iᵢ-1, raise)
            x === nothing && return nothing
            return Complex{T}(zero(x),x)
        else # purely real
            x = tryparse_internal(T, s, i, e, raise)
            x === nothing && return nothing
            return Complex{T}(x)
        end
    end

    if iᵢ < i₊
        raise && throw(ArgumentError("missing imaginary unit"))
        return nothing # no imaginary part
    end

    # parse real part
    re = tryparse_internal(T, s, i, i₊-1, raise)
    re === nothing && return nothing

    # parse imaginary part
    im = tryparse_internal(T, s, i₊+1, iᵢ-1, raise)
    im === nothing && return nothing

    return Complex{T}(re, s[i₊]=='-' ? -im : im)
end

function tryparse_internal(::Type{Bool}, sbuff::DenseStringViewAndSub,
        startpos::Int, endpos::Int, base::Integer, raise::Bool)
    if isempty(sbuff)
        raise && throw(ArgumentError("input string is empty"))
        return nothing
    end

    if isnumeric(sbuff[1])
        intres = tryparse_internal(UInt8, sbuff, startpos, endpos, base, false)
        (intres == 1) && return true
        (intres == 0) && return false
        raise && throw(ArgumentError("invalid Bool representation: $(repr(sbuff))"))
    end

    orig_start = startpos
    orig_end   = endpos

    # Ignore leading and trailing whitespace
    while isspace(sbuff[startpos]) && startpos <= endpos
        startpos = nextind(sbuff, startpos)
    end
    while isspace(sbuff[endpos]) && endpos >= startpos
        endpos = prevind(sbuff, endpos)
    end

    len = endpos - startpos + 1
    p   = pointer(sbuff) + startpos - 1
    GC.@preserve sbuff begin
        (len == 4) && (0 == _memcmp(p, "true", 4)) && (return true)
        (len == 5) && (0 == _memcmp(p, "false", 5)) && (return false)
    end

    if raise
        substr = SubString(sbuff, orig_start, orig_end) # show input string in the error to avoid confusion
        if all(isspace, substr)
            throw(ArgumentError("input string only contains whitespace"))
        else
            throw(ArgumentError("invalid Bool representation: $(repr(substr))"))
        end
    end
    return nothing
end


function init_pseudo_observable_mappings!(m::Model1002)

    pseudo_names = [:y_t, :y_f_t, :NaturalRate, :π_t, :OutputGap, :ExAnteRealRate, :LongRunInflation,
                    :MarginalCost, :Wages, :FlexibleWages, :Hours, :FlexibleHours, :z_t,
                    :Expected10YearRateGap, :NominalFFR, :Expected10YearRate,
                    :Expected10YearNaturalRate,
                    :ExpectedNominalNaturalRate, :NominalRateGap, :LaborProductivityGrowth, :u_t]
                    # :i_f_t, :R_t, :c_f_t, :qk_f_t, :k_f_t, :kbar_f_t, :u_f_t, :rk_f_t, :w_f_t, :L_f_t, :rktil_f_t, :n_f_t, :c_t,
                    # :b_t, :r_f_t]

    if subspec(m) == "ss12"
        to_add = [:g_t, :b_t, :μ_t, :λ_f_t, :λ_w_t, :rm_t, :σ_ω_t, :μ_e_t,
                  :γ_t, :π_star_t, :e_lr_t, :e_tfp_t, :e_gdpdef_t, :e_corepce_t, :e_gdp_t, :e_gdi_t]
        pseudo_names = vcat(pseudo_names, to_add)
    end

    if subspec(m) in ["ss13", "ss14", "ss15", "ss16", "ss17", "ss18", "ss19"]
        push!(pseudo_names, :Sinf_t, :Sinf_w_coef_t, :ι_p, :πtil_t, :πtil_t1, :e_tfp_t)
        if subspec(m) in ["ss14", "ss15", "ss16", "ss18", "ss19"]
            push!(pseudo_names, :e_tfp_t1)
        end
    end

    if haskey(get_settings(m), :add_covid_pseudoobs)
        if get_setting(m, :add_covid_pseudoobs) && subspec(m) in ["ss59", "ss60", "ss61", "ss62", "ss63", "ss64", "ss65", "ss66", "ss67", "ss68", "ss69", "ss70", "ss71", "ss72", "ss73", "ss74", "ss75", "ss76", "ss77", "ss78", "ss79", "ss80", "ss81", "ss82", "ss83", "ss84", "ss85"]
            push!(pseudo_names, :ziid, :varphiiid, :biidc)
        end
    end

    if haskey(get_settings(m), :add_pseudo_gdp)
        if get_setting(m, :add_pseudo_gdp) && subspec(m) in ["ss59", "ss60", "ss61", "ss62", "ss63", "ss64", "ss65", "ss66", "ss67", "ss68", "ss69", "ss70", "ss71", "ss72", "ss73", "ss74", "ss75", "ss76", "ss77", "ss78", "ss79", "ss80", "ss81", "ss82", "ss83", "ss84", "ss85"]
            push!(pseudo_names, :PseudoGDP)
        end
    end

    if haskey(get_settings(m), :add_pseudo_corepce)
        if get_setting(m, :add_pseudo_corepce) && subspec(m) in ["ss59", "ss60", "ss61", "ss62", "ss63", "ss64", "ss65", "ss66", "ss67", "ss68", "ss69", "ss70", "ss71", "ss72", "ss73", "ss74", "ss75", "ss76", "ss77", "ss78", "ss79", "ss80", "ss81", "ss82", "ss83", "ss84", "ss85"]
            push!(pseudo_names, :PseudoCorePCE)
        end
    end

    if haskey(get_settings(m), :add_NominalWageGrowth)
        if get_setting(m, :add_NominalWageGrowth)
            push!(pseudo_names, :NominalWageGrowth)
        end
    end
    if haskey(get_settings(m), :add_ztil)
        if get_setting(m, :add_ztil)
            push!(pseudo_names, :ztil)
        end
    end
    if haskey(get_settings(m), :add_zp)
        if get_setting(m, :add_zp)
            push!(pseudo_names, :zp)
        end
    end
    if haskey(get_settings(m), :add_pgap)
        if get_setting(m, :add_pgap)
            push!(pseudo_names, :pgap)
        end
    end
    if haskey(get_settings(m), :add_ygap)
        if get_setting(m, :add_ygap)
            push!(pseudo_names, :ygap)
        end
    end
    if haskey(get_settings(m), :add_altpolicy_pgap)
        if get_setting(m, :add_altpolicy_pgap)
            push!(pseudo_names, :pgap)
        end
    end
    if haskey(get_settings(m), :add_altpolicy_ygap)
        if get_setting(m, :add_altpolicy_ygap)
            push!(pseudo_names, :ygap)
        end
    end

    if haskey(get_settings(m), :add_urhat)
        if get_setting(m, :add_urhat)
            push!(pseudo_names, :urhat)
        end
    end

    if haskey(get_settings(m), :add_rw)
        if get_setting(m, :add_rw)
            push!(pseudo_names, :rw)
            push!(pseudo_names, :Rref)
        end
    end

    if haskey(get_settings(m), :add_laborshare_measurement)
        if get_setting(m, :add_laborshare_measurement)
            push!(pseudo_names, :laborshare_t)
        end
    end

    if haskey(get_settings(m), :add_laborproductivity_measurement)
        if get_setting(m, :add_laborproductivity_measurement)
            push!(pseudo_names, :laborproductivity)
        end
    end

    if haskey(get_settings(m), :add_nominalgdp_level)
        if get_setting(m, :add_nominalgdp_level)
            push!(pseudo_names, :NominalGDPLevel)
        end
    end

    if haskey(get_settings(m), :add_nominalgdp_growth)
        if get_setting(m, :add_nominalgdp_growth)
            push!(pseudo_names, :NominalGDPGrowth)
        end
    end

    if haskey(get_settings(m), :add_flexible_price_growth)
        if get_setting(m, :add_flexible_price_growth)
            push!(pseudo_names, :FlexibleGDPGrowth, :FlexibleConsumptionGrowth, :FlexibleInvestmentGrowth)
        end
    end

    if haskey(get_settings(m), :add_cumulative)
        if get_setting(m, :add_cumulative)
            push!(pseudo_names, :AccumOutputGap, :GDPLevel,
                  :TechnologyLevel, :FlexibleGDPLevel, :ConsumptionLevel,
                  :FlexibleConsumptionLevel, :InvestmentLevel,
                  :FlexibleInvestmentLevel)
        end
    end

    if haskey(get_settings(m), :add_laborproductivitygrowth_nome_measurement)
        if get_setting(m, :add_laborproductivitygrowth_nome_measurement)
            push!(pseudo_names, :LaborProductivityGrowthNoME)
        end
    end

    if haskey(get_settings(m), :add_Epi_t_measurement)
        if get_setting(m, :add_Epi_t_measurement)
            push!(pseudo_names, :Epi_t)
        end
    end

    # Create PseudoObservable objects
    pseudo = OrderedDict{Symbol,PseudoObservable}()
    for k in pseudo_names
        pseudo[k] = PseudoObservable(k)
    end

    # Fill in names and reverse transforms
    pseudo[:y_t].name = "Output Growth"
    pseudo[:y_t].longname = "Output Growth Per Capita"

    pseudo[:y_f_t].name = "Flexible Output Growth"
    pseudo[:y_f_t].longname = "Output that would prevail in a flexible-price economy."

    pseudo[:NaturalRate].name = "Real Natural Rate"
    pseudo[:NaturalRate].longname = "The real interest rate that would prevail in a flexible-price economy."
    pseudo[:NaturalRate].rev_transform = quartertoannual

    pseudo[:π_t].name = "Inflation"
    pseudo[:π_t].longname = "Inflation"
    pseudo[:π_t].rev_transform = quartertoannual

    pseudo[:OutputGap].name = "Output Gap"
    pseudo[:OutputGap].longname = "Output Gap"

    pseudo[:ExAnteRealRate].name = "Ex Ante Real Rate"
    pseudo[:ExAnteRealRate].longname = "Ex Ante Real Rate"
    pseudo[:ExAnteRealRate].rev_transform = quartertoannual

    pseudo[:LongRunInflation].name = "Long Run Inflation"
    pseudo[:LongRunInflation].longname = "Long Run Inflation"
    pseudo[:LongRunInflation].rev_transform = quartertoannual

    pseudo[:MarginalCost].name = "Marginal Cost"
    pseudo[:MarginalCost].longname = "Marginal Cost"

    pseudo[:Wages].name = "Wages"
    pseudo[:Wages].longname = "Wages"

    pseudo[:FlexibleWages].name = "Flexible Wages"
    pseudo[:FlexibleWages].longname = "Wages that would prevail in a flexible-wage economy"

    pseudo[:Hours].name = "Hours"
    pseudo[:Hours].longname = "Hours"

    pseudo[:FlexibleHours].name     = "Flexible Hours"
    pseudo[:FlexibleHours].longname = "Flexible Hours"

    pseudo[:z_t].name     = "z_t (Technology Growth minus Steady State Growth)"
    pseudo[:z_t].longname = "z_t (Technology Growth minus Steady State Growth)"

    pseudo[:Expected10YearRateGap].name     = "Expected 10-Year Rate Gap"
    pseudo[:Expected10YearRateGap].longname = "Expected 10-Year Rate Gap"
    pseudo[:Expected10YearRateGap].rev_transform = quartertoannual

    pseudo[:NominalFFR].name     = "Nominal FFR"
    pseudo[:NominalFFR].longname = "Nominal FFR at an annual rate"
    pseudo[:NominalFFR].rev_transform = quartertoannual

    pseudo[:Expected10YearRate].name     = "Expected 10-Year Rate"
    pseudo[:Expected10YearRate].longname = "Expected 10-Year Interest Rate"
    pseudo[:Expected10YearRate].rev_transform = quartertoannual

    pseudo[:Expected10YearNaturalRate].name     = "Expected 10-Year Natural Rate"
    pseudo[:Expected10YearNaturalRate].longname = "Expected 10-Year Natural Rate of Interest"
    pseudo[:Expected10YearNaturalRate].rev_transform = quartertoannual

    pseudo[:ExpectedNominalNaturalRate].name     = "Expected Nominal Natural Rate"
    pseudo[:ExpectedNominalNaturalRate].longname = "Natural Rate + Expected Inflation"
    pseudo[:ExpectedNominalNaturalRate].rev_transform = quartertoannual

    pseudo[:NominalRateGap].name     = "Nominal Rate Gap"
    pseudo[:NominalRateGap].longname = "Nominal FFR - Nominal Natural Rate"
    pseudo[:NominalRateGap].rev_transform = quartertoannual

    pseudo[:LaborProductivityGrowth].name     = "Labor Productivity Growth"
    pseudo[:LaborProductivityGrowth].longname = "Labor Productivity Growth Rate"
    pseudo[:LaborProductivityGrowth].rev_transform = quartertoannual

    pseudo[:u_t].name     = "u_t"
    pseudo[:u_t].longname = "u_t"

    # pseudo[:i_f_t].name     = "i_f_t"
    # pseudo[:i_f_t].longname = "i_f_t"

    # pseudo[:R_t].name     = "R_t"
    # pseudo[:R_t].longname = "R_t"

    # pseudo[:c_f_t].name     = "c_f_t"
    # pseudo[:c_f_t].longname = "c_f_t"

    # pseudo[:r_f_t].name     = "r_f_t"
    # pseudo[:r_f_t].longname = "r_f_t"

    # pseudo[:qk_f_t].name     = "qk_f_t"
    # pseudo[:qk_f_t].longname = "qk_f_t"

    # pseudo[:k_f_t].name     = "k_f_t"
    # pseudo[:k_f_t].longname = "k_f_t"

    # pseudo[:kbar_f_t].name     = "kbar_f_t"
    # pseudo[:kbar_f_t].longname = "kbar_f_t"

    # pseudo[:u_f_t].name     = "u_f_t"
    # pseudo[:u_f_t].longname = "u_f_t"

    # pseudo[:rk_f_t].name     = "rk_f_t"
    # pseudo[:rk_f_t].longname = "rk_f_t"

    # pseudo[:w_f_t].name     = "w_f_t"
    # pseudo[:w_f_t].longname = "w_f_t"

    # pseudo[:L_f_t].name     = "L_f_t"
    # pseudo[:L_f_t].longname = "L_f_t"

    # pseudo[:rktil_f_t].name     = "rktil_f_t"
    # pseudo[:rktil_f_t].longname = "rktil_f_t"

    # pseudo[:b_t].name     = "b_t"
    # pseudo[:b_t].longname = "b_t"

    # pseudo[:n_f_t].name     = "n_f_t"
    # pseudo[:n_f_t].longname = "n_f_t"

    if subspec(m) in ["ss13", "ss14", "ss15", "ss16", "ss17", "ss18", "ss19"]
        pseudo[:Sinf_t].name     = "Sinf_t"
        pseudo[:Sinf_t].longname = "Sinf_t, PDV of Emc_t"
        pseudo[:Sinf_w_coef_t].name     = "Sinf_w_coef_t"
        pseudo[:Sinf_w_coef_t].longname = "Sinf_w_coef_t, PDV of Emc_t multiplied by coefficient"
        pseudo[:ι_p].name = "iota_p"
        pseudo[:ι_p].longname = "iota_p"
        pseudo[:πtil_t].name     = "pitil_t"
        pseudo[:πtil_t].longname = "Fundamental Inflation"
        pseudo[:πtil_t1].name     = "pitil_t1"
        pseudo[:πtil_t1].longname = "Fundamental Inflation Lag 1"
        pseudo[:e_tfp_t].name = "e_tfp_t"
        pseudo[:e_tfp_t].longname = "e_tfp_t"
        if subspec(m) in ["ss14", "ss15", "ss16", "ss18", "ss19"]
            pseudo[:e_tfp_t1].name = "e_tfp_t1"
            pseudo[:e_tfp_t1].longname = "e_tfp_t1"
        end
    end

    # Other exogenous processes
    if subspec(m) == "ss12"
        for i in to_add
            pseudo[i].name = string(DSGE.detexify(i))
            pseudo[i].longname = string(DSGE.detexify(i))
        end
    end

    if haskey(get_settings(m), :add_NominalWageGrowth)
        if get_setting(m, :add_NominalWageGrowth)
            pseudo[:NominalWageGrowth].name = "Nominal Wage Growth"
            pseudo[:NominalWageGrowth].longname = "Nominal Wage Growth"
        end
    end

    if haskey(get_settings(m), :add_laborshare_measurement)
        if get_setting(m, :add_laborshare_measurement)
            pseudo[:laborshare_t].name     = "Log Labor Share"
            pseudo[:laborshare_t].longname = "Log Labor Share"
        end
    end

    if haskey(get_settings(m), :add_nominalgdp_level)
        if get_setting(m, :add_nominalgdp_level)
            pseudo[:NominalGDPLevel].name     = "Nominal GDP Level"
            pseudo[:NominalGDPLevel].longname = "Nominal GDP Level"
        end
    end

    if haskey(get_settings(m), :add_nominalgdp_growth)
        if get_setting(m, :add_nominalgdp_growth)
            pseudo[:NominalGDPGrowth].name     = "Nominal GDP Growth"
            pseudo[:NominalGDPGrowth].longname = "Nominal GDP Growth"
        end
    end

    if haskey(get_settings(m), :add_cumulative)
        if get_setting(m, :add_cumulative)
            pseudo[:AccumOutputGap].name     = "Accumulated Output Gap"
            pseudo[:AccumOutputGap].longname = "Accumulated Output Gap"
            pseudo[:GDPLevel].name     = "GDP Level"
            pseudo[:GDPLevel].longname = "GDP Level"
            pseudo[:TechnologyLevel].name     = "Technology (z) Level"
            pseudo[:TechnologyLevel].longname = "Technology (z) Level"
            pseudo[:FlexibleGDPLevel].name     = "Flexible GDP Level"
            pseudo[:FlexibleGDPLevel].longname = "Flexible GDP Level"
            pseudo[:ConsumptionLevel].name     = "Consumption Level"
            pseudo[:ConsumptionLevel].longname = "Consumption Level"
            pseudo[:FlexibleConsumptionLevel].name     = "Flexible Consumption Level"
            pseudo[:FlexibleConsumptionLevel].longname = "Flexible Consumption Level"
            pseudo[:InvestmentLevel].name     = "Investment Level"
            pseudo[:InvestmentLevel].longname = "Investment Level"
            pseudo[:FlexibleInvestmentLevel].name     = "Flexible Investment Level"
            pseudo[:FlexibleInvestmentLevel].longname = "Flexible Investment Level"
        end
    end

    if haskey(get_settings(m), :add_flexible_price_growth)
        if get_setting(m, :add_flexible_price_growth)
            pseudo[:FlexibleGDPGrowth].name     = "Flexible GDP Level"
            pseudo[:FlexibleGDPGrowth].longname = "Flexible GDP Growth"
            pseudo[:FlexibleConsumptionGrowth].name     = "Flexible Consumption Growth"
            pseudo[:FlexibleConsumptionGrowth].longname = "Flexible Consumption Growth"
            pseudo[:FlexibleInvestmentGrowth].name     = "Flexible Investment Growth"
            pseudo[:FlexibleInvestmentGrowth].longname = "Flexible Investment Growth"
        end
    end

    if haskey(get_settings(m), :add_ztil)
        if get_setting(m, :add_ztil)
            pseudo[:ztil].name     = "ztil"
            pseudo[:ztil].longname = "ztil"
        end
    end

    if haskey(get_settings(m), :add_zp)
        if get_setting(m, :add_zp)
            pseudo[:zp].name     = "zp"
            pseudo[:zp].longname = "zp"
        end
    end

    if haskey(get_settings(m), :add_altpolicy_pgap)
        if get_setting(m, :add_altpolicy_pgap)
            pseudo[:pgap].name     = "pgap"
            pseudo[:pgap].longname = "pgap"
        end
    end

    if haskey(get_settings(m), :add_altpolicy_ygap)
        if get_setting(m, :add_altpolicy_ygap)
            pseudo[:ygap].name     = "ygap"
            pseudo[:ygap].longname = "ygap"
        end
    end

    if haskey(get_settings(m), :add_urhat)
        if get_setting(m, :add_urhat)
            pseudo[:urhat].name     = "urhat"
            pseudo[:urhat].longname = "urhat"
        end
    end

    if haskey(get_settings(m), :add_rw)
        if get_setting(m, :add_rw)
            pseudo[:rw].name     = "rw"
            pseudo[:rw].longname = "rw"
            pseudo[:Rref].name     = "Rref"
            pseudo[:Rref].longname = "Rref"
        end
    end

    if haskey(get_settings(m), :add_covid_pseudoobs)
        if get_setting(m, :add_covid_pseudoobs) && subspec(m) in ["ss59", "ss60", "ss61", "ss62", "ss63", "ss64", "ss65", "ss66", "ss67", "ss68", "ss69", "ss70", "ss71", "ss72", "ss73", "ss74", "ss75", "ss76", "ss77", "ss78", "ss79", "ss80", "ss81", "ss82", "ss83", "ss84", "ss85"]
            pseudo[:ziid].name     = "ziid"
            pseudo[:ziid].longname = "ziid"
            pseudo[:biidc].name     = "biidc"
            pseudo[:biidc].longname = "biidc"
            pseudo[:varphiiid].name     = "varphiiid"
            pseudo[:varphiiid].longname = "varphiiid"
        end
    end

    if haskey(get_settings(m), :add_pseudo_gdp)
        if get_setting(m, :add_pseudo_gdp) && subspec(m) in ["ss59", "ss60", "ss61", "ss62", "ss63", "ss64", "ss65", "ss66", "ss67", "ss68", "ss69", "ss70", "ss71", "ss72", "ss73", "ss74", "ss75", "ss76", "ss77", "ss78", "ss79", "ss80", "ss81", "ss82", "ss83", "ss84", "ss85"]
            pseudo[:PseudoGDP].name     = "GDP Growth Pseudo-observable"
            pseudo[:PseudoGDP].longname = "GDP Growth Pseudo-observable"
            pseudo[:PseudoGDP].rev_transform = loggrowthtopct_annualized_percapita
        end
    end

    if haskey(get_settings(m), :add_pseudo_corepce)
        if get_setting(m, :add_pseudo_corepce) && subspec(m) in ["ss59", "ss60", "ss61", "ss62", "ss63", "ss64", "ss65", "ss66", "ss67", "ss68", "ss69", "ss70", "ss71", "ss72", "ss73", "ss74", "ss75", "ss76", "ss77", "ss78", "ss79", "ss80", "ss81", "ss82", "ss83", "ss84", "ss85"]
            pseudo[:PseudoCorePCE].name     = "Core PCE Pseudo-observable"
            pseudo[:PseudoCorePCE].longname = "Core PCE Pseudo-observable"
            pseudo[:PseudoCorePCE].rev_transform = loggrowthtopct_annualized
        end
    end

    if haskey(get_settings(m), :add_laborproductivity_measurement)
        if get_setting(m, :add_laborproductivity_measurement)
            pseudo[:laborproductivity].name     = "Log Labor Productivity"
            pseudo[:laborproductivity].longname = "Log Labor Productivity"
        end
    end

    if haskey(get_settings(m), :add_laborproductivitygrowth_nome_measurement)
        if get_setting(m, :add_laborproductivitygrowth_nome_measurement)
            pseudo[:LaborProductivityGrowthNoME].name     = "Labor Productivity Growth (No ME)"
            pseudo[:LaborProductivityGrowthNoME].longname = "Labor Productivity Growth (No ME)"
        end
    end

    if haskey(get_settings(m), :add_Epi_t_measurement)
        if get_setting(m, :add_Epi_t_measurement)
            pseudo[:Epi_t].name     = "Short-term Inflation Expectations"
            pseudo[:Epi_t].longname = "Short-term Inflation Expectations"
        end
    end

    # Needed to implement pseudo-measurement equation correctly
    m <= Setting(:forward_looking_pseudo_observables, [:Expected10YearRateGap, :Expected10YearRate, :Expected10YearNaturalRate])

    # Add to model object
    m.pseudo_observable_mappings = pseudo
end


module PDBTools

using Parameters
using Printf
using InteractiveUtils
using Formatting
using StaticArrays
using LinearAlgebra: norm

export readPDB, writePDB, getseq, wget, edit!, oneletter, threeletter, residuename
export Atom, printatom, index, index_pdb, name, beta, occup
export Residue, eachresidue, resname, residue, resnum, chain, model, segname
export coor, maxmin, distance, closest
export atomic_number, element, mass, element_name
export formula, stoichiometry
export Sequence

#
# Data structures
#
include("./Atom.jl")
include("./Indexes_mmCIF_fields.jl")
include("./Residue.jl")

#
# Input and output functions
#
include("./all.jl")
include("./oneletter.jl")
include("./threeletter.jl")
include("./residuename.jl")
include("./formula.jl")
include("./check_mmCIF.jl")
include("./parse_int.jl")
include("./read_atom.jl")
include("./write_atom.jl")
include("./readPDB.jl")
include("./edit.jl")
include("./writePDB.jl")
include("./getseq.jl")
include("./coor.jl")
include("./distance.jl")
include("./maxmin.jl")
include("./same_residue.jl")
include("./wget.jl")

# Element properties
include("./elements.jl")
include("./atom_properties.jl")

# For selections
include("./select/select_includes.jl")
include("./alternate_conformation.jl")

end


# Compatibility shims until users upgrade to new NNlib format
function conv!(y::CuArray{T}, x::CuArray{T}, w::CuArray{T}; pad=0, stride=1, flipkernel=0, dilation=1, kwargs...) where {T<:CUDNNFloat}
    cdims = DenseConvDims(x, w; padding=pad, stride=stride, flipkernel=flipkernel, dilation=dilation)
    return conv!(y, x, w, cdims; kwargs...)
end

function ∇conv_filter!(dw::CuArray{T}, dy::CuArray{T}, x::CuArray{T}; pad=0, stride=1, flipkernel=0, dilation=1, kwargs...) where {T<:CUDNNFloat}
    cdims = DenseConvDims(x, dw; padding=pad, stride=stride, flipkernel=flipkernel, dilation=dilation)
    # NOTE!!!  This compat shim re-arranges the argument order!
    return ∇conv_filter!(dw, x, dy, cdims; kwargs...)
end

function maxpool!(y::CuArray{T}, x::CuArray{T}, k; pad=map(_->0,k), stride=k) where {T<:CUDNNFloat}
    pdims = PoolDims(x, k; padding=pad, stride=stride)
    return maxpool!(y, x, pdims)
end

function meanpool!(y::CuArray{T}, x::CuArray{T}, k; pad=map(_->0,k), stride=k) where {T<:CUDNNFloat}
    pdims = PoolDims(x, k; padding=pad, stride=stride)
    return meanpool!(y, x, pdims)
end


"""
Fitness objects are used to provide an empirical means of assessing the
relative quality of one potential solution to another.
"""
module fitness
importall common
using utility, core, distance
export assign, uses, maximise, Fitness, FitnessScheme

"""
Base type used by all fitness values.
"""
abstract Fitness

"""
Base type used by all fitness schemes.
"""
abstract FitnessScheme

# Load each of the fitness schemes.
include("fitness/_scalar.jl")
include("fitness/_pareto.jl")
include("fitness/_aggregate.jl")
include("fitness/_belegundu.jl")
include("fitness/_goldberg.jl")
include("fitness/_moga.jl")

"""
Responsible for generating a fitness value according to a given scheme.
"""
assign(s::FitnessScheme) =
  error("No 'fitness' function defined for this fitness scheme.")

"""
Returns the type of object by this fitness scheme to store fitness values.
"""
uses(s::FitnessScheme) =
  error("No 'uses' function defined for this fitness scheme.")

"""
Returns a flag indicating whether or not fitness values produced by a given
scheme should be maximised or minimised.
"""
maximise(::FitnessScheme) =
  error("No 'maximise' function defined for this fitness scheme.")
end


## GPUArrays interfaces

GPUArrays.device(x::MtlArray) = x.dev

## execution

struct mtlArrayBackend <: AbstractGPUBackend end

struct mtlKernelContext <: AbstractKernelContext end

@inline function GPUArrays.launch_heuristic(::mtlArrayBackend, f::F, args::Vararg{Any,N};
                                             elements::Int, elements_per_thread::Int) where {F,N}
    kernel = @metal launch=false f(mtlKernelContext(), args...)

    # TODO: Be more efficient about creating pipeline states as they're expensive
    # The pipeline state automatically computes occupancy stats, so we just need to parse its fields
    pipeline = MtlComputePipelineState(kernel.fun.lib.device, kernel.fun)
    threads_needed = cld(elements, elements_per_thread)
    # Limit the threadgroup size
    _threads = min(threads_needed, pipeline.maxTotalThreadsPerThreadgroup)
    _blocks  = cld(threads_needed, _threads)
    return (threads=_threads, blocks=_blocks)
end

function GPUArrays.gpu_call(::mtlArrayBackend, f, args, threads::Int, blocks::Int;
                            name::Union{String,Nothing})
    @metal threads=threads grid=blocks name=name f(mtlKernelContext(), args...)
end

## on-device

# indexing
GPUArrays.blockidx(ctx::mtlKernelContext)  = Metal.threadgroup_position_in_grid_1d()
GPUArrays.blockdim(ctx::mtlKernelContext)  = Metal.threads_per_threadgroup_1d()
GPUArrays.threadidx(ctx::mtlKernelContext) = Metal.thread_position_in_threadgroup_1d()
GPUArrays.griddim(ctx::mtlKernelContext)   = Metal.threadgroups_per_grid_1d()

# math

@inline GPUArrays.cos(ctx::mtlKernelContext, x)  = cos(x)
@inline GPUArrays.sin(ctx::mtlKernelContext, x)  = sin(x)
@inline GPUArrays.sqrt(ctx::mtlKernelContext, x) = sqrt(x)
@inline GPUArrays.log(ctx::mtlKernelContext, x)  = log(x)

# memory

# @inline function GPUArrays.LocalMemory(::oneKernelContext, ::Type{T}, ::Val{dims}, ::Val{id}
#                                       ) where {T, dims, id}
#     ptr = oneAPI.emit_localmemory(Val(id), T, Val(prod(dims)))
#     oneDeviceArray(dims, LLVMPtr{T, onePI.AS.Local}(ptr))
# end

# synchronization

@inline GPUArrays.synchronize_threads(::mtlKernelContext) = Metal.threadgroup_barrier()

#
# Host abstractions
#

GPUArrays.backend(::Type{<:MtlArray}) = mtlArrayBackend()

const GLOBAL_RNGs = Dict{MtlDevice,GPUArrays.RNG}()
function GPUArrays.default_rng(::Type{<:MtlArray})
    dev = MtlDevice(1)
    get!(GLOBAL_RNGs, dev) do
        N = 128 # Size of default oneAPI working group with barrier, so should be good for Metal
        state = MtlArray{NTuple{4, UInt32}}(undef, N)
        rng = GPUArrays.RNG(state)
        Random.seed!(rng)
        rng
    end
end