module MPS

import ..FileFormats

import MathOptInterface
const MOI = MathOptInterface

MOI.Utilities.@model(Model,
    (MOI.ZeroOne, MOI.Integer),
    (MOI.EqualTo, MOI.GreaterThan, MOI.LessThan, MOI.Interval),
    (),
    (MOI.SOS1, MOI.SOS2),
    (),
    (MOI.ScalarAffineFunction,),
    (MOI.VectorOfVariables,),
    ()
)

MOI.supports(::Model, ::MOI.ObjectiveFunction{MOI.SingleVariable}) = false

struct Options
    warn::Bool
end

get_options(m::Model) = get(m.ext, :MPS_OPTIONS, Options(false))

"""
    Model(; kwargs...)

Create an empty instance of FileFormats.MPS.Model.

Keyword arguments are:

 - `warn::Bool=false`: print a warning when variables or constraints are renamed.
"""
function Model(;
    warn::Bool = false,
)
    model = Model{Float64}()
    model.ext[:MPS_OPTIONS] = Options(warn)
    return model
end

function Base.show(io::IO, ::Model)
    print(io, "A Mathematical Programming System (MPS) model")
    return
end

@enum(VType, VTYPE_CONTINUOUS, VTYPE_INTEGER, VTYPE_BINARY)

# ==============================================================================
#
#   Base.write
#
# ==============================================================================

"""
    Base.write(io::IO, model::FileFormats.MPS.Model)

Write `model` to `io` in the MPS file format.
"""
function Base.write(io::IO, model::Model)
    options = get_options(model)
    FileFormats.create_unique_names(
        model;
        warn = options.warn,
        replacements = Function[s -> replace(s, ' ' => '_')]
    )
    ordered_names = String[]
    names = Dict{MOI.VariableIndex, String}()
    for x in MOI.get(model, MOI.ListOfVariableIndices())
        n = MOI.get(model, MOI.VariableName(), x)
        push!(ordered_names, n)
        names[x] = n
    end
    write_model_name(io, model)
    write_rows(io, model)
    discovered_columns = write_columns(io, model, ordered_names, names)
    write_rhs(io, model)
    write_ranges(io, model)
    write_bounds(io, model, discovered_columns, ordered_names, names)
    write_sos(io, model, names)
    println(io, "ENDATA")
    return
end

# ==============================================================================
#   Model name
# ==============================================================================

function write_model_name(io::IO, model::Model)
    model_name = MOI.get(model, MOI.Name())
    println(io, "NAME          ", model_name)
    return
end

# ==============================================================================
#   ROWS
# ==============================================================================

const SET_TYPES = (
    (MOI.LessThan{Float64}, 'L'),
    (MOI.GreaterThan{Float64}, 'G'),
    (MOI.EqualTo{Float64}, 'E'),
    (MOI.Interval{Float64}, 'L'),  # See the note in the RANGES section.
)

function _write_rows(io, model, S, sense_char)
    for index in MOI.get(
        model,
        MOI.ListOfConstraintIndices{MOI.ScalarAffineFunction{Float64}, S}()
    )
        row_name = MOI.get(model, MOI.ConstraintName(), index)
        if row_name == ""
            error("Row name is empty: $(index).")
        end
        println(io, " ", sense_char, "  ", row_name)
    end
    return
end

function write_rows(io::IO, model::Model)
    println(io, "ROWS\n N  OBJ")
    for (set_type, sense_char) in SET_TYPES
        _write_rows(io, model, set_type, sense_char)
    end
    return
end

# ==============================================================================
#   COLUMNS
# ==============================================================================

function _list_of_integer_variables(model, names, integer_variables, S)
    for index in MOI.get(
        model, MOI.ListOfConstraintIndices{MOI.SingleVariable, S}()
    )
        v_index = MOI.get(model, MOI.ConstraintFunction(), index)
        push!(integer_variables, names[v_index.variable])
    end
    return
end

function list_of_integer_variables(model::Model, names)
    integer_variables = Set{String}()
    for S in (MOI.ZeroOne, MOI.Integer)
        _list_of_integer_variables(model, names, integer_variables, S)
    end
    return integer_variables
end

function extract_terms(
    v_names::Dict{MOI.VariableIndex, String},
    coefficients::Dict{String, Vector{Tuple{String, Float64}}},
    row_name::String,
    func::MOI.ScalarAffineFunction,
    discovered_columns::Set{String},
    multiplier::Float64 = 1.0,
)
    for term in func.terms
        variable_name = v_names[term.variable_index]
        push!(
            coefficients[variable_name],
            (row_name, multiplier * term.coefficient)
        )
        push!(discovered_columns, variable_name)
    end
    return
end

function _collect_coefficients(
    model,
    S,
    v_names::Dict{MOI.VariableIndex, String},
    coefficients::Dict{String, Vector{Tuple{String, Float64}}},
    discovered_columns::Set{String},
)
    for index in MOI.get(
        model,
        MOI.ListOfConstraintIndices{MOI.ScalarAffineFunction{Float64}, S}(),
    )
        row_name = MOI.get(model, MOI.ConstraintName(), index)
        func = MOI.get(model, MOI.ConstraintFunction(), index)
        extract_terms(
            v_names, coefficients, row_name, func, discovered_columns
        )
    end
    return
end

function write_columns(io::IO, model::Model, ordered_names, names)
    coefficients = Dict{String, Vector{Tuple{String, Float64}}}(
        n => Tuple{String, Float64}[] for n in ordered_names
    )
    # Many MPS readers (e.g., CPLEX and GAMS) will error if a variable (column)
    # appears in the BOUNDS section but did not appear in the COLUMNS section.
    # This is likely because such variables are meaningless - they don't appear
    # in the objective or constraints and so can be trivially removed.
    # To avoid generating MPS files that crash existing readers, we cache the
    # names of all variables seen in COLUMNS into `discovered_columns`, and then
    # pass this set to `write_bounds` so that it can act appropriately.
    discovered_columns = Set{String}()
    # Build constraint coefficients
    for (S, _) in SET_TYPES
        _collect_coefficients(model, S, names, coefficients, discovered_columns)
    end
    # Build objective
    # MPS doesn't support maximization so we flip the sign on the objective
    # coefficients.
    s = MOI.get(model, MOI.ObjectiveSense()) == MOI.MAX_SENSE ? -1.0 : 1.0
    obj_func = MOI.get(
        model, MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}()
    )
    extract_terms(names, coefficients, "OBJ", obj_func, discovered_columns, s)
    integer_variables = list_of_integer_variables(model, names)
    println(io, "COLUMNS")
    int_open = false
    for variable in ordered_names
        is_int = variable in integer_variables
        if is_int && !int_open
            println(io, "    MARKER    'MARKER'                 'INTORG'")
            int_open = true
        elseif !is_int && int_open
            println(io, "    MARKER    'MARKER'                 'INTEND'")
            int_open = false
        end
        for (constraint, coefficient) in coefficients[variable]
            print(io, "     ")
            print(io, rpad(variable, 8))
            print(io, " ")
            print(io, rpad(constraint, 8))
            print(io, " ")
            Base.Grisu.print_shortest(io, coefficient)
            println(io)
        end
    end
    return discovered_columns
end

# ==============================================================================
#   RHS
# ==============================================================================

value(set::MOI.LessThan) = set.upper
value(set::MOI.GreaterThan) = set.lower
value(set::MOI.EqualTo) = set.value
value(set::MOI.Interval) = set.upper  # See the note in the RANGES section.

function _write_rhs(io, model, S, sense_char)
    for index in MOI.get(
        model,
        MOI.ListOfConstraintIndices{MOI.ScalarAffineFunction{Float64}, S}()
    )
        row_name = MOI.get(model, MOI.ConstraintName(), index)
        print(io, "    rhs       ", rpad(row_name, 8), "  ")
        set = MOI.get(model, MOI.ConstraintSet(), index)
        Base.Grisu.print_shortest(io, value(set))
        println(io)
    end
end

function write_rhs(io::IO, model::Model)
    println(io, "RHS")
    for (set_type, sense_char) in SET_TYPES
        _write_rhs(io, model, set_type, sense_char)
    end
    return
end

# ==============================================================================
#  RANGES
#
# Here is how RANGE information is encoded.
#
#     Row type | Range value |  lower bound  |  upper bound
#     ------------------------------------------------------
#         G    |     +/-     |     rhs       | rhs + |range|
#         L    |     +/-     | rhs - |range| |     rhs
#         E    |      +      |     rhs       | rhs + range
#         E    |      -      | rhs + range   |     rhs
#
# We elect to write out ScalarAffineFunction-in-Interval constraints in terms of
# LessThan (L) constraints with a range shift. The RHS term is set to the upper
# bound, and the RANGE term to upper - lower.
# ==============================================================================

function write_ranges(io::IO, model::Model)
    println(io, "RANGES")
    for index in MOI.get(
        model,
        MOI.ListOfConstraintIndices{
            MOI.ScalarAffineFunction{Float64}, MOI.Interval{Float64}
        }()
    )
        row_name = MOI.get(model, MOI.ConstraintName(), index)
        print(io, "    rhs       ", rpad(row_name, 8), "  ")
        set = MOI.get(model, MOI.ConstraintSet(), index)::MOI.Interval{Float64}
        Base.Grisu.print_shortest(io, set.upper - set.lower)
        println(io)
    end
    return
end

# ==============================================================================
#   BOUNDS
# Variables default to [0, ∞).
#     FR    free variable      -∞ < x < ∞
#     FX    fixed variable     x == b
#     MI    lower bound -inf   -∞ < x
#     UP    upper bound             x <= b
#     PL    upper bound +inf        x < ∞
#     LO    lower bound        b <= x
#     BV    binary variable    x = 0 or 1
#
#  Not yet implemented:
#     LI    integer variable   b <= x
#     UI    integer variable        x <= b
#     SC    semi-cont variable x = 0 or l <= x <= b
#           l is the lower bound on the variable. If none set then defaults to 1
# ==============================================================================

function write_single_bound(io::IO, var_name::String, lower, upper)
    name = rpad(var_name, 8)
    if lower == upper
        print(io, " FX bounds    ", name, " ")
        Base.Grisu.print_shortest(io, lower)
        println(io)
    elseif lower == -Inf && upper == Inf
        println(io, " FR bounds    ", var_name)
    else
        if lower == -Inf
            println(io, " MI bounds    ", name)
        else
            print(io, " LO bounds    ", name, " ")
            Base.Grisu.print_shortest(io, lower)
            println(io)
        end
        if upper == Inf
            println(io, " PL bounds    ", name)
        else
            print(io, " UP bounds    ", name, " ")
            Base.Grisu.print_shortest(io, upper)
            println(io)
        end
    end
    return
end

function update_bounds(x::Tuple{Float64, Float64, VType}, set::MOI.GreaterThan)
    return (max(x[1], set.lower), x[2], x[3])
end

function update_bounds(x::Tuple{Float64, Float64, VType}, set::MOI.LessThan)
    return (x[1], min(x[2], set.upper), x[3])
end

function update_bounds(x::Tuple{Float64, Float64, VType}, set::MOI.Interval)
    return (set.lower, set.upper, x[3])
end

function update_bounds(x::Tuple{Float64, Float64, VType}, set::MOI.EqualTo)
    return (set.value, set.value, x[3])
end

function update_bounds(x::Tuple{Float64, Float64, VType}, set::MOI.ZeroOne)
    return (x[1], x[2], VTYPE_BINARY)
end

function _collect_bounds(bounds, model, S, names)
    for index in MOI.get(
        model, MOI.ListOfConstraintIndices{MOI.SingleVariable, S}()
    )
        func = MOI.get(model, MOI.ConstraintFunction(), index)
        set = MOI.get(model, MOI.ConstraintSet(), index)::S
        name = names[func.variable]
        bounds[name] = update_bounds(bounds[name], set)
    end
    return
end

function write_bounds(
    io::IO, model::Model, discovered_columns::Set{String}, ordered_names, names
)
    println(io, "BOUNDS")
    bounds = Dict{String, Tuple{Float64, Float64, VType}}(
        n => (-Inf, Inf, VTYPE_CONTINUOUS) for n in ordered_names
    )
    for S in (
        MOI.LessThan{Float64},
        MOI.GreaterThan{Float64},
        MOI.EqualTo{Float64},
        MOI.Interval{Float64},
        MOI.ZeroOne,
    )
        _collect_bounds(bounds, model, S, names)
    end
    for var_name in ordered_names
        lower, upper, vtype = bounds[var_name]
        if !(var_name in discovered_columns)
            @warn(
                "Variable $var_name is mentioned in BOUNDS, but is not " *
                "mentioned in the COLUMNS section. We are ignoring it."
            )
            continue
        elseif vtype == VTYPE_BINARY
            println(io, " BV bounds    ", var_name)
        end
        write_single_bound(io, var_name, lower, upper)
    end
    return
end

# ==============================================================================
#   SOS
# ==============================================================================

function write_sos_constraint(io::IO, model::Model, index, names)
    func = MOI.get(model, MOI.ConstraintFunction(), index)
    set = MOI.get(model, MOI.ConstraintSet(), index)
    for (variable, weight) in zip(func.variables, set.weights)
        print(io, "    ", rpad(names[variable], 8), "  ")
        Base.Grisu.print_shortest(io, weight)
        println(io)
    end
end

function write_sos(io::IO, model::Model, names)
    sos1_indices = MOI.get(
        model,
        MOI.ListOfConstraintIndices{MOI.VectorOfVariables, MOI.SOS1{Float64}}(),
    )
    sos2_indices = MOI.get(
        model,
        MOI.ListOfConstraintIndices{MOI.VectorOfVariables, MOI.SOS2{Float64}}(),
    )
    if length(sos1_indices) + length(sos2_indices) > 0
        println(io, "SOS")
        idx = 1
        for (sos_type, indices) in enumerate([sos1_indices, sos2_indices])
            for index in indices
                println(io, " S", sos_type, " SOS", idx)
                write_sos_constraint(io, model, index, names)
                idx += 1
            end
        end
    end
    return
end

# ==============================================================================
#
#   Base.read!
#
# Here is a template for an MPS file, reproduced from
# http://lpsolve.sourceforge.net/5.5/mps-format.htm.
#
# Field:    1           2          3         4         5         6
# Columns:  2-3        5-12      15-22     25-36     40-47     50-61
#           NAME   problem name
#           ROWS
#            type     name
#           COLUMNS
#                     column      row       value     row      value
#                     name        name                name
#           RHS
#                     rhs         row       value     row      value
#                     name        name                name
#           RANGES
#                     range       row       value     row      value
#                     name        name                name
#           BOUNDS
#            type     bound       column    value
#                     name        name
#           SOS
#            type     CaseName    SOSName   SOSpriority
#                     CaseName    VarName1  VarWeight1
#                     CaseName    VarName2  VarWeight2
#                     CaseName    VarNameN  VarWeightN
#           ENDATA
# ==============================================================================

@enum(Sense, SENSE_N, SENSE_G, SENSE_L, SENSE_E, SENSE_UNKNOWN)
function Sense(s::String)
    if s == "G"
        return SENSE_G
    elseif s == "L"
        return SENSE_L
    elseif s == "E"
        return SENSE_E
    elseif s == "N"
        return SENSE_N
    else
        return SENSE_UNKNOWN
    end
end
mutable struct TempMPSModel
    name::String
    obj_name::String
    c::Vector{Float64}
    col_lower::Vector{Float64}
    col_upper::Vector{Float64}
    row_lower::Vector{Float64}
    row_upper::Vector{Float64}
    sense::Vector{Sense}
    A::Vector{Vector{Tuple{Int, Float64}}}
    vtype::Vector{VType}
    name_to_col::Dict{String, Int}
    col_to_name::Vector{String}
    name_to_row::Dict{String, Int}
    row_to_name::Vector{String}
    intorg_flag::Bool  # A flag used to parse COLUMNS section.
end

function TempMPSModel()
    return TempMPSModel(
        "",
        "",
        Float64[],  # c
        Float64[],  # col_lower
        Float64[],  # col_upper
        Float64[],  # row_lower
        Float64[],  # row_upper
        Sense[],   # sense
        Vector{Vector{Tuple{Int, Float64}}}[],  # A
        Bool[],
        Dict{String, Int}(),
        String[],
        Dict{String, Int}(),
        String[],
        false,
    )
end

@enum(
    Headers,
    HEADER_NAME,
    HEADER_ROWS,
    HEADER_COLUMNS,
    HEADER_RHS,
    HEADER_RANGES,
    HEADER_BOUNDS,
    HEADER_SOS,
    HEADER_ENDATA,
    HEADER_UNKNOWN,
)

# Headers(s) gets called _alot_ (on every line), so we try very hard to be
# efficient.]
function Headers(s::AbstractString)
    N = length(s)
    if N > 7 || N < 3
        return HEADER_UNKNOWN
    elseif N == 3
        x = first(s)
        if (x == 'R' || x == 'r') && uppercase(s) == "RHS"
            return HEADER_RHS
        elseif (x == 'S' || x == 's') && uppercase(s) == "SOS"
            return HEADER_SOS
        end
    elseif N == 4
        x = first(s)
        if (x == 'R' || x == 'r') && uppercase(s) == "ROWS"
            return HEADER_ROWS
        end
    elseif N == 5
        return HEADER_UNKNOWN
    elseif N == 6
        x = first(s)
        if (x == 'R' || x == 'r') && uppercase(s) == "RANGES"
            return HEADER_RANGES
        elseif (x == 'B' || x == 'b') && uppercase(s) == "BOUNDS"
            return HEADER_BOUNDS
        elseif (x == 'E' || x == 'e') && uppercase(s) == "ENDATA"
            return HEADER_ENDATA
        end
    elseif N == 7
        x = first(s)
        if (x == 'C' || x == 'c') && (uppercase(s) == "COLUMNS")
            return HEADER_COLUMNS
        end
    end
    return HEADER_UNKNOWN
end

function line_to_items(line)
    items = split(line, " "; keepempty = false)
    return String.(items)
end

"""
    Base.read!(io::IO, model::FileFormats.MPS.Model)

Read `io` in the MPS file format and store the result in `model`.
"""
function Base.read!(io::IO, model::Model)
    if !MOI.is_empty(model)
        error("Cannot read in file because model is not empty.")
    end
    data = TempMPSModel()
    header = HEADER_NAME
    multi_objectives = String[]
    while !eof(io)
        line = string(strip(readline(io)))
        if isempty(line) || first(line) == '*'
            continue  # Skip blank lines and comments.
        end
        h = Headers(line)
        if h != HEADER_UNKNOWN
            header = h
            continue
        else
            # Carry on with the previous header
        end
        # TODO: split into hard fields based on column indices.
        items = line_to_items(line)
        if header == HEADER_NAME
            parse_name_line(data, items)
        elseif header == HEADER_ROWS
            multi_obj = parse_rows_line(data, items)
            if multi_obj !== nothing
                push!(multi_objectives, multi_obj)
            end
        elseif header == HEADER_COLUMNS
            parse_columns_line(data, items, multi_objectives)
        elseif header == HEADER_RHS
            parse_rhs_line(data, items)
        elseif header == HEADER_RANGES
            parse_ranges_line(data, items)
        elseif header == HEADER_BOUNDS
            parse_bounds_line(data, items)
        elseif header == HEADER_SOS
            error("TODO(odow): implement parsing of SOS constraints.")
        else
            @assert header == HEADER_ENDATA
            break
        end
    end
    copy_to(model, data)
    return
end

function bounds_to_set(lower, upper)
    if -Inf < lower < upper < Inf
        return MOI.Interval(lower, upper)
    elseif -Inf < lower && upper == Inf
        return MOI.GreaterThan(lower)
    elseif -Inf == lower && upper < Inf
        return MOI.LessThan(upper)
    elseif lower == upper
        return MOI.EqualTo(upper)
    end
    return  # free variable
end

function copy_to(model::Model, data::TempMPSModel)
    MOI.set(model, MOI.Name(), data.name)
    variable_map = Dict{String, MOI.VariableIndex}()
    for (i, name) in enumerate(data.col_to_name)
        _add_variable(model, data, variable_map, i, name)
    end
    _add_objective(model, data, variable_map)
    for (j, c_name) in enumerate(data.row_to_name)
        set = bounds_to_set(data.row_lower[j], data.row_upper[j])
        if set !== nothing
            _add_linear_constraint(model, data, variable_map, j, c_name, set)
        end
        # `else` is a free constraint. Don't add it.
    end
    return
end

function _add_variable(model, data, variable_map, i, name)
    x = MOI.add_variable(model)
    variable_map[name] = x
    MOI.set(model, MOI.VariableName(), x, name)
    set = bounds_to_set(data.col_lower[i], data.col_upper[i])
    if set !== nothing
        MOI.add_constraint(model, MOI.SingleVariable(x), set)
    end
    if data.vtype[i] == VTYPE_INTEGER
        MOI.add_constraint(model, MOI.SingleVariable(x), MOI.Integer())
    elseif data.vtype[i] == VTYPE_BINARY
        MOI.add_constraint(model, MOI.SingleVariable(x), MOI.ZeroOne())
    end
    return
end

function _add_objective(model, data, variable_map)
    MOI.set(model, MOI.ObjectiveSense(), MOI.MIN_SENSE)
    MOI.set(
        model,
        MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(),
        MOI.ScalarAffineFunction(
            [
                MOI.ScalarAffineTerm(data.c[i], variable_map[v])
                for (i, v) in enumerate(data.col_to_name) if !iszero(data.c[i])
            ],
            0.0,
        )
    )
    return
end

function _add_linear_constraint(model, data, variable_map, j, c_name, set)
    terms = MOI.ScalarAffineTerm{Float64}[
        MOI.ScalarAffineTerm(coef, variable_map[data.col_to_name[i]])
        for (i, coef) in data.A[j]
    ]
    c = MOI.add_constraint(model, MOI.ScalarAffineFunction(terms, 0.0), set)
    MOI.set(model, MOI.ConstraintName(), c, c_name)
    return
end

# ==============================================================================
#   NAME
# ==============================================================================

function parse_name_line(data::TempMPSModel, items::Vector{String})
    if !(1 <= length(items) <= 2) || uppercase(items[1]) != "NAME"
        error("Malformed NAME line: $(join(items, " "))")
    end
    if length(items) == 2
        data.name = items[2]
    elseif length(items) == 1
        data.name = ""
    end
    return
end

# ==============================================================================
#   ROWS
# ==============================================================================

function parse_rows_line(data::TempMPSModel, items::Vector{String})
    if length(items) != 2
        error("Malformed ROWS line: $(join(items, " "))")
    end
    sense, name = Sense(items[1]), items[2]
    if haskey(data.name_to_row, name)
        error("Duplicate row encountered: $(join(items, " ")).")
    elseif sense == SENSE_UNKNOWN
        error("Invalid row sense: $(join(items, " "))")
    end
    if sense == SENSE_N
        if data.obj_name != ""
            return name  # Detected a duplicate objective. Skip it.
        end
        data.obj_name = name
        return
    end
    if name == data.obj_name
        error("Found row with same name as objective: $(join(items, " ")).")
    end
    # Add some default bounds for the constraints.
    push!(data.row_to_name, name)
    row = length(data.row_to_name)
    data.name_to_row[name] = row
    push!(data.sense, sense)
    push!(data.A, Tuple{Int, Float64}[])
    if sense == SENSE_G
        push!(data.row_lower, 0.0)
        push!(data.row_upper, Inf)
        data.row_upper[row] = Inf
    elseif sense == SENSE_L
        push!(data.row_lower, -Inf)
        push!(data.row_upper, 0.0)
    else
        @assert sense == SENSE_E
        push!(data.row_lower, 0.0)
        push!(data.row_upper, 0.0)
    end
    return
end

# ==============================================================================
#   COLUMNS
# ==============================================================================

function parse_single_coefficient(
    data, row_name::String, column::Int, value::Float64
)
    if row_name == data.obj_name
        data.c[column] += value
        return
    end
    row = get(data.name_to_row, row_name, nothing)
    if row === nothing
        error("ROW name $(row_name) not recognised. Is it in the ROWS field?")
    end
    push!(data.A[row], (column, value))
    return
end

function _add_new_column(data, column_name)
    if haskey(data.name_to_col, column_name)
        return
    end
    push!(data.col_to_name, column_name)
    data.name_to_col[column_name] = length(data.col_to_name)
    push!(data.c, 0.0)
    push!(data.col_lower, 0.0)
    push!(data.col_upper, Inf)
    push!(data.vtype, VTYPE_CONTINUOUS)
    return
end

function _set_intorg(data, column, column_name)
     if data.intorg_flag
        data.vtype[column] = VTYPE_INTEGER
    elseif data.vtype[column] != VTYPE_CONTINUOUS
        error(
            "Variable $(column_name) appeared in COLUMNS outside an " *
            "`INT` marker after already being declared as integer."
        )
    end
    return
end

function parse_columns_line(
    data::TempMPSModel, items::Vector{String}, multi_objectives::Vector{String}
)
    if length(items) == 3
        # [column name] [row name] [value]
        column_name, row_name, value = items
        if row_name == "'MARKER'"
            if value == "'INTORG'"
                data.intorg_flag = true
                return
            elseif value == "'INTEND'"
                data.intorg_flag = false
                return
            end
        elseif row_name in multi_objectives
            return
        end
        _add_new_column(data, column_name)
        column = data.name_to_col[column_name]
        parse_single_coefficient(data, row_name, column, parse(Float64, value))
        _set_intorg(data, column, column_name)
    elseif length(items) == 5
        # [column name] [row name] [value] [row name 2] [value 2]
        column_name, row_name_1, value_1, row_name_2, value_2 = items
        if row_name_1 in multi_objectives || row_name_2 in multi_objectives
            return
        end
        _add_new_column(data, column_name)
        column = data.name_to_col[column_name]
        parse_single_coefficient(
            data, row_name_1, column, parse(Float64, value_1)
        )
        parse_single_coefficient(
            data, row_name_2, column, parse(Float64, value_2)
        )
        _set_intorg(data, column, column_name)
    else
        error("Malformed COLUMNS line: $(join(items, " "))")
    end
    return
end

# ==============================================================================
#   RHS
# ==============================================================================

function parse_single_rhs(
    data, row_name::String, value::Float64, items::Vector{String}
)
    row = get(data.name_to_row, row_name, nothing)
    if row === nothing
        error("ROW name $(row_name) not recognised. Is it in the ROWS field?")
    end
    if data.sense[row] == SENSE_E
        data.row_upper[row] = value
        data.row_lower[row] = value
    elseif data.sense[row] == SENSE_G
        data.row_lower[row] = value
    elseif data.sense[row] == SENSE_L
        data.row_upper[row] = value
    else
        @assert data.sense[row] == SENSE_N
        error("Cannot have RHS for objective: $(join(items, " "))")
    end
    return
end

# TODO: handle multiple RHS vectors.
function parse_rhs_line(data::TempMPSModel, items::Vector{String})
    if length(items) == 3
        # [rhs name] [row name] [value]
        rhs_name, row_name, value = items
        parse_single_rhs(data, row_name, parse(Float64, value), items)
    elseif length(items) == 5
        # [rhs name] [row name 1] [value 1] [row name 2] [value 2]
        rhs_name, row_name_1, value_1, row_name_2, value_2 = items
        parse_single_rhs(data, row_name_1, parse(Float64, value_1), items)
        parse_single_rhs(data, row_name_2, parse(Float64, value_2), items)
    else
        error("Malformed RHS line: $(join(items, " "))")
    end
    return
end

# ==============================================================================
#  RANGES
#
# Here is how RANGE information is encoded. (We repeat this comment because it
# is so non-trivial.)
#
#     Row type | Range value |  lower bound  |  upper bound
#     ------------------------------------------------------
#         G    |     +/-     |     rhs       | rhs + |range|
#         L    |     +/-     | rhs - |range| |     rhs
#         E    |      +      |     rhs       | rhs + range
#         E    |      -      | rhs + range   |     rhs
# ==============================================================================

function parse_single_range(data, row_name::String, value::Float64)
    row = get(data.name_to_row, row_name, nothing)
    if row === nothing
        error("ROW name $(row_name) not recognised. Is it in the ROWS field?")
    end
    if data.sense[row] == SENSE_G
        data.row_upper[row] = data.row_lower[row] + abs(value)
    elseif data.sense[row] == SENSE_L
        data.row_lower[row] = data.row_upper[row] - abs(value)
    elseif data.sense[row] == SENSE_E
        if value > 0.0
            data.row_upper[row] = data.row_lower[row] + value
        else
            data.row_lower[row] = data.row_upper[row] + value
        end
    end
    return
end

# TODO: handle multiple RANGES vectors.
function parse_ranges_line(data::TempMPSModel, items::Vector{String})
    if length(items) == 3
        # [rhs name] [row name] [value]
        _, row_name, value = items
        parse_single_range(data, row_name, parse(Float64, value))
    elseif length(items) == 5
        # [rhs name] [row name] [value] [row name 2] [value 2]
        _, row_name_1, value_1, row_name_2, value_2 = items
        parse_single_range(data, row_name_1, parse(Float64, value_1))
        parse_single_range(data, row_name_2, parse(Float64, value_2))
    else
        error("Malformed RANGES line: $(join(items, " "))")
    end
    return
end

# ==============================================================================
#   BOUNDS
# ==============================================================================

function _parse_single_bound(data, column_name::String, bound_type::String)
    col = get(data.name_to_col, column_name, nothing)
    if col === nothing
        error("Column name $(column_name) not found.")
    end
    if bound_type == "PL"
        data.col_upper[col] = Inf
    elseif bound_type == "MI"
        data.col_lower[col] = -Inf
    elseif bound_type == "FR"
        data.col_lower[col] = -Inf
        data.col_upper[col] = Inf
    elseif bound_type == "BV"
        data.col_lower[col] = -Inf
        data.col_upper[col] = Inf
        data.vtype[col] = VTYPE_BINARY
    else
        error("Invalid bound type $(bound_type): $(column_name)")
    end
end

function _parse_single_bound(
    data, column_name::String, bound_type::String, value::Float64
)
    col = get(data.name_to_col, column_name, nothing)
    if col === nothing
        error("Column name $(column_name) not found.")
    end
    if bound_type == "FX"
        data.col_lower[col] = value
        data.col_upper[col] = value
    elseif bound_type == "UP"
        data.col_upper[col] = value
    elseif bound_type == "LO"
        data.col_lower[col] = value
    elseif bound_type == "LI"
        data.col_lower[col] = value
        data.vtype[col] = VTYPE_INTEGER
    elseif bound_type == "UI"
        data.col_upper[col] = value
        data.vtype[col] = VTYPE_INTEGER
    elseif bound_type == "FR"
        # So even though FR bounds should be of the form:
        #  FR BOUND1    VARNAME
        # there are cases in MIPLIB2017 (e.g., leo1 and leo2) like so:
        #  FR BOUND1    C0000001       .000000
        # In these situations, just ignore the value.
        data.col_lower[col] = -Inf
        data.col_upper[col] = Inf
    elseif bound_type == "BV"
        # A similar situation happens with BV bounds in leo1 and leo2.
        data.col_lower[col] = -Inf
        data.col_upper[col] = Inf
        data.vtype[col] = VTYPE_BINARY
    else
        error("Invalid bound type $(bound_type): $(column_name)")
    end
end

function parse_bounds_line(data::TempMPSModel, items::Vector{String})
    if length(items) == 3
        bound_type, _, column_name = items
        _parse_single_bound(data, column_name, bound_type)
    elseif length(items) == 4
        bound_type, _, column_name, value = items
        _parse_single_bound(data, column_name, bound_type, parse(Float64, value))
    else
        error("Malformed BOUNDS line: $(join(items, " "))")
    end
    return
end

end


using ClimateMachine.BalanceLaws

abstract type AbstractFluidModel <: BalanceLaw end

"""
    ModelSetup <: AbstractFluidModel
"""
struct ModelSetup{𝒯, 𝒰, 𝒱, 𝒲, 𝒳} <: AbstractFluidModel
    physics::𝒯
    boundary_conditions::𝒰
    initial_conditions::𝒱
    numerics::𝒲
    parameters::𝒳
end

function ModelSetup(; physics, boundary_conditions, initial_conditions, numerics, parameters)
    return ModelSetup(
        physics,
        unpack_boundary_conditions(boundary_conditions),
        initial_conditions,
        numerics,
        parameters,
    )
end

"""
    DryAtmosModel <: AbstractFluidModel

    temporarily use this struct
"""
Base.@kwdef struct DryAtmosModel{𝒯, 𝒰, 𝒱, 𝒲, 𝒳} <: AbstractFluidModel
    physics::𝒯
    boundary_conditions::𝒰
    initial_conditions::𝒱
    numerics::𝒲
    parameters::𝒳
end

"""
    DryAtmosLinearModel <: AbstractFluidModel

    temporarily use this struct
"""
Base.@kwdef struct DryAtmosLinearModel{𝒯, 𝒰, 𝒱, 𝒲, 𝒳} <: AbstractFluidModel
    physics::𝒯
    boundary_conditions::𝒰
    initial_conditions::𝒱
    numerics::𝒲
    parameters::𝒳
end

function unpack_boundary_conditions(bcs)
    # We need to repackage the boundary conditions to match the
    # boundary conditions interface of the Balance Law and DGModel
    boundaries = (:west, :east, :south, :north, :bottom, :top)
    repackaged_bcs = []

    for boundary in boundaries
        fields = get(bcs, boundary, nothing)
        new_bc = isnothing(fields) ? FluidBC() : FluidBC(fields...)
        push!(repackaged_bcs, new_bc)
    end

    return Tuple(repackaged_bcs)
end


module ContactImplicitMPC

using BenchmarkTools
using Colors
using CoordinateTransformations
using FFMPEG
using FileIO
using ForwardDiff
using GeometryBasics
using IfElse
using InteractiveUtils
using JLD2
using LinearAlgebra
using Logging
using QDLDL
using MeshCat
using MeshCatMechanisms
using Meshing
using MeshIO
using Parameters
using Plots
using Random
using Rotations
using SparseArrays
using StaticArrays
using Symbolics
using Test

# Utilities
include("utils.jl") #

# Solver
include("solver/gn.jl")
include("solver/lu.jl")
include("solver/ldl.jl")
include("solver/qr.jl")
include("solver/schur.jl")
include("solver/cones.jl") #

# Environment
include("simulator/environment.jl")

# Dynamics
include("dynamics/model.jl")

# Simulator
include("simulation/index.jl")

# Solver
include("solver/interior_point.jl")

# Simulator
include("simulation/contact_methods.jl")
include("simulation/simulation.jl")
include("simulator/trajectory.jl")

include("dynamics/code_gen_dynamics.jl")
include("dynamics/fast_methods_dynamics.jl")

# Models
include("dynamics/quaternions.jl")
include("dynamics/mrp.jl")
include("dynamics/euler.jl")

include("dynamics/particle_2D/model.jl")
include("dynamics/particle/model.jl")
include("dynamics/hopper_2D/model.jl")
include("dynamics/hopper_3D/model.jl")
include("dynamics/quadruped/model.jl")
include("dynamics/flamingo/model.jl")
include("dynamics/pushbot/model.jl")
include("dynamics/rigidbody/model.jl")

# Simulator
include("simulator/policy.jl")
include("simulator/disturbances.jl")
include("simulator/simulator.jl")

# Simulation
include("simulation/environments/flat.jl")
include("simulation/environments/piecewise.jl")
include("simulation/environments/quadratic.jl")
include("simulation/environments/slope.jl")
include("simulation/environments/sinusoidal.jl")
include("simulation/environments/stairs.jl")

include("simulation/residual_approx.jl")
include("simulation/code_gen_simulation.jl")

# Controller
include("controller/linearized_step.jl")
include("controller/implicit_dynamics.jl")
include("controller/objective.jl")
include("controller/linearized_solver.jl")
include("controller/newton.jl")
include("controller/newton_indices.jl")
include("controller/newton_residual.jl")
include("controller/newton_jacobian.jl")
include("controller/mpc_utils.jl")
include("controller/policy.jl")
include("controller/newton_structure_solver/methods.jl")

# Visuals
include("dynamics/visuals.jl")
include("dynamics/visual_utils.jl")

export
    World,
    LinearizedCone,
    NonlinearCone,
    ContactModel,
    Dimensions,
    BaseMethods,
    DynamicsMethods,
    ResidualMethods,
    Environment,
    environment_2D,
    environment_3D,
    environment_2D_flat,
    environment_3D_flat,
    get_model,
    SparseStructure,
    LinearizedStep,
    bil_addition!,
    r_linearized!,
    rz_linearized!,
    ImplicitTraj,
    linearization!,
    implicit_dynamics!,
    TrackingObjective,
    TrackingVelocityObjective,
    second_order_cone_product,
    generate_base_expressions,
    RLin,
    RZLin,
    RθLin,
    ContactTraj,
    Simulation,
    num_var,
    num_data,
    get_simulation,
    get_trajectory,
    interior_point,
    InteriorPointOptions,
    interior_point_solve!,
    r!,
    rz!,
    rθ,
    generate_base_expressions,
    save_expressions,
    instantiate_base!,
    generate_dynamics_expressions,
    save_expressions,
    instantiate_dynamics!,
    environment_3D_flat,
    friction_dim,
    dim,
    sqrt_quat,
    cayley_map,
    L_multiply,
    R_multiply,
    R2,
    R3,
    FrictionCone,
    rotation,
    module_dir,
    open_loop_disturbances,
    disturbances,
    open_loop_policy,
    policy,
    linearized_mpc_policy,
    NewtonOptions,
    LinearizedMPCOptions,
    SimulatorOptions,
    simulator,
    simulate!,
    generate_residual_expressions,
    instantiate_residual!,
    ϕ_func,
    tracking_error,
    repeat_ref_traj,
    Schur,
    schur_factorize!,
    schur_solve!,
    LinearSolver,
    LUSolver,
    lu_solver,
    factorize!,
    linear_solve!,
    OptimizationSpace,
    OptimizationSpace13,
    index_q2,
    index_γ1,
    index_b1,
    index_ψ1,
    index_s1,
    index_η1,
    index_s2,
    index_q0,
    index_q1,
    index_u1,
    index_w1,
    index_μ,
    index_h,
    index_dyn,
    index_imp,
    index_mdp,
    index_fri,
    index_bimp,
    index_bmdp,
    index_bfri,
    linearization_var_index,
    linearization_term_index,
    index_ort,
    index_soc,
    num_var,
    num_data,
    num_bilinear,
    index_equr,
    index_ortr,
    index_socr

end # module


# Copyright (c) 2013: Iain Dunning, Miles Lubin, and contributors
#
# Use of this source code is governed by an MIT-style license that can be found
# in the LICENSE.md file or at https://opensource.org/licenses/MIT.

mutable struct _ConstraintInfo{F,S}
    func::F
    set::S
    dual_start::Union{Nothing,Float64}
end

_ConstraintInfo(func, set) = _ConstraintInfo(func, set, nothing)

"""
    Optimizer()

Create a new Ipopt optimizer.
"""
mutable struct Optimizer <: MOI.AbstractOptimizer
    inner::Union{Nothing,IpoptProblem}
    invalid_model::Bool
    variables::MOI.Utilities.VariablesContainer{Float64}
    variable_primal_start::Vector{Union{Nothing,Float64}}
    variable_lower_start::Vector{Union{Nothing,Float64}}
    variable_upper_start::Vector{Union{Nothing,Float64}}
    nlp_data::MOI.NLPBlockData
    sense::MOI.OptimizationSense
    objective::Union{
        Nothing,
        MOI.VariableIndex,
        MOI.ScalarAffineFunction{Float64},
        MOI.ScalarQuadraticFunction{Float64},
    }
    linear_le_constraints::Vector{
        _ConstraintInfo{
            MOI.ScalarAffineFunction{Float64},
            MOI.LessThan{Float64},
        },
    }
    linear_ge_constraints::Vector{
        _ConstraintInfo{
            MOI.ScalarAffineFunction{Float64},
            MOI.GreaterThan{Float64},
        },
    }
    linear_eq_constraints::Vector{
        _ConstraintInfo{MOI.ScalarAffineFunction{Float64},MOI.EqualTo{Float64}},
    }
    quadratic_le_constraints::Vector{
        _ConstraintInfo{
            MOI.ScalarQuadraticFunction{Float64},
            MOI.LessThan{Float64},
        },
    }
    quadratic_ge_constraints::Vector{
        _ConstraintInfo{
            MOI.ScalarQuadraticFunction{Float64},
            MOI.GreaterThan{Float64},
        },
    }
    quadratic_eq_constraints::Vector{
        _ConstraintInfo{
            MOI.ScalarQuadraticFunction{Float64},
            MOI.EqualTo{Float64},
        },
    }
    nlp_dual_start::Union{Nothing,Vector{Float64}}
    silent::Bool
    options::Dict{String,Any}
    solve_time::Float64
    callback::Union{Nothing,Function}

    function Optimizer()
        return new(
            nothing,
            false,
            MOI.Utilities.VariablesContainer{Float64}(),
            Union{Nothing,Float64}[],
            Union{Nothing,Float64}[],
            Union{Nothing,Float64}[],
            MOI.NLPBlockData([], _EmptyNLPEvaluator(), false),
            MOI.FEASIBILITY_SENSE,
            nothing,
            _ConstraintInfo{
                MOI.ScalarAffineFunction{Float64},
                MOI.LessThan{Float64},
            }[],
            _ConstraintInfo{
                MOI.ScalarAffineFunction{Float64},
                MOI.GreaterThan{Float64},
            }[],
            _ConstraintInfo{
                MOI.ScalarAffineFunction{Float64},
                MOI.EqualTo{Float64},
            }[],
            _ConstraintInfo{
                MOI.ScalarQuadraticFunction{Float64},
                MOI.LessThan{Float64},
            }[],
            _ConstraintInfo{
                MOI.ScalarQuadraticFunction{Float64},
                MOI.GreaterThan{Float64},
            }[],
            _ConstraintInfo{
                MOI.ScalarQuadraticFunction{Float64},
                MOI.EqualTo{Float64},
            }[],
            nothing,
            false,
            Dict{String,Any}(),
            NaN,
            nothing,
        )
    end
end

MOI.get(::Optimizer, ::MOI.SolverVersion) = "3.14.4"

### _EmptyNLPEvaluator

struct _EmptyNLPEvaluator <: MOI.AbstractNLPEvaluator end

MOI.features_available(::_EmptyNLPEvaluator) = [:Grad, :Jac, :Hess]
MOI.initialize(::_EmptyNLPEvaluator, ::Any) = nothing
MOI.eval_constraint(::_EmptyNLPEvaluator, g, x) = nothing
MOI.jacobian_structure(::_EmptyNLPEvaluator) = Tuple{Int64,Int64}[]
MOI.hessian_lagrangian_structure(::_EmptyNLPEvaluator) = Tuple{Int64,Int64}[]
MOI.eval_constraint_jacobian(::_EmptyNLPEvaluator, J, x) = nothing
MOI.eval_hessian_lagrangian(::_EmptyNLPEvaluator, H, x, σ, μ) = nothing

function MOI.empty!(model::Optimizer)
    model.inner = nothing
    model.invalid_model = false
    MOI.empty!(model.variables)
    empty!(model.variable_primal_start)
    empty!(model.variable_lower_start)
    empty!(model.variable_upper_start)
    model.nlp_data = MOI.NLPBlockData([], _EmptyNLPEvaluator(), false)
    model.sense = MOI.FEASIBILITY_SENSE
    model.objective = nothing
    empty!(model.linear_le_constraints)
    empty!(model.linear_ge_constraints)
    empty!(model.linear_eq_constraints)
    empty!(model.quadratic_le_constraints)
    empty!(model.quadratic_ge_constraints)
    empty!(model.quadratic_eq_constraints)
    model.nlp_dual_start = nothing
    return
end

function MOI.is_empty(model::Optimizer)
    return MOI.is_empty(model.variables) &&
           isempty(model.variable_primal_start) &&
           isempty(model.variable_lower_start) &&
           isempty(model.variable_upper_start) &&
           model.nlp_data.evaluator isa _EmptyNLPEvaluator &&
           model.sense == MOI.FEASIBILITY_SENSE &&
           isempty(model.linear_le_constraints) &&
           isempty(model.linear_ge_constraints) &&
           isempty(model.linear_eq_constraints) &&
           isempty(model.quadratic_le_constraints) &&
           isempty(model.quadratic_ge_constraints) &&
           isempty(model.quadratic_eq_constraints)
end

MOI.supports_incremental_interface(::Optimizer) = true

function MOI.copy_to(model::Optimizer, src::MOI.ModelLike)
    return MOI.Utilities.default_copy_to(model, src)
end

MOI.get(::Optimizer, ::MOI.SolverName) = "Ipopt"

function MOI.supports_constraint(
    ::Optimizer,
    ::Type{
        <:Union{
            MOI.VariableIndex,
            MOI.ScalarAffineFunction{Float64},
            MOI.ScalarQuadraticFunction{Float64},
        },
    },
    ::Type{
        <:Union{
            MOI.LessThan{Float64},
            MOI.GreaterThan{Float64},
            MOI.EqualTo{Float64},
        },
    },
)
    return true
end

function MOI.get(model::Optimizer, ::MOI.ListOfConstraintTypesPresent)
    ret = MOI.get(model.variables, MOI.ListOfConstraintTypesPresent())
    constraints = Set{Tuple{Type,Type}}()
    for F in (
        MOI.ScalarAffineFunction{Float64},
        MOI.ScalarQuadraticFunction{Float64},
    )
        for S in (
            MOI.LessThan{Float64},
            MOI.GreaterThan{Float64},
            MOI.EqualTo{Float64},
        )
            if !isempty(_constraints(model, F, S))
                push!(constraints, (F, S))
            end
        end
    end
    return append!(ret, collect(constraints))
end

### MOI.Silent

MOI.supports(::Optimizer, ::MOI.Silent) = true

function MOI.set(model::Optimizer, ::MOI.Silent, value)
    model.silent = value
    return
end

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

### MOI.TimeLimitSec

MOI.supports(::Optimizer, ::MOI.TimeLimitSec) = true

function MOI.set(model::Optimizer, ::MOI.TimeLimitSec, value::Real)
    MOI.set(model, MOI.RawOptimizerAttribute("max_cpu_time"), Float64(value))
    return
end

function MOI.set(model::Optimizer, ::MOI.TimeLimitSec, ::Nothing)
    delete!(model.options, "max_cpu_time")
    return
end

function MOI.get(model::Optimizer, ::MOI.TimeLimitSec)
    return get(model.options, "max_cpu_time", nothing)
end

### MOI.RawOptimizerAttribute

MOI.supports(::Optimizer, ::MOI.RawOptimizerAttribute) = true

function MOI.set(model::Optimizer, p::MOI.RawOptimizerAttribute, value)
    model.options[p.name] = value
    return
end

function MOI.get(model::Optimizer, p::MOI.RawOptimizerAttribute)
    if !haskey(model.options, p.name)
        error("RawParameter with name $(p.name) is not set.")
    end
    return model.options[p.name]
end

### Variables

"""
    column(x::MOI.VariableIndex)

Return the column associated with a variable.
"""
column(x::MOI.VariableIndex) = x.value

function MOI.add_variable(model::Optimizer)
    push!(model.variable_primal_start, nothing)
    push!(model.variable_lower_start, nothing)
    push!(model.variable_upper_start, nothing)
    return MOI.add_variable(model.variables)
end

function MOI.is_valid(model::Optimizer, x::MOI.VariableIndex)
    return MOI.is_valid(model.variables, x)
end

function MOI.get(
    model::Optimizer,
    attr::Union{MOI.NumberOfVariables,MOI.ListOfVariableIndices},
)
    return MOI.get(model.variables, attr)
end

function MOI.is_valid(
    model::Optimizer,
    ci::MOI.ConstraintIndex{MOI.VariableIndex,S},
) where {S<:Union{MOI.LessThan,MOI.GreaterThan,MOI.EqualTo}}
    return MOI.is_valid(model.variables, ci)
end

function MOI.get(
    model::Optimizer,
    attr::Union{
        MOI.NumberOfConstraints{MOI.VariableIndex,S},
        MOI.ListOfConstraintIndices{MOI.VariableIndex,S},
    },
) where {S<:Union{MOI.LessThan,MOI.GreaterThan,MOI.EqualTo}}
    return MOI.get(model.variables, attr)
end

function MOI.get(
    model::Optimizer,
    attr::Union{MOI.ConstraintFunction,MOI.ConstraintSet},
    c::MOI.ConstraintIndex{MOI.VariableIndex,S},
) where {S<:Union{MOI.LessThan,MOI.GreaterThan,MOI.EqualTo}}
    return MOI.get(model.variables, attr, c)
end

function MOI.add_constraint(
    model::Optimizer,
    x::MOI.VariableIndex,
    set::Union{
        MOI.LessThan{Float64},
        MOI.GreaterThan{Float64},
        MOI.EqualTo{Float64},
    },
)
    return MOI.add_constraint(model.variables, x, set)
end

function MOI.set(
    model::Optimizer,
    ::MOI.ConstraintSet,
    ci::MOI.ConstraintIndex{MOI.VariableIndex,S},
    set::S,
) where {S<:Union{MOI.LessThan,MOI.GreaterThan,MOI.EqualTo}}
    MOI.set(model.variables, MOI.ConstraintSet(), ci, set)
    return
end

function MOI.delete(
    model::Optimizer,
    ci::MOI.ConstraintIndex{MOI.VariableIndex,S},
) where {S<:Union{MOI.LessThan,MOI.GreaterThan,MOI.EqualTo}}
    MOI.delete(model.variables, ci)
    return
end

### ScalarAffineFunction and ScalarQuadraticFunction constraints

function MOI.is_valid(
    model::Optimizer,
    ci::MOI.ConstraintIndex{F,S},
) where {
    F<:Union{
        MOI.ScalarAffineFunction{Float64},
        MOI.ScalarQuadraticFunction{Float64},
    },
    S<:Union{MOI.LessThan,MOI.GreaterThan,MOI.EqualTo},
}
    return 1 <= ci.value <= length(_constraints(model, F, S))
end

function _constraints(
    model::Optimizer,
    ::Type{MOI.ScalarAffineFunction{Float64}},
    ::Type{MOI.LessThan{Float64}},
)
    return model.linear_le_constraints
end

function _constraints(
    model::Optimizer,
    ::Type{MOI.ScalarAffineFunction{Float64}},
    ::Type{MOI.GreaterThan{Float64}},
)
    return model.linear_ge_constraints
end

function _constraints(
    model::Optimizer,
    ::Type{MOI.ScalarAffineFunction{Float64}},
    ::Type{MOI.EqualTo{Float64}},
)
    return model.linear_eq_constraints
end

function _constraints(
    model::Optimizer,
    ::Type{MOI.ScalarQuadraticFunction{Float64}},
    ::Type{MOI.LessThan{Float64}},
)
    return model.quadratic_le_constraints
end

function _constraints(
    model::Optimizer,
    ::Type{MOI.ScalarQuadraticFunction{Float64}},
    ::Type{MOI.GreaterThan{Float64}},
)
    return model.quadratic_ge_constraints
end

function _constraints(
    model::Optimizer,
    ::Type{MOI.ScalarQuadraticFunction{Float64}},
    ::Type{MOI.EqualTo{Float64}},
)
    return model.quadratic_eq_constraints
end

function _check_inbounds(model::Optimizer, var::MOI.VariableIndex)
    MOI.throw_if_not_valid(model, var)
    return
end

function _check_inbounds(model::Optimizer, aff::MOI.ScalarAffineFunction)
    for term in aff.terms
        MOI.throw_if_not_valid(model, term.variable)
    end
    return
end

function _check_inbounds(model::Optimizer, quad::MOI.ScalarQuadraticFunction)
    for term in quad.affine_terms
        MOI.throw_if_not_valid(model, term.variable)
    end
    for term in quad.quadratic_terms
        MOI.throw_if_not_valid(model, term.variable_1)
        MOI.throw_if_not_valid(model, term.variable_2)
    end
    return
end

function MOI.add_constraint(
    model::Optimizer,
    func::F,
    set::S,
) where {
    F<:Union{
        MOI.ScalarAffineFunction{Float64},
        MOI.ScalarQuadraticFunction{Float64},
    },
    S<:MOI.AbstractScalarSet,
}
    _check_inbounds(model, func)
    constraints = _constraints(model, F, S)
    push!(constraints, _ConstraintInfo(func, set))
    return MOI.ConstraintIndex{F,S}(length(constraints))
end

function MOI.get(
    model::Optimizer,
    ::MOI.NumberOfConstraints{F,S},
) where {
    F<:Union{
        MOI.ScalarAffineFunction{Float64},
        MOI.ScalarQuadraticFunction{Float64},
    },
    S,
}
    return length(_constraints(model, F, S))
end

function MOI.get(
    model::Optimizer,
    ::MOI.ListOfConstraintIndices{F,S},
) where {
    F<:Union{
        MOI.ScalarAffineFunction{Float64},
        MOI.ScalarQuadraticFunction{Float64},
    },
    S,
}
    return MOI.ConstraintIndex{F,S}[
        MOI.ConstraintIndex{F,S}(i) for
        i in eachindex(_constraints(model, F, S))
    ]
end

function MOI.get(
    model::Optimizer,
    ::MOI.ConstraintFunction,
    c::MOI.ConstraintIndex{F,S},
) where {
    F<:Union{
        MOI.ScalarAffineFunction{Float64},
        MOI.ScalarQuadraticFunction{Float64},
    },
    S,
}
    return _constraints(model, F, S)[c.value].func
end

function MOI.get(
    model::Optimizer,
    ::MOI.ConstraintSet,
    c::MOI.ConstraintIndex{F,S},
) where {
    F<:Union{
        MOI.ScalarAffineFunction{Float64},
        MOI.ScalarQuadraticFunction{Float64},
    },
    S,
}
    return _constraints(model, F, S)[c.value].set
end

function MOI.supports(
    ::Optimizer,
    ::MOI.ConstraintDualStart,
    ::Type{MOI.ConstraintIndex{F,S}},
) where {
    F<:Union{
        MOI.ScalarAffineFunction{Float64},
        MOI.ScalarQuadraticFunction{Float64},
    },
    S,
}
    return true
end

function MOI.set(
    model::Optimizer,
    ::MOI.ConstraintDualStart,
    ci::MOI.ConstraintIndex{F,S},
    value::Union{Real,Nothing},
) where {
    F<:Union{
        MOI.ScalarAffineFunction{Float64},
        MOI.ScalarQuadraticFunction{Float64},
    },
    S,
}
    MOI.throw_if_not_valid(model, ci)
    constraints = _constraints(model, F, S)
    constraints[ci.value].dual_start = value
    return
end

function MOI.get(
    model::Optimizer,
    ::MOI.ConstraintDualStart,
    ci::MOI.ConstraintIndex{F,S},
) where {
    F<:Union{
        MOI.ScalarAffineFunction{Float64},
        MOI.ScalarQuadraticFunction{Float64},
    },
    S,
}
    MOI.throw_if_not_valid(model, ci)
    constraints = _constraints(model, F, S)
    return constraints[ci.value].dual_start
end

### MOI.VariablePrimalStart

function MOI.supports(
    ::Optimizer,
    ::MOI.VariablePrimalStart,
    ::Type{MOI.VariableIndex},
)
    return true
end

function MOI.set(
    model::Optimizer,
    ::MOI.VariablePrimalStart,
    vi::MOI.VariableIndex,
    value::Union{Real,Nothing},
)
    MOI.throw_if_not_valid(model, vi)
    model.variable_primal_start[column(vi)] = value
    return
end

### MOI.ConstraintDualStart

_dual_start(::Optimizer, ::Nothing, ::Int = 1) = 0.0

function _dual_start(model::Optimizer, value::Real, scale::Int = 1)
    return _dual_multiplier(model) * value * scale
end

function MOI.supports(
    ::Optimizer,
    ::MOI.ConstraintDualStart,
    ::Type{
        MOI.ConstraintIndex{
            MOI.VariableIndex,
            <:Union{MOI.GreaterThan,MOI.LessThan,MOI.EqualTo},
        },
    },
)
    return true
end

function MOI.set(
    model::Optimizer,
    ::MOI.ConstraintDualStart,
    ci::MOI.ConstraintIndex{MOI.VariableIndex,MOI.GreaterThan{Float64}},
    value::Union{Real,Nothing},
)
    MOI.throw_if_not_valid(model, ci)
    model.variable_lower_start[ci.value] = value
    return
end

function MOI.get(
    model::Optimizer,
    ::MOI.ConstraintDualStart,
    ci::MOI.ConstraintIndex{MOI.VariableIndex,MOI.GreaterThan{Float64}},
)
    MOI.throw_if_not_valid(model, ci)
    return model.variable_lower_start[ci.value]
end

function MOI.set(
    model::Optimizer,
    ::MOI.ConstraintDualStart,
    ci::MOI.ConstraintIndex{MOI.VariableIndex,MOI.LessThan{Float64}},
    value::Union{Real,Nothing},
)
    MOI.throw_if_not_valid(model, ci)
    model.variable_upper_start[ci.value] = value
    return
end

function MOI.get(
    model::Optimizer,
    ::MOI.ConstraintDualStart,
    ci::MOI.ConstraintIndex{MOI.VariableIndex,MOI.LessThan{Float64}},
)
    MOI.throw_if_not_valid(model, ci)
    return model.variable_upper_start[ci.value]
end

function MOI.set(
    model::Optimizer,
    ::MOI.ConstraintDualStart,
    ci::MOI.ConstraintIndex{MOI.VariableIndex,MOI.EqualTo{Float64}},
    value::Union{Real,Nothing},
)
    MOI.throw_if_not_valid(model, ci)
    if value === nothing
        model.variable_lower_start[ci.value] = nothing
        model.variable_upper_start[ci.value] = nothing
    elseif value >= 0.0
        model.variable_lower_start[ci.value] = value
        model.variable_upper_start[ci.value] = 0.0
    else
        model.variable_lower_start[ci.value] = 0.0
        model.variable_upper_start[ci.value] = value
    end
    return
end

function MOI.get(
    model::Optimizer,
    ::MOI.ConstraintDualStart,
    ci::MOI.ConstraintIndex{MOI.VariableIndex,MOI.EqualTo{Float64}},
)
    MOI.throw_if_not_valid(model, ci)
    l = model.variable_lower_start[ci.value]
    u = model.variable_upper_start[ci.value]
    return (l === u === nothing) ? nothing : (l + u)
end

### MOI.NLPBlockDualStart

MOI.supports(::Optimizer, ::MOI.NLPBlockDualStart) = true

function MOI.set(
    model::Optimizer,
    ::MOI.NLPBlockDualStart,
    values::Union{Nothing,Vector},
)
    model.nlp_dual_start = values
    return
end

MOI.get(model::Optimizer, ::MOI.NLPBlockDualStart) = model.nlp_dual_start

### MOI.NLPBlock

MOI.supports(::Optimizer, ::MOI.NLPBlock) = true

function MOI.set(model::Optimizer, ::MOI.NLPBlock, nlp_data::MOI.NLPBlockData)
    model.nlp_data = nlp_data
    return
end

### ObjectiveSense

MOI.supports(::Optimizer, ::MOI.ObjectiveSense) = true

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

MOI.get(model::Optimizer, ::MOI.ObjectiveSense) = model.sense

### ObjectiveFunction

MOI.get(model::Optimizer, ::MOI.ObjectiveFunctionType) = typeof(model.objective)

function MOI.get(model::Optimizer, ::MOI.ObjectiveFunction{F}) where {F}
    return convert(F, model.objective)::F
end

function MOI.supports(
    ::Optimizer,
    ::MOI.ObjectiveFunction{
        <:Union{
            MOI.VariableIndex,
            MOI.ScalarAffineFunction{Float64},
            MOI.ScalarQuadraticFunction{Float64},
        },
    },
)
    return true
end

function MOI.set(
    model::Optimizer,
    ::MOI.ObjectiveFunction{F},
    func::F,
) where {
    F<:Union{
        MOI.VariableIndex,
        MOI.ScalarAffineFunction{Float64},
        MOI.ScalarQuadraticFunction{Float64},
    },
}
    _check_inbounds(model, func)
    model.objective = func
    return
end

### Ipopt callback functions
### In setting up the data for Ipopt, we order the constraints as follows:
### - linear_le_constraints
### - linear_ge_constraints
### - linear_eq_constraints
### - quadratic_le_constraints
### - quadratic_ge_constraints
### - quadratic_eq_constraints
### - nonlinear constraints from nlp_data

const _CONSTRAINT_ORDERING = (
    :linear_le_constraints,
    :linear_ge_constraints,
    :linear_eq_constraints,
    :quadratic_le_constraints,
    :quadratic_ge_constraints,
    :quadratic_eq_constraints,
)

function _offset(
    ::Optimizer,
    ::Type{<:MOI.ScalarAffineFunction},
    ::Type{<:MOI.LessThan},
)
    return 0
end

function _offset(
    model::Optimizer,
    ::Type{<:MOI.ScalarAffineFunction},
    ::Type{<:MOI.GreaterThan},
)
    return length(model.linear_le_constraints)
end

function _offset(
    model::Optimizer,
    F::Type{<:MOI.ScalarAffineFunction},
    ::Type{<:MOI.EqualTo},
)
    return _offset(model, F, MOI.GreaterThan{Float64}) +
           length(model.linear_ge_constraints)
end

function _offset(
    model::Optimizer,
    ::Type{<:MOI.ScalarQuadraticFunction},
    ::Type{<:MOI.LessThan},
)
    x = _offset(model, MOI.ScalarAffineFunction{Float64}, MOI.EqualTo{Float64})
    return x + length(model.linear_eq_constraints)
end

function _offset(
    model::Optimizer,
    F::Type{<:MOI.ScalarQuadraticFunction},
    ::Type{<:MOI.GreaterThan},
)
    return _offset(model, F, MOI.LessThan{Float64}) +
           length(model.quadratic_le_constraints)
end

function _offset(
    model::Optimizer,
    F::Type{<:MOI.ScalarQuadraticFunction},
    ::Type{<:MOI.EqualTo},
)
    return _offset(model, F, MOI.GreaterThan{Float64}) +
           length(model.quadratic_ge_constraints)
end

function _nlp_constraint_offset(model::Optimizer)
    x = _offset(
        model,
        MOI.ScalarQuadraticFunction{Float64},
        MOI.EqualTo{Float64},
    )
    return x + length(model.quadratic_eq_constraints)
end

_eval_function(::Nothing, ::Any) = 0.0

_eval_function(f, x) = MOI.Utilities.eval_variables(xi -> x[xi.value], f)

### Eval_F_CB

function _eval_objective(model::Optimizer, x)
    if model.nlp_data.has_objective
        return MOI.eval_objective(model.nlp_data.evaluator, x)
    end
    return _eval_function(model.objective, x)
end

### Eval_Grad_F_CB

_fill_gradient(::Any, ::Any, ::Nothing) = nothing

function _fill_gradient(grad, ::Vector, f::MOI.VariableIndex)
    grad[f.value] = 1.0
    return
end

function _fill_gradient(grad, ::Vector, f::MOI.ScalarAffineFunction{Float64})
    for term in f.terms
        grad[term.variable.value] += term.coefficient
    end
    return
end

function _fill_gradient(
    grad,
    x::Vector,
    quad::MOI.ScalarQuadraticFunction{Float64},
)
    for term in quad.affine_terms
        grad[term.variable.value] += term.coefficient
    end
    for term in quad.quadratic_terms
        row_idx = term.variable_1
        col_idx = term.variable_2
        if row_idx == col_idx
            grad[row_idx.value] += term.coefficient * x[row_idx.value]
        else
            grad[row_idx.value] += term.coefficient * x[col_idx.value]
            grad[col_idx.value] += term.coefficient * x[row_idx.value]
        end
    end
    return
end

function _eval_objective_gradient(model::Optimizer, grad, x)
    if model.nlp_data.has_objective
        MOI.eval_objective_gradient(model.nlp_data.evaluator, grad, x)
    else
        fill!(grad, 0.0)
        _fill_gradient(grad, x, model.objective)
    end
    return
end

### Eval_G_CB

function _eval_constraint(model::Optimizer, g, x)
    row = 1
    for key in _CONSTRAINT_ORDERING
        for info in getfield(model, key)
            g[row] = _eval_function(info.func, x)
            row += 1
        end
    end
    nlp_g = view(g, row:length(g))
    MOI.eval_constraint(model.nlp_data.evaluator, nlp_g, x)
    return
end

### Eval_Jac_G_CB

function _append_to_jacobian_sparsity(J, f::MOI.ScalarAffineFunction, row)
    for term in f.terms
        push!(J, (row, term.variable.value))
    end
    return
end

function _append_to_jacobian_sparsity(J, f::MOI.ScalarQuadraticFunction, row)
    for term in f.affine_terms
        push!(J, (row, term.variable.value))
    end
    for term in f.quadratic_terms
        row_idx = term.variable_1
        col_idx = term.variable_2
        if row_idx == col_idx
            push!(J, (row, row_idx.value))
        else
            push!(J, (row, row_idx.value))
            push!(J, (row, col_idx.value))
        end
    end
    return
end

function _jacobian_structure(model::Optimizer)
    J = Tuple{Int64,Int64}[]
    row = 1
    for key in _CONSTRAINT_ORDERING
        for info in getfield(model, key)
            _append_to_jacobian_sparsity(J, info.func, row)
            row += 1
        end
    end
    if length(model.nlp_data.constraint_bounds) > 0
        for (nlp_row, col) in MOI.jacobian_structure(model.nlp_data.evaluator)
            push!(J, (nlp_row + row - 1, col))
        end
    end
    return J
end

function _fill_constraint_jacobian(
    values,
    offset,
    ::Vector,
    f::MOI.ScalarAffineFunction,
)
    num_coefficients = length(f.terms)
    for i in 1:num_coefficients
        values[offset+i] = f.terms[i].coefficient
    end
    return num_coefficients
end

function _fill_constraint_jacobian(
    values,
    offset,
    x,
    f::MOI.ScalarQuadraticFunction,
)
    nterms = 0
    for term in f.affine_terms
        nterms += 1
        values[offset+nterms] = term.coefficient
    end
    for term in f.quadratic_terms
        row_idx = term.variable_1
        col_idx = term.variable_2
        if row_idx == col_idx
            nterms += 1
            values[offset+nterms] = term.coefficient * x[col_idx.value]
        else
            # Note that the order matches the Jacobian sparsity pattern.
            nterms += 2
            values[offset+nterms-1] = term.coefficient * x[col_idx.value]
            values[offset+nterms] = term.coefficient * x[row_idx.value]
        end
    end
    return nterms
end

function _eval_constraint_jacobian(model::Optimizer, values, x)
    offset = 0
    for key in _CONSTRAINT_ORDERING
        for info in getfield(model, key)
            offset += _fill_constraint_jacobian(values, offset, x, info.func)
        end
    end
    nlp_values = view(values, (1+offset):length(values))
    MOI.eval_constraint_jacobian(model.nlp_data.evaluator, nlp_values, x)
    return
end

### Eval_H_CB

_append_to_hessian_sparsity(::Any, ::Any) = nothing

function _append_to_hessian_sparsity(H, f::MOI.ScalarQuadraticFunction)
    for term in f.quadratic_terms
        push!(H, (term.variable_1.value, term.variable_2.value))
    end
    return
end

function _append_hessian_lagrangian_structure(H, model::Optimizer)
    if !model.nlp_data.has_objective
        _append_to_hessian_sparsity(H, model.objective)
    end
    for info in model.quadratic_le_constraints
        _append_to_hessian_sparsity(H, info.func)
    end
    for info in model.quadratic_ge_constraints
        _append_to_hessian_sparsity(H, info.func)
    end
    for info in model.quadratic_eq_constraints
        _append_to_hessian_sparsity(H, info.func)
    end
    append!(H, MOI.hessian_lagrangian_structure(model.nlp_data.evaluator))
    return
end

_fill_hessian_lagrangian(::Any, ::Any, ::Any, ::Any) = 0

function _fill_hessian_lagrangian(H, offset, λ, f::MOI.ScalarQuadraticFunction)
    for term in f.quadratic_terms
        H[offset+1] = λ * term.coefficient
        offset += 1
    end
    return length(f.quadratic_terms)
end

function _eval_hessian_lagrangian(
    ::Type{S},
    model::Optimizer,
    H,
    μ,
    offset,
) where {S}
    F = MOI.ScalarQuadraticFunction{Float64}
    offset_start = _offset(model, F, S)
    for (i, info) in enumerate(_constraints(model, F, S))
        offset +=
            _fill_hessian_lagrangian(H, offset, μ[offset_start+i], info.func)
    end
    return offset
end

function _eval_hessian_lagrangian(model::Optimizer, H, x, σ, μ)
    offset = 0
    if !model.nlp_data.has_objective
        offset += _fill_hessian_lagrangian(H, 0, σ, model.objective)
    end
    # Handles any quadratic constraints that are present. The order matters.
    offset =
        _eval_hessian_lagrangian(MOI.LessThan{Float64}, model, H, μ, offset)
    offset =
        _eval_hessian_lagrangian(MOI.GreaterThan{Float64}, model, H, μ, offset)
    offset = _eval_hessian_lagrangian(MOI.EqualTo{Float64}, model, H, μ, offset)
    # Handles the Hessian in the nonlinear block
    MOI.eval_hessian_lagrangian(
        model.nlp_data.evaluator,
        view(H, 1+offset:length(H)),
        x,
        σ,
        view(μ, 1+_nlp_constraint_offset(model):length(μ)),
    )
    return
end

### MOI.optimize!

_bounds(s::MOI.LessThan) = (-Inf, s.upper)
_bounds(s::MOI.GreaterThan) = (s.lower, Inf)
_bounds(s::MOI.EqualTo) = (s.value, s.value)

function MOI.optimize!(model::Optimizer)
    # TODO: Reuse model.inner for incremental solves if possible.
    num_quadratic_constraints =
        length(model.quadratic_le_constraints) +
        length(model.quadratic_ge_constraints) +
        length(model.quadratic_eq_constraints)
    num_nlp_constraints = length(model.nlp_data.constraint_bounds)
    has_hessian = :Hess in MOI.features_available(model.nlp_data.evaluator)
    init_feat = [:Grad]
    if has_hessian
        push!(init_feat, :Hess)
    end
    if num_nlp_constraints > 0
        push!(init_feat, :Jac)
    end
    MOI.initialize(model.nlp_data.evaluator, init_feat)
    jacobian_sparsity = _jacobian_structure(model)
    hessian_sparsity = Tuple{Int,Int}[]
    if has_hessian
        _append_hessian_lagrangian_structure(hessian_sparsity, model)
    end
    function eval_f_cb(x)
        # TODO(odow): FEASIBILITY_SENSE could produce confusing solver output if
        # a nonzero objective is set.
        if model.sense == MOI.FEASIBILITY_SENSE
            return 0.0
        end
        return _eval_objective(model, x)
    end
    function eval_grad_f_cb(x, grad_f)
        if model.sense == MOI.FEASIBILITY_SENSE
            grad_f .= zero(eltype(grad_f))
        else
            _eval_objective_gradient(model, grad_f, x)
        end
        return
    end
    eval_g_cb(x, g) = _eval_constraint(model, g, x)
    function eval_jac_g_cb(x, rows, cols, values)
        if values === nothing
            for i in 1:length(jacobian_sparsity)
                rows[i], cols[i] = jacobian_sparsity[i]
            end
        else
            _eval_constraint_jacobian(model, values, x)
        end
        return
    end
    function eval_h_cb(x, rows, cols, obj_factor, lambda, values)
        if values === nothing
            for i in 1:length(hessian_sparsity)
                rows[i], cols[i] = hessian_sparsity[i]
            end
        else
            _eval_hessian_lagrangian(model, values, x, obj_factor, lambda)
        end
        return
    end
    g_L, g_U = Float64[], Float64[]
    for key in _CONSTRAINT_ORDERING
        for info in getfield(model, key)
            l, u = _bounds(info.set)
            push!(g_L, l)
            push!(g_U, u)
        end
    end
    for bound in model.nlp_data.constraint_bounds
        push!(g_L, bound.lower)
        push!(g_U, bound.upper)
    end
    start_time = time()
    if length(model.variables.lower) == 0
        # Don't attempt to create a problem because Ipopt will error.
        model.invalid_model = true
        return
    end
    model.inner = CreateIpoptProblem(
        length(model.variables.lower),
        model.variables.lower,
        model.variables.upper,
        length(g_L),
        g_L,
        g_U,
        length(jacobian_sparsity),
        length(hessian_sparsity),
        eval_f_cb,
        eval_g_cb,
        eval_grad_f_cb,
        eval_jac_g_cb,
        has_hessian ? eval_h_cb : nothing,
    )
    if model.sense == MOI.MIN_SENSE
        AddIpoptNumOption(model.inner, "obj_scaling_factor", 1.0)
    elseif model.sense == MOI.MAX_SENSE
        AddIpoptNumOption(model.inner, "obj_scaling_factor", -1.0)
    end
    # Ipopt crashes by default if NaN/Inf values are returned from the
    # evaluation callbacks. This option tells Ipopt to explicitly check for them
    # and return Invalid_Number_Detected instead. This setting may result in a
    # minor performance loss and can be overwritten by specifying
    # check_derivatives_for_naninf="no".
    AddIpoptStrOption(model.inner, "check_derivatives_for_naninf", "yes")
    if !has_hessian
        AddIpoptStrOption(
            model.inner,
            "hessian_approximation",
            "limited-memory",
        )
    end
    if num_nlp_constraints == 0 && num_quadratic_constraints == 0
        AddIpoptStrOption(model.inner, "jac_c_constant", "yes")
        AddIpoptStrOption(model.inner, "jac_d_constant", "yes")
        if !model.nlp_data.has_objective
            # We turn on this option if all constraints are linear and the
            # objective is linear or quadratic. From the documentation, it's
            # unclear if it may also apply if the constraints are at most
            # quadratic.
            AddIpoptStrOption(model.inner, "hessian_constant", "yes")
        end
    end
    if model.silent
        AddIpoptIntOption(model.inner, "print_level", 0)
    end
    # Other misc options that over-ride the ones set above.
    for (name, value) in model.options
        if value isa String
            AddIpoptStrOption(model.inner, name, value)
        elseif value isa Integer
            AddIpoptIntOption(model.inner, name, value)
        else
            @assert value isa Float64
            AddIpoptNumOption(model.inner, name, value)
        end
    end
    # Initialize the starting point, projecting variables from 0 onto their
    # bounds if VariablePrimalStart  is not provided.
    for (i, v) in enumerate(model.variable_primal_start)
        if v !== nothing
            model.inner.x[i] = v
        else
            model.inner.x[i] = max(0.0, model.variables.lower[i])
            model.inner.x[i] = min(model.inner.x[i], model.variables.upper[i])
        end
    end
    # Initialize the dual start to 0.0 if NLPBlockDualStart is not provided.
    if model.nlp_dual_start === nothing
        model.nlp_dual_start = zeros(Float64, num_nlp_constraints)
    end
    # ConstraintDualStart
    row = 1
    for key in _CONSTRAINT_ORDERING
        for info in getfield(model, key)
            model.inner.mult_g[row] = _dual_start(model, info.dual_start, -1)
            row += 1
        end
    end
    for dual_start in model.nlp_dual_start
        model.inner.mult_g[row] = _dual_start(model, dual_start, -1)
        row += 1
    end
    # ConstraintDualStart for variable bounds
    for i in 1:length(model.inner.n)
        model.inner.mult_x_L[i] =
            _dual_start(model, model.variable_lower_start[i])
        model.inner.mult_x_U[i] =
            _dual_start(model, model.variable_upper_start[i], -1)
    end
    # Initialize CallbackFunction.
    if model.callback !== nothing
        SetIntermediateCallback(model.inner, model.callback)
    end
    IpoptSolve(model.inner)
    # Store SolveTimeSec.
    model.solve_time = time() - start_time
    return
end

# From Ipopt/src/Interfaces/IpReturnCodes_inc.h
const _STATUS_CODES = Dict(
    0 => :Solve_Succeeded,
    1 => :Solved_To_Acceptable_Level,
    2 => :Infeasible_Problem_Detected,
    3 => :Search_Direction_Becomes_Too_Small,
    4 => :Diverging_Iterates,
    5 => :User_Requested_Stop,
    6 => :Feasible_Point_Found,
    -1 => :Maximum_Iterations_Exceeded,
    -2 => :Restoration_Failed,
    -3 => :Error_In_Step_Computation,
    -4 => :Maximum_CpuTime_Exceeded,
    -5 => :Maximum_WallTime_Exceeded,
    -10 => :Not_Enough_Degrees_Of_Freedom,
    -11 => :Invalid_Problem_Definition,
    -12 => :Invalid_Option,
    -13 => :Invalid_Number_Detected,
    -100 => :Unrecoverable_Exception,
    -101 => :NonIpopt_Exception_Thrown,
    -102 => :Insufficient_Memory,
    -199 => :Internal_Error,
)

### MOI.ResultCount

# Ipopt always has an iterate available.
function MOI.get(model::Optimizer, ::MOI.ResultCount)
    return (model.inner !== nothing) ? 1 : 0
end

### MOI.TerminationStatus

function MOI.get(model::Optimizer, ::MOI.TerminationStatus)
    if model.invalid_model
        return MOI.INVALID_MODEL
    elseif model.inner === nothing
        return MOI.OPTIMIZE_NOT_CALLED
    end
    status = _STATUS_CODES[model.inner.status]
    if status == :Solve_Succeeded || status == :Feasible_Point_Found
        return MOI.LOCALLY_SOLVED
    elseif status == :Infeasible_Problem_Detected
        return MOI.LOCALLY_INFEASIBLE
    elseif status == :Solved_To_Acceptable_Level
        return MOI.ALMOST_LOCALLY_SOLVED
    elseif status == :Search_Direction_Becomes_Too_Small
        return MOI.NUMERICAL_ERROR
    elseif status == :Diverging_Iterates
        return MOI.NORM_LIMIT
    elseif status == :User_Requested_Stop
        return MOI.INTERRUPTED
    elseif status == :Maximum_Iterations_Exceeded
        return MOI.ITERATION_LIMIT
    elseif status == :Maximum_CpuTime_Exceeded
        return MOI.TIME_LIMIT
    elseif status == :Maximum_WallTime_Exceeded
        return MOI.TIME_LIMIT
    elseif status == :Restoration_Failed
        return MOI.NUMERICAL_ERROR
    elseif status == :Error_In_Step_Computation
        return MOI.NUMERICAL_ERROR
    elseif status == :Invalid_Option
        return MOI.INVALID_OPTION
    elseif status == :Not_Enough_Degrees_Of_Freedom
        return MOI.INVALID_MODEL
    elseif status == :Invalid_Problem_Definition
        return MOI.INVALID_MODEL
    elseif status == :Invalid_Number_Detected
        return MOI.INVALID_MODEL
    elseif status == :Unrecoverable_Exception
        return MOI.OTHER_ERROR
    elseif status == :NonIpopt_Exception_Thrown
        return MOI.OTHER_ERROR
    else
        @assert status == :Insufficient_Memory
        return MOI.MEMORY_LIMIT
    end
end

### MOI.RawStatusString

function MOI.get(model::Optimizer, ::MOI.RawStatusString)
    if model.invalid_model
        return "The model has no variable"
    elseif model.inner === nothing
        return "Optimize not called"
    else
        return string(_STATUS_CODES[model.inner.status])
    end
end

### MOI.PrimalStatus

function MOI.get(model::Optimizer, attr::MOI.PrimalStatus)
    if !(1 <= attr.result_index <= MOI.get(model, MOI.ResultCount()))
        return MOI.NO_SOLUTION
    end
    status = _STATUS_CODES[model.inner.status]
    if status == :Solve_Succeeded
        return MOI.FEASIBLE_POINT
    elseif status == :Feasible_Point_Found
        return MOI.FEASIBLE_POINT
    elseif status == :Solved_To_Acceptable_Level
        # Solutions are only guaranteed to satisfy the "acceptable" convergence
        # tolerances.
        return MOI.NEARLY_FEASIBLE_POINT
    elseif status == :Infeasible_Problem_Detected
        return MOI.INFEASIBLE_POINT
    else
        return MOI.UNKNOWN_RESULT_STATUS
    end
end

### MOI.DualStatus

function MOI.get(model::Optimizer, attr::MOI.DualStatus)
    if !(1 <= attr.result_index <= MOI.get(model, MOI.ResultCount()))
        return MOI.NO_SOLUTION
    end
    status = _STATUS_CODES[model.inner.status]
    if status == :Solve_Succeeded
        return MOI.FEASIBLE_POINT
    elseif status == :Feasible_Point_Found
        return MOI.FEASIBLE_POINT
    elseif status == :Solved_To_Acceptable_Level
        # Solutions are only guaranteed to satisfy the "acceptable" convergence
        # tolerances.
        return MOI.NEARLY_FEASIBLE_POINT
    else
        return MOI.UNKNOWN_RESULT_STATUS
    end
end

### MOI.SolveTimeSec

MOI.get(model::Optimizer, ::MOI.SolveTimeSec) = model.solve_time

### MOI.ObjectiveValue

function MOI.get(model::Optimizer, attr::MOI.ObjectiveValue)
    MOI.check_result_index_bounds(model, attr)
    return model.inner.obj_val
end

### MOI.VariablePrimal

function MOI.get(
    model::Optimizer,
    attr::MOI.VariablePrimal,
    vi::MOI.VariableIndex,
)
    MOI.check_result_index_bounds(model, attr)
    MOI.throw_if_not_valid(model, vi)
    return model.inner.x[column(vi)]
end

### MOI.ConstraintPrimal

function MOI.get(
    model::Optimizer,
    attr::MOI.ConstraintPrimal,
    ci::MOI.ConstraintIndex{F,S},
) where {
    F<:Union{
        MOI.ScalarAffineFunction{Float64},
        MOI.ScalarQuadraticFunction{Float64},
    },
    S,
}
    MOI.check_result_index_bounds(model, attr)
    MOI.throw_if_not_valid(model, ci)
    return model.inner.g[_offset(model, F, S)+ci.value]
end

function MOI.get(
    model::Optimizer,
    attr::MOI.ConstraintPrimal,
    ci::MOI.ConstraintIndex{
        MOI.VariableIndex,
        <:Union{
            MOI.LessThan{Float64},
            MOI.GreaterThan{Float64},
            MOI.EqualTo{Float64},
        },
    },
)
    MOI.check_result_index_bounds(model, attr)
    MOI.throw_if_not_valid(model, ci)
    return model.inner.x[ci.value]
end

### MOI.ConstraintDual

_dual_multiplier(model::Optimizer) = model.sense == MOI.MIN_SENSE ? 1.0 : -1.0

function MOI.get(
    model::Optimizer,
    attr::MOI.ConstraintDual,
    ci::MOI.ConstraintIndex{F,S},
) where {
    F<:Union{
        MOI.ScalarAffineFunction{Float64},
        MOI.ScalarQuadraticFunction{Float64},
    },
    S,
}
    MOI.check_result_index_bounds(model, attr)
    MOI.throw_if_not_valid(model, ci)
    s = -_dual_multiplier(model)
    return s * model.inner.mult_g[_offset(model, F, S)+ci.value]
end

function MOI.get(
    model::Optimizer,
    attr::MOI.ConstraintDual,
    ci::MOI.ConstraintIndex{MOI.VariableIndex,MOI.LessThan{Float64}},
)
    MOI.check_result_index_bounds(model, attr)
    MOI.throw_if_not_valid(model, ci)
    rc = model.inner.mult_x_L[ci.value] - model.inner.mult_x_U[ci.value]
    return min(0.0, _dual_multiplier(model) * rc)
end

function MOI.get(
    model::Optimizer,
    attr::MOI.ConstraintDual,
    ci::MOI.ConstraintIndex{MOI.VariableIndex,MOI.GreaterThan{Float64}},
)
    MOI.check_result_index_bounds(model, attr)
    MOI.throw_if_not_valid(model, ci)
    rc = model.inner.mult_x_L[ci.value] - model.inner.mult_x_U[ci.value]
    return max(0.0, _dual_multiplier(model) * rc)
end

function MOI.get(
    model::Optimizer,
    attr::MOI.ConstraintDual,
    ci::MOI.ConstraintIndex{MOI.VariableIndex,MOI.EqualTo{Float64}},
)
    MOI.check_result_index_bounds(model, attr)
    MOI.throw_if_not_valid(model, ci)
    rc = model.inner.mult_x_L[ci.value] - model.inner.mult_x_U[ci.value]
    return _dual_multiplier(model) * rc
end

### MOI.NLPBlockDual

function MOI.get(model::Optimizer, attr::MOI.NLPBlockDual)
    MOI.check_result_index_bounds(model, attr)
    s = -_dual_multiplier(model)
    return s .* model.inner.mult_g[(1+_nlp_constraint_offset(model)):end]
end

### Ipopt.CallbackFunction

"""
    CallbackFunction()

A solver-dependent callback for Ipopt's IntermediateCallback.

The callback should be a function like the following:
```julia
function my_intermediate_callback(
    alg_mod::Cint,
    iter_count::Cint,
    obj_value::Float64,
    inf_pr::Float64,
    inf_du::Float64,
    mu::Float64,
    d_norm::Float64,
    regularization_size::Float64,
    alpha_du::Float64,
    alpha_pr::Float64,
    ls_trials::Cint,
)
    # ... user code ...
    return true # or `return false` to terminate the solve.
end
```

The arguments are defined in the Ipopt documentation:
https://coin-or.github.io/Ipopt/OUTPUT.html

Note: Calling `SetIntermediateCallback` will over-write this callback! Don't
call both.
"""
struct CallbackFunction <: MOI.AbstractCallback end

function MOI.set(model::Optimizer, ::CallbackFunction, f::Function)
    model.callback = f
    return
end

function MOI.get(
    model::Optimizer,
    ::MOI.CallbackVariablePrimal,
    x::MOI.VariableIndex,
)
    return model.inner.x[column(x)]
end


module Variational

using ..Core, ..Core.RandomVariables, ..Utilities
using Distributions, Bijectors
using ProgressMeter, LinearAlgebra
using ..Turing: PROGRESS
using ..Turing: Model, SampleFromPrior, SampleFromUniform
using ..Turing: Turing
using Random: AbstractRNG

using ForwardDiff
using Tracker

import ..Core: getchunksize, getADtype

using Requires
function __init__()
    @require Flux="587475ba-b771-5e3f-ad9e-33799f191a9c" apply!(o, x, Δ) = Flux.Optimise.apply!(o, x, Δ)
end

export
    vi,
    ADVI,
    ELBO,
    TruncatedADAGrad

abstract type VariationalInference{AD} end

getchunksize(::T) where {T <: VariationalInference} = getchunksize(T)
getchunksize(::Type{<:VariationalInference{AD}}) where AD = getchunksize(AD)
getADtype(alg::VariationalInference) = getADtype(typeof(alg))
getADtype(::Type{<: VariationalInference{AD}}) where {AD} = AD

abstract type VariationalObjective end

const VariationalPosterior = Distribution{Multivariate, Continuous}

"""
    rand(vi::VariationalInference, num_samples)

Produces `num_samples` samples for the given VI method using number of samples equal
to `num_samples`.
"""
function rand(vi::VariationalPosterior, num_samples) end

"""
    objective(vi::VariationalInference, q::VariationalPosterior, model::Model, args...)

Computes the variational objective to be optimized for a given VI method.
"""
function objective(
    vi::VariationalInference,
    q::VariationalPosterior,
    model::Model,
    num_samples)
end

"""
    grad!(vo, vi::VariationalInference, q::VariationalPosterior, model::Model, θ, out, args...)

Computes the gradients used in `optimize!`. Default implementation is provided for
`VariationalInference{AD}` where `AD` is either `ForwardDiffAD` or `TrackerAD`.
This implicitly also gives a default implementation of `optimize!`.

Variance reduction techniques, e.g. control variates, should be implemented in this function.
"""
function grad!(
    vo, vi::VariationalInference{AD},
    q::VariationalPosterior,
    model::Model,
    θ,
    out,
    args...
) where AD
    error("Turing.Variational.grad!: unmanaged variational inference algorithm: "
          * "$(typeof(alg))")
end

"""
    vi(model::Model, alg::VariationalInference)
    vi(model::Model, alg::VariationalInference, q::VariationalPosterior)

Constructs the variational posterior from the `model` and performs the optimization
following the configuration of the given `VariationalInference` instance.
"""
function vi(model::Model, alg::VariationalInference)
    error("Turing.Variational.vi: variational inference algorithm $(typeof(alg)) "
          * "is not implemented")
end
function vi(model::Model, alg::VariationalInference, q::VariationalPosterior)
    error("Turing.Variational.vi: variational inference algorithm $(typeof(alg)) "
          * "is not implemented")
end

# default implementations
function grad!(
    vo,
    alg::VariationalInference{AD},
    q::VariationalPosterior,
    model::Model,
    θ::AbstractVector{T},
    out::DiffResults.MutableDiffResult,
    args...
) where {T <: Real, AD <: ForwardDiffAD}
    # TODO: this probably slows down executation quite a bit; exists a better way
    # of doing this?
    f(θ_) = - vo(alg, q, model, θ_, args...)

    chunk_size = getchunksize(alg)
    # Set chunk size and do ForwardMode.
    chunk = ForwardDiff.Chunk(min(length(θ), chunk_size))
    config = ForwardDiff.GradientConfig(f, θ, chunk)
    ForwardDiff.gradient!(out, f, θ, config)
end

function grad!(
    vo,
    alg::VariationalInference{AD},
    q::VariationalPosterior,
    model::Model,
    θ::AbstractVector{T},
    out::DiffResults.MutableDiffResult,
    args...
) where {T <: Real, AD <: TrackerAD}
    θ_tracked = [Tracker.param(θ[i]) for i ∈ eachindex(θ)]
    y = - vo(alg, q, model, θ_tracked, args...)
    Tracker.back!(y, 1.0)

    DiffResults.value!(out, Tracker.data(y))
    DiffResults.gradient!(out, [Tracker.grad(θ_tracked[i]) for i ∈ eachindex(θ_tracked)])
end

"""
    optimize!(vo, alg::VariationalInference{AD}, q::VariationalPosterior, model::Model, θ; optimizer = TruncatedADAGrad())

Iteratively updates parameters by calling `grad!` and using the given `optimizer` to compute
the steps.
"""
function optimize!(
    vo,
    alg::VariationalInference{AD},
    q::VariationalPosterior,
    model::Model,
    θ;
    optimizer = TruncatedADAGrad()
) where AD
    # TODO: should we always assume `samples_per_step` and `max_iters` for all algos?
    alg_name = alg_str(alg)
    samples_per_step = alg.samples_per_step
    max_iters = alg.max_iters

    # number of previous gradients to use to compute `s` in adaGrad
    stepsize_num_prev = 10

    num_params = length(q)

    # TODO: really need a better way to warn the user about potentially
    # not using the correct accumulator
    if (optimizer isa TruncatedADAGrad) && (θ ∉ keys(optimizer.acc))
        # this message should only occurr once in the optimization process
        @info "[$alg_name] Should only be seen once: optimizer created for θ" objectid(θ)
    end

    diff_result = DiffResults.GradientResult(θ)

    i = 0
    prog = if PROGRESS[]
        ProgressMeter.Progress(max_iters, 1, "[$alg_name] Optimizing...", 0)
    else
        0
    end

    # add criterion? A running mean maybe?
    time_elapsed = @elapsed while (i < max_iters) # & converged
        grad!(vo, alg, q, model, θ, diff_result, samples_per_step)

        # apply update rule
        Δ = DiffResults.gradient(diff_result)
        Δ = apply!(optimizer, θ, Δ)
        @. θ = θ - Δ

        Turing.DEBUG && @debug "Step $i" Δ DiffResults.value(diff_result)
        PROGRESS[] && (ProgressMeter.next!(prog))

        i += 1
    end

    return θ
end

# distributions
include("distributions.jl")

# objectives
include("objectives.jl")

# optimisers
include("optimisers.jl")

# VI algorithms
include("advi.jl")

end


################################################################################
#
# This file simulates an adiabatic computing
#
################################################################################

include("Base.jl")
include("Hamiltonian.jl")
include("Operator.jl")


using NeXLSpectrum
using Test

@testset "XRF Spectrum" begin
    fn1 = joinpath(@__DIR__, "Other", "20170510_D1_Mo_50kv_100mA_acryl 0.spx")
    sp1a = loadspectrum(fn1)
    #sp1b = FileIO.load(fn1)
    @test sp1a[:RealTime] == 310.610
    @test sp1a[:LiveTime] == 300.294
    @test sp1a[:BrukerThroughput] == 130000
    @test sp1a[:DetectorModel] == "Bruker XFlash 430"
    @test sp1a[:DetectorSerialNumber] == "11881_0239"
    @test sp1a[:DetectorThickness] == 0.1 * 0.45 # cm
    @test sp1a[:DeadLayerThickness] == 0.1 * 0.029 # ???
    @test sp1a[:BeamEnergy] == 1000.0 * 50.0
    @test sp1a[:XRFTubeAnode] == n"Mo"
    @test sp1a[:ProbeCurrent] == 1000.0 * 99
    @test sp1a[:XRFTubeIncidentAngle] == deg2rad(84.0)
    @test sp1a[:XRFTubeTakeOffAngle] == deg2rad(6.0)
    @test sp1a[:XRFExcitationAngle] == deg2rad(50.0)
    @test sp1a[:Elevation] == deg2rad(50.0)
    @test sp1a[:TakeOffAngle] == deg2rad(50.0)
    @test sp1a[:XRFExcitationPathLength] == 0.1 * 10.0
    @test sp1a[:XRFDetectionPathLength] == 0.1 * 20.0
    @test sp1a[:DetectorSolidAngle] == 0.0065
    @test sp1a[:ChamberPressure] == 20.0
    @test sp1a[:ChamberAtmosphere] == "Air"
    @test sp1a[:XRFSampleTilt] == 0.0
    f1 = sp1a[:XRFTubeWindow]
    @test isequal(f1.material, pure(n"Be"))
    @test isequal(f1.thickness, 200.0e-4)
end

@testset "SEM/EDS" begin
    fn2 = joinpath(@__DIR__, "Other", "Albite_ChMxd_20kV0p5nA130kHzTC_5p9kHzOCR_4ks.spx")
    sp2a = loadspectrum(fn2)
    #sp2b = FileIO.load(fn2)
    @test sp2a[:RealTime] == 4028.702
    @test sp2a[:LiveTime] == 4000.005
    @test sp2a[:BrukerThroughput] == 130000
    @test sp2a[:DetectorModel] == "Bruker XFlash 6|30"
    @test sp2a[:DetectorSerialNumber] == "13773"
    @test sp2a[:DetectorThickness] == 0.1 * 0.45
    @test sp2a[:DeadLayerThickness] == 0.1 * 0.029
    @test sp2a[:Elevation] == deg2rad(35.0)
    win = sp2a[:Window]
    @test length(win) == 5
    @test isapprox(win[1], Film(pure(n"B"), 1.3E1 * 1.0e-7))
    @test isapprox(win[2], Film(pure(n"C"), 1.45E1 * 1.0e-7))
    @test isapprox(win[3], Film(pure(n"N"), 4.5 * 1.0e-7))
    @test isapprox(win[4], Film(pure(n"O"), 8.5 * 1.0e-7))
    @test isapprox(win[5], Film(pure(n"Al"), 3.5 * 1.0e-7))
    @test sp2a[:BeamEnergy] == 20.0e3

    @test n"C" in sp2a[:Elements]
    @test n"O" in sp2a[:Elements]
    @test n"Na" in sp2a[:Elements]
    @test n"Al" in sp2a[:Elements]
    @test n"Si" in sp2a[:Elements]
    @test sp2a[:TakeOffAngle] == deg2rad(35.0)
    @test sp2a[:WorkingDistance] == 1.1
end


# General case - to be used with bc::NoBC or m::PrescribedXModels
function soil_boundary_flux!(
    nf,
    bc::AbstractBoundaryConditions,
    m::AbstractSoilComponentModel,
    land::LandModel,
    state⁺::Vars,
    diff⁺::Vars,
    aux⁺::Vars,
    nM,
    state⁻::Vars,
    diff⁻::Vars,
    aux⁻::Vars,
    t,
    _...,
) end

function soil_boundary_state!(
    nf,
    bc::AbstractBoundaryConditions,
    m::AbstractSoilComponentModel,
    land::LandModel,
    state⁺::Vars,
    aux⁺::Vars,
    nM,
    state⁻::Vars,
    aux⁻::Vars,
    t,
    _...,
)

end

# Dirichlet Methods for SoilHeat and SoilWater

"""
    function soil_boundary_state!(
        nf,
        bc::Dirichlet,
        water::SoilWaterModel,
        land::LandModel,
        state⁺::Vars,
        aux⁺::Vars,
        nM,
        state⁻::Vars,
        aux⁻::Vars,
        t,
    )

The Dirichlet-type method for `soil_boundary_state!` for the
`SoilWaterModel`.
"""
function soil_boundary_state!(
    nf,
    bc::Dirichlet,
    water::SoilWaterModel,
    land::LandModel,
    state⁺::Vars,
    aux⁺::Vars,
    nM,
    state⁻::Vars,
    aux⁻::Vars,
    t,
    _...,
)
    bc_function = bc.state_bc
    state⁺.soil.water.ϑ_l = bc_function(aux⁻, t)

end

"""
    function soil_boundary_state!(
        nf,
        bc::Dirichlet,
        heat::SoilHeatModel,
        land::LandModel,
        state⁺::Vars,
        aux⁺::Vars,
        nM,
        state⁻::Vars,
        aux⁻::Vars,
        t,
    )

The Dirichlet-type method for `soil_boundary_state!` for the
`SoilHeatModel`.
"""
function soil_boundary_state!(
    nf,
    bc::Dirichlet,
    heat::SoilHeatModel,
    land::LandModel,
    state⁺::Vars,
    aux⁺::Vars,
    nM,
    state⁻::Vars,
    aux⁻::Vars,
    t,
    _...,
)
    bc_function = bc.state_bc
    param_set = parameter_set(land)
    ϑ_l, θ_i = get_water_content(land.soil.water, aux⁻, state⁻, t)
    θ_l = volumetric_liquid_fraction(ϑ_l, land.soil.param_functions.porosity)
    ρc_s = volumetric_heat_capacity(
        θ_l,
        θ_i,
        land.soil.param_functions.ρc_ds,
        param_set,
    )

    ρe_int_bc =
        volumetric_internal_energy(θ_i, ρc_s, bc_function(aux⁻, t), param_set)

    state⁺.soil.heat.ρe_int = ρe_int_bc
end

# Neumann conditions for SoilHeat and SoilWater

"""
    function soil_boundary_flux!(
        nf,
        bc::Neumann,
        water::SoilWaterModel,
        land::LandModel,
        state⁺::Vars,
        diff⁺::Vars,
        aux⁺::Vars,
        n̂,
        state⁻::Vars,
        diff⁻::Vars,
        aux⁻::Vars,
        t,
    )

The Neumann method for `soil_boundary_flux!` for the
`SoilWaterModel`.
"""
function soil_boundary_flux!(
    nf,
    bc::Neumann,
    water::SoilWaterModel,
    land::LandModel,
    state⁺::Vars,
    diff⁺::Vars,
    aux⁺::Vars,
    n̂,
    state⁻::Vars,
    diff⁻::Vars,
    aux⁻::Vars,
    t,
    _...,
)
    bc_function = bc.scalar_flux_bc
    # on faces at the top and bottom of the domain,
    # the vector n̂ can point in either ± ẑ, depending on which side of the domain
    # we are on. The user supplies a scalar flux F, such that the flux at the boundary
    # is assumed to be F⃗ = F ẑ - i.e. the user doesn't need to worry about the normal vector.
    # so, we take the absolute value of n̂ here.
    # The minus sign is because the condition is applied on minus the flux.
    # the same argument applies for the other directions.
    diff⁺.soil.water.K∇h = abs.(n̂) * (-bc_function(aux⁻, t))
end

"""
    function soil_boundary_flux!(
        nf,
        bc::Neumann,
        heat::SoilHeatModel,
        land::LandModel,
        state⁺::Vars,
        diff⁺::Vars,
        aux⁺::Vars,
        n̂,
        state⁻::Vars,
        diff⁻::Vars,
        aux⁻::Vars,
        t,
    )

The Neumann method for `soil_boundary_flux!` for the
`SoilHeatModel`.
"""
function soil_boundary_flux!(
    nf,
    bc::Neumann,
    heat::SoilHeatModel,
    land::LandModel,
    state⁺::Vars,
    diff⁺::Vars,
    aux⁺::Vars,
    n̂,
    state⁻::Vars,
    diff⁻::Vars,
    aux⁻::Vars,
    t,
    _...,
)
    bc_function = bc.scalar_flux_bc
    # on faces at the top and bottom of the domain,
    # the vector n̂ can point in either ± ẑ, depending on which side of the domain
    # we are on. The user supplies a scalar flux F, such that the flux at the boundary
    # is assumed to be F⃗ = F ẑ - i.e. the user doesn't need to worry about the normal vector.
    # so, we take the absolute value of n̂ here.
    # The minus sign is because the condition is applied on minus the flux.
    # the same argument applies for the other directions.
    diff⁺.soil.heat.κ∇T = abs.(n̂) * (-bc_function(aux⁻, t))
end

# SurfaceDriven conditions for for SoilHeat and SoilWater

"""
    function soil_boundary_flux!(
        nf,
        bc::SurfaceDrivenWaterBoundaryConditions,
        water::SoilWaterModel,
        land::LandModel,
        state⁺::Vars,
        diff⁺::Vars,
        aux⁺::Vars,
        n̂,
        state⁻::Vars,
        diff⁻::Vars,
        aux⁻::Vars,
        t,
    )

The Surface Driven BC method for `soil_boundary_flux!` for the
`SoilWaterModel`.
"""
function soil_boundary_flux!(
    nf,
    bc::SurfaceDrivenWaterBoundaryConditions,
    water::SoilWaterModel,
    land::LandModel,
    state⁺::Vars,
    diff⁺::Vars,
    aux⁺::Vars,
    n̂,
    state⁻::Vars,
    diff⁻::Vars,
    aux⁻::Vars,
    t,
    _...,
)

    diff⁺.soil.water.K∇h = compute_surface_grad_bc(
        land.soil,
        bc.runoff_model,
        bc.precip_model,
        n̂,
        state⁻,
        diff⁻,
        aux⁻,
        t,
    )

end

"""
    soil_boundary_flux!(
        nf,
        bc::SurfaceDrivenHeatBoundaryConditions,
        heat::SoilHeatModel,
        land::LandModel,
        state⁺::Vars,
        diff⁺::Vars,
        aux⁺::Vars,
        n̂,
        state⁻::Vars,
        diff⁻::Vars,
        aux⁻::Vars,
        t,
    )

The Surface Driven BC method for `soil_boundary_flux!` for the
`SoilHeatModel`.
"""
function soil_boundary_flux!(
    nf,
    bc::SurfaceDrivenHeatBoundaryConditions,
    heat::SoilHeatModel,
    land::LandModel,
    state⁺::Vars,
    diff⁺::Vars,
    aux⁺::Vars,
    n̂,
    state⁻::Vars,
    diff⁻::Vars,
    aux⁻::Vars,
    t,
    _...,
)

    net_surface_flux = compute_net_radiative_energy_flux(bc.nswf_model, t)
    diff⁺.soil.heat.κ∇T = n̂ * (-net_surface_flux)
end


function test_simulate()
    d = 2
    lims = repeat([-1.0 1.0], d, 1)
    cdist = UniformContext(lims)
    rdist = GaussianRewards([x -> x[1], x -> -x[1]], 1.0)
    policy = FixedPolicy(1)
    traj = Trajectory(d)
    simulate!(10, traj, cdist, rdist, policy; verbose=false)
    simulate!(10, traj, cdist, rdist, policy; verbose=false)
    @test traj.t == 20
    @test size(traj.X) == (20, d)
    @test traj.a == fill(1, 20)

    # Test matrix extension
    # traj = Trajectory(d)
    # ContextualBandits.add_memory!(traj, 10)
    # X = rand(10, d)
    # a = ones(Int, 10)
    # r = rand(10)
    # total_regret = zeros(10)
    # ContextualBandits.extend!(traj, X, a, r, total_regret)
    # @test traj.X == X
    # @test traj.r == r
    # @test traj.a == a
    # @test traj.total_regret == total_regret
    # @test traj.t == 10
end

test_simulate()