# The following is largely inspired from JuMP/test/JuMPExtension.jl

@enum Level BOTH LOWER_ONLY UPPER_ONLY DUAL_OF_LOWER

abstract type AbstractBilevelModel <: JuMP.AbstractModel end

Base.broadcastable(model::AbstractBilevelModel) = Ref(model)

mutable struct ConstraintInfo{T<:Union{Float64, Vector{Float64}}}
    start::T
    upper::T
    lower::T
    function ConstraintInfo{Float64}()
        new(NaN, NaN, NaN)
    end
    function ConstraintInfo{Vector{Float64}}(N::Integer)
        new(fill(NaN, N), fill(NaN, N), fill(NaN, N))
    end
end

mutable struct VariableInfo{T<:Union{Float64, Vector{Float64}}}
    upper::T
    lower::T
    function VariableInfo{Float64}()
        new(NaN, NaN)
    end
    function VariableInfo{Vector{Float64}}(N::Integer)
        new(fill(NaN, N), fill(NaN, N))
    end
end

mutable struct BilevelModel <: AbstractBilevelModel
    # Structured data
    upper::JuMP.AbstractModel
    lower::JuMP.AbstractModel

    # Model data
    nextvaridx::Int                                 # Next variable index is nextvaridx+1
    variables::Dict{Int, JuMP.AbstractVariable}     # Map varidx -> variable
    varnames::Dict{Int, String}                     # Map varidx -> name
    var_level::Dict{Int, Level}
    var_upper::Dict{Int, JuMP.AbstractVariableRef}
    var_lower::Dict{Int, JuMP.AbstractVariableRef}
    var_info::Dict{Int, VariableInfo}

    # upper level decisions that are "parameters" of the second level
    upper_to_lower_link::Dict{JuMP.AbstractVariableRef, JuMP.AbstractVariableRef}
    # lower level decisions that are input to upper
    lower_to_upper_link::Dict{JuMP.AbstractVariableRef, JuMP.AbstractVariableRef}
    # lower level decisions that are input to upper
    upper_var_to_lower_ctr_link::Dict{JuMP.AbstractVariableRef, JuMP.ConstraintRef}
    # joint link
    link::Dict{JuMP.AbstractVariableRef, JuMP.AbstractVariableRef}

    nextconidx::Int                                 # Next constraint index is nextconidx+1
    constraints::Dict{Int, JuMP.AbstractConstraint} # Map conidx -> variable
    connames::Dict{Int, String}                     # Map varidx -> name
    ctr_level::Dict{Int, Level}
    ctr_upper::Dict{Int, JuMP.ConstraintRef}
    ctr_lower::Dict{Int, JuMP.ConstraintRef}
    ctr_info::Dict{Int, ConstraintInfo}

    upper_objective_sense::MOI.OptimizationSense
    upper_objective_function::JuMP.AbstractJuMPScalar

    lower_objective_sense::MOI.OptimizationSense
    lower_objective_function::JuMP.AbstractJuMPScalar

    # solution data
    solver#::MOI.ModelLike
    upper_to_sblm
    lower_to_sblm
    lower_dual_to_sblm
    sblm_to_solver
    lower_primal_dual_map

    objdict::Dict{Symbol, Any}    # Same that JuMP.Model's field `objdict`

    function BilevelModel()

        model = new(
            JuMP.Model(),
            JuMP.Model(),

            # var
            0, Dict{Int, JuMP.AbstractVariable}(),   Dict{Int, String}(),    # Model Variables
            Dict{Int, Level}(), Dict{Int, JuMP.AbstractVariable}(), Dict{Int, JuMP.AbstractVariable}(),
            Dict{Int, VariableInfo}(),
            # links
            Dict{JuMP.AbstractVariable, JuMP.AbstractVariable}(), Dict{JuMP.AbstractVariable, JuMP.AbstractVariable}(),
            Dict{JuMP.AbstractVariable, JuMP.ConstraintRef}(),
            Dict{JuMP.AbstractVariable, JuMP.AbstractVariable}(),
            #ctr
            0, Dict{Int, JuMP.AbstractConstraint}(), Dict{Int, String}(),    # Model Constraints
            Dict{Int, Level}(), Dict{Int, JuMP.AbstractConstraint}(), Dict{Int, JuMP.AbstractConstraint}(),
            Dict{Int, ConstraintInfo}(),
            #obj
            MOI.FEASIBILITY_SENSE,
            zero(JuMP.GenericAffExpr{Float64, BilevelVariableRef}), # Model objective
            MOI.FEASIBILITY_SENSE,
            zero(JuMP.GenericAffExpr{Float64, BilevelVariableRef}), # Model objective

            nothing,
            nothing,
            nothing,
            nothing,
            nothing,
            nothing,
            Dict{Symbol, Any}(),
            )

        return model
    end
end

abstract type InnerBilevelModel <: AbstractBilevelModel end
struct UpperModel <: InnerBilevelModel
    m::BilevelModel
end
Upper(m::BilevelModel) = UpperModel(m)
UpperToLower(m::BilevelModel) = UpperModel(m)
struct LowerModel <: InnerBilevelModel
    m::BilevelModel
end
Lower(m::BilevelModel) = LowerModel(m)
LowerToUpper(m::BilevelModel) = LowerModel(m)
bilevel_model(m::InnerBilevelModel) = m.m
mylevel_model(m::UpperModel) = bilevel_model(m).upper
mylevel_model(m::LowerModel) = bilevel_model(m).lower
level(m::LowerModel) = LOWER_ONLY
level(m::UpperModel) = UPPER_ONLY
mylevel_ctr_list(m::LowerModel) = bilevel_model(m).ctr_lower
mylevel_ctr_list(m::UpperModel) = bilevel_model(m).ctr_upper
mylevel_var_list(m::LowerModel) = bilevel_model(m).var_lower
mylevel_var_list(m::UpperModel) = bilevel_model(m).var_upper

# obj

mylevel_obj_sense(m::LowerModel) = bilevel_model(m).lower_objective_sense
mylevel_obj_function(m::LowerModel) = bilevel_model(m).lower_objective_function
mylevel_obj_sense(m::UpperModel) = bilevel_model(m).upper_objective_sense
mylevel_obj_function(m::UpperModel) = bilevel_model(m).upper_objective_function

set_mylevel_obj_sense(m::LowerModel, val) = bilevel_model(m).lower_objective_sense = val
set_mylevel_obj_function(m::LowerModel, val) = bilevel_model(m).lower_objective_function = val
set_mylevel_obj_sense(m::UpperModel, val) = bilevel_model(m).upper_objective_sense = val
set_mylevel_obj_function(m::UpperModel, val) = bilevel_model(m).upper_objective_function = val

function set_link!(m::UpperModel, upper::JuMP.AbstractVariableRef, lower::JuMP.AbstractVariableRef)
    bilevel_model(m).upper_to_lower_link[upper] = lower
    bilevel_model(m).link[upper] = lower
    nothing
end
function set_link!(m::LowerModel, upper::JuMP.AbstractVariableRef, lower::JuMP.AbstractVariableRef)
    bilevel_model(m).lower_to_upper_link[lower] = upper
    bilevel_model(m).link[upper] = lower
    nothing
end

# Models to deal with variables that are not exchanged between models
abstract type SingleBilevelModel <: AbstractBilevelModel end
struct UpperOnlyModel <: SingleBilevelModel
    m::BilevelModel
end
UpperOnly(m::BilevelModel) = UpperOnlyModel(m)
struct LowerOnlyModel <: SingleBilevelModel
    m::BilevelModel
end
LowerOnly(m::BilevelModel) = LowerOnlyModel(m)

bilevel_model(m::SingleBilevelModel) = m.m
mylevel_model(m::UpperOnlyModel) = bilevel_model(m).upper
mylevel_model(m::LowerOnlyModel) = bilevel_model(m).lower
level(m::LowerOnlyModel) = LOWER_ONLY
level(m::UpperOnlyModel) = UPPER_ONLY
mylevel_var_list(m::LowerOnlyModel) = bilevel_model(m).var_lower
mylevel_var_list(m::UpperOnlyModel) = bilevel_model(m).var_upper

in_upper(l::Level) = l == BOTH || l == UPPER_ONLY || l == DUAL_OF_LOWER
in_lower(l::Level) = l == BOTH || l == LOWER_ONLY

#### Model ####

# Variables
struct BilevelVariableRef <: JuMP.AbstractVariableRef
    model::BilevelModel # `model` owning the variable
    idx::Int       # Index in `model.variables`
    level::Level
end
mylevel(v::BilevelVariableRef) = v.level
in_level(v::BilevelVariableRef, level::Level) = (
    v.level === BOTH ||
    v.level === level ||
    (v.level === DUAL_OF_LOWER && level === UPPER_ONLY))

in_upper(v::BilevelVariableRef) = in_upper(mylevel(v))
in_lower(v::BilevelVariableRef) = in_lower(mylevel(v))
upper_ref(v::BilevelVariableRef) = v.model.var_upper[v.idx]
lower_ref(v::BilevelVariableRef) = v.model.var_lower[v.idx]

function bound_ref(v::BilevelVariableRef)
    if mylevel(v) == LOWER_ONLY
        return lower_ref(v)
    elseif mylevel(v) == UPPER_ONLY || mylevel(v) == DUAL_OF_LOWER
        return upper_ref(v)
    elseif mylevel(v) == BOTH
        m = v.model
        i = v.idx
        upper_ref = m.var_upper[i]
        # this checks actual owner:
        if haskey(m.upper_to_lower_link, upper_ref)
            return upper_ref
        else
            return m.var_lower[i]
        end
    else
        error("Unknown level")
    end
end

function solver_ref(v::BilevelVariableRef)
    m = v.model
    if mylevel(v) == LOWER_ONLY
        return m.sblm_to_solver[
            m.lower_to_sblm[JuMP.index(lower_ref(v))]]
    else
        return m.sblm_to_solver[
            m.upper_to_sblm[JuMP.index(upper_ref(v))]]
    end
end

Base.broadcastable(v::BilevelVariableRef) = Ref(v)
Base.copy(v::BilevelVariableRef) = v
# TODO
# Base.copy(v::BilevelVariableRef, new_model::BilevelModel) = BilevelVariableRef(new_model, v.idx, v.level) # also requires a map
Base.:(==)(v::BilevelVariableRef, w::BilevelVariableRef) =
    v.model === w.model && v.idx == w.idx && v.level == w.level
JuMP.owner_model(v::BilevelVariableRef) = v.model
JuMP.isequal_canonical(v::BilevelVariableRef, w::BilevelVariableRef) = v == w
JuMP.variable_type(::AbstractBilevelModel) = BilevelVariableRef
# add in BOTH levels
function JuMP.add_variable(inner::InnerBilevelModel, v::JuMP.AbstractVariable, name::String="")
    m = bilevel_model(inner)
    m.nextvaridx += 1
    vref = BilevelVariableRef(m, m.nextvaridx, BOTH)
    # break info so that bounds go to correct level
    var_upper, var_lower = split_variable(inner, v)
    v_upper = JuMP.add_variable(m.upper, var_upper, name)
    m.var_upper[vref.idx] = v_upper
    v_lower = JuMP.add_variable(m.lower, var_lower, name)
    m.var_lower[vref.idx] = v_lower
    m.var_level[vref.idx] = BOTH
    set_link!(inner, v_upper, v_lower)
    m.variables[vref.idx] = v # save complete data
    JuMP.set_name(vref, name)
    m.var_info[vref.idx] = empty_info(v)
    vref
end
function JuMP.add_variable(single::SingleBilevelModel, v::JuMP.AbstractVariable, name::String="")
    m = bilevel_model(single)
    m.nextvaridx += 1
    vref = BilevelVariableRef(m, m.nextvaridx, level(single))
    v_level = JuMP.add_variable(mylevel_model(single), v, name)
    mylevel_var_list(single)[vref.idx] = v_level
    m.var_level[vref.idx] = level(single)
    m.variables[vref.idx] = v
    JuMP.set_name(vref, name)
    m.var_info[vref.idx] = empty_info(v)
    vref
end
function MOI.delete!(m::AbstractBilevelModel, vref::BilevelVariableRef)
    error("No deletion on bilevel models")
    delete!(m.variables, vref.idx)
    delete!(m.varnames, vref.idx)
end
MOI.is_valid(m::BilevelModel, vref::BilevelVariableRef) = vref.idx in keys(m.variables)
JuMP.num_variables(m::BilevelModel) = length(m.variables)
JuMP.num_variables(m::InnerBilevelModel) = JuMP.num_variables(bilevel_model(m))
function empty_info(::JuMP.AbstractVariable)
    return VariableInfo{Float64}()
end

"""
Split variable because actual owner of the variable should be the one holding the bounds.
"""
function split_variable(::UpperModel, v::JuMP.AbstractVariable)
    var_upper = v
    var_lower = JuMP.ScalarVariable(JuMP.VariableInfo(
        false, NaN,
        false, NaN,
        false, NaN,
        v.info.has_start, v.info.start,
        false, false))
    return var_upper, var_lower
end
function split_variable(::LowerModel, v::JuMP.AbstractVariable)
    var_lower = v
    var_upper = JuMP.ScalarVariable(JuMP.VariableInfo(
        false, NaN,
        false, NaN,
        false, NaN,
        v.info.has_start, v.info.start,
        false, false))
    return var_upper, var_lower
end

# Internal function
variable_info(vref::BilevelVariableRef) = vref.model.variables[vref.idx].info
function update_variable_info(vref::BilevelVariableRef, info::JuMP.VariableInfo)
    vref.model.variables[vref.idx] = JuMP.ScalarVariable(info)
end

function JuMP.has_lower_bound(vref::BilevelVariableRef)
    if mylevel(vref) == DUAL_OF_LOWER
        return !isnan(get_dual_lower_bound(get_constrain_ref(vref)))
    end
    return variable_info(vref).has_lb
end
function JuMP.lower_bound(vref::BilevelVariableRef)
    if mylevel(vref) == DUAL_OF_LOWER
        return get_dual_lower_bound(get_constrain_ref(vref))
    end
    @assert !JuMP.is_fixed(vref)
    variable_info(vref).lower_bound
end
function JuMP.set_lower_bound(vref::BilevelVariableRef, lower)
    if mylevel(vref) == DUAL_OF_LOWER
        set_dual_lower_bound(get_constrain_ref(vref), lower)
        return vref.model.variables[vref.idx]
    end
    info = variable_info(vref)
    JuMP.set_lower_bound(bound_ref(vref), lower)
    update_variable_info(vref,
                         JuMP.VariableInfo(true, lower,
                                           info.has_ub, info.upper_bound,
                                           info.has_fix, info.fixed_value,
                                           info.has_start, info.start,
                                           info.binary, info.integer))
end
function JuMP.delete_lower_bound(vref::BilevelVariableRef)
    if mylevel(vref) == DUAL_OF_LOWER
        set_dual_lower_bound(get_constrain_ref(vref), NaN)
        return vref.model.variables[vref.idx]
    end
    info = variable_info(vref)
    JuMP.delete_lower_bound(bound_ref(vref))
    update_variable_info(vref,
                         JuMP.VariableInfo(false, info.lower_bound,
                                           info.has_ub, info.upper_bound,
                                           info.has_fix, info.fixed_value,
                                           info.has_start, info.start,
                                           info.binary, info.integer))
end
function JuMP.has_upper_bound(vref::BilevelVariableRef)
    if mylevel(vref) == DUAL_OF_LOWER
        return !isnan(get_dual_upper_bound(get_constrain_ref(vref)))
    end
    return variable_info(vref).has_ub
end
function JuMP.upper_bound(vref::BilevelVariableRef)
    if mylevel(vref) == DUAL_OF_LOWER
        return get_dual_upper_bound(get_constrain_ref(vref))
    end
    @assert !JuMP.is_fixed(vref)
    variable_info(vref).upper_bound
end
function JuMP.set_upper_bound(vref::BilevelVariableRef, upper)
    if mylevel(vref) == DUAL_OF_LOWER
        set_dual_upper_bound(get_constrain_ref(vref), upper)
        return vref.model.variables[vref.idx]
    end
    info = variable_info(vref)
    JuMP.set_upper_bound(bound_ref(vref), upper)
    update_variable_info(vref,
                         JuMP.VariableInfo(info.has_lb, info.lower_bound,
                                           true, upper,
                                           info.has_fix, info.fixed_value,
                                           info.has_start, info.start,
                                           info.binary, info.integer))
end
function JuMP.delete_upper_bound(vref::BilevelVariableRef)
    if mylevel(vref) == DUAL_OF_LOWER
        set_dual_upper_bound(get_constrain_ref(vref), NaN)
        return vref.model.variables[vref.idx]
    end
    info = variable_info(vref)
    JuMP.delete_upper_bound(bound_ref(vref))
    update_variable_info(vref,
                         JuMP.VariableInfo(info.has_lb, info.lower_bound,
                                           false, info.upper_bound,
                                           info.has_fix, info.fixed_value,
                                           info.has_start, info.start,
                                           info.binary, info.integer))
end
JuMP.is_fixed(vref::BilevelVariableRef) = variable_info(vref).has_fix
JuMP.fix_value(vref::BilevelVariableRef) = variable_info(vref).fixed_value
function JuMP.fix(vref::BilevelVariableRef, value)
    if mylevel(vref) == DUAL_OF_LOWER
        error("Dual variable cannot be fixed.")
    end
    info = variable_info(vref)
    JuMP.fix(bound_ref(vref), value)
    update_variable_info(vref,
                         JuMP.VariableInfo(info.has_lb, info.lower_bound,
                                           info.has_ub, info.upper_bound,
                                           true, value,
                                           info.has_start, info.start,
                                           info.binary, info.integer))
end
function JuMP.unfix(vref::BilevelVariableRef)
    if mylevel(vref) == DUAL_OF_LOWER
        error("Dual variable cannot be fixed.")
    end
    info = variable_info(vref)
    JuMP.unfix(bound_ref(vref))
    update_variable_info(vref,
                         JuMP.VariableInfo(info.has_lb, info.lower_bound,
                                           info.has_ub, info.upper_bound,
                                           false, info.fixed_value,
                                           info.has_start, info.start,
                                           info.binary, info.integer))
end
function JuMP.start_value(vref::BilevelVariableRef)
    if mylevel(vref) == DUAL_OF_LOWER
        return get_dual_start(get_constrain_ref(vref))
    end
    variable_info(vref).start
end
function JuMP.set_start_value(vref::BilevelVariableRef, start)
    if mylevel(vref) == DUAL_OF_LOWER
        set_dual_start(get_constrain_ref(vref), start)
        return vref.model.variables[vref.idx]
    end
    info = variable_info(vref)
    in_upper(vref) && JuMP.set_start_value(upper_ref(vref), start)
    in_lower(vref) && JuMP.set_start_value(lower_ref(vref), start)
    update_variable_info(vref,
                         JuMP.VariableInfo(info.has_lb, info.lower_bound,
                                           info.has_ub, info.upper_bound,
                                           info.has_fix, info.fixed_value,
                                           true, start,
                                           info.binary, info.integer))
end
JuMP.is_binary(vref::BilevelVariableRef) = variable_info(vref).binary
function JuMP.set_binary(vref::BilevelVariableRef)
    if mylevel(vref) == DUAL_OF_LOWER
        error("Dual variable cannot be binary.")
    end
    @assert !JuMP.is_integer(vref)
    info = variable_info(vref)
    JuMP.set_binary(bound_ref(vref))
    update_variable_info(vref,
                         JuMP.VariableInfo(info.has_lb, info.lower_bound,
                                           info.has_ub, info.upper_bound,
                                           info.has_fix, info.fixed_value,
                                           info.has_start, info.start,
                                           true, info.integer))
end
function JuMP.unset_binary(vref::BilevelVariableRef)
    if mylevel(vref) == DUAL_OF_LOWER
        error("Dual variable cannot be binary.")
    end
    info = variable_info(vref)
    JuMP.unset_binary(bound_ref(vref))
    update_variable_info(vref,
                         JuMP.VariableInfo(info.has_lb, info.lower_bound,
                                           info.has_ub, info.upper_bound,
                                           info.has_fix, info.fixed_value,
                                           info.has_start, info.start,
                                           false, info.integer))
end
JuMP.is_integer(vref::BilevelVariableRef) = variable_info(vref).integer
function JuMP.set_integer(vref::BilevelVariableRef)
    if mylevel(vref) == DUAL_OF_LOWER
        error("Dual variable cannot be integer.")
    end
    @assert !JuMP.is_binary(vref)
    info = variable_info(vref)
    JuMP.set_integer(bound_ref(vref))
    update_variable_info(vref,
                         JuMP.VariableInfo(info.has_lb, info.lower_bound,
                                           info.has_ub, info.upper_bound,
                                           info.has_fix, info.fixed_value,
                                           info.has_start, info.start,
                                           info.binary, true))
end
function JuMP.unset_integer(vref::BilevelVariableRef)
    if mylevel(vref) == DUAL_OF_LOWER
        error("Dual variable cannot be integer.")
    end
    info = variable_info(vref)
    JuMP.unset_integer(bound_ref(vref))
    update_variable_info(vref,
                         JuMP.VariableInfo(info.has_lb, info.lower_bound,
                                           info.has_ub, info.upper_bound,
                                           info.has_fix, info.fixed_value,
                                           info.has_start, info.start,
                                           info.binary, false))
end

# Constraints
struct BilevelConstraintRef
    model::BilevelModel # `model` owning the constraint
    idx::Int       # Index in `model.constraints`
    level::Level
end
JuMP.constraint_type(::AbstractBilevelModel) = BilevelConstraintRef
my_level(cref::BilevelConstraintRef) = cref.level
function JuMP.add_constraint(m::BilevelModel, c::JuMP.AbstractConstraint, name::String="")
    error(
        "Can't add constraint directly to the bilevel model `m`, "*
        "attach the constraint to the upper or lower model "*
        "with @constraint(Upper(m), ...) or @constraint(Lower(m), ...)")
end
# function constraint_object(con_ref::ConstraintRef{Model, _MOICON{FuncType, SetType}}) where
#     {FuncType <: MOI.AbstractScalarFunction, SetType <: MOI.AbstractScalarSet}
#     model = con_ref.model
#     f = MOI.get(model, MOI.ConstraintFunction(), con_ref)::FuncType
#     s = MOI.get(model, MOI.ConstraintSet(), con_ref)::SetType
#     return ScalarConstraint(jump_function(model, f), s)
# end
# JuMP.add_constraint(m::UpperModel, c::JuMP.VectorConstraint, name::String="") =
#     error("no vec ctr")
function JuMP.add_constraint(m::InnerBilevelModel, c::Union{JuMP.ScalarConstraint{F,S},JuMP.VectorConstraint{F,S}}, name::String="") where {F,S}
    blm = bilevel_model(m)
    blm.nextconidx += 1
    cref = BilevelConstraintRef(blm, blm.nextconidx, level(m))
    func = JuMP.jump_function(c)
    # @show c
    # @show func
    level_func = replace_variables(func, bilevel_model(m), mylevel_model(m), mylevel_var_list(m), level(m))
    level_c = JuMP.build_constraint(error, level_func, c.set)
    level_cref = JuMP.add_constraint(mylevel_model(m), level_c, name)
    blm.ctr_level[cref.idx] = level(m)
    mylevel_ctr_list(m)[cref.idx] = level_cref
    blm.constraints[cref.idx] = c
    blm.ctr_info[cref.idx] = empty_info(c)
    JuMP.set_name(cref, name)
    cref
end
function MOI.delete!(m::AbstractBilevelModel, cref::BilevelConstraintRef)
    error("can't delete")
    delete!(m.constraints, cref.idx)
    delete!(m.connames, cref.idx)
end
MOI.is_valid(m::BilevelModel, cref::BilevelConstraintRef) = cref.idx in keys(m.constraints)
MOI.is_valid(m::InnerBilevelModel, cref::BilevelConstraintRef) =
    MOI.is_valid(bilevel_model(m), cref) && bilevel_model(m).ctr_level[cref.idx] == level(m)
function JuMP.constraint_object(cref::BilevelConstraintRef, F::Type, S::Type)
    c = cref.model.constraints[cref.idx]
    # `TypeError` should be thrown is `F` and `S` are not correct
    # This is needed for the tests in `constraints.jl`
    c.func::F
    c.set::S
    c
end
function empty_info(c::JuMP.ScalarConstraint{F,S}) where {F,S}
    return ConstraintInfo{Float64}()
end
function empty_info(c::JuMP.VectorConstraint{F,S}) where {F,S}
    return ConstraintInfo{Vector{Float64}}(MOI.dimension(c.set))
end

function set_dual_start(cref::BilevelConstraintRef, value::T) where T<:Number
    cref.model.ctr_info[cref.idx].start = value
end
function set_dual_start(cref::BilevelConstraintRef, value::T) where T<:Vector{S} where S
    array = cref.model.ctr_info[cref.idx].start
    @assert length(array) == length(value)
    copyto!(array, value)
end
function get_dual_start(cref::BilevelConstraintRef)
    cref.model.ctr_info[cref.idx].start
end
function set_dual_upper_bound(cref::BilevelConstraintRef, value::T) where T<:Number
    cref.model.ctr_info[cref.idx].upper = value
end
function set_dual_upper_bound(cref::BilevelConstraintRef, value::T) where T<:Vector{S} where S
    array = cref.model.ctr_info[cref.idx].upper
    @assert length(array) == length(value)
    copyto!(array, value)
end
function get_dual_upper_bound(cref::BilevelConstraintRef)
    cref.model.ctr_info[cref.idx].upper
end
function set_dual_lower_bound(cref::BilevelConstraintRef, value::T) where T<:Number
    cref.model.ctr_info[cref.idx].lower = value
end
function set_dual_lower_bound(cref::BilevelConstraintRef, value::T) where T<:Vector{S} where S
    array = cref.model.ctr_info[cref.idx].lower
    @assert length(array) == length(value)
    copyto!(array, value)
end
function get_dual_lower_bound(cref::BilevelConstraintRef)
    cref.model.ctr_info[cref.idx].lower
end

function set_primal_upper_bound_hint(vref::BilevelVariableRef, value::T) where T<:Number
    vref.model.var_info[vref.idx].upper = value
end
function get_primal_upper_bound_hint(vref::BilevelVariableRef)
    vref.model.var_info[vref.idx].upper
end
function set_primal_lower_bound_hint(vref::BilevelVariableRef, value::T) where T<:Number
    vref.model.var_info[vref.idx].lower = value
end
function get_primal_lower_bound_hint(vref::BilevelVariableRef)
    vref.model.var_info[vref.idx].lower
end

function JuMP.value(cref::BilevelConstraintRef)
    if my_level(cref) == BilevelJuMP.LOWER_ONLY
        # Constraint index on the lower model
        con_lower_idx = cref.model.ctr_lower[cref.idx].index
        # Single bilevel model constraint associated with the lower level constraint
        con_sblm_idx = cref.model.lower_to_sblm[con_lower_idx]
    else
        # Constraint index on the lower model
        con_upper_idx = cref.model.ctr_upper[cref.idx].index
        # Single bilevel model constraint associated with the lower level constraint
        con_sblm_idx = cref.model.upper_to_sblm[con_upper_idx]
    end
    # Solver constraint associated with the single bilevel model constraint
    con_solver_idx = cref.model.sblm_to_solver[con_sblm_idx]
    return MOI.get(cref.model.solver, MOI.ConstraintPrimal(), con_solver_idx)
end
# variables again (duals)
# code for using dual variables associated with lower level constraints
# in the upper level

struct DualOf
    ci::BilevelConstraintRef
end
struct DualVariableInfo
    info::JuMP.VariableInfo
    ci::BilevelConstraintRef
end
function JuMP.build_variable(
    _error::Function,
    info::JuMP.VariableInfo,
    dual_of::DualOf;
    extra_kw_args...,
)
    for (kwarg, _) in extra_kw_args
        _error("Unrecognized keyword argument $kwarg")
    end

    if info.has_lb
        set_dual_lower_bound(dual_of.ci, info.lower_bound)
        # info.has_lb = false
        # info.lower_bound = NaN
    end
    if info.has_ub
        set_dual_upper_bound(dual_of.ci, info.upper_bound)
        # info.has_ub = false
        # info.upper_bound = NaN
    end
    info.has_fix   && _error("Dual variable does not support fixing")
    if info.has_start
        set_dual_start(dual_of.ci, info.start)
        # info.has_start = false
        # info.start = NaN
    end
    info.binary    && _error("Dual variable cannot be binary")
    info.integer   && _error("Dual variable cannot be integer")

    info = JuMP.VariableInfo(false, NaN, false, NaN,
                             false, NaN, false, NaN,
                             false, false)

    if dual_of.ci.level != LOWER_ONLY
        error("Variables can only be tied to LOWER level constraints, got $(dual_of.ci.level) level")
    end

    return DualVariableInfo(
        info,
        dual_of.ci
    )
end
function JuMP.add_variable(inner::UpperModel, dual_info::DualVariableInfo, name::String="")
    # TODO vector version
    m = bilevel_model(inner)
    m.nextvaridx += 1
    vref = BilevelVariableRef(m, m.nextvaridx, DUAL_OF_LOWER)
    v_upper = JuMP.add_variable(m.upper, JuMP.ScalarVariable(dual_info.info), name)
    m.var_upper[vref.idx] = v_upper
    m.var_level[vref.idx] = DUAL_OF_LOWER
    m.upper_var_to_lower_ctr_link[v_upper] = m.ctr_lower[dual_info.ci.idx] # TODO improve this
    m.variables[vref.idx] = JuMP.ScalarVariable(dual_info.info)
    JuMP.set_name(vref, name)
    m.var_info[vref.idx] = empty_info(dual_info)
    vref
end
function empty_info(::DualVariableInfo)
    return VariableInfo{Float64}()
end

function get_constrain_ref(vref::BilevelVariableRef)
    model = vref.model
    ctr_ref = model.upper_var_to_lower_ctr_link[model.var_upper[vref.idx]]
    idx = -1
    for (ind, ref) in model.ctr_lower
        if ref == ctr_ref
            idx = ind
        end
    end
    @assert idx != -1
    return BilevelConstraintRef(model, idx, DUAL_OF_LOWER)
end

# Objective
function JuMP.set_objective(m::InnerBilevelModel, sense::MOI.OptimizationSense,
                            f::JuMP.AbstractJuMPScalar)
    set_mylevel_obj_sense(m, sense)
    set_mylevel_obj_function(m, f)
    level_f = replace_variables(f, bilevel_model(m), mylevel_model(m), mylevel_var_list(m), level(m))
    JuMP.set_objective(mylevel_model(m), sense, level_f)
end
JuMP.objective_sense(m::InnerBilevelModel) = mylevel_obj_sense(m)
JuMP.objective_function_type(m::InnerBilevelModel) = typeof(mylevel_obj_function(m))
JuMP.objective_function(m::InnerBilevelModel) = mylevel_obj_function(m)
function JuMP.objective_function(m::InnerBilevelModel, FT::Type)
    mylevel_obj_function(m) isa FT || error("The objective function is not of type $FT")
    mylevel_obj_function(m)
end

# todo remove
JuMP.objective_sense(m::AbstractBilevelModel) = MOI.FEASIBILITY_SENSE
# end todo remove
JuMP.num_variables(m::AbstractBilevelModel) = JuMP.num_variables(bilevel_model(m))
JuMP.show_constraints_summary(::Any, ::AbstractBilevelModel) = "no summary"
JuMP.show_backend_summary(::Any, ::AbstractBilevelModel) = "no summary"
JuMP.object_dictionary(m::BilevelModel) = m.objdict
JuMP.object_dictionary(m::AbstractBilevelModel) = JuMP.object_dictionary(bilevel_model(m))
JuMP.show_objective_function_summary(::IO, ::AbstractBilevelModel) = "no summary"
# TODO
# function JuMP.constraints_string(print_mode, model::MyModel)

bileve_obj_error() = error("There is no objective for BilevelModel use Upper(.) and Lower(.)")

function JuMP.set_objective(m::BilevelModel, sense::MOI.OptimizationSense,
    f::JuMP.AbstractJuMPScalar)
    bileve_obj_error()
end
JuMP.objective_sense(m::BilevelModel) = JuMP.objective_sense(m.upper)#bileve_obj_error()
JuMP.objective_function_type(model::BilevelModel) = bileve_obj_error()
JuMP.objective_function(model::BilevelModel) = bileve_obj_error()
function JuMP.objective_function(model::BilevelModel, FT::Type)
    bileve_obj_error()
end

# Names
JuMP.name(vref::BilevelVariableRef) = vref.model.varnames[vref.idx]
function JuMP.set_name(vref::BilevelVariableRef, name::String)
    vref.model.varnames[vref.idx] = name
end
JuMP.name(cref::BilevelConstraintRef) = cref.model.connames[cref.idx]
function JuMP.set_name(cref::BilevelConstraintRef, name::String)
    cref.model.connames[cref.idx] = name
end
# TODO
# function JuMP.variable_by_name(model::MyModel, name::String)
# see jump extensions
# function JuMP.constraint_by_name(model::MyModel, name::String)

# replace variables
function replace_variables(var::BilevelVariableRef,
    model::BilevelModel,
    inner::JuMP.AbstractModel,
    variable_map::Dict{Int, V},
    level::Level) where {V<:JuMP.AbstractVariableRef}
    if var.model === model && in_level(var, level)
        return variable_map[var.idx]
    elseif var.model === model
        error("Variable $(var) belonging Only to $(var.level) level, was added in the $(level) level.")
    else
        error("A BilevelModel cannot have expression using variables of a BilevelModel different from itself")
    end
end
function replace_variables(aff::JuMP.GenericAffExpr{C, BilevelVariableRef},
    model::BilevelModel,
    inner::JuMP.AbstractModel,
    variable_map::Dict{Int, V},
    level::Level) where {C,V<:JuMP.AbstractVariableRef}
    result = JuMP.GenericAffExpr{C, JuMP.VariableRef}(0.0)#zero(aff)
    result.constant = aff.constant
    for (coef, var) in JuMP.linear_terms(aff)
        JuMP.add_to_expression!(result,
        coef,
        replace_variables(var, model, model, variable_map, level))
    end
    return result
end
function replace_variables(quad::JuMP.GenericQuadExpr{C, BilevelVariableRef},
    model::BilevelModel,
    inner::JuMP.AbstractModel,
    variable_map::Dict{Int, V},
    level::Level) where {C,V<:JuMP.AbstractVariableRef}
    aff = replace_variables(quad.aff, model, model, variable_map, level)
    quadv = JuMP.GenericQuadExpr{C, JuMP.VariableRef}(aff)
    for (coef, var1, var2) in JuMP.quad_terms(quad)
        JuMP.add_to_expression!(quadv,
        coef,
        replace_variables(var1, model, model, variable_map, level),
        replace_variables(var2, model, model, variable_map, level))
    end
    return quadv
end
replace_variables(funcs::Vector, args...) = map(f -> replace_variables(f, args...), funcs)

function print_lp(m, name)
    dest = MOI.FileFormats.Model(format = MOI.FileFormats.FORMAT_MOF)
    MOI.copy_to(dest, m)
    MOI.write_to_file(dest, name)
end

JuMP.optimize!(::T) where {T<:AbstractBilevelModel} =
    error("cant solve a model of type: $T ")
function JuMP.optimize!(model::BilevelModel, optimizer, mode::BilevelSolverMode{T} = SOS1Mode();
    lower_prob = "", upper_prob = "", bilevel_prob = "", solver_prob = "") where T

    if !MOI.is_empty(optimizer)
        error("An empty optimizer must be provided")
    end

    upper = JuMP.backend(model.upper)
    lower = JuMP.backend(model.lower)

    if length(lower_prob) > 0
        print_lp(lower, lower_prob)
    end
    if length(upper_prob) > 0
        print_lp(upper, upper_prob)
    end

    moi_upper = JuMP.index.(
        collect(values(model.upper_to_lower_link)))
    moi_link = JuMP.index(model.link)
    moi_link2 = index2(model.upper_var_to_lower_ctr_link)

    single_blm, upper_to_sblm, lower_to_sblm, lower_primal_dual_map, lower_dual_to_sblm =
    build_bilevel(upper, lower, moi_link, moi_upper, mode, moi_link2)

    # pass lower level dual variables info (start, upper, lower)
    for (idx, info) in model.ctr_info
        if haskey(model.ctr_lower, idx)
            ctr = model.ctr_lower[idx]
            pre_duals = lower_primal_dual_map.primal_con_dual_var[JuMP.index(ctr)] # vector
            duals = map(x->lower_dual_to_sblm[x], pre_duals)
            pass_dual_info(single_blm, duals, info)
        end
    end
    for (idx, info) in model.var_info
        if haskey(model.var_lower, idx)
            var = lower_to_sblm[JuMP.index(model.var_lower[idx])]
        elseif haskey(model.var_upper, idx)
            var = upper_to_sblm[JuMP.index(model.var_upper[idx])]
        else
            continue
        end
        pass_primal_info(single_blm, var, info)
    end

    if length(bilevel_prob) > 0
        print_lp(single_blm, bilevel_prob)
    end
    # print_lp(single_blm, "bilevel_orig.mof")

    solver = optimizer#MOI.Bridges.full_bridge_optimizer(optimizer, Float64)
    sblm_to_solver = MOI.copy_to(solver, single_blm, copy_names = true)
    # print_lp(solver, "bilevel_bridge.mof")
    # print_lp(solver.model, "bilevel_cache.mof")

    if length(solver_prob) > 0
        print_lp(solver, solver_prob)
    end

    MOI.optimize!(solver)

    model.solver  = solver#::MOI.ModelLike
    model.upper_to_sblm = upper_to_sblm
    model.lower_to_sblm = lower_to_sblm
    model.lower_dual_to_sblm = lower_dual_to_sblm
    model.lower_primal_dual_map = lower_primal_dual_map
    model.sblm_to_solver = sblm_to_solver

    return nothing
end
function pass_dual_info(single_blm, dual, info::ConstraintInfo{Float64})
    if !isnan(info.start)
        MOI.set(single_blm, MOI.VariablePrimalStart(), dual[], info.start)
    end
    if !isnan(info.upper) &&
        !MOI.is_valid(single_blm, CI{SVF,LT{Float64}}(dual[].value))
        MOI.add_constraint(single_blm,
            SVF(dual[]), LT{Float64}(info.upper))
    end
    if !isnan(info.lower) &&
        !MOI.is_valid(single_blm, CI{SVF,GT{Float64}}(dual[].value))
        MOI.add_constraint(single_blm,
            SVF(dual[]), GT{Float64}(info.lower))
    end
    return
end
function pass_dual_info(single_blm, dual, info::ConstraintInfo{Vector{Float64}})
    for i in eachindex(dual)
        if !isnan(info.start[i])
            MOI.set(single_blm, MOI.VariablePrimalStart(), dual[i], info.start[i])
        end
        if !isnan(info.upper[i]) &&
            !MOI.is_valid(single_blm, CI{SVF,LT{Float64}}(dual[i].value))
            MOI.add_constraint(single_blm,
                SVF(dual[i]), LT{Float64}(info.upper[i]))
        end
        if !isnan(info.lower[i]) &&
            !MOI.is_valid(single_blm, CI{SVF,GT{Float64}}(dual[i].value))
            MOI.add_constraint(single_blm,
                SVF(dual[i]), MOI.GreaterThan{Float64}(info.lower[i]))
        end
    end
    return
end

function pass_primal_info(single_blm, primal, info::VariableInfo{Float64})
    if !isnan(info.upper) &&
        !MOI.is_valid(single_blm, CI{SVF,LT{Float64}}(primal.value))
        MOI.add_constraint(single_blm,
            SVF(primal), LT{Float64}(info.upper))
    end
    if !isnan(info.lower) &&
        !MOI.is_valid(single_blm, CI{SVF,GT{Float64}}(primal.value))
        MOI.add_constraint(single_blm,
            SVF(primal), GT{Float64}(info.lower))
    end
    return
end

function JuMP.index(d::Dict)
    ret = Dict{VI,VI}()
    # sizehint!(ret, length(d))
    for (k,v) in d
        ret[JuMP.index(k)] = JuMP.index(v)
    end
    return ret
end
function index2(d::Dict)
    ret = Dict{VI,CI}()
    # sizehint!(ret, length(d))
    for (k,v) in d
        ret[JuMP.index(k)] = JuMP.index(v)
    end
    return ret
end

function JuMP.value(v::BilevelVariableRef)::Float64
    m = owner_model(v)
    solver = m.solver
    ref = solver_ref(v)
    return MOI.get(solver, MOI.VariablePrimal(), ref)
end

function JuMP.dual(cref::BilevelConstraintRef)
    # Right now this code assumes there is no possibility for vectorized constraints
    if my_level(cref) == BilevelJuMP.LOWER_ONLY
        # Constraint index on the lower model
        con_lower_idx = cref.model.ctr_lower[cref.idx].index
        # Dual variable associated with constraint index
        model_var_idxs = cref.model.lower_primal_dual_map.primal_con_dual_var[con_lower_idx]
        # Single bilevel model variable associated with the dual variable
        sblm_var_idxs = MOI.VariableIndex[]
        for vi in model_var_idxs
            push!(sblm_var_idxs, cref.model.lower_dual_to_sblm[vi])
        end
        # Solver variable associated withe the sblm model
        solver_var_idxs = MOI.VariableIndex[]
        for vi in sblm_var_idxs
            push!(solver_var_idxs, cref.model.sblm_to_solver[vi])
        end
        return MOI.get(cref.model.solver, MOI.VariablePrimal(), solver_var_idxs)
    else
        error("Dual solutions of upper level constraints are not available")
    end
end
function JuMP.primal_status(model::BilevelModel)
    return MOI.get(model.solver, MOI.PrimalStatus())
end
JuMP.dual_status(model::BilevelModel) = error("dual status cant be queried for BilevelModel")
function JuMP.termination_status(model::BilevelModel)
    return MOI.get(model.solver, MOI.TerminationStatus())
end
function JuMP.raw_status(model::BilevelModel)
    return MOI.get(model.solver, MOI.RawStatusString())
end
function JuMP.objective_value(model::BilevelModel)
    return MOI.get(model.solver, MOI.ObjectiveValue())
end
function lower_objective_value(model::BilevelModel)
    # Create a dict with lower variables in the lower objective
    vals = Dict()
    for v in model.lower_objective_function.terms.keys
        vals[MOI.VariableIndex(v.idx)] = JuMP.value(v)
    end
    # Evaluate the lower objective expression
    return MOIU.eval_variables(vi -> vals[vi],
        model.lower.moi_backend.model_cache.model.objective)
end


# Boundary condition functions

for f in [:prescribed_surface_jump!,:prescribed_surface_average!]
  @eval function $f(dfb,t::Real,sys::ILMSystem)
          @unpack base_cache, bc, phys_params, motions = sys
          $f(dfb,t,base_cache,bc,phys_params,motions)
        end
  @eval function $f(dfb,sys::ILMSystem)
          @unpack base_cache, bc, phys_params = sys
          $f(dfb,base_cache,bc,phys_params)
        end
end

function prescribed_surface_jump!(dfb,t::Real,base_cache::BasicILMCache{N},bc,phys_params,motions) where N
    dfb .= bc["exterior"](t,base_cache,phys_params,motions)
    dfb .-= bc["interior"](t,base_cache,phys_params,motions)
    return dfb
end

function prescribed_surface_jump!(dfb,base_cache::BasicILMCache{N},bc,phys_params) where N
    dfb .= bc["exterior"](base_cache,phys_params)
    dfb .-= bc["interior"](base_cache,phys_params)
    return dfb
end

function prescribed_surface_jump!(dfb,t::Real,base_cache::BasicILMCache{0},bc,phys_params,motions)
    fill!(dfb,0.0)
    return dfb
end

function prescribed_surface_jump!(dfb,base_cache::BasicILMCache{0},bc,phys_params)
    fill!(dfb,0.0)
    return dfb
end

function prescribed_surface_average!(fb,t::Real,base_cache::BasicILMCache{N},bc,phys_params,motions) where N
    fb .= 0.5*bc["exterior"](t,base_cache,phys_params,motions)
    fb .+= 0.5*bc["interior"](t,base_cache,phys_params,motions)
    return fb
end

function prescribed_surface_average!(fb,t::Real,base_cache::BasicILMCache{0},bc,phys_params,motions)
    fill!(fb,0.0)
    return fb
end

function prescribed_surface_average!(fb,base_cache::BasicILMCache{N},bc,phys_params) where N
    fb .= 0.5*bc["exterior"](base_cache,phys_params)
    fb .+= 0.5*bc["interior"](base_cache,phys_params)
    return fb
end

function prescribed_surface_average!(fb,base_cache::BasicILMCache{0},bc,phys_params)
    fill!(fb,0.0)
    return fb
end


"""
    $(TYPEDEF)

Make a two-dimensional bipedal robot walk forward as fast as possible. The robot model is
based on the work T. Erez, Y. Tassa, and E. Todorov: ["Infinite Horizon Model Predictive
Control for Nonlinear Periodic Tasks"](https://homes.cs.washington.edu/~todorov/papers/ErezRSS11.pdf), 2011.

* **State: (29, )**
* **Action: (6, )**
* **Observation: (7, )**
"""
mutable struct Walker2DV2{Sim, SSpace, OSpace} <: AbstractMuJoCoEnvironment
    sim::Sim
    statespace::SSpace
    observationspace::OSpace
    last_torso_x::Float64
    randreset_distribution::Uniform{Float64}
    function Walker2DV2(sim::MJSim)
        sspace = MultiShape(
            simstate=statespace(sim),
            last_torso_x=ScalarShape(Float64)
        )
        ospace = MultiShape(
            cropped_qpos = VectorShape(Float64, sim.m.nq - 1),
            qvel = VectorShape(Float64, sim.m.nv)
        )
        env = new{typeof(sim), typeof(sspace), typeof(ospace)}(
            sim, sspace, ospace, 0, Uniform(-0.005, 0.005)
        )
        reset!(env)
    end
end

Walker2DV2() = first(tconstruct(Walker2DV2, 1))

function LyceumBase.tconstruct(::Type{Walker2DV2}, n::Integer)
    modelpath = joinpath(@__DIR__, "walker2d.xml")
    Tuple(Walker2DV2(s) for s in tconstruct(MJSim, n, modelpath, skip=4))
end

@inline getsim(env::Walker2DV2) = env.sim

@inline LyceumBase.statespace(env::Walker2DV2) = env.statespace

function LyceumBase.getstate!(state, env::Walker2DV2)
    checkaxes(statespace(env), state)
    @uviews state begin
        shaped = statespace(env)(state)
        getstate!(shaped.simstate, env.sim)
        shaped.last_torso_x = env.last_torso_x
    end
    state
end

function LyceumBase.setstate!(env::Walker2DV2, state)
    checkaxes(statespace(env), state)
    @uviews state begin
        shaped = statespace(env)(state)
        setstate!(env.sim, shaped.simstate)
        env.last_torso_x = shaped.last_torso_x
    end
    env
end

@inline LyceumBase.observationspace(env::Walker2DV2) = env.observationspace

function LyceumBase.getobservation!(obs, env::Walker2DV2)
    checkaxes(observationspace(env), obs)
    qpos = env.sim.d.qpos
    @views @uviews qpos obs begin
        shaped = observationspace(env)(obs)
        copyto!(shaped.cropped_qpos, qpos[2:end])
        copyto!(shaped.qvel, env.sim.d.qvel)
        clamp!(shaped.qvel, -10, 10)
    end
    obs
end

function LyceumBase.getreward(state, action, ::Any, env::Walker2DV2)
    checkaxes(statespace(env), state)
    checkaxes(actionspace(env), action)
    @uviews state begin
        shapedstate = statespace(env)(state)
        alive_bonus = 1.0
        reward = (_torso_x(shapedstate, env) - shapedstate.last_torso_x) / timestep(env)
        reward += alive_bonus
        reward -= 1e-3 * sum(x->x^2, action)
        reward
    end
end

function LyceumBase.reset!(env::Walker2DV2)
    reset!(env.sim)
    env.last_torso_x = _torso_x(env)
    env
end

function LyceumBase.randreset!(rng::AbstractRNG, env::Walker2DV2)
    reset_nofwd!(env.sim)
    perturb!(rng, env.randreset_distribution, env.sim.d.qpos)
    perturb!(rng, env.randreset_distribution, env.sim.d.qvel)
    forward!(env.sim)
    env.last_torso_x = _torso_x(env)
    env
end

function LyceumBase.step!(env::Walker2DV2)
    env.last_torso_x = _torso_x(env)
    step!(env.sim)
    env
end

function LyceumBase.isdone(state, o, env::Walker2DV2)
    checkaxes(statespace(env), state)
    @uviews state begin
        shapedstate = statespace(env)(state)
        torso_x = _torso_x(shapedstate, env)
        height = _torso_height(shapedstate, env)
        torso_ang = _torso_ang(shapedstate, env)
        qpos = shapedstate.simstate.qpos
        qvel = shapedstate.simstate.qvel

        done = !(
            height > 0.8
            && height < 2.0
            && abs(torso_ang) < 1.0
        )
        done
    end
end

@inline _torso_x(shapedstate::ShapedView, ::Walker2DV2) = shapedstate.simstate.qpos[1]
@inline _torso_x(env::Walker2DV2) = env.sim.d.qpos[1]
@inline _torso_height(shapedstate::ShapedView, ::Walker2DV2) = shapedstate.simstate.qpos[2]
@inline _torso_ang(shapedstate::ShapedView, ::Walker2DV2) = shapedstate.simstate.qpos[3]


# This file is a part of BAT.jl, licensed under the MIT License (MIT).

using BAT
using Test

using ValueShapes, Distributions

@testset "variate_transform" begin
    @testset "identity_vt" begin
        @test @inferred(BAT.IdentityVT(ScalarShape{Real}())) isa BAT.VariateTransform{<:ScalarShape,<:ScalarShape}
        @test @inferred(BAT.IdentityVT(ArrayShape{Int}(2))) isa BAT.VariateTransform{<:ArrayShape,<:ArrayShape}
        @test @inferred(BAT.IdentityVT(ArrayShape{Float32}(2, 3))) isa BAT.VariateTransform{<:ArrayShape,<:ArrayShape}

        ntshape = NamedTupleShape(a = ScalarShape{Real}(), b = ArrayShape{Int}(2))
        @test @inferred(BAT.IdentityVT(ntshape)) isa BAT.VariateTransform{<:ValueShapes.NamedTupleShape{(:a, :b)},<:ValueShapes.NamedTupleShape{(:a, :b)}}

        nttrafo = BAT.IdentityVT(ntshape)
        ntvalue = (a = 4.2, b = [5, 7])

        @test @inferred(varshape(nttrafo)) == ntshape
        @test @inferred(inv(nttrafo)) === nttrafo

        @test @inferred(BAT.apply_vartrafo(nttrafo, ntvalue, 0)).v === ntvalue
        @test BAT.apply_vartrafo(nttrafo, ntvalue, 0.79).ladj == 0.79
        @test ismissing(BAT.apply_vartrafo(nttrafo, ntvalue, missing).ladj)
        @test @inferred(BAT.apply_vartrafo(inv(nttrafo), ntvalue, 0.79)).v === ntvalue
        @test BAT.apply_vartrafo(inv(nttrafo), ntvalue, 0.79).ladj == 0.79
        @test ismissing(BAT.apply_vartrafo(inv(nttrafo), ntvalue, missing).ladj)

        @test @inferred(nttrafo ∘ nttrafo) === nttrafo

        arrayshape_trafo = @inferred(BAT.IdentityVT(ArrayShape{Real}(1)))
        hoeldertable_trafo = @inferred((arrayshape_trafo)(HoelderTable(), volcorr=Val(false)))

        @test hoeldertable_trafo.volcorr == @inferred(BAT.TDNoCorr())
        @test hoeldertable_trafo.trafo isa BAT.IdentityVT
    end
end


function SubPlane(arg0::Hyperplane, arg1::Region)
    return SubPlane((Hyperplane, Region), arg0, arg1)
end

function apply_transform(obj::AbstractSubHyperplane, arg0::Transform)
    return jcall(obj, "applyTransform", AbstractSubHyperplane, (Transform,), arg0)
end

function copy_self(obj::AbstractSubHyperplane)
    return jcall(obj, "copySelf", AbstractSubHyperplane, ())
end

function get_hyperplane(obj::AbstractSubHyperplane)
    return jcall(obj, "getHyperplane", Hyperplane, ())
end

function get_remaining_region(obj::AbstractSubHyperplane)
    return jcall(obj, "getRemainingRegion", Region, ())
end

function get_size(obj::AbstractSubHyperplane)
    return jcall(obj, "getSize", jdouble, ())
end

function is_empty(obj::AbstractSubHyperplane)
    return jcall(obj, "isEmpty", jboolean, ())
end

function reunite(obj::AbstractSubHyperplane, arg0::SubHyperplane)
    return jcall(obj, "reunite", AbstractSubHyperplane, (SubHyperplane,), arg0)
end

function split(obj::SubPlane, arg0::Hyperplane)
    return jcall(obj, "split", SubHyperplane_SplitSubHyperplane, (Hyperplane,), arg0)
end


"""
    visualize(mechanism, storage; vis, build, show_contact, animation, color, name)

    visualize mechanism using trajectory from storage

    mechanism: Mechanism
    storage: Storage
    vis: Visualizer
    build: flag to construct mechanism visuals (only needs to be built once)
    show_contact: flag to show contact locations on system
    color: RGBA
    name: unique identifier for mechanism
"""
function visualize(mechanism::Mechanism, storage::Storage{T,N}; vis::Visualizer=Visualizer(),
    build::Bool=true,
    show_joint=false,
    joint_radius=0.1,
    show_contact=false,
    animation=nothing,
    color=nothing,
    name::Symbol=:robot) where {T,N}

    storage = deepcopy(storage)
    bodies = mechanism.bodies
    origin = mechanism.origin

    # Build robot in the visualizer
    build && build_robot(mechanism,
        vis=vis,
        show_joint=show_joint,
        show_contact=show_contact,
        color=color,
        name=name)

    # Create animations
    framerate = Int64(round(1/mechanism.timestep))
    (animation == nothing) && (animation =
        MeshCat.Animation(Dict{MeshCat.SceneTrees.Path,MeshCat.AnimationClip}(), framerate))

    # Bodies and Contacts
    for (id, body) in enumerate(bodies)
        shape = body.shape
        visshape = convert_shape(shape)
        subvisshape = nothing
        showshape = false
        if visshape !== nothing
            subvisshape = vis[name][:bodies][Symbol(body.name, "__id_$id")]
            showshape = true
        end

        animate_node!(storage, id, shape, animation, subvisshape, showshape)

        if show_joint
            for (jd, joint) in enumerate(mechanism.joints)
                if joint.child_id == body.id
                    radius = 0.1
                    joint_shape = Sphere(radius,
                        position_offset=joint.translational.vertices[2],
                        color=RGBA(0.0, 0.0, 1.0, 0.5))
                    visshape = convert_shape(joint_shape)
                    subvisshape = nothing
                    showshape = false
                    if visshape !== nothing
                        subvisshape = vis[name][:joints][Symbol(joint.name, "__id_$(jd)")]
                        showshape = true
                    end
                    animate_node!(storage, id, joint_shape, animation, subvisshape, showshape)
                end
            end
        end

        if show_contact
            for (jd, contact) in enumerate(mechanism.contacts)
                if contact.parent_id == body.id
                    radius = abs(contact.model.collision.contact_radius)
                    (radius == 0.0) && (radius = 0.01)
                    contact_shape = Sphere(radius,
                        position_offset=contact.model.collision.contact_origin, #TODO: generalize for collision checking
                        axis_offset=one(Quaternion),
                        color=RGBA(1.0, 0.0, 0.0, 0.5))
                    visshape = convert_shape(contact_shape)
                    subvisshape = nothing
                    showshape = false
                    if visshape !== nothing
                        subvisshape = vis[name][:contacts][Symbol(contact.name, "__id_$(jd)")]
                        showshape = true
                    end
                    animate_node!(storage, id, contact_shape, animation, subvisshape, showshape)
                end
            end
        end
    end

    # Origin
    id = origin.id
    shape = origin.shape
    visshape = convert_shape(shape)
    if visshape !== nothing
        subvisshape = vis[name][:bodies][Symbol(:origin, "_id")]
        shapetransform = transform(szeros(T,3), one(Quaternion{T}), shape)
        settransform!(subvisshape, shapetransform)
    end

    setanimation!(vis, animation)
    return vis, animation
end

"""
    build_robot(mechanism; vis, show_contact, name, color)

    construct visuals for mechanism

    mechanism: Mechanism
    vis: Visualizer
    show_contact: flag to show contact locations on mechanism
    name: unique identifier
    color: RGBA
"""
function build_robot(mechanism::Mechanism;
    vis::Visualizer=Visualizer(),
    show_joint=false,
    joint_radius=0.1,
    show_contact=false,
    name::Symbol=:robot,
    color=nothing) where {T,N}

    bodies = mechanism.bodies
    origin = mechanism.origin
    set_background!(vis)
    set_light!(vis)
    set_floor!(vis)

    # Bodies and Contacts
    for (id, body) in enumerate(bodies)
        if color !== nothing
            shape = deepcopy(body.shape)
            set_color!(shape, color)
        else
            shape = body.shape
        end
        visshape = convert_shape(shape)
        subvisshape = nothing
        if visshape !== nothing
            subvisshape = vis[name][:bodies][Symbol(body.name, "__id_$id")]
            setobject!(subvisshape, visshape, shape,
                transparent=(show_joint || show_contact))
        end

        if show_joint
            for (jd, joint) in enumerate(mechanism.joints)
                if joint.child_id == body.id
                    radius = joint_radius
                    joint_shape = Sphere(radius,
                        position_offset=joint.translational.vertices[2],
                        color=RGBA(0.0, 0.0, 1.0, 0.5))
                    visshape = convert_shape(joint_shape)
                    subvisshape = nothing
                    if visshape !== nothing
                        subvisshape = vis[name][:joints][Symbol(joint.name, "__id_$(jd)")]
                        setobject!(subvisshape, visshape, joint_shape,
                            transparent=false)
                    end
                end
            end
        end

        if show_contact
            for (jd, contact) in enumerate(mechanism.contacts)
                if contact.parent_id == body.id
                    radius = abs(contact.model.collision.contact_radius)
                    (radius == 0.0) && (radius = 0.01)
                    contact_shape = Sphere(radius,
                        position_offset=(contact.model.collision.contact_origin),
                        axis_offset=one(Quaternion), color=RGBA(1.0, 0.0, 0.0, 0.5))
                    visshape = convert_shape(contact_shape)
                    subvisshape = nothing
                    if visshape !== nothing
                        subvisshape = vis[name][:contacts][Symbol(contact.name, "__id_$(jd)")]
                        setobject!(subvisshape,visshape,contact_shape,transparent=false)
                    end
                end
            end
        end
    end

    # Origin
    id = origin.id
    shape = origin.shape
    visshape = convert_shape(shape)
    if visshape !== nothing
        subvisshape = vis[name][:bodies][Symbol(:origin, "_id")]
        setobject!(subvisshape,visshape,shape,transparent=(show_joint || show_contact))
    end
    return vis
end

"""
    set_robot(vis, mechanism, z; show_contact, name)

    visualze mechanism configuration from maximal representation

    vis: Visualizer
    mechanism: Mechanism
    z: maximal state
    show_contact: flag to show contact locations on mechanism
    name: unique identifier
"""
function set_robot(vis::Visualizer, mechanism::Mechanism, z::Vector{T};
    show_joint::Bool=false,
    joint_radius=0.1,
    show_contact::Bool=true,
    name::Symbol=:robot) where {T,N}

    (length(z) == minimal_dimension(mechanism)) && (z = minimal_to_maximal(mechanism, z))
    bodies = mechanism.bodies
    origin = mechanism.origin

    # Bodies and Contacts
    for (id, body) in enumerate(bodies)
        x, _, q, _ = unpack_maximal_state(z, id)
        shape = body.shape
        visshape = convert_shape(shape)
        subvisshape = nothing
        showshape = false
        if visshape !== nothing
            subvisshape = vis[name][:bodies][Symbol(body.name, "__id_$id")]
            showshape = true
        end

        set_node!(x, q, id, shape, subvisshape, showshape)

        if show_joint
            for (jd, joint) in enumerate(mechanism.joints)
                if joint.child_id == body.id
                    radius = joint_radius
                    joint_shape = Sphere(radius,
                        position_offset=joint.translational.vertices[2],
                        color=RGBA(0.0, 0.0, 1.0, 0.5))
                    visshape = convert_shape(joint_shape)
                    subvisshape = nothing
                    showshape = false
                    if visshape !== nothing
                        subvisshape = vis[name][:joints][Symbol(joint.name, "__id_$(jd)")]
                        showshape = true
                    end
                    set_node!(x, q, id, joint_shape, subvisshape, showshape)
                end
            end
        end

        if show_contact
            for (jd, contact) in enumerate(mechanism.contacts)
                if contact.parent_id == body.id
                    radius = abs(contact.model.collision.contact_radius)
                    (radius == 0.0) && (radius = 0.01)
                    contact_shape = Sphere(radius,
                        position_offset=(contact.model.collision.contact_origin),
                        axis_offset=one(Quaternion), color=RGBA(1.0, 0.0, 0.0, 0.5))
                    visshape = convert_shape(contact_shape)
                    subvisshape = nothing
                    showshape = false
                    if visshape !== nothing
                        subvisshape = vis[name][:contacts][Symbol(contact.name, "__id_$(jd)")]
                        showshape = true
                    end
                    set_node!(x, q, id, contact_shape, subvisshape, showshape)
                end
            end
        end
    end

    # Origin
    id = origin.id
    shape = origin.shape
    visshape = convert_shape(shape)
    if visshape !== nothing
        subvisshape = vis[name][:bodies][Symbol(:origin, "_id")]
        shapetransform = transform(szeros(T,3), one(Quaternion{T}), shape)
        settransform!(subvisshape, shapetransform)
    end
    return vis
end

function transform(x, q, shape)
    scale_transform = MeshCat.LinearMap(diagm(shape.scale))
    x_transform = MeshCat.Translation(x + vector_rotate(shape.position_offset, q))
    q_transform = MeshCat.LinearMap(q * shape.axis_offset)
    return MeshCat.compose(x_transform, q_transform, scale_transform)
end

MeshCat.js_scaling(s::AbstractVector) = s
MeshCat.js_position(p::AbstractVector) = p

function set_node!(x, q, id, shape, shapevisualizer, showshape) where {T,N}
    if showshape
        # TODO currently setting props directly because MeshCat/Rotations doesn't convert scaled rotation properly.
        # If this changes, do similarily to origin
        setprop!(shapevisualizer, "scale", MeshCat.js_scaling(shape.scale))
        setprop!(shapevisualizer, "position", MeshCat.js_position(x + vector_rotate(shape.position_offset, q)))
        setprop!(shapevisualizer, "quaternion", MeshCat.js_quaternion(q * shape.axis_offset))
    end
    return
end

function animate_node!(storage::Storage{T,N}, id, shape, animation, shapevisualizer, showshape) where {T,N}
    for i=1:N
        x = storage.x[id][i]
        q = storage.q[id][i]
        atframe(animation, i) do
            set_node!(x, q, id, shape, shapevisualizer, showshape)
        end
    end
    return
end

function MeshCat.setobject!(subvisshape, visshape, shapes::Shapes;
    transparent=false)
    for (i, s) in enumerate(shapes.shape)
        v = subvisshape["node_$i"]
        setobject!(v, visshape[i], s, transparent=transparent)
        scale_transform = MeshCat.LinearMap(diagm(s.scale))
        x_transform = MeshCat.Translation(s.position_offset)
        q_transform = MeshCat.LinearMap(s.axis_offset)
        t = MeshCat.compose(x_transform, q_transform, scale_transform)
        settransform!(v, t)
    end
end

function MeshCat.setobject!(subvisshape, visshape, shape::Shape;
    transparent=false)
    setobject!(subvisshape, visshape, MeshPhongMaterial(color=(transparent ? RGBA(0.75, 0.75, 0.75, 0.5) : shape.color)))
end

function MeshCat.setobject!(subvisshape, visshape::Vector, shape::Capsule; transparent=false)
    setobject!(subvisshape["cylinder"], visshape[1], MeshPhongMaterial(color=(transparent ? RGBA(0.75, 0.75, 0.75, 0.5) : shape.color)))
    setobject!(subvisshape["cap1"], visshape[2], MeshPhongMaterial(color=(transparent ? RGBA(0.75, 0.75, 0.75, 0.5) : shape.color)))
    setobject!(subvisshape["cap2"], visshape[3], MeshPhongMaterial(color=(transparent ? RGBA(0.75, 0.75, 0.75, 0.5) : shape.color)))
end

function MeshCat.setobject!(subvisshape, visshape, shape::Mesh; transparent=false)
    if visshape.mtl_library == ""
        visshape = MeshFileGeometry(visshape.contents, visshape.format)
        setobject!(subvisshape, visshape, MeshPhongMaterial(color=shape.color))
    else
        setobject!(subvisshape, visshape)
    end
end


# So we can access "binsize" before a `CompactHeap` is instantiated.
struct PowerOfTwo
    val::Int
end

poweroftwo(x::PowerOfTwo) = x
poweroftwo(x::Integer) = PowerOfTwo(x)

getbin(x::PowerOfTwo, sz) = ((sz - 1) >> x.val) + 1
binsize(x::PowerOfTwo, bin) = bin << x.val

#####
##### Iterate over a vector of `Blocks` as a vector of `FreelistPtr`.
#####

struct FreelistIterator{T}
    lists::Vector{T}
end

Base.length(x::FreelistIterator) = length(x.lists)
Base.pointer(x::FreelistIterator, i) = pointer(x.lists, i)
function Base.iterate(x::FreelistIterator, i = 1)
    i > length(x) && return nothing
    return FreelistPtr(pointer(x, i)), i + 1
end

#####
##### CompactHeap
#####

mutable struct CompactHeap <: AbstractHeap
    minallocation::PowerOfTwo

    # Where is and what kind of memory do we have?
    base::Ptr{Nothing}
    len::Int
    allocated::Int
    pool::Pool

    # Maintaining freelist status
    status::FindNextTree
    freelists::Vector{Block}
    canalloc::Bool

    # Allocator token.
    # Can be used to preserve `base` as long as the heap is still alive.
    token::Any
end

topointer(ptr::Ptr) = ptr
topointer(A::AbstractArray) = convert(Ptr{Nothing}, pointer(A))
getstate(heap::CompactHeap) = (heap.allocated, heap.len)

getfreelist(heap::CompactHeap, i) = FreelistPtr(pointer(heap.freelists, i))
freelists(heap::CompactHeap) = FreelistIterator(heap.freelists)

function CompactHeap(
    allocator::AbstractAllocator,
    sz;
    pool = Local,
    # Power of two
    minallocation = PowerOfTwo(21),
)
    minallocation = poweroftwo(minallocation)

    # Round down size to the nearest multiple of `minallocation`
    sz = (sz >> minallocation.val) << minallocation.val

    # Allocate the memory managed by this heap
    base, token = allocate(allocator, sz)
    allocated = 0

    numbins = getbin(minallocation, sz)
    status = FindNextTree(numbins)
    freelists = [Block() for _ in Base.OneTo(numbins)]

    heap = CompactHeap(
        minallocation,
        base,
        sz,
        allocated,
        pool,
        status,
        freelists,
        true,
        token,
    )

    finalizer(heap) do x
        free(x.base, x.token)
    end

    # Add an entry for the biggest freelist.
    block = Block(base)
    block.size = sz
    block.backsize = 0
    block.free = true
    block.pool = pool
    block.next = Block()
    block.previous = Block()

    push!(heap, block)
    return heap
end

getbin(heap::CompactHeap, sz) = getbin(heap.minallocation, sz)
getbin(heap::CompactHeap, block::Block) = getbin(heap, block.size)

binsize(heap::CompactHeap, bin) = binsize(heap.minallocation, bin)

basepointer(heap::CompactHeap) = heap.base
Base.sizeof(heap::CompactHeap) = heap.len
baseaddress(heap::CompactHeap) = convert(UInt, heap.base)
numbins(heap::CompactHeap) = length(heap.freelists)

# The CompactHeap maintains backsizes, so we can also walk backwards.
function walkprevious(::CompactHeap, block)
    # Invariant, the first block will always have a backsize of zero.
    backsize = block.backsize
    if iszero(backsize)
        return nothing
    else
        return Block(pointer(block) - backsize)
    end
end

"""
    sortfreelists!(heap)

Sort all of the freelists in `heap` in order of increasing address.
"""
function sortfreelists!(heap)
    for freelist in freelists(heap)
        sort!(freelist)
    end
end

function Base.push!(heap::CompactHeap, block::Block)
    bin = getbin(heap, block)
    list = @inbounds(getfreelist(heap, bin))

    isempty(list) && @inbounds(setentry!(heap.status, bin))
    push!(list, block)
    return nothing
end

popfreelist!(heap::CompactHeap, bin::Integer) = popfreelist!(heap, getfreelist(heap, bin), bin)
Base.@propagate_inbounds function popfreelist!(
    heap::CompactHeap,
    list::AbstractFreelist,
    bin::Integer,
)
    block = pop!(list)
    isempty(list) && clearentry!(heap.status, bin)
    return block
end

# Pop a block for this bin.
# If a freelist doesn't exist for this bin, split the next largest block.
#
# Assume `sz` is a multiple of the heap granularity
function Base.pop!(heap::CompactHeap, bin)
    if isset(heap.status, bin)
        return @inbounds(popfreelist!(heap, bin))
    end
    # list = @inbounds getfreelist(heap, bin)
    # if !isempty(list)
    #     return @inbounds(popfreelist!(heap, list, bin))
    # end

    # Okay, we don't have a block.
    # We need to get the next highest available block and split it.
    nextbin = @inbounds(findnext(heap.status, bin))

    # No bigger block available.
    # This means we can't fulfill this request.
    # Return `nothing` to indicate this.
    nextbin === nothing && return nothing

    # Otherwise, take this block and split it.
    # TODO: Maybe set a limit where we don't split if the over-allocation amount is
    # less than this limit.
    block = @inbounds(popfreelist!(heap, nextbin))

    # Now that we have the block, split it down to the size we want and put the remainder
    # back in the heap.
    sz = binsize(heap, bin)
    remainder = Block(pointer(block) + sz)

    remainder.free = true
    remainder.size = block.size - sz
    remainder.next = Block()
    remainder.previous = Block()

    # Since we split this block, we need to maintain the backsize pointers for efficient
    # merging.
    block.size = sz
    remainder.backsize = block.size
    next = walknext(heap, remainder)
    if next !== nothing
        next.backsize = remainder.size
    end

    push!(heap, remainder)
    return block
end

function canpop(heap::CompactHeap, bin)
    isset(heap.status, bin) && return true

    # Okay, we don't have a block.
    # We need to get the next highest available block and split it.
    nextbin = @inbounds(findnext(heap.status, bin))

    # No bigger block available - cannot allocate.
    return nextbin === nothing ? false : true
end

function remove!(heap::CompactHeap, block::Block)
    bin = getbin(heap, block)
    list = getfreelist(heap, bin)
    remove!(list, block)
    isempty(list) && @inbounds(clearentry!(heap.status, bin))
    return nothing
end

function putback!(heap::CompactHeap, block::Block)
    @check isfree(block)

    # Check the block before and after are free.
    #
    # If so, merge the blocks
    previous = walkprevious(heap, block)
    next = walknext(heap, block)

    if previous !== nothing && isfree(previous)
        remove!(heap, previous)
        previous.size += block.size

        # Purge metadata for this block to try and catch errors earlier.
        zerometa!(block)
        block = previous
    end

    if next !== nothing
        # Similar to the case with "block.previous", if the next block is being evicted,
        # we don't want to try to merge since that will break eviction logic.
        if isfree(next)
            remove!(heap, next)
            block.size += next.size

            # Set the backedge for the nextnext block
            nextnext = walknext(heap, next)
            if nextnext !== nothing
                nextnext.backsize = block.size
            end

            # Purge metadata for the next block to try and catch errors earlier.
            zerometa!(next)
        else
            next.backsize = block.size
        end
    end

    push!(heap, block)
    return nothing
end

#####
##### High level API
#####

const CAN_ALLOC_MESSAGE = """
Something went wrong. Trying to allocate from a heap that has been temporarily marked
as "no-allocate". It's likely something went wrong during eviction.
"""

function alloc(heap::CompactHeap, bytes::Integer, id::UInt)
    heap.canalloc || error(CAN_ALLOC_MESSAGE)
    iszero(bytes) && return nothing

    needed_size = max(bytes + headersize(), one(bytes) << heap.minallocation.val)
    needed_size > sizeof(heap) && return nothing
    bin = getbin(heap, needed_size)
    block = pop!(heap, bin)
    block === nothing && return nothing

    # Setup the default state for the block
    heap.allocated += sizeof(block)
    clearbits!(block)
    block.id = id
    block.pool = heap.pool

    # Preemptively mark as dirty.
    # Up to callee's to determine if this is actually clean or not.
    block.dirty = true
    ptr = pointer(block) + headersize()
    return ptr
end

@inline function canalloc(heap::CompactHeap, bytes::Integer)
    iszero(bytes) && return false

    needed_size = max(bytes + headersize(), one(bytes) << heap.minallocation.val)
    needed_size > sizeof(heap) && return false
    bin = getbin(heap, needed_size)
    return canpop(heap, bin)
end

function free(heap::CompactHeap, ptr::Ptr{Nothing})
    block = unsafe_block(ptr)
    if block.evicting == true
        block.orphaned = true
        return nothing
    end
    block.queued = false
    block.free = true
    heap.allocated -= sizeof(block)
    putback!(heap, block)
    return nothing
end

#####
##### Defrag API
#####

function nextfree(heap, block = baseblock(heap))
    while block !== nothing && !isfree(block)
        block = walknext(heap, block)
    end
    return block
end

"""
    defrag!(f::F, heap::CompactHeap; [nthreads], [queued_callback])

Defragment `heap`. Callback function `f` is given an ID and a new block.
If a block is marked as `queued`, then, by default, we can't touch it.
However, the caller may provide a zero argument `queued_callback` which must free this
block, in which case defragmentation may continue.
"""
function defrag!(
    f::F,
    heap::CompactHeap;
    nthreads = Threads.nthreads(),
    queued_callback = nothing,
) where {F}
    @checknothing freeblock = nextfree(heap)
    @checknothing block = walknext(heap, freeblock)

    while true
        # By the invariants we keep on the heap, the next block should be not-free.
        # However, it MIGHT be queued for deletion, in which case, we can't move it.
        if block.queued
            if queued_callback !== nothing
                queued_callback()
                # Recurse to begin the walk again since it's possible that blocks
                # we've already visited have now been cleaned up.
                # println("REWALKING!")
                defrag!(f, heap; nthreads, queued_callback)
                return nothing
            else
                @checknothing freeblock = nextfree(heap, block)
                @checknothing block = walknext(heap, freeblock)
            end
            continue
        end

        # The following operations need to be atomic as far as the GC is concerned.
        # This takes care of if the GC runs while we're moving data and the block we
        # are moving suddenly becomes free ...

        # Now, we need to move "block" to "freeblock" and then fixup the heap.
        id = block.id
        # safeprint("Defrag: Moving $id")
        pool = block.pool
        block.backsize = freeblock.backsize
        freeblock_size = freeblock.size
        remove!(heap, freeblock)

        # After the "memcpy", the new location for "block" is "freeblock".
        # Thus, we need to perform the callback so that references can be updated.
        aliases = pointer(freeblock) + sizeof(block) > pointer(block)
        f(id, freeblock, block)
        if aliases
            _memcpy!(pointer(freeblock), pointer(block), sizeof(block); nthreads = nothing)
        else
            _memcpy!(pointer(freeblock), pointer(block), sizeof(block); nthreads)
        end

        @check freeblock.id == id

        # Update siblings
        sibling = getsibling(freeblock)
        if sibling !== nothing
            sibling.sibling = freeblock
        end

        # Since we haven't changed the size of the block, "walknext" should still work,
        # though the actual block returned will be garbage.
        block = freeblock
        freeblock = walknext(heap, block)

        freeblock.backsize = block.size
        freeblock.size = freeblock_size
        freeblock.free = true

        # TODO: In all likelyhood, we're just going to be removing the freeblock again.
        # However, dealing with queued blocks makes this tricky so do the conservative
        # thing for now.
        putback!(heap, freeblock)

        @checknothing block = walknext(heap, freeblock)
        block.backsize = freeblock.size
    end
    return nothing
end

#####
##### Eviction API
#####

function unsafe_free!(f::F, block::Block) where {F}
    @check !isfree(block)

    # Mark this block as being evicted.
    # This will short-circuit any calls to `free` the result from the callback while
    # still keeping the block as `not-free` to be compatible with the other heap logic.
    block.evicting = true

    # Do the callback on the block and then mark the blocked as freed.
    return f(block)
end

function evict!(f::F, heap::CompactHeap, block::Block) where {F}
    if isfree(block)
        remove!(heap, block)
        block.free = false
        heap.allocated += sizeof(block)
        return nothing
    else
        return unsafe_free!(f, block)
    end
end

# Helper method to get a block from a pointer.
function evictfrom!(heap::CompactHeap, ptr::Ptr{Nothing}, sz; kw...)
    evictfrom!(heap, unsafe_block(ptr), sz; kw...)
end

function evictfrom!(heap::CompactHeap, block::Block, sz; cb = donothing)
    heap.canalloc = false
    sizefreed = 0

    last = block
    current = block
    aborted = false

    # NOTE ON ABORTING.
    # Control flow must always pass through the bottom of this function to ensure
    # that even if an eviction is aborted early that we maintain the heap in a
    # consistent state.

    # Walk forward until either we have freed enough space or reach the end of the heap.
    while true
        abort = evict!(cb, heap, current)

        # Check for early abort
        if (abort === true)
            # Cleanup eviction status
            current.evicting = false

            # Aborted on the first block.
            # If this block was NOT orphaned as a side-effect of the eviction callback,
            # then we need to rollback the current block to the last successfully
            # processed block.
            #
            # However, if this block was recaimed, then it is free but hasn't been
            # put back on the heap.
            #
            # As a result, we need to include it in the freed frontier.
            if current.orphaned
                sizefreed += current.size
            else
                if iszero(sizefreed)
                    heap.canalloc = true
                    return nothing
                end
                current = last
            end

            @goto cleanup
        end
        sizefreed += current.size

        # Update state
        last = current
        current = walknext(heap, current)
        # If we've either freed enough memory or walked off the end of the array, exit.
        (sizefreed >= sz || current === nothing) && break
    end

    # Walk backwards if we have to in order to free enough space.
    while sizefreed < sz
        candidate = walkprevious(heap, block)
        candidate === nothing && break
        # Intenionally overwrite the argument "block".
        block = candidate
        abort = evict!(cb, heap, block)
        if (abort === true)
            block.evicting = false
            if !block.orphaned
                # Roll back block to last successful eviction.
                block = walknext(heap, block)
            else
                sizefreed += block.size
            end
            @goto cleanup
        end
        sizefreed += block.size
    end

    # Resize our block and put it back in the heap.
    @label cleanup
    block.free = true
    block.orphaned = false
    block.size = sizefreed
    heap.allocated -= sizefreed

    putback!(heap, block)
    heap.canalloc = true
    return nothing
end

#####
##### Materializer
#####

# Touch all the pages in an allocation.
function materialize_os_pages!(heap::CompactHeap)
    # First, we must make sure that NOTHING is allocated in the heap - otherwise, we might
    # over write some existing data.
    if length(heap) != 1 || !isfree(first(heap))
        throw(error("Can only materialize OS pages for an empty heap!"))
    end

    # Now, we threaded touch all pages in the heap
    #
    # Take steps of 4096, which is the smallest possible page size.
    ptr = convert(Ptr{UInt8}, datapointer(baseblock(heap)))
    Polyester.@batch per = core for i = 1:4096:sizeof(heap)
        unsafe_store!(ptr, one(UInt8), i)
    end
    return nothing
end

function touchheap(heap::CompactHeap)
    v = zeros(UInt8, Threads.nthreads())
    ptr = convert(Ptr{UInt8}, datapointer(baseblock(heap)))
    Polyester.@batch per = core for i = 1:64:sizeof(heap)
        v[Threads.threadid()] += unsafe_load(ptr, i)
    end
    return sum(v)
end

#####
##### Checker
#####

# Check the various invariants of the CompactHeap
#
# 1. First Block should always have zero backsize
# 2. Any freelist entry in the freelist dict must have an entry in the `status` tree.
# 3. All non-entries in the `status` tree must not have entry in the freelist dict.
# 4. Forward and backsizes must be consistent.
# 5. Free blocks found by walking the heap must be consistent with free-blocks found walking
#    the freelists

function first_block_invariant(heap::CompactHeap)
    block = baseblock(heap)
    if !iszero(block.backsize)
        println("Base block has nonzero backsize: $(block.backsize)")
        return false
    end
    return true
end

function freelist_status_invariant(heap::CompactHeap)
    maxbin = getbin(heap, sizeof(heap))
    passed = true

    # Walk the `status` tree and correlate the entries with the freelist dict.
    for (i, mask) in enumerate(last(heap.status.runs))
        for j = 1:64
            bin = 64 * (i - 1) + j
            bin > maxbin && return passed

            if hasentryat(mask, j)
                # There must be a key in the `freelists` dictionary and it must be non-empty
                if isempty(getfreelist(heap, bin))
                    println(
                        "Heap has a `status` entry for bin $bin but no corresponding freelist bound",
                    )
                    passed = false
                end
            else
                # No entry here - there must not be an entry in the dict.
                if !isempty(getfreelist(heap, bin))
                    println("Heap as no `status` entry for $bin but found a freelist")
                    passed = false
                end
            end
        end
    end
    return passed
end

function size_invariant(heap::CompactHeap)
    passed = true

    # This certainly better not be nothing!
    # Heaps should always have at least one entry
    y = iterate(heap)
    if y === nothing
        println("Empty Heap!!")
        passed = false
    end
    block, state = y

    count = 1
    while true
        y = iterate(heap, state)
        y === nothing && break
        next, state = y

        if next.backsize != block.size
            println(
                "Block $(count + 1) has backsize $(next.backsize). Expected $(block.size)",
            )
            passed = false
        end

        if block.free && next.free
            println("Unmerged free blocks in position $(count) and $(count + 1)")
            passed = false
        end
        count += 1
        block = next
    end
    return passed
end

function free_invariant(heap::CompactHeap)
    passed = true
    counts = Dict{Int,Int}()

    for block in heap
        if isfree(block)
            bin = getbin(heap, block)
            get!(counts, bin, 0)
            counts[bin] += 1
        end
    end

    for (bin, freelist) in enumerate(freelists(heap))
        if !isempty(freelist)
            if !haskey(counts, bin)
                println("Found a freelist for bin $bin but no block of this size found")
                passed = false
            else
                if length(freelist) != counts[bin]
                    println(
                        "Length of freelist is $(length(freelist)) but counted $(counts[bin]) blocks",
                    )
                    passed = false
                end
                delete!(counts, bin)
            end
        end
    end

    # Did we miss any
    if !isempty(counts)
        println("Unaccounted blocks")
        passed = false
        for (bin, count) in counts
            println("    Bin: $bin - Count: $count")
        end
    end

    return passed
end

function check(heap::CompactHeap)
    passed = true
    if !first_block_invariant(heap)
        printstyled(stdout, "Heap failed First Block Invariant\n"; color = :red)
        passed = false
    end

    if !freelist_status_invariant(heap)
        printstyled(stdout, "Heap failed First List Invariant\n"; color = :red)
        passed = false
    end

    if !size_invariant(heap)
        printstyled(stdout, "Heap failed size invariant\n"; color = :red)
        passed = false
    end

    if !free_invariant(heap)
        printstyled(stdout, "Heap failed free invariant\n"; color = :red)
        passed = false
    end
    return passed
end


export RainbowLearner

using Flux
using Zygote
using StatsBase
using Random

mutable struct RainbowLearner{
    Tq<:AbstractApproximator,
    Tt<:AbstractApproximator,
    Tf,
    Ts,
    Tss<:Union{Int,Nothing},
    R<:AbstractRNG,
} <: AbstractLearner
    approximator::Tq
    target_approximator::Tt
    loss_func::Tf
    Vₘₐₓ::Float32
    Vₘᵢₙ::Float32
    n_actions::Int
    n_atoms::Int
    support::Ts
    stack_size::Tss
    delta_z::Float32
    γ::Float32
    batch_size::Int
    update_horizon::Int
    min_replay_history::Int
    update_freq::Int
    target_update_freq::Int
    update_step::Int
    default_priority::Float64
    rng::R
end

function RainbowLearner(;
    approximator,
    target_approximator,
    loss_func,
    Vₘₐₓ,
    Vₘᵢₙ,
    n_actions,
    n_atoms = 51,
    support = collect(range(Float32(-Vₘₐₓ), Float32(Vₘₐₓ), length = n_atoms)),
    stack_size = 4,
    delta_z = Float32(support[2] - support[1]),
    γ = 0.99,
    batch_size = 32,
    update_horizon = 1,
    min_replay_history = 32,
    update_freq = 1,
    target_update_freq = 500,
    update_step = 0,
    default_priority = 100.0,
    seed = nothing,
)
    copyto!(approximator, target_approximator)  # force sync
    support = send_to_device(device(approximator), support)
    rng = MersenneTwister(seed)
    RainbowLearner(
        approximator,
        target_approximator,
        loss_func,
        Vₘₐₓ,
        Vₘᵢₙ,
        n_actions,
        n_atoms,
        support,
        stack_size,
        delta_z,
        γ,
        batch_size,
        update_horizon,
        min_replay_history,
        update_freq,
        target_update_freq,
        update_step,
        default_priority,
        rng,
    )
end

function (learner::RainbowLearner)(obs)
    state = send_to_device(device(learner.approximator), get_state(obs))
    logits = batch_estimate(learner.approximator, state)
    q = learner.support .* softmax(reshape(logits, :, learner.n_actions))
    # probs = vec(sum(q, dims=1)) .+ legal_action
    vec(sum(q, dims = 1)) |> send_to_host
end

function RLBase.update!(learner::RainbowLearner, batch)
    learner.update_step += 1
    learner.update_step % learner.update_freq == 0 || return nothing

    Q, Qₜ, γ, loss_func, n_atoms, n_actions, support, delta_z, update_horizon, batch_size =
        learner.approximator,
        learner.target_approximator,
        learner.γ,
        learner.loss_func,
        learner.n_atoms,
        learner.n_actions,
        learner.support,
        learner.delta_z,
        learner.update_horizon,
        learner.batch_size

    states, rewards, terminals, next_states = map(
        x -> send_to_device(device(Q), x),
        (batch.states, batch.rewards, batch.terminals, batch.next_states),
    )
    actions = CartesianIndex.(batch.actions, 1:batch_size)
    target_support =
        reshape(rewards, 1, :) .+
        (reshape(support, :, 1) * reshape((γ^update_horizon) .* (1 .- terminals), 1, :))

    updated_priorities = send_to_device(device(Q), Vector{Float32}())

    next_logits = batch_estimate(Qₜ, next_states)
    next_probs = reshape(softmax(reshape(next_logits, n_atoms, :)), n_atoms, n_actions, :)
    next_q = reshape(sum(support .* next_probs, dims = 1), n_actions, :)
    # next_q_argmax = argmax(cpu(next_q .+ next_legal_actions), dims=1)
    next_prob_select = select_best_probs(next_probs, next_q)

    target_distribution = project_distribution(
        target_support,
        next_prob_select,
        support,
        delta_z,
        learner.Vₘᵢₙ,
        learner.Vₘₐₓ,
    )

    gs = gradient(Flux.params(Q)) do
        logits = reshape(batch_estimate(Q, states), n_atoms, n_actions, :)
        select_logits = logits[:, actions]
        batch_losses = loss_func(select_logits, target_distribution)

        updated_priorities = vec(clamp.(sqrt.(batch_losses .+ 1f-10), 1.f0, 1.f2))
        target_priorities = 1.0f0 ./ sqrt.(updated_priorities .+ 1f-10)
        normalized_target_priorities = target_priorities ./ maximum(target_priorities)

        mean(Zygote.dropgrad(normalized_target_priorities) .* batch_losses)
    end

    update!(Q, gs)

    if learner.update_step % learner.target_update_freq == 0
        copyto!(Qₜ, Q)
    end

    updated_priorities |> send_to_host
end

@inline function select_best_probs(probs, q)
    q_argmax = argmax(q, dims = 1)
    prob_select = @inbounds probs[:, q_argmax] # !!! without @inbounds it would be really slow
    reshape(prob_select, :, length(q_argmax))
end

function project_distribution(supports, weights, target_support, delta_z, vmin, vmax)
    batch_size, n_atoms = size(supports, 2), length(target_support)
    clampped_support = clamp.(supports, vmin, vmax)
    tiled_support = reshape(
        repeat(clampped_support; outer = (n_atoms, 1)),
        n_atoms,
        n_atoms,
        batch_size,
    )

    projection =
        clamp.(
            1 .- abs.(tiled_support .- reshape(target_support, 1, :)) ./ delta_z,
            0,
            1,
        ) .* reshape(weights, n_atoms, 1, batch_size)
    reshape(sum(projection, dims = 1), n_atoms, batch_size)
end

function RLBase.extract_experience(t::AbstractTrajectory, learner::RainbowLearner)
    s = learner.stack_size
    h = learner.update_horizon
    n = learner.batch_size
    γ = learner.γ
    if length(t) > learner.min_replay_history
        # 1. sample indices based on priority
        inds = Vector{Int}(undef, n)
        valid_ind_range = isnothing(s) ? (1:(length(t)-h)) : (s:(1:(length(t)-h)))
        for i in 1:n
            ind, p = sample(learner.rng, get_trace(t, :priority))
            while ind ∉ valid_ind_range
                ind, p = sample(learner.rng, get_trace(t, :priority))
            end
            inds[i] = ind
        end

        # 2. extract SARTS
        states = consecutive_view(get_trace(t, :state), inds; n_stack = s)
        actions = consecutive_view(get_trace(t, :action), inds)
        next_states = consecutive_view(get_trace(t, :state), inds .+ h; n_stack = s)
        consecutive_rewards = consecutive_view(get_trace(t, :reward), inds; n_horizon = h)
        consecutive_terminals =
            consecutive_view(get_trace(t, :terminal), inds; n_horizon = h)
        rewards, terminals = zeros(Float32, n), fill(false, n)

        # 3. make sure that we only consider experiences in current episode for rewards and terminals
        for i in 1:n
            m = findfirst(view(consecutive_terminals, :, i))
            if isnothing(m)
                terminals[i] = false
                rewards[i] = discount_rewards_reduced(view(consecutive_rewards, :, i), γ)
            else
                terminals[i] = true
                rewards[i] = discount_rewards_reduced(view(consecutive_rewards, 1:m, i), γ)
            end
        end
        inds,
        (
            states = states,
            actions = actions,
            rewards = rewards,
            terminals = terminals,
            next_states = next_states,
        )
    else
        nothing
    end
end

function RLBase.update!(p::QBasedPolicy{<:RainbowLearner}, t::AbstractTrajectory)
    indexed_experience = extract_experience(t, p)
    if !isnothing(indexed_experience)
        inds, experience = indexed_experience
        priorities = update!(p.learner, experience)
        if !isnothing(priorities)
            get_trace(t, :priority)[inds] .= priorities
        end
    end
end

function (
    agent::Agent{<:QBasedPolicy{<:RainbowLearner},<:CircularCompactPSARTSATrajectory}
)(
    ::PostActStage,
    obs,
)
    push!(
        agent.trajectory;
        reward = get_reward(obs),
        terminal = get_terminal(obs),
        priority = agent.policy.learner.default_priority,
    )
    nothing
end


function get_trigger_limits(config_paths)
    triggers = Dict()
    triggers["BUY"] = Dict()
    triggers["SELL"] = Dict()
    for (symbol, paths) in config_paths
        !haskey(triggers["BUY"],  symbol) ? triggers["BUY"][symbol]  = Dict() : ""
        !haskey(triggers["SELL"], symbol) ? triggers["SELL"][symbol] = Dict() : ""
        for bot_path in paths
            open_orders = JSON.parsefile(joinpath(bot_path, "open_orders.json"))
            for (k,v) in open_orders["BUY"]
                if !haskey(triggers["BUY"][symbol], bot_path)
                    if haskey(v, "price")
                        triggers["BUY"][symbol][bot_path] = parse(Float64, v["price"])
                    end
                else
                    if parse(Float64, v["price"]) < triggers["BUY"][symbol][bot_path]
                        if haskey(v, "price")
                            triggers["BUY"][symbol][bot_path] = parse(Float64, v["price"])
                        end
                    end
                end
            end
            for (k,v) in open_orders["SELL"]
                if !haskey(triggers["SELL"][symbol], bot_path)
                    if haskey(v, "price")
                        triggers["SELL"][symbol][bot_path] = parse(Float64, v["price"])
                    end
                else
                    if parse(Float64, v["price"]) < triggers["SELL"][symbol][bot_path]
                        if haskey(v, "price")
                            triggers["BUY"][symbol][bot_path] = parse(Float64, v["price"])
                        end
                    end
                end
            end
        end
    end
    return triggers
end
function place_triggers(config_paths, symbol, price)
    if typeof(price) == String
        price = parse(Float64, price)
    end
    triggers = get_trigger_limits(config_paths)
    if haskey(triggers["SELL"], symbol)
        for (k,v) in triggers["SELL"][symbol]
            if v * 0.9985 <= price
                if !isfile(joinpath(k, "trigger.json"))
                    open(joinpath(k, "trigger.json"), "w") do f
                        JSON.print(f, Dict(), 4)
                    end
                end
            end
        end
    end
    if haskey(triggers["BUY"], symbol)
        for (k,v) in triggers["BUY"][symbol]
            if v * 1.0015 >= price
                if !isfile(joinpath(k, "trigger.json"))
                    open(joinpath(k, "trigger.json"), "w") do f
                        JSON.print(f, Dict(), 4)
                    end
                end
            end
        end
    end
end