moi_wrapper.jl

# Copyright (c) 2020: Akshay Sharma 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.

"""
    function diff_optimizer(
        optimizer_constructor;
        with_bridge_type = Float64,
        with_cache_type = Float64,
        with_outer_cache = !isnothing(with_bridge_type),
        allow_parametric_opt_interface = true,
    )

Creates a `DiffOpt.Optimizer`, which is an MOI layer with an internal optimizer
and other utility methods. Results (primal, dual and slack values) are obtained
by querying the internal optimizer instantiated using the
`optimizer_constructor`. These values are required for find jacobians with respect to problem data.

The inner optimizer is instantiated with
`MOI.instantiate(optimizer_constructor; with_bridge_type, with_cache_type)`;
see the docs of `MOI.instantiate`.

If `allow_parametric_opt_interface` is `true` and the inner optimizer does not
*natively* (in the sense, without the bridge layer) supports `ParameterSet`, then
a `ParametricOptInterface.Optimizer` layer is added.

If `with_outer_cache` is `true`, an additional layer of cache is added.

One define a differentiable model by using any solver of choice. Example:

```julia
julia> import DiffOpt, HiGHS

julia> model = DiffOpt.diff_optimizer(HiGHS.Optimizer)
julia> set_attribute(model, DiffOpt.ModelConstructor, DiffOpt.QuadraticProgram.Model) # optional selection of diff method
julia> x = model.add_variable(model)
julia> model.add_constraint(model, ...)
```
"""
function diff_optimizer(
    optimizer_constructor;
    with_bridge_type = Float64,
    with_cache_type = Float64,
    with_outer_cache = !isnothing(with_bridge_type),
    allow_parametric_opt_interface = true,
)
    optimizer = MOI.instantiate(
        optimizer_constructor;
        with_bridge_type,
        with_cache_type,
    )
    if allow_parametric_opt_interface &&
       !MOI.supports_add_constrained_variable(
        isnothing(with_bridge_type) ? optimizer : optimizer.model,
        MOI.Parameter{Float64},
    )
        optimizer = POI.Optimizer(optimizer)
    end
    # When we do `MOI.copy_to(diff, optimizer)` we need to efficiently `MOI.get`
    # the model information from `optimizer`. However, 1) `optimizer` may not
    # implement some getters or it may be inefficient and 2) the getters may be
    # unimplemented or inefficient through some bridges.
    # For this reason we add a cache layer, the same cache JuMP adds.
    if with_outer_cache
        optimizer = MOI.Utilities.CachingOptimizer(
            MOI.Utilities.UniversalFallback(
                MOI.Utilities.Model{
                    something(with_bridge_type, with_cache_type),
                }(),
            ),
            optimizer,
        )
    end
    return Optimizer(optimizer)
end

mutable struct Optimizer{OT<:MOI.ModelLike} <: MOI.AbstractOptimizer
    # main optimizer responsible for caching the optimization problem data
    # and for solving the optimization problem
    optimizer::OT

    # list of differentiation backends for automatic differentiation
    # without mode selection
    model_constructors::Vector{Any}
    # used to select a single differentiation backend
    # if not `nothing`, it is used to select the model constructor
    # and the above is ignored
    model_constructor::Any

    # instantiated differentiation backend from the options above
    diff::Any

    # mapping between the `optimizer` and the `diff` models
    index_map::Union{Nothing,MOI.Utilities.IndexMap}

    # sensitivity input cache using MOI-like sparse format
    input_cache::InputCache

    function Optimizer(optimizer::OT) where {OT<:MOI.ModelLike}
        output =
            new{OT}(optimizer, Any[], nothing, nothing, nothing, InputCache())
        add_all_model_constructors(output)
        add_default_factorization(output)
        return output
    end
end

"""
   add_model_constructor(optimizer::Optimizer, model_constructor)

Add the constructor of [`AbstractModel`](@ref) for `optimizer` to choose
from when trying to differentiate.
"""
function add_model_constructor(optimizer::Optimizer, model_constructor)
    push!(optimizer.model_constructors, model_constructor)
    return
end

function MOI.add_variable(model::Optimizer)
    model.diff = nothing
    return MOI.add_variable(model.optimizer)
end

function MOI.add_variables(model::Optimizer, N::Int)
    model.diff = nothing
    return MOI.VariableIndex[MOI.add_variable(model) for i in 1:N]
end

function MOI.add_constraint(
    model::Optimizer,
    f::MOI.AbstractFunction,
    s::MOI.AbstractSet,
)
    model.diff = nothing
    return MOI.add_constraint(model.optimizer, f, s)
end

function MOI.add_constraints(
    model::Optimizer,
    f::AbstractVector{F},
    s::AbstractVector{S},
) where {F<:MOI.AbstractFunction,S<:MOI.AbstractSet}
    model.diff = nothing
    return MOI.ConstraintIndex{F,S}[
        MOI.add_constraint(model, f[i], s[i]) for i in eachindex(f)
    ]
end

function MOI.set(
    model::Optimizer,
    attr::MOI.ObjectiveFunction{F},
    f::F,
) where {F<:MOI.AbstractFunction}
    model.diff = nothing
    return MOI.set(model.optimizer, attr, f)
end

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

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

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

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

function MOI.get(
    model::Optimizer,
    attr::MOI.ConstraintSet,
    ci::MOI.ConstraintIndex,
)
    return MOI.get(model.optimizer, attr, ci)
end

function MOI.set(
    model::Optimizer,
    attr::MOI.ConstraintSet,
    ci::MOI.ConstraintIndex{F,S},
    s::S,
) where {F,S}
    model.diff = nothing
    return MOI.set(model.optimizer, attr, ci, s)
end

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

function MOI.set(
    model::Optimizer,
    attr::MOI.ConstraintFunction,
    ci::MOI.ConstraintIndex{F,S},
    f::F,
) where {F,S}
    model.diff = nothing
    return MOI.set(model.optimizer, attr, ci, f)
end

# `MOI.supports` methods

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

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

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

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

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

function MOI.empty!(model::Optimizer)
    MOI.empty!(model.optimizer)
    model.diff = nothing
    model.index_map = nothing
    empty!(model.input_cache)
    return
end

function MOI.is_empty(model::Optimizer)
    return MOI.is_empty(model.optimizer) && model.diff === nothing
end

function MOI.supports_incremental_interface(model::Optimizer)
    if !MOI.supports_incremental_interface(model.optimizer)
        error(
            "DiffOpt requires a solver that " *
            "`MOI.supports_incremental_interface`, which is not the case for " *
            "$(model.optimizer)",
        )
    end
    return true
end

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

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

function MOI.set(
    model::Optimizer,
    ::MOI.VariablePrimalStart,
    vi::MOI.VariableIndex,
    value::Float64,
)
    model.diff = nothing
    MOI.set(model.optimizer, MOI.VariablePrimalStart(), vi, value)
    return
end

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

function MOI.set(
    model::Optimizer,
    attr::MOI.AbstractVariableAttribute,
    v::MOI.VariableIndex,
    value,
)
    MOI.set(model.optimizer, attr, v, value)
    return
end

function MOI.get(
    model::Optimizer,
    attr::MOI.AbstractVariableAttribute,
    vi::MOI.VariableIndex,
)
    return MOI.get(model.optimizer, attr, vi)
end

function MOI.delete(model::Optimizer, ci::MOI.ConstraintIndex{F,S}) where {F,S}
    model.diff = nothing
    MOI.delete(model.optimizer, ci)
    return
end

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

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

function MOI.is_valid(model::Optimizer, con::MOI.ConstraintIndex)
    return MOI.is_valid(model.optimizer, con)
end

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

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

# helper methods to check if a constraint contains a Variable
function _constraint_contains(
    model::Optimizer,
    v::MOI.VariableIndex,
    ci::MOI.ConstraintIndex{MOI.VariableIndex},
)
    return v == MOI.get(model, MOI.ConstraintFunction(), ci)
end

function _constraint_contains(
    model::Optimizer,
    v::MOI.VariableIndex,
    ci::MOI.ConstraintIndex{MOI.ScalarAffineFunction{Float64}},
)
    func = MOI.get(model, MOI.ConstraintFunction(), ci)
    return any(term -> v == term.variable, func.terms)
end

function _constraint_contains(
    model::Optimizer,
    v::MOI.VariableIndex,
    ci::MOI.ConstraintIndex{MOI.VectorOfVariables},
)
    func = MOI.get(model, MOI.ConstraintFunction(), ci)
    return v in func.variables
end

function _constraint_contains(
    model::Optimizer,
    v::MOI.VariableIndex,
    ci::MOI.ConstraintIndex{MOI.VectorAffineFunction{Float64}},
)
    func = MOI.get(model, MOI.ConstraintFunction(), ci)
    return any(term -> v == term.scalar_term.variable, func.terms)
end

function MOI.delete(model::Optimizer, v::MOI.VariableIndex)
    model.diff = nothing
    MOI.delete(model.optimizer, v)
    return
end

# for array deletion
function MOI.delete(model::Optimizer, indices::Vector{MOI.VariableIndex})
    model.diff = nothing
    for i in indices
        MOI.delete(model, i)
    end
    return
end

function MOI.modify(
    model::Optimizer,
    ::MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}},
    chg::MOI.AbstractFunctionModification,
)
    model.diff = nothing
    MOI.modify(
        model.optimizer,
        MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(),
        chg,
    )
    return
end

function MOI.modify(
    model::Optimizer,
    ci::MOI.ConstraintIndex,
    chg::MOI.AbstractFunctionModification,
)
    model.diff = nothing
    MOI.modify(model.optimizer, ci, chg)
    return
end

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

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

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

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

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

function MOI.set(model::Optimizer, ::MOI.Name, name::String)
    MOI.set(model.optimizer, MOI.Name(), name)
    return
end

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

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

function MOI.set(
    model::Optimizer,
    ::MOI.TimeLimitSec,
    value::Union{Real,Nothing},
)
    MOI.set(model.optimizer, MOI.TimeLimitSec(), value)
    return
end

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

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

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

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

function MOI.get(
    model::Optimizer,
    attr::A,
) where {A<:Union{MOI.SolverName,MOI.SolverVersion}}
    return MOI.get(model.optimizer, attr)
end

function MOI.optimize!(model::Optimizer)
    model.diff = nothing
    MOI.optimize!(model.optimizer)
    return
end

"""
    ModelConstructor <: MOI.AbstractOptimizerAttribute

Determines which subtype of [`DiffOpt.AbstractModel`](@ref) to use for
differentiation. When set to `nothing`, the first one out of
`model.model_constructors` that support the problem is used.

Examples:

```julia
julia> MOI.set(model, DiffOpt.ModelConstructor(), DiffOpt.QuadraticProgram.Model)

julia> MOI.set(model, DiffOpt.ModelConstructor(), DiffOpt.ConicProgram.Model)

julia> MOI.set(model, DiffOpt.ModelConstructor(), DiffOpt.NonlinearProgram.Model)
```
"""
struct ModelConstructor <: MOI.AbstractOptimizerAttribute end

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

function MOI.set(model::Optimizer, ::ModelConstructor, model_constructor)
    model.diff = nothing
    model.model_constructor = model_constructor
    return
end

MOI.get(model::Optimizer, ::ModelConstructor) = model.model_constructor

function reverse_differentiate!(model::Optimizer)
    st = MOI.get(model.optimizer, MOI.TerminationStatus())
    if !in(st, (MOI.LOCALLY_SOLVED, MOI.OPTIMAL))
        error(
            "Trying to compute the reverse differentiation on a model with termination status $(st)",
        )
    end
    if !iszero(model.input_cache.dobj) &&
       (!isempty(model.input_cache.dx) || !isempty(model.input_cache.dy))
        if !MOI.get(model, AllowObjectiveAndSolutionInput())
            @warn "Computing reverse differentiation with both solution sensitivities and objective sensitivities. " *
                  "Set `DiffOpt.AllowObjectiveAndSolutionInput()` to `true` to silence this warning."
        end
    end
    diff = _diff(model)
    MOI.set(
        diff,
        NonLinearKKTJacobianFactorization(),
        model.input_cache.factorization,
    )
    MOI.set(
        diff,
        AllowObjectiveAndSolutionInput(),
        model.input_cache.allow_objective_and_solution_input,
    )
    for (vi, value) in model.input_cache.dx
        MOI.set(diff, ReverseVariablePrimal(), model.index_map[vi], value)
    end
    for (vi, value) in model.input_cache.dy
        MOI.set(diff, ReverseConstraintDual(), model.index_map[vi], value)
    end
    if !iszero(model.input_cache.dobj)
        try
            MOI.set(diff, ReverseObjectiveSensitivity(), model.input_cache.dobj)
        catch e
            if e isa MOI.UnsupportedAttribute
                _fallback_set_reverse_objective_sensitivity(
                    model,
                    model.input_cache.dobj,
                )
            else
                rethrow(e)
            end
        end
    end
    return reverse_differentiate!(diff)
end

# Gradient evaluation functions for objective sensitivity fallbacks
function _eval_gradient(::Optimizer, ::Number)
    return Dict{MOI.VariableIndex,Float64}()
end

function _eval_gradient(::Optimizer, f::MOI.VariableIndex)
    return Dict{MOI.VariableIndex,Float64}(f => 1.0)
end

function _eval_gradient(::Optimizer, f::MOI.ScalarAffineFunction{Float64})
    grad = Dict{MOI.VariableIndex,Float64}()
    for term in f.terms
        grad[term.variable] = get(grad, term.variable, 0.0) + term.coefficient
    end
    return grad
end

function _eval_gradient(
    model::Optimizer,
    f::MOI.ScalarQuadraticFunction{Float64},
)
    grad = Dict{MOI.VariableIndex,Float64}()
    for term in f.affine_terms
        grad[term.variable] = get(grad, term.variable, 0.0) + term.coefficient
    end
    # MOI convention: function is 0.5 * x' * Q * x, so derivative of diagonal
    # term 0.5 * coef * xi^2 is coef * xi (not 2 * coef * xi)
    for term in f.quadratic_terms
        xi, xj = term.variable_1, term.variable_2
        coef = term.coefficient
        xi_val = MOI.get(model, MOI.VariablePrimal(), xi)
        xj_val = MOI.get(model, MOI.VariablePrimal(), xj)
        if xi == xj
            grad[xi] = get(grad, xi, 0.0) + coef * xi_val
        else
            grad[xi] = get(grad, xi, 0.0) + coef * xj_val
            grad[xj] = get(grad, xj, 0.0) + coef * xi_val
        end
    end
    return grad
end

function _fallback_set_reverse_objective_sensitivity(model::Optimizer, val)
    diff = _diff(model)
    obj_type = MOI.get(model, MOI.ObjectiveFunctionType())
    obj_func = MOI.get(model, MOI.ObjectiveFunction{obj_type}())
    grad = _eval_gradient(model, obj_func)
    for (xi, df_dxi) in grad
        MOI.set(
            diff,
            ReverseVariablePrimal(),
            model.index_map[xi],
            df_dxi * val,
        )
    end
    return
end

function _copy_forward_in_constraint(diff, index_map, con_map, constraints)
    for (index, value) in constraints
        MOI.set(
            diff,
            ForwardConstraintFunction(),
            con_map[index],
            MOI.Utilities.map_indices(index_map, value),
        )
    end
    return
end

function forward_differentiate!(model::Optimizer)
    st = MOI.get(model.optimizer, MOI.TerminationStatus())
    if !in(st, (MOI.LOCALLY_SOLVED, MOI.OPTIMAL))
        error(
            "Trying to compute the forward differentiation on a model with termination status $(st)",
        )
    end
    diff = _diff(model)
    MOI.set(
        diff,
        NonLinearKKTJacobianFactorization(),
        model.input_cache.factorization,
    )
    MOI.set(
        diff,
        AllowObjectiveAndSolutionInput(),
        model.input_cache.allow_objective_and_solution_input,
    )
    T = Float64
    list = MOI.get(
        model,
        MOI.ListOfConstraintIndices{MOI.VariableIndex,MOI.Parameter{T}}(),
    )
    parametric_diff = !isempty(list)
    if parametric_diff # MOI.supports_constraint(model, MOI.VariableIndex, MOI.Parameter{T})
        # @show "param mode"
        for (vi, value) in model.input_cache.parameter_constraints
            MOI.set(
                diff,
                ForwardConstraintSet(),
                model.index_map[vi],
                MOI.Parameter(value),
            )
        end
        return forward_differentiate!(diff)
    end
    # @show "func mode"
    if model.input_cache.objective !== nothing
        MOI.set(
            diff,
            ForwardObjectiveFunction(),
            MOI.Utilities.map_indices(
                model.index_map,
                model.input_cache.objective,
            ),
        )
    end
    for (F, S) in MOI.Utilities.DoubleDicts.nonempty_outer_keys(
        model.input_cache.scalar_constraints,
    )
        _copy_forward_in_constraint(
            diff,
            model.index_map,
            model.index_map.con_map[F, S],
            model.input_cache.scalar_constraints[F, S],
        )
    end
    for (F, S) in MOI.Utilities.DoubleDicts.nonempty_outer_keys(
        model.input_cache.vector_constraints,
    )
        _copy_forward_in_constraint(
            diff,
            model.index_map,
            model.index_map.con_map[F, S],
            model.input_cache.vector_constraints[F, S],
        )
    end
    return forward_differentiate!(diff)
end

function empty_input_sensitivities!(model::Optimizer)
    empty!(model.input_cache)
    if model.diff !== nothing
        empty_input_sensitivities!(model.diff)
    end
    return
end

function _add_bridges(instantiated_model)
    model = MOI.Bridges.LazyBridgeOptimizer(instantiated_model)
    # We don't add any variable bridge here because:
    # 1) If `ZerosBridge` is used, `MOI.Bridges.unbridged_function` does not work.
    #    This is in fact expected: since `ZerosBridge` drops the variable, we dont
    #    compute the derivative of the value of this variable as a function of its fixed value.
    #    This could be easily determined as the same as the derivative of the value but
    #    since the variable was also dropped from other constraints, we would ignore its impact on the other constraints.
    # 2) For affine variable bridges, `bridged_function` and `unbridged_function` don't treat the function as a derivative hence they will add constants
    MOI.Bridges.Constraint.add_all_bridges(model, Float64)
    MOI.Bridges.Objective.add_all_bridges(model, Float64)
    return model
end

function _instantiate_diff(model::Optimizer, constructor)
    # parametric_diff = MOI.supports_constraint(
    #     model,
    #     MOI.VariableIndex,
    #     MOI.Parameter{Float64},
    # )
    list = MOI.get(
        model,
        MOI.ListOfConstraintIndices{MOI.VariableIndex,MOI.Parameter{Float64}}(),
    )
    parametric_diff = !isempty(list)
    model_instance = MOI.instantiate(constructor)
    needs_poi =
        !MOI.supports_add_constrained_variable(
            model_instance,
            MOI.Parameter{Float64},
        )
    model_bridged = _add_bridges(model_instance)
    if needs_poi && parametric_diff
        return POI.Optimizer(model_bridged)
    end
    return model_bridged
end

function _diff(model::Optimizer)
    if model.diff === nothing
        _check_termination_status(model)
        model_constructor = MOI.get(model, ModelConstructor())
        if isnothing(model_constructor)
            for constructor in model.model_constructors
                model.diff = _instantiate_diff(model, constructor)
                try
                    model.index_map = MOI.copy_to(model.diff, model.optimizer)
                catch err
                    if err isa MOI.UnsupportedConstraint ||
                       err isa MOI.UnsupportedAttribute
                        model.diff = nothing
                    else
                        rethrow(err)
                    end
                end
                if !isnothing(model.diff)
                    break
                end
            end
            if isnothing(model.diff)
                error(
                    "No differentiation model supports the problem. If you " *
                    "believe it should be supported, say by " *
                    "`DiffOpt.QuadraticProgram.Model`, use " *
                    "`MOI.set(model, DiffOpt.ModelConstructor, DiffOpt.QuadraticProgram.Model)`" *
                    "and try again to see an error indicating why it is not supported.",
                )
            end
        else
            model.diff = _instantiate_diff(model, model_constructor)
            model.index_map = MOI.copy_to(model.diff, model.optimizer)
        end
        _copy_dual(model.diff, model.optimizer, model.index_map)
    end
    return model.diff
end

function _check_termination_status(model::Optimizer)
    if !in(
        MOI.get(model, MOI.TerminationStatus()),
        (MOI.LOCALLY_SOLVED, MOI.OPTIMAL),
    )
        error(
            "problem status: ",
            MOI.get(model.optimizer, MOI.TerminationStatus()),
        )
    end
    return
end

# DiffOpt attributes redirected to `diff`

function _checked_diff(model::Optimizer, attr::MOI.AnyAttribute, call)
    if model.diff === nothing
        error("Cannot get attribute `$attr`. First call `DiffOpt.$call`.")
    end
    return model.diff
end

function MOI.get(model::Optimizer, attr::ReverseObjectiveFunction)
    return IndexMappedFunction(
        MOI.get(_checked_diff(model, attr, :reverse_differentiate!), attr),
        model.index_map,
    )
end

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

function MOI.set(model::Optimizer, ::ForwardObjectiveFunction, objective)
    T = Float64
    list = MOI.get(
        model,
        MOI.ListOfConstraintIndices{MOI.VariableIndex,MOI.Parameter{T}}(),
    )
    parametric_diff = !isempty(list)
    if parametric_diff
        error(
            "Cannot set forward objective function for a model with parameters. " *
            "Use `MOI.set(model, ForwardConstraintSet(), ParameterRef(vi), Parameter(val))` instead.",
        )
    end
    model.input_cache.objective = objective
    return
end

function MOI.get(model::Optimizer, ::ForwardObjectiveFunction)
    return model.input_cache.objective
end

function MOI.get(
    model::Optimizer,
    attr::ForwardVariablePrimal,
    vi::MOI.VariableIndex,
)
    return MOI.get(
        _checked_diff(model, attr, :forward_differentiate!),
        attr,
        model.index_map[vi],
    )
end

function MOI.get(
    model::Optimizer,
    attr::ReverseConstraintSet,
    ci::MOI.ConstraintIndex,
)
    return MOI.get(
        _checked_diff(model, attr, :reverse_differentiate!),
        attr,
        model.index_map[ci],
    )
end

function MOI.get(
    model::Optimizer,
    attr::ForwardConstraintDual,
    ci::MOI.ConstraintIndex,
)
    return MOI.get(
        _checked_diff(model, attr, :reverse_differentiate!),
        attr,
        model.index_map[ci],
    )
end

function MOI.get(model::Optimizer, attr::ForwardObjectiveSensitivity)
    diff_model = _checked_diff(model, attr, :forward_differentiate!)
    val = 0.0
    try
        val = MOI.get(diff_model, attr)
    catch e
        if e isa MOI.UnsupportedAttribute
            val = _fallback_get_forward_objective_sensitivity(model)
        else
            rethrow(e)
        end
    end
    return val
end

function _fallback_get_forward_objective_sensitivity(model::Optimizer)
    obj_type = MOI.get(model, MOI.ObjectiveFunctionType())
    obj_func = MOI.get(model, MOI.ObjectiveFunction{obj_type}())
    grad = _eval_gradient(model, obj_func)
    ret = 0.0
    for (xi, df_dxi) in grad
        dx_dp = MOI.get(model, ForwardVariablePrimal(), xi)
        ret += df_dxi * dx_dp
    end
    return ret
end

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

function MOI.set(
    model::Optimizer,
    ::ReverseVariablePrimal,
    vi::MOI.VariableIndex,
    val,
)
    T = Float64
    is_param = MOI.is_valid(
        model,
        MOI.ConstraintIndex{MOI.VariableIndex,MOI.Parameter{T}}(vi.value),
    )
    if is_param
        error("Trying to set a backward variable sensitivity for a parameter")
    end
    model.input_cache.dx[vi] = val
    return
end

function MOI.get(
    model::Optimizer,
    ::ReverseVariablePrimal,
    vi::MOI.VariableIndex,
)
    return get(model.input_cache.dx, vi, 0.0)
end

function MOI.supports(
    ::Optimizer,
    ::ReverseConstraintDual,
    ::Type{MOI.ConstraintIndex},
)
    return true
end

function MOI.set(
    model::Optimizer,
    ::ReverseConstraintDual,
    ci::MOI.ConstraintIndex,
    val,
)
    model.input_cache.dy[ci] = val
    return
end

function MOI.set(model::Optimizer, ::ReverseObjectiveSensitivity, val)
    model.input_cache.dobj = val
    return
end

function MOI.get(
    model::Optimizer,
    ::ReverseConstraintDual,
    ci::MOI.ConstraintIndex,
    val,
)
    return return get(model.input_cache.dy, ci, 0.0)
end

function MOI.get(
    model::Optimizer,
    attr::ReverseConstraintFunction,
    ci::MOI.ConstraintIndex,
)
    return IndexMappedFunction(
        MOI.get(
            _checked_diff(model, attr, :reverse_differentiate!),
            attr,
            model.index_map[ci],
        ),
        model.index_map,
    )
end

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

function MOI.set(
    model::Optimizer,
    ::ForwardConstraintFunction,
    ci::MOI.ConstraintIndex{MOI.ScalarAffineFunction{T},S},
    func::MOI.ScalarAffineFunction{T},
) where {T,S}
    list = MOI.get(
        model,
        MOI.ListOfConstraintIndices{MOI.VariableIndex,MOI.Parameter{T}}(),
    )
    parametric_diff = !isempty(list)
    if parametric_diff
        error(
            "Cannot set forward constraint function for a model with parameters. " *
            "Use `MOI.set(model, ForwardConstraintSet(), ParameterRef(vi), Parameter(val))` instead.",
        )
    end
    model.input_cache.scalar_constraints[ci] = func
    return
end

function MOI.set(
    model::Optimizer,
    ::ForwardConstraintFunction,
    ci::MOI.ConstraintIndex{MOI.VectorAffineFunction{T},S},
    func::MOI.VectorAffineFunction{T},
) where {T,S}
    list = MOI.get(
        model,
        MOI.ListOfConstraintIndices{MOI.VariableIndex,MOI.Parameter{T}}(),
    )
    parametric_diff = !isempty(list)
    if parametric_diff
        error(
            "Cannot set forward constraint function for a model with parameters. " *
            "Use `MOI.set(model, ForwardConstraintSet(), ParameterRef(vi), Parameter(val))` instead.",
        )
    end
    model.input_cache.vector_constraints[ci] = func
    return
end

function MOI.get(
    model::Optimizer,
    ::ForwardConstraintFunction,
    ci::MOI.ConstraintIndex{MOI.ScalarAffineFunction{T},S},
) where {T,S}
    return get(
        model.input_cache.scalar_constraints,
        ci,
        zero(MOI.ScalarAffineFunction{T}),
    )
end

function MOI.get(
    model::Optimizer,
    ::ForwardConstraintFunction,
    ci::MOI.ConstraintIndex{MOI.VectorAffineFunction{T},S},
) where {T,S}
    func = get(model.input_cache.vector_constraints, ci, nothing)
    if func === nothing
        set = MOI.get(model, MOI.ConstraintSet(), ci)
        dim = MOI.dimension(set)
        return MOI.Utilities.zero_with_output_dimension(
            MOI.VectorAffineFunction{T},
            dim,
        )
    else
        return func
    end
end

function MOI.supports(
    ::Optimizer,
    ::ForwardConstraintSet,
    ::Type{MOI.ConstraintIndex{MOI.VariableIndex,MOI.Parameter{T}}},
) where {T}
    return true
end

function MOI.get(
    model::Optimizer,
    ::ForwardConstraintSet,
    ci::MOI.ConstraintIndex{MOI.VariableIndex,MOI.Parameter{T}},
) where {T}
    set = get(model.input_cache.parameter_constraints, ci, zero(T))
    return MOI.Parameter{T}(set)
end

function MOI.set(
    model::Optimizer,
    ::ForwardConstraintSet,
    ci::MOI.ConstraintIndex{MOI.VariableIndex,MOI.Parameter{T}},
    set::MOI.Parameter,
) where {T}
    model.input_cache.parameter_constraints[ci] = set.value
    return
end

function MOI.get(model::Optimizer, attr::DifferentiateTimeSec)
    return MOI.get(model.diff, attr)
end

function MOI.supports(
    ::Optimizer,
    ::NonLinearKKTJacobianFactorization,
    ::Function,
)
    return true
end

function MOI.supports(::Optimizer, ::AllowObjectiveAndSolutionInput, ::Bool)
    return true
end

function MOI.set(
    model::Optimizer,
    ::NonLinearKKTJacobianFactorization,
    factorization,
)
    model.input_cache.factorization = factorization
    return
end

function MOI.set(model::Optimizer, ::AllowObjectiveAndSolutionInput, allow)
    model.input_cache.allow_objective_and_solution_input = allow
    return
end

function MOI.get(model::Optimizer, ::NonLinearKKTJacobianFactorization)
    return model.input_cache.factorization
end

function MOI.get(model::Optimizer, ::AllowObjectiveAndSolutionInput)
    return model.input_cache.allow_objective_and_solution_input
end

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