module ParametricOptInterface

using MathOptInterface
using DataStructures: OrderedDict

const MOI = MathOptInterface

const PARAMETER_INDEX_THRESHOLD = 1_000_000_000_000_000_000
const SUPPORTED_SETS =
    (MOI.LessThan{Float64}, MOI.EqualTo{Float64}, MOI.GreaterThan{Float64})
const SUPPORTED_VECTOR_SETS = (MOI.Nonnegatives, MOI.SecondOrderCone)

"""
    Parameter(::Float64)

The `Parameter` structure stores the numerical value associated to a given
parameter.

## Example

```julia-repl
julia> ParametricOptInterface.Parameter(5)
ParametricOptInterface.Parameter(5)
```
"""
struct Parameter <: MOI.AbstractScalarSet
    val::Float64
end

"""
    Optimizer{T, OT <: MOI.ModelLike} <: MOI.AbstractOptimizer

Declares a `Optimizer`, which allows the handling of parameters in a
optimization model.

## Example

```julia-repl
julia> ParametricOptInterface.Optimizer(GLPK.Optimizer())
ParametricOptInterface.Optimizer{Float64,GLPK.Optimizer}
```
"""
mutable struct Optimizer{T,OT<:MOI.ModelLike} <: MOI.AbstractOptimizer
    optimizer::OT
    parameters::Dict{MOI.VariableIndex,T}
    parameters_name::Dict{MOI.VariableIndex,String}
    updated_parameters::Dict{MOI.VariableIndex,T}
    variables::Dict{MOI.VariableIndex,MOI.VariableIndex}
    last_variable_index_added::Int64
    last_parameter_index_added::Int64
    # Store reference to affine constraints with parameters: v + p >= 0
    affine_constraint_cache::MOI.Utilities.DoubleDicts.DoubleDict{
        Vector{MOI.ScalarAffineTerm{Float64}},
    }
    # Store reference to parameter * variable constraints: p*v >= 0.0
    quadratic_constraint_cache_pv::Dict{
        MOI.ConstraintIndex,
        Vector{MOI.ScalarQuadraticTerm{Float64}},
    }
    # Store reference to parameter * parameter constraints: p*p >= var
    quadratic_constraint_cache_pp::Dict{
        MOI.ConstraintIndex,
        Vector{MOI.ScalarQuadraticTerm{Float64}},
    }
    # Store reference to constraints quad_variable_term + affine_with_parameters: v*v + p >= 0
    quadratic_constraint_cache_pc::MOI.Utilities.DoubleDicts.DoubleDict{
        Vector{MOI.ScalarAffineTerm{Float64}},
    }
    # Probably for QPs p*v*v (?)
    quadratic_constraint_variables_associated_to_parameters_cache::Dict{
        MOI.ConstraintIndex,
        Vector{MOI.ScalarAffineTerm{T}},
    }
    # Store the map between quadratic terms add as var * parameter to the resulting shape when implemented in the solver
    # for instance p*p + var -> ScalarAffine(var)
    quadratic_added_cache::OrderedDict{MOI.ConstraintIndex,MOI.ConstraintIndex}
    last_quad_add_added::Int64
    vector_constraint_cache::MOI.Utilities.DoubleDicts.DoubleDict{
        Vector{MOI.VectorAffineTerm{Float64}},
    }
    # Ditto from the constraint caches but for the objective function
    affine_objective_cache::Vector{MOI.ScalarAffineTerm{T}}
    quadratic_objective_cache_pv::Vector{MOI.ScalarQuadraticTerm{T}}
    quadratic_objective_cache_pp::Vector{MOI.ScalarQuadraticTerm{T}}
    quadratic_objective_cache_pc::Vector{MOI.ScalarAffineTerm{T}}
    quadratic_objective_variables_associated_to_parameters_cache::Vector{
        MOI.ScalarAffineTerm{T},
    }
    multiplicative_parameters::Set{Int64}
    dual_value_of_parameters::Vector{Float64}
    evaluate_duals::Bool
    number_of_parameters_in_model::Int64
    function Optimizer(optimizer::OT; evaluate_duals::Bool = true) where {OT}
        return new{Float64,OT}(
            optimizer,
            Dict{MOI.VariableIndex,Float64}(),
            Dict{MOI.VariableIndex,String}(),
            Dict{MOI.VariableIndex,Float64}(),
            Dict{MOI.VariableIndex,MOI.VariableIndex}(),
            0,
            PARAMETER_INDEX_THRESHOLD,
            MOI.Utilities.DoubleDicts.DoubleDict{
                Vector{MOI.ScalarAffineTerm{Float64}},
            }(),
            Dict{MOI.ConstraintIndex,Vector{MOI.ScalarQuadraticTerm{Float64}}}(),
            Dict{MOI.ConstraintIndex,Vector{MOI.ScalarQuadraticTerm{Float64}}}(),
            MOI.Utilities.DoubleDicts.DoubleDict{
                Vector{MOI.ScalarAffineTerm{Float64}},
            }(),
            Dict{MOI.ConstraintIndex,Vector{MOI.ScalarAffineTerm{Float64}}}(),
            OrderedDict{MOI.ConstraintIndex,MOI.ConstraintIndex}(),
            0,
            MOI.Utilities.DoubleDicts.DoubleDict{
                Vector{MOI.VectorAffineTerm{Float64}},
            }(),
            Vector{MOI.ScalarAffineTerm{Float64}}(),
            Vector{MOI.ScalarQuadraticTerm{Float64}}(),
            Vector{MOI.ScalarQuadraticTerm{Float64}}(),
            Vector{MOI.ScalarAffineTerm{Float64}}(),
            Vector{MOI.ScalarAffineTerm{Float64}}(),
            Set{Int64}(),
            Vector{Float64}(),
            evaluate_duals,
            0,
        )
    end
end

function ParametricOptimizer(args...; kwargs...)
    @warn("ParametricOptimizer is deprecated. Use `POI.Optimizer` instead.")
    return Optimizer(args...; kwargs...)
end

include("utils.jl")
include("duals.jl")
include("update_parameters.jl")

function MOI.is_empty(model::Optimizer)
    return MOI.is_empty(model.optimizer) &&
           isempty(model.parameters) &&
           isempty(model.parameters_name) &&
           isempty(model.variables) &&
           isempty(model.updated_parameters) &&
           isempty(model.variables) &&
           model.last_variable_index_added == 0 &&
           model.last_parameter_index_added == PARAMETER_INDEX_THRESHOLD &&
           isempty(model.affine_constraint_cache) &&
           isempty(model.quadratic_constraint_cache_pv) &&
           isempty(model.quadratic_constraint_cache_pp) &&
           isempty(model.quadratic_constraint_cache_pc) &&
           isempty(
               model.quadratic_constraint_variables_associated_to_parameters_cache,
           ) &&
           isempty(model.quadratic_added_cache) &&
           model.last_quad_add_added == 0 &&
           isempty(model.affine_objective_cache) &&
           isempty(model.quadratic_objective_cache_pv) &&
           isempty(model.quadratic_objective_cache_pp) &&
           isempty(model.quadratic_objective_cache_pc) &&
           isempty(
               model.quadratic_objective_variables_associated_to_parameters_cache,
           ) &&
           isempty(model.dual_value_of_parameters) &&
           model.number_of_parameters_in_model == 0
end

function MOI.supports_constraint(
    model::Optimizer,
    F::Union{
        Type{MOI.SingleVariable},
        Type{MOI.ScalarAffineFunction{T}},
        Type{MOI.VectorOfVariables},
        Type{MOI.VectorAffineFunction{T}},
    },
    S::Type{<:MOI.AbstractSet},
) where {T}
    return MOI.supports_constraint(model.optimizer, F, S)
end

function MOI.supports_constraint(
    model::Optimizer,
    ::Type{MOI.ScalarQuadraticFunction{T}},
    S::Type{<:MOI.AbstractSet},
) where {T}
    return MOI.supports_constraint(
        model.optimizer,
        MOI.ScalarAffineFunction{T},
        S,
    )
end

function MOI.supports_constraint(
    model::Optimizer,
    ::Type{MOI.VectorQuadraticFunction{T}},
    S::Type{<:MOI.AbstractSet},
) where {T}
    return MOI.supports_constraint(
        model.optimizer,
        MOI.VectorAffineFunction{T},
        S,
    )
end

function MOI.supports(
    model::Optimizer,
    attr::Union{
        MOI.ObjectiveSense,
        MOI.ObjectiveFunction{MOI.SingleVariable},
        MOI.ObjectiveFunction{MOI.ScalarAffineFunction{T}},
        MOI.ObjectiveFunction{MOI.ScalarQuadraticFunction{T}},
    },
) where {T}
    return MOI.supports(model.optimizer, attr)
end

function MOI.Utilities.supports_default_copy_to(model::Optimizer, bool::Bool)
    return MOI.Utilities.supports_default_copy_to(model.optimizer, bool)
end
function MOI.supports(model::Optimizer, ::MOI.Name)
    return MOI.supports(model.optimizer, MOI.Name())
end
function MOI.get(model::Optimizer, ::MOI.ListOfModelAttributesSet)
    return MOI.get(model.optimizer, MOI.ListOfModelAttributesSet())
end
MOI.get(model::Optimizer, ::MOI.Name) = MOI.get(model.optimizer, MOI.Name())
function MOI.set(model::Optimizer, ::MOI.Name, name::String)
    return MOI.set(model.optimizer, MOI.Name(), name)
end
function MOI.get(model::Optimizer, ::MOI.ListOfVariableIndices)
    return MOI.get(model.optimizer, MOI.ListOfVariableIndices())
end
function MOI.get(model::Optimizer, ::MOI.ListOfVariableAttributesSet)
    return MOI.get(model.optimizer, MOI.ListOfVariableAttributesSet())
end
function MOI.get(model::Optimizer, ::MOI.ListOfConstraintAttributesSet)
    return MOI.get(model.optimizer, MOI.ListOfConstraintAttributesSet())
end
function MOI.set(
    model::Optimizer,
    ::MOI.ConstraintFunction,
    c::MOI.ConstraintIndex{F,S},
    f::F,
) where {F,S}
    MOI.set(model.optimizer, MOI.ConstraintFunction(), c, f)
    return
end
function MOI.set(
    model::Optimizer,
    ::MOI.ConstraintSet,
    c::MOI.ConstraintIndex{F,S},
    s::S,
) where {F,S}
    MOI.set(model.optimizer, MOI.ConstraintSet(), c, s)
    return
end
function MOI.modify(
    model::Optimizer,
    c::MOI.ConstraintIndex{F,S},
    chg::MOI.ScalarCoefficientChange{Float64},
) where {F,S}
    MOI.modify(model.optimizer, c, chg)
    return
end
function MOI.modify(
    model::Optimizer,
    c::MOI.ObjectiveFunction{F},
    chg::MOI.ScalarCoefficientChange{Float64},
) where {F<:MathOptInterface.AbstractScalarFunction}
    MOI.modify(model.optimizer, c, chg)
    return
end
function MOI.modify(
    model::Optimizer,
    c::MOI.ObjectiveFunction{F},
    chg::MOI.ScalarConstantChange{Float64},
) where {F<:MathOptInterface.AbstractScalarFunction}
    MOI.modify(model.optimizer, c, chg)
    return
end

function MOI.empty!(model::Optimizer{T}) where {T}
    MOI.empty!(model.optimizer)
    empty!(model.parameters)
    empty!(model.parameters_name)
    empty!(model.updated_parameters)
    empty!(model.variables)
    model.last_variable_index_added = 0
    model.last_parameter_index_added = PARAMETER_INDEX_THRESHOLD
    empty!(model.affine_constraint_cache)
    empty!(model.quadratic_constraint_cache_pv)
    empty!(model.quadratic_constraint_cache_pp)
    empty!(model.quadratic_constraint_cache_pc)
    empty!(model.quadratic_constraint_variables_associated_to_parameters_cache)
    empty!(model.quadratic_added_cache)
    model.last_quad_add_added = 0
    empty!(model.vector_constraint_cache)
    empty!(model.affine_objective_cache)
    empty!(model.quadratic_objective_cache_pv)
    empty!(model.quadratic_objective_cache_pp)
    empty!(model.quadratic_objective_cache_pc)
    empty!(model.quadratic_objective_variables_associated_to_parameters_cache)
    empty!(model.dual_value_of_parameters)
    model.number_of_parameters_in_model = 0
    return
end

function MOI.set(
    model::Optimizer,
    attr::MOI.VariableName,
    v::MOI.VariableIndex,
    name::String,
)
    if is_parameter_in_model(model, v)
        model.parameters_name[v] = name
    else
        MOI.set(model.optimizer, attr, v, name)
    end
    return
end

function MOI.get(model::Optimizer, attr::MOI.VariableName, v::MOI.VariableIndex)
    if is_parameter_in_model(model, v)
        return get(model.parameters_name, v, "")
    else
        return MOI.get(model.optimizer, attr, v)
    end
end

function MOI.get(model::Optimizer, tp::Type{MOI.VariableIndex}, attr::String)
    return MOI.get(model.optimizer, tp, attr)
end

function MOI.get(model::Optimizer, attr::MOI.ObjectiveBound)
    return MOI.get(model.optimizer, attr)
end

function MOI.supports(
    model::Optimizer,
    attr::MOI.VariableName,
    tp::Type{MOI.VariableIndex},
)
    return MOI.supports(model.optimizer, attr, tp)
end

function MOI.set(
    model::Optimizer,
    attr::MOI.ConstraintName,
    c::MOI.ConstraintIndex,
    name::String,
)
    MOI.set(model.optimizer, attr, c, name)
    return
end

function MOI.get(
    model::Optimizer,
    attr::MOI.ConstraintName,
    c::MOI.ConstraintIndex,
)
    return MOI.get(model.optimizer, attr, c)
end

function MOI.get(model::Optimizer, ::MOI.NumberOfVariables)
    return length(model.parameters) + length(model.variables)
end

function MOI.supports(
    model::Optimizer,
    attr::MOI.ConstraintName,
    tp::Type{<:MOI.ConstraintIndex},
)
    return MOI.supports(model.optimizer, attr, tp)
end

function MOI.get(
    model::Optimizer,
    attr::MOI.ConstraintFunction,
    ci::MOI.ConstraintIndex{F,S},
) where {F,S}
    # Check if the index was cached as Affine expression
    if haskey(model.affine_constraint_cache, ci)
        return model.affine_constraint_cache[ci]
        # Check if the index was cached as a Quadratic
    elseif haskey(model.quadratic_added_cache, ci)
        return model.quadratic_added_cache[ci]
    else
        return MOI.get(model.optimizer, attr, ci)
    end
end

function MOI.get(
    model::Optimizer,
    tp::Type{MOI.ConstraintIndex{F,S}},
    attr::String,
) where {F,S}
    return MOI.get(model.optimizer, tp, attr)
end

function MOI.get(model::Optimizer, tp::Type{MOI.ConstraintIndex}, attr::String)
    return MOI.get(model.optimizer, tp, attr)
end

function MOI.is_valid(model::Optimizer, vi::MOI.VariableIndex)
    return MOI.is_valid(model.optimizer, vi)
end

function MOI.supports(model::Optimizer, ::MOI.NumberOfThreads)
    return MOI.supports(model.optimizer, MOI.NumberOfThreads())
end

function MOI.supports(model::Optimizer, ::MOI.TimeLimitSec)
    return MOI.supports(model.optimizer, MOI.TimeLimitSec())
end

function MOI.set(model::Optimizer, ::MOI.TimeLimitSec, value)
    MOI.set(model.optimizer, MOI.TimeLimitSec(), value)
    return
end

function MOI.get(model::Optimizer, ::MOI.TimeLimitSec)
    return MOI.get(model.optimizer, MOI.TimeLimitSec())
end

function MOI.get(model::Optimizer, ::MOI.SolveTime)
    return MOI.get(model.optimizer, MOI.SolveTime())
end

function MOI.supports(model::Optimizer, ::MOI.Silent)
    return MOI.supports(model.optimizer, MOI.Silent())
end

function MOI.set(model::Optimizer, ::MOI.Silent, value::Bool)
    MOI.set(model.optimizer, MOI.Silent(), value)
    return
end

MOI.get(model::Optimizer, ::MOI.Silent) = MOI.get(model.optimizer, MOI.Silent())

function MOI.get(model::Optimizer, ::MOI.RawStatusString)
    return MOI.get(model.optimizer, MOI.RawStatusString())
end

function MOI.get(model::Optimizer, ::MOI.NumberOfConstraints{F,S}) where {F,S}
    return length(MOI.get(model, MOI.ListOfConstraintIndices{F,S}()))
end

function MOI.get(
    model::Optimizer,
    attr::MOI.ConstraintSet,
    ci::MOI.ConstraintIndex{F,S},
) where {F,S}
    MOI.throw_if_not_valid(model, ci)
    return MOI.get(model.optimizer, attr, ci)
end

function MOI.get(
    model::Optimizer,
    attr::MOI.ConstraintSet,
    ci::MOI.ConstraintIndex{F,S},
) where {
    F<:Union{MOI.SingleVariable,MOI.VectorOfVariables,MOI.VectorAffineFunction},
    S<:MOI.AbstractSet,
}
    return MOI.get(model.optimizer, attr, ci)
end

function MOI.get(model::Optimizer, attr::MOI.ObjectiveSense)
    return MOI.get(model.optimizer, attr)
end

function MOI.get(model::Optimizer, attr::MOI.ObjectiveFunctionType)
    if !isempty(model.quadratic_constraint_cache_pv) ||
       !isempty(model.quadratic_constraint_cache_pc)
        return MOI.ScalarQuadraticFunction{Float64}
    end
    return MOI.get(model.optimizer, attr)
end

function MOI.get(
    model::Optimizer,
    attr::MOI.ObjectiveFunction{F},
) where {
    F<:Union{
        MOI.SingleVariable,
        MOI.ScalarAffineFunction{T},
        MOI.ScalarQuadraticFunction{T},
    },
} where {T}

    # Check if the index was cached as Affine expression
    if !isempty(model.affine_objective_cache)
        return model.affine_objective_cache
        # Check if the index was cached as a Quadratic
    elseif !isempty(model.quadratic_objective_cache_pv)
        return model.quadratic_objective_cache_pv
    elseif !isempty(model.quadratic_objective_cache_pp)
        return model.quadratic_objective_cache_pp
    elseif !isempty(model.quadratic_objective_cache_pc)
        return model.quadratic_objective_cache_pc
    else
        return MOI.get(model.optimizer, attr)
    end
end

function MOI.get(model::Optimizer, attr::MOI.ResultCount)
    return MOI.get(model.optimizer, attr)
end

# In the AbstractBridgeOptimizer, we collect all the possible constraint types and them filter with NumberOfConstraints.
# If NumberOfConstraints is zero then we remove it from the list.
# Here, you can look over keys(quadratic_added_cache) and add the F-S types of all the keys in constraints.
# To implement NumberOfConstraints, you call NumberOfConstraints to the inner optimizer.
# Then you remove the number of constraints of that that in values(quadratic_added_cache)
function MOI.get(model::Optimizer, ::MOI.ListOfConstraints)
    inner_ctrs = MOI.get(model.optimizer, MOI.ListOfConstraints())
    if !has_quadratic_constraint_caches(model)
        return inner_ctrs
    end

    cache_keys = collect(keys(model.quadratic_added_cache))
    constraints = Set{Tuple{DataType,DataType}}()

    for (F, S) in inner_ctrs
        inner_index =
            MOI.get(model.optimizer, MOI.ListOfConstraintIndices{F,S}())
        cache_map_check =
            quadratic_constraint_cache_map_check(mode, inner_index)
        push!(constraints, typeof.(cache_keys[cache_map])...)
        # If not all the constraints are chached then also push the original type
        if !all(cache_map_check)
            push!(constraints, (F, S))
        end
    end

    return collect(constraints)
end

function MOI.get(
    model::Optimizer,
    attr::MOI.ListOfConstraintIndices{F,S},
) where {S,F<:Union{MOI.VectorOfVariables,MOI.SingleVariable}}
    return MOI.get(model.optimizer, attr)
end

function MOI.get(
    model::Optimizer,
    ::MOI.ListOfConstraintIndices{F,S},
) where {
    S<:MOI.AbstractSet,
    F<:Union{MOI.ScalarAffineFunction{T},MOI.VectorAffineFunction{T}},
} where {T}
    inner_index = MOI.get(model.optimizer, MOI.ListOfConstraintIndices{F,S}())
    if !has_quadratic_constraint_caches(model)
        return inner_index
    end

    cache_map_check = quadratic_constraint_cache_map_check(mode, inner_index)
    return inner_index[cache_map_check]
end

function MOI.get(
    model::Optimizer,
    ::MOI.ListOfConstraintIndices{F,S},
) where {S<:MOI.AbstractSet,F<:MOI.ScalarQuadraticFunction{T}} where {T}
    inner_index = MOI.ConstraintIndex{F,S}[]
    if MOI.supports_constraint(model.optimizer, F, S)
        inner_index =
            MOI.get(model.optimizer, MOI.ListOfConstraintIndices{F,S}())
        if !has_quadratic_constraint_caches(model)
            return inner_index
        end
    end

    quadratic_caches = [
        :quadratic_constraint_cache_pc,
        :quadratic_constraint_cache_pp,
        :quadratic_constraint_cache_pv,
        # JD: Check if this applies here
        # :quadratic_constraint_variables_associated_to_parameters_cache
    ]

    for field in quadratic_caches
        cache = getfield(model, field)
        push!(inner_index, keys(cache)...)
    end
    return inner_index
end

function MOI.supports_add_constrained_variable(::Optimizer, ::Type{Parameter})
    return true
end

function MOI.supports_add_constrained_variables(
    model::Optimizer,
    ::Type{MOI.Reals},
)
    return MOI.supports_add_constrained_variables(model.optimizer, MOI.Reals)
end

function MOI.add_variable(model::Optimizer)
    next_variable_index!(model)
    v_p = MOI.VariableIndex(model.last_variable_index_added)
    v = MOI.add_variable(model.optimizer)
    model.variables[v_p] = v
    return v_p
end

function MOI.add_constrained_variable(model::Optimizer, set::Parameter)
    next_parameter_index!(model)
    p = MOI.VariableIndex(model.last_parameter_index_added)
    model.parameters[p] = set.val
    cp = MOI.ConstraintIndex{MOI.SingleVariable,Parameter}(
        model.last_parameter_index_added,
    )
    update_number_of_parameters!(model)
    return p, cp
end

function MOI.add_constraint(
    model::Optimizer,
    f::MOI.SingleVariable,
    set::MOI.AbstractScalarSet,
)
    if is_parameter_in_model(model, f.variable)
        error("Cannot constrain a parameter")
    elseif !is_variable_in_model(model, f.variable)
        error("Variable not in the model")
    end
    return MOI.add_constraint(model.optimizer, f, set)
end

function add_constraint_with_parameters_on_function(
    model::Optimizer,
    f::MOI.ScalarAffineFunction{T},
    set::MOI.AbstractScalarSet,
) where {T}
    vars, params, param_constant =
        separate_possible_terms_and_calculate_parameter_constant(model, f.terms)
    ci = MOI.Utilities.normalize_and_add_constraint(
        model.optimizer,
        MOI.ScalarAffineFunction(vars, f.constant + param_constant),
        set,
    )
    model.affine_constraint_cache[ci] = params
    return ci
end

function MOI.add_constraint(
    model::Optimizer,
    f::MOI.ScalarAffineFunction{T},
    set::MOI.AbstractScalarSet,
) where {T}
    if !function_has_parameters(model, f)
        return MOI.add_constraint(model.optimizer, f, set)
    else
        return add_constraint_with_parameters_on_function(model, f, set)
    end
end

function MOI.get(
    model::Optimizer,
    attr::MOI.VariablePrimal,
    v::MOI.VariableIndex,
)
    if is_parameter_in_model(model, v)
        return model.parameters[v]
    elseif is_variable_in_model(model, v)
        return MOI.get(model.optimizer, attr, model.variables[v])
    else
        error("Variable not in the model")
    end
end

function MOI.set(
    model::Optimizer,
    ::MOI.ConstraintSet,
    cp::MOI.ConstraintIndex{MOI.SingleVariable,Parameter},
    set::Parameter,
)
    p = MOI.VariableIndex(cp.value)
    if !is_parameter_in_model(model, p)
        error("Parameter not in the model")
    end
    return model.updated_parameters[p] = set.val
end

struct ParameterValue <: MOI.AbstractVariableAttribute end

"""
    MOI.set(
        model::Optimizer,
        ::ParameterValue,
        x::MOI.VariableIndex,
        value::Float64,
    )

Sets the parameter to a given value, using its `MOI.VariableIndex` as reference.

# Example

```julia-repl
julia> MOI.set(model, ParameterValue(), w, 2.0)
2.0
```
"""
function MOI.set(
    model::Optimizer,
    ::ParameterValue,
    vi::MOI.VariableIndex,
    val::Float64,
)
    if !is_parameter_in_model(model, vi)
        error("Parameter not in the model")
    end
    return model.updated_parameters[vi] = val
end

function MOI.set(
    model::Optimizer,
    ::ParameterValue,
    vi::MOI.VariableIndex,
    val::Real,
)
    return MOI.set(model, ParameterValue(), vi, convert(Float64, val))
end

function empty_objective_function_caches!(model::Optimizer)
    empty!(model.affine_objective_cache)
    empty!(model.quadratic_objective_cache_pv)
    empty!(model.quadratic_objective_cache_pp)
    empty!(model.quadratic_objective_cache_pc)
    return
end

function MOI.set(
    model::Optimizer,
    attr::MOI.ObjectiveFunction,
    f::MOI.ScalarAffineFunction{T},
) where {T}
    # clear previously defined objetive function cache
    empty_objective_function_caches!(model)
    if !function_has_parameters(model, f)
        MOI.set(model.optimizer, attr, f)
    else
        vars, params, param_constant =
            separate_possible_terms_and_calculate_parameter_constant(
                model,
                f.terms,
            )
        MOI.set(
            model.optimizer,
            attr,
            MOI.ScalarAffineFunction(vars, f.constant + param_constant),
        )
        model.affine_objective_cache = params
    end
    return
end

function MOI.set(
    model::Optimizer,
    attr::MOI.ObjectiveFunction,
    v::MOI.SingleVariable,
)
    if haskey(model.parameters, v)
        error("Cannot use a parameter as objective function alone")
    elseif !haskey(model.variables, v)
        error("Variable not in the model")
    end
    return MOI.set(
        model.optimizer,
        attr,
        MOI.SingleVariable(model.variables[v.variable]),
    )
end

function MOI.set(
    model::Optimizer,
    attr::MOI.ObjectiveSense,
    sense::MOI.OptimizationSense,
)
    MOI.set(model.optimizer, attr, sense)
    return
end

function MOI.get(
    model::Optimizer,
    attr::T,
) where {
    T<:Union{
        MOI.TerminationStatus,
        MOI.ObjectiveValue,
        MOI.DualObjectiveValue,
        MOI.PrimalStatus,
        MOI.DualStatus,
    },
}
    return MOI.get(model.optimizer, attr)
end

function MOI.get(
    model::Optimizer,
    attr::T,
    c::MOI.ConstraintIndex,
) where {T<:Union{MOI.ConstraintPrimal,MOI.ConstraintDual}}
    return MOI.get(model.optimizer, attr, c)
end

function MOI.set(model::Optimizer, attr::MOI.RawParameter, val::Any)
    MOI.set(model.optimizer, attr, val)
    return
end

# TODO(odow): remove this.
function MOI.set(model::Optimizer, attr::String, val::Any)
    return MOI.set(model.optimizer, MOI.RawParameter(attr), val)
end

function MOI.get(model::Optimizer, ::MOI.SolverName)
    name = MOI.get(model.optimizer, MOI.SolverName())
    return "Parametric Optimizer with $(name) attached"
end

function MOI.add_constraint(
    model::Optimizer,
    f::MOI.VectorOfVariables,
    set::MOI.AbstractVectorSet,
) where {T}
    if function_has_parameters(model, f)
        error("VectorOfVariables does not allow parameters")
    end
    return MOI.add_constraint(model.optimizer, f, set)
end

function MOI.add_constraint(
    model::Optimizer,
    f::MOI.VectorAffineFunction{T},
    set::MOI.AbstractVectorSet,
) where {T}
    if !function_has_parameters(model, f)
        return MOI.add_constraint(model.optimizer, f, set)
    else
        return add_constraint_with_parameters_on_function(model, f, set)
    end
end

function add_constraint_with_parameters_on_function(
    model::Optimizer,
    f::MOI.VectorAffineFunction{T},
    set::MOI.AbstractVectorSet,
) where {T}
    vars, params, param_constants =
        separate_possible_terms_and_calculate_parameter_constant(model, f, set)
    ci = MOI.add_constraint(
        model.optimizer,
        MOI.VectorAffineFunction(vars, f.constants + param_constants),
        set,
    )
    model.vector_constraint_cache[ci] = params
    return ci
end

function add_constraint_with_parameters_on_function(
    model::Optimizer,
    f::MOI.ScalarQuadraticFunction{T},
    set::S,
) where {T,S<:MOI.AbstractScalarSet}
    if function_quadratic_terms_has_parameters(model, f.quadratic_terms)
        (
            quad_vars,
            quad_aff_vars,
            quad_params,
            aff_terms,
            variables_associated_to_parameters,
            quad_param_constant,
        ) = separate_possible_terms_and_calculate_parameter_constant(
            model,
            f.quadratic_terms,
        )
    else
        quad_vars = f.quadratic_terms
        quad_aff_vars = MOI.ScalarQuadraticTerm{T}[]
        quad_params = MOI.ScalarQuadraticTerm{T}[]
        aff_terms = MOI.ScalarAffineTerm{T}[]
        quad_param_constant = zero(T)
        variables_associated_to_parameters = MOI.VariableIndex[]
    end
    if function_affine_terms_has_parameters(model, f.affine_terms)
        (
            aff_vars,
            aff_params,
            terms_with_variables_associated_to_parameters,
            aff_param_constant,
        ) = separate_possible_terms_and_calculate_parameter_constant(
            model,
            f.affine_terms,
            variables_associated_to_parameters,
        )
    else
        aff_vars = f.affine_terms
        aff_params = MOI.ScalarAffineTerm{T}[]
        terms_with_variables_associated_to_parameters =
            MOI.ScalarAffineTerm{T}[]
        aff_param_constant = zero(T)
    end
    aff_terms = vcat(aff_terms, aff_vars)
    const_term = f.constant + aff_param_constant + quad_param_constant
    quad_terms = quad_vars
    f_quad = if !isempty(quad_vars)
        MOI.ScalarQuadraticFunction(aff_terms, quad_terms, const_term)
    else
        MOI.ScalarAffineFunction(aff_terms, const_term)
    end
    model.last_quad_add_added += 1
    ci =
        MOI.Utilities.normalize_and_add_constraint(model.optimizer, f_quad, set)
    # This part is used to remember that ci came from a quadratic function
    # It is particularly useful because sometimes the constraint mutates
    new_ci = MOI.ConstraintIndex{MOI.ScalarQuadraticFunction{T},S}(
        model.last_quad_add_added,
    )
    model.quadratic_added_cache[new_ci] = ci
    fill_quadratic_constraint_caches!(
        model,
        new_ci,
        quad_aff_vars,
        quad_params,
        aff_params,
        terms_with_variables_associated_to_parameters,
    )
    return new_ci
end

function MOI.add_constraint(
    model::Optimizer,
    f::MOI.ScalarQuadraticFunction{T},
    set::MOI.AbstractScalarSet,
) where {T}
    if !function_has_parameters(model, f)
        return MOI.add_constraint(model.optimizer, f, set)
    else
        return add_constraint_with_parameters_on_function(model, f, set)
    end
end

function MOI.delete(model::Optimizer, v::MOI.VariableIndex)
    pop!(model.variables, v, 0) # 0 is a default value if v is not found
    MOI.delete(model.optimizer, v)
    return
end

function MOI.delete(
    model::Optimizer,
    c::MOI.ConstraintIndex{F,S},
) where {F<:MOI.ScalarAffineFunction,S<:MOI.AbstractSet}
    if haskey(model.affine_constraint_cache, c)
        delete!(model.affine_constraint_cache, c)
    end
    MOI.delete(model.optimizer, c)
    return
end

function MOI.delete(
    model::Optimizer,
    c::MOI.ConstraintIndex{F,S},
) where {
    F<:Union{MOI.SingleVariable,MOI.VectorOfVariables,MOI.VectorAffineFunction},
    S<:MOI.AbstractSet,
}
    MOI.delete(model.optimizer, c)
    return
end

function MOI.is_valid(
    model::Optimizer,
    c::MOI.ConstraintIndex{F,S},
) where {
    F<:Union{MOI.SingleVariable,MOI.VectorOfVariables,MOI.VectorAffineFunction},
    S<:MOI.AbstractSet,
}
    return MOI.is_valid(model.optimizer, c)
end

function MOI.is_valid(
    model::Optimizer,
    c::MOI.ConstraintIndex{F,S},
) where {F<:MOI.ScalarAffineFunction,S<:MOI.AbstractSet}
    if haskey(model.affine_constraint_cache, c)
        # TODO: how to check the cached constraint is valid (?)
        return function_has_parameters(
            model,
            model.affine_constraint_cache[c],
        ) && MOI.is_valid(model.optimizer, c)
    end
    return MOI.is_valid(model.optimizer, c)
end

function MOI.set(
    model::Optimizer,
    attr::MOI.ObjectiveFunction,
    f::MOI.ScalarQuadraticFunction{T},
) where {T}
    # clear previously defined objetive function cache
    empty_objective_function_caches!(model)
    if !function_has_parameters(model, f)
        MOI.set(model.optimizer, attr, f)
        return
    end
    if function_quadratic_terms_has_parameters(model, f.quadratic_terms)
        (
            quad_vars,
            quad_aff_vars,
            quad_params,
            aff_terms,
            variables_associated_to_parameters,
            quad_param_constant,
        ) = separate_possible_terms_and_calculate_parameter_constant(
            model,
            f.quadratic_terms,
        )
    else
        quad_vars = f.quadratic_terms
        quad_aff_vars = MOI.ScalarQuadraticTerm{T}[]
        quad_params = MOI.ScalarQuadraticTerm{T}[]
        aff_terms = MOI.ScalarAffineTerm{T}[]
        quad_param_constant = zero(T)
        variables_associated_to_parameters = MOI.VariableIndex[]
    end

    if function_affine_terms_has_parameters(model, f.affine_terms)
        (
            aff_vars,
            aff_params,
            terms_with_variables_associated_to_parameters,
            aff_param_constant,
        ) = separate_possible_terms_and_calculate_parameter_constant(
            model,
            f.affine_terms,
            variables_associated_to_parameters,
        )
    else
        aff_vars = f.affine_terms
        aff_params = MOI.ScalarAffineTerm{T}[]
        terms_with_variables_associated_to_parameters =
            MOI.ScalarAffineTerm{T}[]
        aff_param_constant = zero(T)
    end

    aff_terms = vcat(aff_terms, aff_vars)
    const_term = f.constant + aff_param_constant + quad_param_constant
    quad_terms = quad_vars

    if !isempty(quad_vars)
        MOI.set(
            model.optimizer,
            attr,
            MOI.ScalarQuadraticFunction(aff_terms, quad_terms, const_term),
        )
    else
        MOI.set(
            model.optimizer,
            MOI.ObjectiveFunction{MOI.ScalarAffineFunction{T}}(),
            MOI.ScalarAffineFunction(aff_terms, const_term),
        )
    end

    if !isempty(quad_terms)
        f_quad = MOI.ScalarQuadraticFunction(aff_terms, quad_terms, const_term)

        MOI.set(model.optimizer, attr, f_quad)
    else
        f_quad = MOI.ScalarAffineFunction(aff_terms, const_term)

        MOI.set(
            model.optimizer,
            MOI.ObjectiveFunction{MOI.ScalarAffineFunction{T}}(),
            f_quad,
        )
    end

    model.quadratic_objective_cache_pv = quad_aff_vars
    model.quadratic_objective_cache_pp = quad_params
    model.quadratic_objective_cache_pc = aff_params
    model.quadratic_objective_variables_associated_to_parameters_cache =
        terms_with_variables_associated_to_parameters

    return
end

function MOI.optimize!(model::Optimizer)
    if !isempty(model.updated_parameters)
        update_parameters!(model)
    end
    MOI.optimize!(model.optimizer)
    if model.evaluate_duals
        calculate_dual_of_parameters(model)
    end
    return
end

end # module


# This file provides most of the implementation of multi-robot target *tracking*
# problems.
#
# The tracking problems involve estimation of target positions and optimization
# via maximizing mutual information.

using LinearAlgebra

export MultiRobotTargetTrackingConfigs, MultiRobotTargetTrackingProblem

# Automatically switch to sparse filters for large numbers of robots
const num_robots_sparse_filtering_threshold = 15
const sparsity_threshold=1e-3

# Stores configuration variables for multi-robot target tracking
struct MultiRobotTargetTrackingConfigs
  grid::Grid
  sensor::RangingSensor
  horizon::Int64

  # Configuration for the MCTS solver
  solver_iterations::Int64
  solver_information_samples::Int64

  # We will generally use a larger number of samples when querying actual
  # solution values
  objective_information_samples::Int64

  robot_target_range_limit::Float64

  function MultiRobotTargetTrackingConfigs(;
                                           grid,
                                           horizon=default_horizon,
                                           sensor=RangingSensor(),
                                           solver_iterations =
                                           default_num_iterations[horizon],
                                           solver_information_samples =
                                             default_solver_information_samples,
                                           objective_information_samples =
                                             default_num_information_samples,
                                           robot_target_range_limit = Inf
                                          )

    new(grid, sensor, horizon, solver_iterations, solver_information_samples,
        objective_information_samples, robot_target_range_limit)
  end
end

# Constructor that can infer parameters based on the number of robots
#
# Note that inferred parameters can still be overriden, but we probably don't
# want to do so
function MultiRobotTargetTrackingConfigs(num_robots;
                                         kwargs...)
  grid=Grid(num_robots=num_robots)
  MultiRobotTargetTrackingConfigs(;grid=grid, kwargs...)
end

# Copy constructor with overrides
function
  MultiRobotTargetTrackingConfigs(configs::MultiRobotTargetTrackingConfigs;
                                  kwargs...)

  fields = fieldnames(MultiRobotTargetTrackingConfigs)
  old_settings = Dict(field => getfield(configs, field) for field in fields)
  settings = merge(old_settings, kwargs)

  MultiRobotTargetTrackingConfigs(; settings...)
end

# Solution elements consist of the robot index and the associated trajectory
# We construct solutions as such because the output may not have the same
# ordering as the robots
#
# Warning: this "problem" object may become invalid if any of the underlying
# objects change and may copy some but not of the inputs.
struct MultiRobotTargetTrackingProblem{F<:AbstractFilter} <: AbstractMultiRobotProblem
  # Target tracking problems are defined by vectors of robot states
  partition_matroid::Vector{State}

  target_filters::Vector{F}
  filter_means::Vector{Tuple{Float64, Float64}}

  configs::MultiRobotTargetTrackingConfigs

  function MultiRobotTargetTrackingProblem(robot_states::Vector{State},
                 target_filters::Vector{F},
                 configs::MultiRobotTargetTrackingConfigs) where F <: AbstractFilter

    if length(robot_states) > num_robots_sparse_filtering_threshold
      sparse_filters = map(x->SparseFilter(x, threshold=sparsity_threshold),
                           target_filters)

      filter_means = compute_filter_means(sparse_filters)

      new{SparseFilter{Int64}}(robot_states, sparse_filters, filter_means,
                               configs)
    else
      filter_means = compute_filter_means(target_filters)

      new{F}(robot_states, target_filters, filter_means, configs)
    end
  end
end

# Construct a target tracking problem with configs. I may or may not need this
function MultiRobotTargetTrackingProblem(robot_states::Vector{State},
                                         target_filters::Vector{<:AbstractFilter};
                                         kwargs...)
  configs = MultiRobotTargetTrackingConfigs(;kwargs...)
  MultiRobotTargetTrackingProblem(robot_states, target_filters, configs)
end

function objective(p::MultiRobotTargetTrackingProblem, X)
  configs = p.configs

  trajectories = map(last, X)

  sum(p.target_filters) do filter
    finite_horizon_information(configs.grid, filter, configs.sensor,
                               trajectories;
                               num_samples=configs.objective_information_samples
                              ).reward
  end
end

# Returns distance from robot state to target mean
# (due to convexity and Jensen's inequality, this is a lower bound on the range)
function target_mean_distance(p::MultiRobotTargetTrackingProblem,
                     robot_state::State, target_index::Integer)
  norm(robot_state .- p.filter_means[target_index])
end
function target_mean_distance(p::MultiRobotTargetTrackingProblem,
                     robot_index::Integer, target_index::Integer)
  target_mean_distance(p, p.partition_matroid[robot_index], target_index)
end

function filter_targets_in_range(p::MultiRobotTargetTrackingProblem,
                                 robot_index::Integer;
                                 range_limit=p.configs.robot_target_range_limit
                                )

  # filter by range, producing tuples
  in_bounds = filter(1:length(p.filter_means)) do target_index
    target_mean_distance(p, robot_index, target_index) < range_limit
  end

  map(x->p.target_filters[x[1]], in_bounds)
end

function solve_block(p::MultiRobotTargetTrackingProblem, block::Integer,
                     selections::Vector)
  configs = p.configs

  trajectories = map(last, selections)

  # Select filters in range
  local targets_in_range
  if p.configs.robot_target_range_limit == Inf
    targets_in_range = p.target_filters
  else
    targets_in_range = filter_targets_in_range(p, block)
  end

  problem = SingleRobotTargetTrackingProblem(configs.grid, configs.sensor,
                                             configs.horizon, targets_in_range,
                                             prior_trajectories=trajectories,
                                             num_information_samples=
                                             configs.solver_information_samples
                                            )

  state = get_state(p, block)

  solution = solve_single_robot(problem, state,
                                n_iterations=configs.solver_iterations)

  (block, solution.trajectory)
end

function compute_filter_means(target_filters::Vector{<:AbstractFilter})
  map(target_filters) do x
    Tuple{Float64, Float64}(weighted_average(x))
  end
end


export
    LocalSimulation,
    simulate
"""
    LocalSimulation

Describe a Local-BPS sampling
"""
struct LocalSimulation
    fg::FactorGraph
    x0::Vector{AllowedVarType}  # Starting point
    v0::Vector{AllowedVarType}  # Starting velocity
    T::Real                     # Simulation time
    maxnevents::Int             # Maximum number of events
    lambdaref::Real             # Refreshment rate
    # -- implicit
    nvars::Int                  # Number of nodes
    nfactors::Int               # Number of factors
    function LocalSimulation(fg, x0, v0, T, maxnevents, lambdaref)
        # check none of the default named arguments went through
        cond =  !(  fg == :undefined ||
                    x0 == :undefined ||
                    v0 == :undefined ||
                    T  == :undefined )
        @assert cond "One or more essential argument undefined"
        cond = (length(x0) == length(v0) == fg.structure.nvars > 0)
        @assert cond "inconsistent dims"
        cond = (T > 0.0 && maxnevents > 0 && lambdaref > 0.0)
        @assert cond "inconsistent params"
        # creating the object
        new( fg, x0, v0, T,
             maxnevents,
             lambdaref,
             length(x0),
             length(fg.factors) )
    end
end
# Constructor with named arguments
function LocalSimulation(;
            factorgraph = :undefined,
            x0          = :undefined,
            v0          = :undefined,
            T           = :undefined,
            maxnevents  = 50000,
            lambdaref   = 0.5)::LocalSimulation
    # calling the unnamed constructor
    LocalSimulation( factorgraph,
                     x0, v0, T,
                     maxnevents,
                     lambdaref )
end

"""
    simulate(sim)

Run a local BPS simulation following the specifications of `sim`.
"""
function simulate(sim::LocalSimulation)::Tuple{AllEventList, Dict}

    (start, all_evlist, pq, tref) = ls_init(sim)

    # counters for the first branch (bounce) and the second branch (refresh)
    ev_firstbranch  = 0
    ev_secondbranch = 0
    nevents         = 0 # sum of above two
    # global clock
    globalclock     = 0.0

    # global iteration
    prog = Progress(sim.maxnevents, 1)

    for evnum in 1:sim.maxnevents
        # get bounce and dequeue
        (fidx, tbounce) = peek(pq)
        dequeue!(pq)
        t = tbounce
        globalclock += t
        # ------------
        # FIRST BRANCH (t<tref => )
        if t < tref
            ev_firstbranch += 1
            ls_firstbranch!(sim.fg, fidx, all_evlist, pq, t)
        # -------------
        # SECOND BRANCH (t>=tref => refreshment)
        else
            ev_secondbranch += 1
            pq    = ls_refreshment(sim.fg, tref, all_evlist)
            tref += randexp()/sim.lambdaref
        end
        # if use while loop
        # nevents += 1
        next!(prog)
    end

    details = Dict(
        "clocktime"       => time()-start,
        "nevents"         => nevents,
        "ev_firstbranch"  => ev_firstbranch,
        "ev_secondbranch" => ev_secondbranch,
        "globalclock"     => globalclock
    )

    (all_evlist, details)
end

# =============================================================================
# NOTE: The functions with a `ls_` prefix are helper functions exclusively
# designed for the local simulation. They aren't expected to be used elsewhere.

"""
    ls_init(sim)

Initialise a local simulation: store the starting clock time, create a list of
all eventlists and push the original event, create the priorityqueue and fill
it with the initial bounce times
"""
function ls_init(sim::LocalSimulation)
    # instantiate wall clock
    start = time()
    # initialisation of the eventlists for every node in the graph
    all_evlist = AllEventList([typeof(sim.x0[i]) for i in 1:sim.nvars])
    for i in 1:sim.nvars
        evi = Event(sim.x0[i], sim.v0[i], 0.0)
        pushevent!(all_evlist.evl[i], evi)
    end
    # initialisation of the priority queue and the refreshment time
    pq   = PriorityQueue{Int, Real}()
    tref = randexp()/sim.lambdaref
    # filling of the priority queue with initial position
    for fidx in 1:sim.fg.structure.nfactors
        (xf, vf, g) = ls_retrieve(sim.fg, fidx, all_evlist, 0.0)
        ls_updatepq!(pq, sim.fg, fidx, xf, vf, g, 0.0)
    end
    return (start, all_evlist, pq, tref)
end

"""
    ls_firstbranch(fg, fidx, all_evlist, pq, t)

Operation corresponding to the first branch of events in simulate.
"""
function ls_firstbranch!(fg::FactorGraph, fidx::Int, all_evlist::AllEventList,
                         pq::PriorityQueue, t::Real)
    # retrieve xf, vf corresponding to factor
    (xf, vf, g, vars) = ls_retrieve(fg, fidx, all_evlist, t, true)
    ls_saveupdate!(all_evlist, vars, xf, vf, t)
    ls_updatepq!(pq, fg, fidx, xf, vf, g, t)
    # same story for linked factors (fp)
    for fpidx in linkedfactors(fg, fidx)
        # we don't need to retrieve `vars` here
        (xfp, vfp, gp) = ls_retrieve(fg, fpidx, all_evlist, t)
        ls_updatepq!(pq, fg, fpidx, xfp, vfp, gp, t)
    end
    return (all_evlist, pq)
end

"""
    ls_reshape(u, v)

(Local Simulation, helper function) Reshape a single vector `u` to a vector of
AllowedVarType following the structure of `v`. For example, let
`v = [0., [0., 0.]]` and `u=[1., 2., 3.]` then `ls_reshape(u,v)` will return
`[1., [2., 3.]]`.
"""
function ls_reshape(u::Vector{<:Real}, v::V) where V <: Vector{AllowedVarType}
    w = similar(v)
    # offsets to know where to look for next block of information
    offset = 0
    for i ∈ 1:length(v)
        tmpi = length(v[i])
        # if length is 1, return only single point, otherwise vector
        @inbounds w[i] = (tmpi==1) ? u[offset+1] : u[offset .+ (1:tmpi)]
        offset += tmpi
    end
    return w
end

"""
    ls_retrieve(fg, fidx, all_evlist, t, doreflect)

(Local simulation, helper function) Retrieve the structure corresponding to all
positions of the nodes attached to a factor (`xf`) and the velocities (`vf`) as
well as the gradient at `xf` and the list of indices of the attached variables.
"""
function ls_retrieve(fg::FactorGraph, fidx::Int,
                    all_evlist::AllEventList,
                    t::Real, doreflect::Bool=false)
    # indices of the variable associated with factor fidx
    vars = assocvariables(fg, fidx)
    # allocate xf, vf, note the different variables
    # don't necessarily have the same types
    vf = Vector{AllowedVarType}(undef, length(vars))
    xf = similar(vf)
    # shortcut to initial event
    if t == 0.0
        for (i, k) ∈ enumerate(vars)
            ev = getevent(all_evlist.evl[k], 1)
            xf[i], vf[i] = ev.x, ev.v
        end
    # retrieving later than initial event (standard case)
    else
        # retrieve vf, xf(t)
        for (i, k) ∈ enumerate(vars)
            # get the eventlist corresponding to variable k
            ev = getlastevent(all_evlist.evl[k])
            # store the components and extrapolate from ev.x as needed
            vf[i] = ev.v
            xf[i] = ev.x + (t-ev.t)*ev.v
        end
    end
    # compute the gradient of the factor's log likelihood at xf
    g  = fg.factors[fidx].gll(vcat(xf...))
    # compute the reflection (if required)
    vf = doreflect ? ls_reshape(reflect_bps!(g, vcat(vf...)), vf) : vf
    (xf, vf, g, vars)
end

"""
    ls_saveupdate!(all_evlist, vars, xf, vf, t)

(Local simulation, helper function) Append the information contained in `xf`,
`vf` in the appropriate EventLists corresponding to the relevant variables
indexed in `vars` (variables associated with the factor `f`, `xf` and `vf` are
the recovered positions and velocities obtained when considering factor `f`).
"""
function ls_saveupdate!(all_evlist::AllEventList, vars::Vector{Int},
                        xf::Vector{AllowedVarType}, vf::Vector{AllowedVarType},
                        t::Real)
    # Add the new sample to the list
    for (i, k) ∈ enumerate(vars)
        pushevent!(all_evlist.evl[k], Event(xf[i], vf[i], t))
    end
    return all_evlist
end

"""
    ls_updatepq!(pq, fg, fidx, xf, vf, g, t)

(Local simulation, helper function) Compute the bounce time (possibly via
thinning depending on the factor's implementation) associated with a factor and
push it to the PriorityQueue `pq`.
"""
function ls_updatepq!(pq::PriorityQueue, fg::FactorGraph, fidx::Int,
                      xf::Vector{AllowedVarType}, vf::Vector{AllowedVarType},
                      g::Vector{<:Real}, t::Real)
    # useful temporary variables
    vcxf, vcvf = vcat(xf...), vcat(vf...)
    # Update time in Priority Queue for the current factor
    acc  = false
    tauf = 0.0
    while !acc
        bounce = fg.factors[fidx].nextevent(vcxf, vcvf)
        acc, tauf = bounce.dobounce(g, vcvf), bounce.tau
    end
    # add the time to current priorityQueue
    pq[fidx] = t+tauf
    return pq
end

"""
    ls_random(evl)

(Local simulation, helper function) Generate N(0,1) random numbers of dimension
corresponding to the passed `EventList`.
"""
function ls_random(evl::EventList)
    l = length(evl.xs[end])
    return l > 1 ? randn(l) : randn()
end

"""
    ls_refreshment(fg, t, all_evlist)

(Local simulation, helper function) Generate a new PriorityQueue by drawing
new velocities for each node and storing it (cf. NewQueue algorithm in the BPS
paper).
"""
function ls_refreshment(fg::FactorGraph, t::Real, all_evlist::AllEventList)
    # draw a new bunch of velocities for each node
    v = Vector{AllowedVarType}(undef, fg.structure.nvars)
    for i ∈ 1:fg.structure.nvars
        @inbounds v[i] = ls_random(all_evlist.evl[i])
    end
    # Instantiate a new priority queue
    pq = PriorityQueue{Int, Real}()
    for fidx ∈ 1:fg.structure.nfactors
        # retrieve xf, vf corresponding to factor
        (xf, vf, g, vars) = ls_retrieve(fg, fidx, all_evlist, t)
        # update the priority queue, note v[vars] and not vf
        # since we are using the refreshed velocity
        ls_updatepq!(pq, fg, fidx, xf, v[vars], g, t)
    end
    return pq
end


"""
    Channel

Struct containing asymptotic kernels for a channel.
* `q1   :: Matrix{Float64}`   : kernel with both fermionic frequencies -> Inf
* `q2_1 :: Array{Float64, 3}` : kernel with second fermionic frequency -> Inf
* `q2_2 :: Array{Float64, 3}` : kernel with first fermionic frequency -> Inf
* `q3   :: Array{Float64, 4}` : full channel
"""
struct Channel
    q1   :: Matrix{Float64}
    q2_1 :: Array{Float64, 3}
    q2_2 :: Array{Float64, 3}
    q3   :: Array{Float64, 4}
end

# generate channel dummy
function get_channel_empty(
    r :: Reduced_lattice,
    m :: Mesh
    ) :: Channel

    num_sites = length(r.sites)

    # init kernels
    q1   = zeros(Float64, num_sites, m.num_Ω)
    q2_1 = zeros(Float64, num_sites, m.num_Ω, m.num_ν)
    q2_2 = zeros(Float64, num_sites, m.num_Ω, m.num_ν)
    q3   = zeros(Float64, num_sites, m.num_Ω, m.num_ν, m.num_ν)

    # build channel
    ch = Channel(q1, q2_1, q2_2, q3)

    return ch
end

# get interpolated value of q1
function get_q1(
    site :: Int64,
    p    :: Param,
    ch   :: Channel
    )    :: Float64

    val  = p.lower_weight * ch.q1[site, p.lower_index]
    val += p.upper_weight * ch.q1[site, p.upper_index]

    return val
end

# get interpolated value of q1 on all lattice sites
function get_q1_avx!(
    r        :: Reduced_lattice,
    p        :: Param,
    ch       :: Channel,
    temp     :: SubArray{Float64, 1, Array{Float64, 3}},
    exchange :: Bool,
    sgn      :: Float64
    )        :: Nothing

    # deref channel
    q1 = ch.q1

    # deref param
    lower_index  = p.lower_index
    upper_index  = p.upper_index
    lower_weight = p.lower_weight
    upper_weight = p.upper_weight

    # check for site exchange
    if exchange
        indices = r.exchange

        @turbo unroll = 1 for i in eachindex(temp)
            val      = lower_weight * q1[indices[i], lower_index]
            val     += upper_weight * q1[indices[i], upper_index]
            temp[i] += sgn * val
        end
    else
        @turbo unroll = 1 for i in eachindex(temp)
            val      = lower_weight * q1[i, lower_index]
            val     += upper_weight * q1[i, upper_index]
            temp[i] += sgn * val
        end
    end

    return nothing
end

# get interpolated value of q2_1
function get_q2_1(
    site :: Int64,
    p1   :: Param,
    p2   :: Param,
    ch   :: Channel
    )    :: Float64

    val  = p1.lower_weight * p2.lower_weight * ch.q2_1[site, p1.lower_index, p2.lower_index]
    val += p1.upper_weight * p2.lower_weight * ch.q2_1[site, p1.upper_index, p2.lower_index]
    val += p1.lower_weight * p2.upper_weight * ch.q2_1[site, p1.lower_index, p2.upper_index]
    val += p1.upper_weight * p2.upper_weight * ch.q2_1[site, p1.upper_index, p2.upper_index]

    return val
end

# get interpolated value of q2_1 on all lattice sites
function get_q2_1_avx!(
    r        :: Reduced_lattice,
    p1       :: Param,
    p2       :: Param,
    ch       :: Channel,
    temp     :: SubArray{Float64, 1, Array{Float64, 3}},
    exchange :: Bool,
    sgn      :: Float64
    )        :: Nothing

    # deref channel
    q2_1 = ch.q2_1

    # deref param
    lower_index1  = p1.lower_index
    upper_index1  = p1.upper_index
    lower_weight1 = p1.lower_weight
    upper_weight1 = p1.upper_weight

    lower_index2  = p2.lower_index
    upper_index2  = p2.upper_index
    lower_weight2 = p2.lower_weight
    upper_weight2 = p2.upper_weight

    # check for site exchange
    if exchange
        indices = r.exchange

        @turbo unroll = 1 for i in eachindex(temp)
            val      = lower_weight1 * lower_weight2 * q2_1[indices[i], lower_index1, lower_index2]
            val     += upper_weight1 * lower_weight2 * q2_1[indices[i], upper_index1, lower_index2]
            val     += lower_weight1 * upper_weight2 * q2_1[indices[i], lower_index1, upper_index2]
            val     += upper_weight1 * upper_weight2 * q2_1[indices[i], upper_index1, upper_index2]
            temp[i] += sgn * val
        end
    else
        @turbo unroll = 1 for i in eachindex(temp)
            val      = lower_weight1 * lower_weight2 * q2_1[i, lower_index1, lower_index2]
            val     += upper_weight1 * lower_weight2 * q2_1[i, upper_index1, lower_index2]
            val     += lower_weight1 * upper_weight2 * q2_1[i, lower_index1, upper_index2]
            val     += upper_weight1 * upper_weight2 * q2_1[i, upper_index1, upper_index2]
            temp[i] += sgn * val
        end
    end

    return nothing
end

# get interpolated value of q2_2
function get_q2_2(
    site :: Int64,
    p1   :: Param,
    p2   :: Param,
    ch   :: Channel
    )    :: Float64

    val  = p1.lower_weight * p2.lower_weight * ch.q2_2[site, p1.lower_index, p2.lower_index]
    val += p1.upper_weight * p2.lower_weight * ch.q2_2[site, p1.upper_index, p2.lower_index]
    val += p1.lower_weight * p2.upper_weight * ch.q2_2[site, p1.lower_index, p2.upper_index]
    val += p1.upper_weight * p2.upper_weight * ch.q2_2[site, p1.upper_index, p2.upper_index]

    return val
end

# get interpolated value of q2_2 on all lattice sites
function get_q2_2_avx!(
    r        :: Reduced_lattice,
    p1       :: Param,
    p2       :: Param,
    ch       :: Channel,
    temp     :: SubArray{Float64, 1, Array{Float64, 3}},
    exchange :: Bool,
    sgn      :: Float64
    )        :: Nothing

    # deref channel
    q2_2 = ch.q2_2

    # deref param
    lower_index1  = p1.lower_index
    upper_index1  = p1.upper_index
    lower_weight1 = p1.lower_weight
    upper_weight1 = p1.upper_weight

    lower_index2  = p2.lower_index
    upper_index2  = p2.upper_index
    lower_weight2 = p2.lower_weight
    upper_weight2 = p2.upper_weight

    # check for site exchange
    if exchange
        indices = r.exchange

        @turbo unroll = 1 for i in eachindex(temp)
            val      = lower_weight1 * lower_weight2 * q2_2[indices[i], lower_index1, lower_index2]
            val     += upper_weight1 * lower_weight2 * q2_2[indices[i], upper_index1, lower_index2]
            val     += lower_weight1 * upper_weight2 * q2_2[indices[i], lower_index1, upper_index2]
            val     += upper_weight1 * upper_weight2 * q2_2[indices[i], upper_index1, upper_index2]
            temp[i] += sgn * val
        end
    else
        @turbo unroll = 1 for i in eachindex(temp)
            val      = lower_weight1 * lower_weight2 * q2_2[i, lower_index1, lower_index2]
            val     += upper_weight1 * lower_weight2 * q2_2[i, upper_index1, lower_index2]
            val     += lower_weight1 * upper_weight2 * q2_2[i, lower_index1, upper_index2]
            val     += upper_weight1 * upper_weight2 * q2_2[i, upper_index1, upper_index2]
            temp[i] += sgn * val
        end
    end

    return nothing
end

# get interpolated value of q3
function get_q3(
    site :: Int64,
    p1   :: Param,
    p2   :: Param,
    p3   :: Param,
    ch   :: Channel
    )    :: Float64

    val  = p1.lower_weight * p2.lower_weight * p3.lower_weight * ch.q3[site, p1.lower_index, p2.lower_index, p3.lower_index]
    val += p1.upper_weight * p2.lower_weight * p3.lower_weight * ch.q3[site, p1.upper_index, p2.lower_index, p3.lower_index]
    val += p1.lower_weight * p2.upper_weight * p3.lower_weight * ch.q3[site, p1.lower_index, p2.upper_index, p3.lower_index]
    val += p1.upper_weight * p2.upper_weight * p3.lower_weight * ch.q3[site, p1.upper_index, p2.upper_index, p3.lower_index]
    val += p1.lower_weight * p2.lower_weight * p3.upper_weight * ch.q3[site, p1.lower_index, p2.lower_index, p3.upper_index]
    val += p1.upper_weight * p2.lower_weight * p3.upper_weight * ch.q3[site, p1.upper_index, p2.lower_index, p3.upper_index]
    val += p1.lower_weight * p2.upper_weight * p3.upper_weight * ch.q3[site, p1.lower_index, p2.upper_index, p3.upper_index]
    val += p1.upper_weight * p2.upper_weight * p3.upper_weight * ch.q3[site, p1.upper_index, p2.upper_index, p3.upper_index]

    return val
end

# get interpolated value of q3 on all lattice sites
function get_q3_avx!(
    r        :: Reduced_lattice,
    p1       :: Param,
    p2       :: Param,
    p3       :: Param,
    ch       :: Channel,
    temp     :: SubArray{Float64, 1, Array{Float64, 3}},
    exchange :: Bool,
    sgn      :: Float64
    )        :: Nothing

    # deref channel
    q3 = ch.q3

    # deref param
    lower_index1  = p1.lower_index
    upper_index1  = p1.upper_index
    lower_weight1 = p1.lower_weight
    upper_weight1 = p1.upper_weight

    lower_index2  = p2.lower_index
    upper_index2  = p2.upper_index
    lower_weight2 = p2.lower_weight
    upper_weight2 = p2.upper_weight

    lower_index3  = p3.lower_index
    upper_index3  = p3.upper_index
    lower_weight3 = p3.lower_weight
    upper_weight3 = p3.upper_weight

    # check for site exchange
    if exchange
        indices = r.exchange

        @turbo unroll = 1 for i in eachindex(temp)
            val      = lower_weight1 * lower_weight2 * lower_weight3 * q3[indices[i], lower_index1, lower_index2, lower_index3]
            val     += upper_weight1 * lower_weight2 * lower_weight3 * q3[indices[i], upper_index1, lower_index2, lower_index3]
            val     += lower_weight1 * upper_weight2 * lower_weight3 * q3[indices[i], lower_index1, upper_index2, lower_index3]
            val     += upper_weight1 * upper_weight2 * lower_weight3 * q3[indices[i], upper_index1, upper_index2, lower_index3]
            val     += lower_weight1 * lower_weight2 * upper_weight3 * q3[indices[i], lower_index1, lower_index2, upper_index3]
            val     += upper_weight1 * lower_weight2 * upper_weight3 * q3[indices[i], upper_index1, lower_index2, upper_index3]
            val     += lower_weight1 * upper_weight2 * upper_weight3 * q3[indices[i], lower_index1, upper_index2, upper_index3]
            val     += upper_weight1 * upper_weight2 * upper_weight3 * q3[indices[i], upper_index1, upper_index2, upper_index3]
            temp[i] += sgn * val
        end
    else
        @turbo unroll = 1 for i in eachindex(temp)
            val      = lower_weight1 * lower_weight2 * lower_weight3 * q3[i, lower_index1, lower_index2, lower_index3]
            val     += upper_weight1 * lower_weight2 * lower_weight3 * q3[i, upper_index1, lower_index2, lower_index3]
            val     += lower_weight1 * upper_weight2 * lower_weight3 * q3[i, lower_index1, upper_index2, lower_index3]
            val     += upper_weight1 * upper_weight2 * lower_weight3 * q3[i, upper_index1, upper_index2, lower_index3]
            val     += lower_weight1 * lower_weight2 * upper_weight3 * q3[i, lower_index1, lower_index2, upper_index3]
            val     += upper_weight1 * lower_weight2 * upper_weight3 * q3[i, upper_index1, lower_index2, upper_index3]
            val     += lower_weight1 * upper_weight2 * upper_weight3 * q3[i, lower_index1, upper_index2, upper_index3]
            val     += upper_weight1 * upper_weight2 * upper_weight3 * q3[i, upper_index1, upper_index2, upper_index3]
            temp[i] += sgn * val
        end
    end

    return nothing
end

# get interpolated value of channel for a given frequency buffer
function get_channel(
    site :: Int64,
    b    :: Buffer,
    ch   :: Channel
    )    :: Float64

    val = 0.0

    if b.kernel == 1
        val = get_q1(site, b.p1, ch)
    elseif b.kernel == 2
        val = get_q2_1(site, b.p1, b.p2, ch)
    elseif b.kernel == 3
        val = get_q2_2(site, b.p1, b.p3, ch)
    elseif b.kernel == 4
        val = get_q3(site, b.p1, b.p2, b.p3, ch)
    end

    return val
end

# get interpolated value of channel for a given frequency buffer on all lattice sites
function get_channel_avx!(
    r        :: Reduced_lattice,
    b        :: Buffer,
    ch       :: Channel,
    temp     :: SubArray{Float64, 1, Array{Float64, 3}},
    exchange :: Bool,
    sgn      :: Float64
    )        :: Nothing

    if b.kernel == 1
        get_q1_avx!(r, b.p1, ch, temp, exchange, sgn)
    elseif b.kernel == 2
        get_q2_1_avx!(r, b.p1, b.p2, ch, temp, exchange, sgn)
    elseif b.kernel == 3
        get_q2_2_avx!(r, b.p1, b.p3, ch, temp, exchange, sgn)
    elseif b.kernel == 4
        get_q3_avx!(r, b.p1, b.p2, b.p3, ch, temp, exchange, sgn)
    end

    return nothing
end

# replace channel with another channel
function replace_with!(
    ch1 :: Channel,
    ch2 :: Channel
    )   :: Nothing

    @turbo ch1.q1   .= ch2.q1
    @turbo ch1.q2_1 .= ch2.q2_1
    @turbo ch1.q2_2 .= ch2.q2_2
    @turbo ch1.q3   .= ch2.q3

    return nothing
end

# multiply channel with factor
function mult_with!(
    ch  :: Channel,
    fac :: Float64
    )   :: Nothing

    @turbo ch.q1   .*= fac
    @turbo ch.q2_1 .*= fac
    @turbo ch.q2_2 .*= fac
    @turbo ch.q3   .*= fac

    return nothing
end

# multiply channel with factor and add to other channel
function mult_with_add_to!(
    ch2 :: Channel,
    fac :: Float64,
    ch1 :: Channel
    )   :: Nothing

    @turbo ch1.q1   .+= fac .* ch2.q1
    @turbo ch1.q2_1 .+= fac .* ch2.q2_1
    @turbo ch1.q2_2 .+= fac .* ch2.q2_2
    @turbo ch1.q3   .+= fac .* ch2.q3

    return nothing
end

"""
    get_abs_max(
        ch :: Channel
        )  :: Float64

Returns maximum absolute value of a channel.
"""
function get_abs_max(
    ch :: Channel
    )  :: Float64

    max_ch = maximum(abs.(ch.q3))

    return max_ch
end

# resample a channel to new meshes via trilinear interpolation
function resample_from_to!(
    Ω_old  :: Vector{Float64},
    ν_old  :: Vector{Float64},
    ch_old :: Channel,
    Ω_new  :: Vector{Float64},
    ν_new  :: Vector{Float64},
    ch_new :: Channel
    )      :: Nothing

    # get dimensions
    num_sites = size(ch_old.q2_1, 1)
    num_Ω     = size(ch_old.q2_1, 2)
    num_ν     = size(ch_old.q2_1, 3)

    # resample q1
    for w in 1 : num_Ω
        # get interpolation parameters
        p1 = get_param(Ω_new[w], Ω_old)

        for site in 1 : num_sites
            ch_new.q1[site, w] = get_q1(site, p1, ch_old)
        end
    end

    # resample q2_1 and q2_2
    for v in 1 : num_ν
        for w in 1 : num_Ω
            # get interpolation parameters
            p1 = get_param(Ω_new[w], Ω_old)
            p2 = get_param(ν_new[v], ν_old)

            for site in 1 : num_sites
                ch_new.q2_1[site, w, v] = get_q2_1(site, p1, p2, ch_old)
                ch_new.q2_2[site, w, v] = get_q2_2(site, p1, p2, ch_old)
            end
        end
    end

    # resample q3
    for vp in 1 : num_ν
        for v in 1 : num_ν
            for w in 1 : num_Ω
                # get interpolation parameters
                p1 = get_param(Ω_new[w],  Ω_old)
                p2 = get_param(ν_new[v],  ν_old)
                p3 = get_param(ν_new[vp], ν_old)

                for site in 1 : num_sites
                    ch_new.q3[site, w, v, vp] = get_q3(site, p1, p2, p3, ch_old)
                end
            end
        end
    end

    return nothing
end


function add_variable_cost!(
    container::OptimizationContainer,
    ::U,
    devices::IS.FlattenIteratorWrapper{T},
    ::V,
) where {T <: PSY.Component, U <: VariableType, V <: AbstractDeviceFormulation}
    for d in devices
        op_cost_data = PSY.get_operation_cost(d)
        _add_variable_cost_to_objective!(container, U(), d, op_cost_data, V())
    end
    return
end

function add_variable_cost!(
    container::OptimizationContainer,
    ::U,
    service::T,
    ::V,
) where {T <: PSY.ReserveDemandCurve, U <: VariableType, V <: StepwiseCostReserve}
    _add_variable_cost_to_objective!(container, U(), service, V())
    return
end

function add_shut_down_cost!(
    container::OptimizationContainer,
    ::U,
    devices::IS.FlattenIteratorWrapper{T},
    ::V,
) where {T <: PSY.Component, U <: VariableType, V <: AbstractDeviceFormulation}
    multiplier = objective_function_multiplier(U(), V())
    for d in devices
        op_cost_data = PSY.get_operation_cost(d)
        cost_term = shut_down_cost(op_cost_data, d, V())
        iszero(cost_term) && continue
        for t in get_time_steps(container)
            _add_proportional_term!(container, U(), d, cost_term * multiplier, t)
        end
    end
    return
end

function add_proportional_cost!(
    container::OptimizationContainer,
    ::U,
    devices::IS.FlattenIteratorWrapper{T},
    ::V,
) where {T <: PSY.Component, U <: VariableType, V <: AbstractDeviceFormulation}
    multiplier = objective_function_multiplier(U(), V())
    for d in devices
        op_cost_data = PSY.get_operation_cost(d)
        cost_term = proportional_cost(op_cost_data, U(), d, V())
        iszero(cost_term) && continue
        for t in get_time_steps(container)
            _add_proportional_term!(container, U(), d, cost_term * multiplier, t)
        end
    end
    return
end

function add_proportional_cost!(
    container::OptimizationContainer,
    ::U,
    devices::IS.FlattenIteratorWrapper{T},
    ::V,
) where {T <: PSY.ThermalGen, U <: OnVariable, V <: AbstractCompactUnitCommitment}
    multiplier = objective_function_multiplier(U(), V())
    for d in devices
        op_cost_data = PSY.get_operation_cost(d)
        cost_term = proportional_cost(op_cost_data, U(), d, V())
        iszero(cost_term) && continue
        for t in get_time_steps(container)
            exp = _add_proportional_term!(container, U(), d, cost_term * multiplier, t)
            add_to_expression!(container, ProductionCostExpression, exp, d, t)
        end
    end
    return
end

function add_proportional_cost!(
    container::OptimizationContainer,
    ::U,
    devices::IS.FlattenIteratorWrapper{T},
    ::V,
) where {
    T <: PSY.Storage,
    U <: Union{ActivePowerInVariable, ActivePowerOutVariable},
    V <: AbstractDeviceFormulation,
}
    multiplier = objective_function_multiplier(U(), V())
    for d in devices
        for t in get_time_steps(container)
            _add_proportional_term!(container, U(), d, COST_EPSILON * multiplier, t)
        end
    end
    return
end

function add_proportional_cost!(
    container::OptimizationContainer,
    ::U,
    devices::IS.FlattenIteratorWrapper{T},
    ::V,
) where {
    T <: PSY.Component,
    U <: Union{EnergySurplusVariable, EnergyShortageVariable},
    V <: AbstractDeviceFormulation,
}
    base_p = get_base_power(container)
    multiplier = objective_function_multiplier(U(), V())
    for d in devices
        op_cost_data = PSY.get_operation_cost(d)
        cost_term = proportional_cost(op_cost_data, U(), d, V())
        iszero(cost_term) && continue
        for t in get_time_steps(container)
            _add_proportional_term!(container, U(), d, cost_term * multiplier * base_p, t)
        end
    end
    return
end

function add_proportional_cost!(
    container::OptimizationContainer,
    ::U,
    service::T,
    ::V,
) where {
    T <: Union{PSY.Reserve, PSY.ReserveNonSpinning},
    U <: ActivePowerReserveVariable,
    V <: AbstractReservesFormulation,
}
    base_p = get_base_power(container)
    reserve_variable = get_variable(container, U(), T, PSY.get_name(service))
    for index in Iterators.product(axes(reserve_variable)...)
        add_to_objective_invariant_expression!(
            container,
            DEFAULT_RESERVE_COST / base_p * reserve_variable[index...],
        )
    end
    return
end

function add_proportional_cost!(
    container::OptimizationContainer,
    ::U,
    areas::IS.FlattenIteratorWrapper{T},
    ::PIDSmoothACE,
) where {T <: PSY.Area, U <: LiftVariable}
    lift_variable = get_variable(container, U(), T)
    for index in Iterators.product(axes(lift_variable)...)
        add_to_objective_invariant_expression!(
            container,
            SERVICES_SLACK_COST * lift_variable[index...],
        )
    end
    return
end

function _add_variable_cost_to_objective!(
    container::OptimizationContainer,
    ::T,
    component::PSY.Component,
    op_cost::PSY.OperationalCost,
    ::U,
) where {T <: VariableType, U <: AbstractDeviceFormulation}
    variable_cost_data = variable_cost(op_cost, T(), component, U())
    _add_variable_cost_to_objective!(container, T(), component, variable_cost_data, U())
    return
end

function _add_variable_cost_to_objective!(
    container::OptimizationContainer,
    ::T,
    component::PSY.Component,
    op_cost::PSY.MarketBidCost,
    ::U,
) where {T <: VariableType, U <: AbstractDeviceFormulation}
    component_name = PSY.get_name(component)
    @debug "Market Bid" _group = LOG_GROUP_COST_FUNCTIONS component_name
    time_steps = get_time_steps(container)
    initial_time = get_initial_time(container)
    variable_cost_forecast = PSY.get_variable_cost(
        component,
        op_cost;
        start_time=initial_time,
        len=length(time_steps),
    )
    variable_cost_forecast_values = TimeSeries.values(variable_cost_forecast)
    parameter_container = _get_cost_function_parameter_container(
        container,
        CostFunctionParameter(),
        component,
        T(),
        U(),
        eltype(variable_cost_forecast_values),
    )
    pwl_cost_expressions =
        _add_pwl_term!(container, component, variable_cost_forecast_values, T(), U())
    jump_model = get_jump_model(container)
    for t in time_steps
        set_parameter!(
            parameter_container,
            jump_model,
            PSY.get_cost(variable_cost_forecast_values[t]),
            # Using 1.0 here since we want to reuse the existing code that adds the mulitpler
            #  of base power times the time delta.
            1.0,
            component_name,
            t,
        )
        add_to_expression!(
            container,
            ProductionCostExpression,
            pwl_cost_expressions[t],
            component,
            t,
        )
        add_to_objective_variant_expression!(container, pwl_cost_expressions[t])
    end

    # Service Cost Bid
    ancillary_services = PSY.get_ancillary_services(op_cost)
    for service in ancillary_services
        _add_service_bid_cost!(container, component, service)
    end
    return
end

function _add_variable_cost_to_objective!(
    container::OptimizationContainer,
    ::T,
    component::PSY.Reserve,
    ::U,
) where {T <: VariableType, U <: StepwiseCostReserve}
    component_name = PSY.get_name(component)
    @debug "PWL Variable Cost" _group = LOG_GROUP_COST_FUNCTIONS component_name
    # If array is full of tuples with zeros return 0.0
    time_steps = get_time_steps(container)
    initial_time = get_initial_time(container)
    variable_cost_forecast = get_time_series(container, component, "variable_cost")
    variable_cost_forecast_values = TimeSeries.values(variable_cost_forecast)
    variable_cost_forecast_values = map(PSY.VariableCost, variable_cost_forecast_values)
    parameter_container = _get_cost_function_parameter_container(
        container,
        CostFunctionParameter(),
        component,
        T(),
        U(),
        eltype(variable_cost_forecast_values),
    )
    pwl_cost_expressions =
        _add_pwl_term!(container, component, variable_cost_forecast_values, T(), U())
    jump_model = get_jump_model(container)
    for t in time_steps
        set_parameter!(
            parameter_container,
            jump_model,
            PSY.get_cost(variable_cost_forecast_values[t]),
            # Using 1.0 here since we want to reuse the existing code that adds the mulitpler
            #  of base power times the time delta.
            1.0,
            component_name,
            t,
        )
        add_to_objective_variant_expression!(container, pwl_cost_expressions[t])
    end
    return
end

function add_start_up_cost!(
    container::OptimizationContainer,
    ::U,
    devices::IS.FlattenIteratorWrapper{T},
    ::V,
) where {T <: PSY.Component, U <: VariableType, V <: AbstractDeviceFormulation}
    for d in devices
        op_cost_data = PSY.get_operation_cost(d)
        _add_start_up_cost_to_objective!(container, U(), d, op_cost_data, V())
    end
    return
end

function _add_start_up_cost_to_objective!(
    container::OptimizationContainer,
    ::T,
    component::PSY.Component,
    op_cost::PSY.OperationalCost,
    ::U,
) where {T <: VariableType, U <: AbstractDeviceFormulation}
    cost_term = start_up_cost(op_cost, component, U())
    iszero(cost_term) && return
    multiplier = objective_function_multiplier(T(), U())
    for t in get_time_steps(container)
        _add_proportional_term!(container, T(), component, cost_term * multiplier, t)
    end
    return
end

const MULTI_START_COST_MAP = Dict{DataType, Int}(
    HotStartVariable => 1,
    WarmStartVariable => 2,
    ColdStartVariable => 3,
)

function _add_start_up_cost_to_objective!(
    container::OptimizationContainer,
    ::T,
    component::PSY.Component,
    op_cost::Union{PSY.MultiStartCost, PSY.MarketBidCost},
    ::U,
) where {T <: VariableType, U <: ThermalMultiStartUnitCommitment}
    cost_terms = start_up_cost(op_cost, component, U())
    cost_term = cost_terms[MULTI_START_COST_MAP[T]]
    iszero(cost_term) && return
    multiplier = objective_function_multiplier(T(), U())
    for t in get_time_steps(container)
        _add_proportional_term!(container, T(), component, cost_term * multiplier, t)
    end
    return
end

_get_cost_function_data_type(::Type{PSY.VariableCost{T}}) where {T} = T

function _get_cost_function_parameter_container(
    container::OptimizationContainer,
    ::S,
    component::T,
    ::U,
    ::V,
    cost_type::DataType,
) where {
    S <: ObjectiveFunctionParameter,
    T <: PSY.Component,
    U <: VariableType,
    V <: Union{AbstractDeviceFormulation, AbstractServiceFormulation},
}
    if has_container_key(container, S, T)
        return get_parameter(container, S(), T)
    else
        container_axes = axes(get_variable(container, U(), T))
        if has_container_key(container, OnStatusParameter, T)
            sos_val = SOSStatusVariable.PARAMETER
        else
            sos_val = sos_status(component, V())
        end
        return add_param_container!(
            container,
            S(),
            T,
            sos_val,
            U,
            uses_compact_power(component, V()),
            _get_cost_function_data_type(cost_type),
            container_axes...,
        )
    end
end

function _add_service_bid_cost!(
    container::OptimizationContainer,
    component::PSY.Component,
    service::PSY.Reserve{T},
) where {T <: PSY.ReserveDirection}
    time_steps = get_time_steps(container)
    initial_time = get_initial_time(container)
    base_power = get_base_power(container)
    forecast_data = PSY.get_services_bid(
        component,
        PSY.get_operation_cost(component),
        service;
        start_time=initial_time,
        len=length(time_steps),
    )
    forecast_data_values = PSY.get_cost.(TimeSeries.values(forecast_data)) .* base_power
    reserve_variable = get_variable(container, U(), T, PSY.get_name(service))
    component_name = PSY.get_name(component)
    for t in time_steps
        add_to_objective_invariant_expression!(
            container,
            forecast_data_values[t] * reserve_variable[component_name, t],
        )
    end
end

function _add_service_bid_cost!(::OptimizationContainer, ::PSY.Component, ::PSY.Service) end

function _add_service_bid_cost!(
    ::OptimizationContainer,
    ::PSY.Component,
    service::PSY.ReserveDemandCurve{T},
) where {T <: PSY.ReserveDirection}
    error(
        "The Current version doesn't supports cost bid for ReserveDemandCurve services, \\
        please change the forecast data for $(PSY.get_name(service)) \\
        and open a feature request",
    )
    return
end

"""
Adds to the cost function cost terms for sum of variables with common factor to be used for cost expression for optimization_container model.

# Arguments

  - container::OptimizationContainer : the optimization_container model built in PowerSimulations
  - var_key::VariableKey: The variable name
  - component_name::String: The component_name of the variable container
  - cost_component::PSY.VariableCost{Float64} : container for cost to be associated with variable
"""
function _add_variable_cost_to_objective!(
    container::OptimizationContainer,
    ::T,
    component::PSY.Component,
    cost_component::PSY.VariableCost{Float64},
    ::U,
) where {T <: VariableType, U <: AbstractDeviceFormulation}
    multiplier = objective_function_multiplier(T(), U())
    base_power = get_base_power(container)
    cost_data = PSY.get_cost(cost_component)
    resolution = get_resolution(container)
    dt = Dates.value(Dates.Second(resolution)) / SECONDS_IN_HOUR
    for time_period in get_time_steps(container)
        linear_cost = _add_proportional_term!(
            container,
            T(),
            component,
            cost_data * multiplier * base_power * dt,
            time_period,
        )
        add_to_expression!(
            container,
            ProductionCostExpression,
            linear_cost,
            component,
            time_period,
        )
    end
    return
end

@doc raw"""
Adds to the cost function cost terms for sum of variables with common factor to be used for cost expression for optimization_container model.

# Equation

``` gen_cost = dt*sign*(sum(variable.^2)*cost_data[1] + sum(variable)*cost_data[2]) ```

# LaTeX

`` cost = dt\times sign (sum_{i\in I} c_1 v_i^2 + sum_{i\in I} c_2 v_i ) ``

for quadratic factor large enough. If the first term of the quadratic objective is 0.0, adds a
linear cost term `sum(variable)*cost_data[2]`

# Arguments

* container::OptimizationContainer : the optimization_container model built in PowerSimulations
* var_key::VariableKey: The variable name
* component_name::String: The component_name of the variable container
* cost_component::PSY.VariableCost{NTuple{2, Float64}} : container for quadratic and linear factors
"""
function _add_variable_cost_to_objective!(
    container::OptimizationContainer,
    ::T,
    component::PSY.Component,
    cost_component::PSY.VariableCost{NTuple{2, Float64}},
    ::U,
) where {T <: VariableType, U <: AbstractDeviceFormulation}
    multiplier = objective_function_multiplier(T(), U())
    base_power = get_base_power(container)
    cost_data = PSY.get_cost(cost_component)
    resolution = get_resolution(container)
    dt = Dates.value(Dates.Second(resolution)) / SECONDS_IN_HOUR
    for time_period in get_time_steps(container)
        if cost_data[1] >= eps()
            cost_term = _add_quadratic_term!(
                container,
                T(),
                component,
                cost_data,
                base_power,
                multiplier * dt,
                time_period,
            )
        else
            cost_term = _add_proportional_term!(
                container,
                T(),
                component,
                cost_data[2] * multiplier * base_power * dt,
                time_period,
            )
        end
        add_to_expression!(
            container,
            ProductionCostExpression,
            cost_term,
            component,
            time_period,
        )
    end
    return
end

"""
Creates piecewise linear cost function using a sum of variables and expression with sign and time step included.

# Arguments

  - container::OptimizationContainer : the optimization_container model built in PowerSimulations
  - var_key::VariableKey: The variable name
  - component_name::String: The component_name of the variable container
  - cost_component::PSY.VariableCost{Vector{NTuple{2, Float64}}}
"""
function _add_variable_cost_to_objective!(
    container::OptimizationContainer,
    ::T,
    component::PSY.Component,
    cost_component::PSY.VariableCost{Vector{NTuple{2, Float64}}},
    ::U,
) where {T <: VariableType, U <: AbstractDeviceFormulation}
    component_name = PSY.get_name(component)
    @debug "PWL Variable Cost" _group = LOG_GROUP_COST_FUNCTIONS component_name
    # If array is full of tuples with zeros return 0.0
    cost_data = PSY.get_cost(cost_component)
    if all(iszero.(last.(cost_data)))
        @debug "All cost terms for component $(component_name) are 0.0" _group =
            LOG_GROUP_COST_FUNCTIONS
        return
    end
    pwl_cost_expressions = _add_pwl_term!(container, component, cost_data, T(), U())
    for t in get_time_steps(container)
        add_to_expression!(
            container,
            ProductionCostExpression,
            pwl_cost_expressions[t],
            component,
            t,
        )
        add_to_objective_invariant_expression!(container, pwl_cost_expressions[t])
    end
    return
end

"""
Returns True/False depending on compatibility of the cost data with the convex implementation method
"""
function _slope_convexity_check(slopes::Vector{Float64})
    flag = true
    for ix in 1:(length(slopes) - 1)
        if slopes[ix] > slopes[ix + 1]
            @debug slopes _group = LOG_GROUP_COST_FUNCTIONS
            return flag = false
        end
    end
    return flag
end

function _check_pwl_compact_data(
    min::Float64,
    max::Float64,
    data::Vector{Tuple{Float64, Float64}},
    base_power::Float64,
)
    return isapprox(max - min, data[end][2] / base_power) && iszero(data[1][2])
end

function _check_pwl_compact_data(
    d::PSY.Component,
    data::Vector{Tuple{Float64, Float64}},
    base_power::Float64,
)
    min = PSY.get_active_power_limits(d).min
    max = PSY.get_active_power_limits(d).max
    return _check_pwl_compact_data(min, max, data, base_power)
end

function _get_sos_value(
    container::OptimizationContainer,
    ::Type{V},
    component::T,
) where {T <: PSY.Component, V <: AbstractDeviceFormulation}
    if has_container_key(container, OnStatusParameter, T)
        sos_val = SOSStatusVariable.PARAMETER
    else
        sos_val = sos_status(component, V())
    end
    return sos_val
end

function _get_sos_value(
    container::OptimizationContainer,
    ::Type{V},
    component::T,
) where {T <: PSY.Component, V <: AbstractServiceFormulation}
    return SOSStatusVariable.NO_VARIABLE
end

function _add_pwl_term!(
    container::OptimizationContainer,
    component::T,
    cost_data::Vector{PSY.VariableCost{Float64}},
    ::U,
    ::V,
) where {T <: PSY.Component, U <: VariableType, V <: AbstractDeviceFormulation}
    multiplier = objective_function_multiplier(U(), V())
    resolution = get_resolution(container)
    dt = Dates.value(Dates.Second(resolution)) / SECONDS_IN_HOUR
    base_power = get_base_power(container)
    # Re-scale breakpoints by Basepower
    time_steps = get_time_steps(container)
    cost_expressions = Vector{JuMP.AffExpr}(undef, time_steps[end])
    for t in time_steps
        proportial_value = PSY.get_cost(cost_data[t]) * multiplier * base_power * dt
        cost_expressions[t] =
            _add_proportional_term!(container, U(), component, proportial_value, t)
    end
    return cost_expressions
end

"""
Add PWL cost terms for data coming from the MarketBidCost
"""
function _add_pwl_term!(
    container::OptimizationContainer,
    component::T,
    cost_data::Vector{PSY.VariableCost{Vector{Tuple{Float64, Float64}}}},
    ::U,
    ::V,
) where {T <: PSY.Component, U <: VariableType, V <: AbstractDeviceFormulation}
    multiplier = objective_function_multiplier(U(), V())
    resolution = get_resolution(container)
    dt = Dates.value(Dates.Second(resolution)) / SECONDS_IN_HOUR
    base_power = get_base_power(container)
    # Re-scale breakpoints by Basepower
    name = PSY.get_name(component)
    time_steps = get_time_steps(container)
    pwl_cost_expressions = Vector{JuMP.AffExpr}(undef, time_steps[end])
    sos_val = _get_sos_value(container, V, component)
    for t in time_steps
        data = PSY.get_cost(cost_data[t])
        is_power_data_compact = _check_pwl_compact_data(component, data, base_power)
        if !uses_compact_power(component, V()) && is_power_data_compact
            error(
                "The data provided is not compatible with formulation $V. Use a formulation compatible with Compact Cost Functions",
            )
            # data = _convert_to_full_variable_cost(data, component)
        elseif uses_compact_power(component, V()) && !is_power_data_compact
            data = _convert_to_compact_variable_cost(data)
        else
            @debug uses_compact_power(component, V()) name T V
            @debug is_power_data_compact name T V
        end
        slopes = PSY.get_slopes(data)
        # First element of the return is the average cost at P_min.
        # Shouldn't be passed for convexity check
        is_convex = _slope_convexity_check(slopes[2:end])
        break_points = map(x -> last(x), data) ./ base_power
        _add_pwl_variables!(container, T, name, t, data)
        _add_pwl_constraint!(container, component, U(), break_points, sos_val, t)
        if !is_convex
            _add_pwl_sos_constraint!(container, component, U(), break_points, sos_val, t)
        end
        pwl_cost = _get_pwl_cost_expression(container, component, t, data, multiplier * dt)
        pwl_cost_expressions[t] = pwl_cost
    end
    return pwl_cost_expressions
end

function _add_pwl_term!(
    container::OptimizationContainer,
    component::T,
    cost_data::Vector{PSY.VariableCost{Vector{Tuple{Float64, Float64}}}},
    ::U,
    ::V,
) where {T <: PSY.Component, U <: VariableType, V <: AbstractServiceFormulation}
    multiplier = objective_function_multiplier(U(), V())
    resolution = get_resolution(container)
    dt = Dates.value(Dates.Second(resolution)) / SECONDS_IN_HOUR
    base_power = get_base_power(container)
    # Re-scale breakpoints by Basepower
    name = PSY.get_name(component)
    time_steps = get_time_steps(container)
    pwl_cost_expressions = Vector{JuMP.AffExpr}(undef, time_steps[end])
    sos_val = _get_sos_value(container, V, component)
    for t in time_steps
        data = PSY.get_cost(cost_data[t])
        slopes = PSY.get_slopes(data)
        # Shouldn't be passed for convexity check
        is_convex = false
        break_points = map(x -> last(x), data) ./ base_power
        _add_pwl_variables!(container, T, name, t, data)
        _add_pwl_constraint!(container, component, U(), break_points, sos_val, t)
        if !is_convex
            _add_pwl_sos_constraint!(container, component, U(), break_points, sos_val, t)
        end
        pwl_cost = _get_pwl_cost_expression(container, component, t, data, multiplier * dt)
        pwl_cost_expressions[t] = pwl_cost
    end
    return pwl_cost_expressions
end

"""
Add PWL cost terms for data coming from a constant PWL cost function
"""
function _add_pwl_term!(
    container::OptimizationContainer,
    component::T,
    data::Vector{NTuple{2, Float64}},
    ::U,
    ::V,
) where {T <: PSY.Component, U <: VariableType, V <: AbstractDeviceFormulation}
    multiplier = objective_function_multiplier(U(), V())
    resolution = get_resolution(container)
    dt = Dates.value(Dates.Second(resolution)) / SECONDS_IN_HOUR
    base_power = get_base_power(container)
    # Re-scale breakpoints by Basepower
    name = PSY.get_name(component)

    is_power_data_compact = _check_pwl_compact_data(component, data, base_power)

    if !uses_compact_power(component, V()) && is_power_data_compact
        error(
            "The data provided is not compatible with formulation $V. Use a formulation compatible with Compact Cost Functions",
        )
        # data = _convert_to_full_variable_cost(data, component)
    elseif uses_compact_power(component, V()) && !is_power_data_compact
        data = _convert_to_compact_variable_cost(data)
    else
        @debug uses_compact_power(component, V()) name T V
        @debug is_power_data_compact name T V
    end

    slopes = PSY.get_slopes(data)
    # First element of the return is the average cost at P_min.
    # Shouldn't be passed for convexity check
    is_convex = _slope_convexity_check(slopes[2:end])
    time_steps = get_time_steps(container)
    pwl_cost_expressions = Vector{JuMP.AffExpr}(undef, time_steps[end])
    break_points = map(x -> last(x), data) ./ base_power
    sos_val = _get_sos_value(container, V, component)
    for t in time_steps
        _add_pwl_variables!(container, T, name, t, data)
        _add_pwl_constraint!(container, component, U(), break_points, sos_val, t)
        if !is_convex
            _add_pwl_sos_constraint!(container, component, U(), break_points, sos_val, t)
        end
        pwl_cost = _get_pwl_cost_expression(container, component, t, data, multiplier * dt)
        pwl_cost_expressions[t] = pwl_cost
    end
    return pwl_cost_expressions
end

function _add_pwl_term!(
    container::OptimizationContainer,
    component::T,
    data::Vector{NTuple{2, Float64}},
    ::U,
    ::V,
) where {T <: PSY.ThermalGen, U <: VariableType, V <: ThermalDispatchNoMin}
    multiplier = objective_function_multiplier(U(), V())
    resolution = get_resolution(container)
    dt = Dates.value(Dates.Second(resolution)) / SECONDS_IN_HOUR
    component_name = PSY.get_name(component)
    @debug "PWL cost function detected for device $(component_name) using $V"
    base_power = get_base_power(container)
    slopes = PSY.get_slopes(data)
    if any(slopes .< 0) || !_slope_convexity_check(slopes[2:end])
        throw(
            IS.InvalidValue(
                "The PWL cost data provided for generator $(component_name) is not compatible with $U.",
            ),
        )
    end

    if _check_pwl_compact_data(component, data, base_power)
        error("The data provided is not compatible with formulation $V. \\
              Use a formulation compatible with Compact Cost Functions")
    end

    if slopes[1] != 0.0
        @debug "PWL has no 0.0 intercept for generator $(component_name)"
        # adds a first intercept a x = 0.0 and Y below the intercept of the first tuple to make convex equivalent
        first_pair = data[1]
        intercept_point = (0.0, first_pair[2] - COST_EPSILON)
        data = vcat(intercept_point, data)
        @assert _slope_convexity_check(slopes)
    end

    time_steps = get_time_steps(container)
    pwl_cost_expressions = Vector{JuMP.AffExpr}(undef, time_steps[end])
    break_points = map(x -> last(x), data) ./ base_power
    sos_val = _get_sos_value(container, V, component)
    for t in time_steps
        _add_pwl_variables!(container, T, component_name, t, data)
        _add_pwl_constraint!(container, component, U(), break_points, sos_val, t)
        pwl_cost = _get_pwl_cost_expression(container, component, t, data, multiplier * dt)
        pwl_cost_expressions[t] = pwl_cost
    end
    return pwl_cost_expressions
end

function _add_pwl_variables!(
    container::OptimizationContainer,
    ::Type{T},
    component_name::String,
    time_period::Int,
    cost_data::Vector{NTuple{2, Float64}},
) where {T <: PSY.Component}
    var_container = lazy_container_addition!(container, PieceWiseLinearCostVariable(), T)
    pwlvars = Array{JuMP.VariableRef}(undef, length(cost_data))
    for i in 1:length(cost_data)
        pwlvars[i] =
            var_container[(component_name, i, time_period)] = JuMP.@variable(
                get_jump_model(container),
                base_name = "PieceWiseLinearCostVariable_$(component_name)_{pwl_$(i), $time_period}",
                lower_bound = 0.0,
                upper_bound = 1.0
            )
    end
    return pwlvars
end

function _add_pwl_constraint!(
    container::OptimizationContainer,
    component::T,
    ::U,
    break_points::Vector{Float64},
    sos_status::SOSStatusVariable,
    period::Int,
) where {T <: PSY.Component, U <: VariableType}
    variables = get_variable(container, U(), T)
    const_container = lazy_container_addition!(
        container,
        PieceWiseLinearCostConstraint(),
        T,
        axes(variables)...,
    )
    len_cost_data = length(break_points)
    jump_model = get_jump_model(container)
    pwl_vars = get_variable(container, PieceWiseLinearCostVariable(), T)
    name = PSY.get_name(component)
    const_container[name, period] = JuMP.@constraint(
        jump_model,
        variables[name, period] ==
        sum(pwl_vars[name, ix, period] * break_points[ix] for ix in 1:len_cost_data)
    )

    if sos_status == SOSStatusVariable.NO_VARIABLE
        bin = 1.0
        @debug "Using Piecewise Linear cost function but no variable/parameter ref for ON status is passed. Default status will be set to online (1.0)" _group =
            LOG_GROUP_COST_FUNCTIONS

    elseif sos_status == SOSStatusVariable.PARAMETER
        bin = get_parameter(container, OnStatusParameter(), T).parameter_array[name, period]
        @debug "Using Piecewise Linear cost function with parameter OnStatusParameter, $T" _group =
            LOG_GROUP_COST_FUNCTIONS
    elseif sos_status == SOSStatusVariable.VARIABLE
        bin = get_variable(container, OnVariable(), T)[name, period]
        @debug "Using Piecewise Linear cost function with variable OnVariable $T" _group =
            LOG_GROUP_COST_FUNCTIONS
    else
        @assert false
    end

    JuMP.@constraint(
        jump_model,
        sum(pwl_vars[name, i, period] for i in 1:len_cost_data) == bin
    )
    return
end

function _add_pwl_sos_constraint!(
    container::OptimizationContainer,
    component::T,
    ::U,
    break_points::Vector{Float64},
    sos_status::SOSStatusVariable,
    period::Int,
) where {T <: PSY.Component, U <: VariableType}
    name = PSY.get_name(component)
    @warn(
        "The cost function provided for $(name) is not compatible with a linear PWL cost function.
  An SOS-2 formulation will be added to the model. This will result in additional binary variables."
    )

    jump_model = get_jump_model(container)
    pwl_vars = get_variable(container, PieceWiseLinearCostVariable(), T)
    bp_count = length(break_points)
    pwl_vars_subset = [pwl_vars[name, i, period] for i in 1:bp_count]
    JuMP.@constraint(jump_model, pwl_vars_subset in MOI.SOS2(collect(1:bp_count)))
    return
end

function _get_pwl_cost_expression(
    container::OptimizationContainer,
    component::T,
    time_period::Int,
    cost_data::Vector{NTuple{2, Float64}},
    multiplier::Float64,
) where {T <: PSY.Component}
    name = PSY.get_name(component)
    pwl_var_container = get_variable(container, PieceWiseLinearCostVariable(), T)
    gen_cost = JuMP.AffExpr(0.0)
    slopes = PSY.get_slopes(cost_data)
    upb = PSY.get_breakpoint_upperbounds(cost_data)
    for i in 1:length(cost_data)
        JuMP.add_to_expression!(
            gen_cost,
            cost_data[i][1] * multiplier * pwl_var_container[(name, i, time_period)],
        )
    end
    return gen_cost
end

function _get_no_load_cost(
    component::T,
    ::V,
    ::U,
) where {T <: PSY.Component, U <: VariableType, V <: AbstractDeviceFormulation}
    return no_load_cost(PSY.get_operation_cost(component), U(), component, V())
end

function _convert_to_compact_variable_cost(
    var_cost::Vector{NTuple{2, Float64}},
    no_load_cost::Float64,
    p_min::Float64,
)
    return [(c - no_load_cost, pp - p_min) for (c, pp) in var_cost]
end

function _convert_to_compact_variable_cost(var_cost::Vector{NTuple{2, Float64}})
    no_load_cost, p_min = var_cost[1]
    return _convert_to_compact_variable_cost(var_cost, no_load_cost, p_min)
end

function _add_proportional_term!(
    container::OptimizationContainer,
    ::T,
    component::U,
    linear_term::Float64,
    time_period::Int,
) where {T <: VariableType, U <: PSY.Component}
    component_name = PSY.get_name(component)
    @debug "Linear Variable Cost" _group = LOG_GROUP_COST_FUNCTIONS component_name
    variable = get_variable(container, T(), U)[component_name, time_period]
    lin_cost = variable * linear_term
    add_to_objective_invariant_expression!(container, lin_cost)
    return lin_cost
end

function _add_quadratic_term!(
    container::OptimizationContainer,
    ::T,
    component::U,
    q_terms::NTuple{2, Float64},
    var_multiplier::Float64,
    expression_multiplier::Float64,
    time_period::Int,
) where {T <: VariableType, U <: PSY.Component}
    component_name = PSY.get_name(component)
    @debug "$component_name Quadratic Variable Cost" _group = LOG_GROUP_COST_FUNCTIONS component_name
    var = get_variable(container, T(), U)[component_name, time_period]
    q_cost_ = (var * var_multiplier) .^ 2 * q_terms[1] + var * var_multiplier * q_terms[2]
    q_cost = q_cost_ * expression_multiplier
    add_to_objective_invariant_expression!(container, q_cost)
    return q_cost
end

function _add_quadratic_term!(
    container::OptimizationContainer,
    ::T,
    component::U,
    q_terms::NTuple{2, Float64},
    var_multiplier::Float64,
    expression_multiplier::Float64,
    time_period::Int,
) where {T <: PowerAboveMinimumVariable, U <: PSY.ThermalGen}
    component_name = PSY.get_name(component)
    p_min = PSY.get_active_power_limits(component).min
    @debug "$component_name Quadratic Variable Cost" _group = LOG_GROUP_COST_FUNCTIONS component_name
    var = get_variable(container, T(), U)[component_name, time_period]
    q_cost_ =
        (var * var_multiplier) .^ 2 * q_terms[1] +
        var * var_multiplier * (q_terms[2] + 2 * q_terms[1] * p_min)
    q_cost = q_cost_ * expression_multiplier
    add_to_objective_invariant_expression!(container, q_cost)
    return q_cost
end