# Constructs a Differentiable Optimizer model from a MOI Optimizer.
# Supports `forward_differentiate!` and `reverse_differentiate!` methods for solving and differentiating the model respectectively.

# ## Note
# Currently supports differentiating linear and quadratic programs only.

const MOIDD = MOI.Utilities.DoubleDicts

Base.@kwdef mutable struct InputCache
    dx::Dict{VI, Float64} = Dict{VI, Float64}()# dz for QP
    # ds
    # dy #= [d\lambda, d\nu] for QP
    # FIXME Would it be possible to have a DoubleDict where the value depends
    #       on the function type ? Here, we need to create two dicts to have
    #       concrete value types.
    # `scalar_constraints` and `vector_constraints` includes `A` and `b` for CPs
    # or `G` and `h` for QPs
    scalar_constraints::MOIDD.DoubleDict{MOI.ScalarAffineFunction{Float64}} = MOIDD.DoubleDict{MOI.ScalarAffineFunction{Float64}}() # also includes G for QPs
    vector_constraints::MOIDD.DoubleDict{MOI.VectorAffineFunction{Float64}} = MOIDD.DoubleDict{MOI.VectorAffineFunction{Float64}}() # also includes G for QPs
    objective::Union{Nothing,MOI.AbstractScalarFunction} = nothing
end

function Base.empty!(cache::InputCache)
    empty!(cache.dx)
    empty!(cache.scalar_constraints)
    empty!(cache.vector_constraints)
    cache.objective = nothing
    return
end

"""
    reverse_differentiate!(model::MOI.ModelLike)

Wrapper method for the backward pass / reverse differentiation.
This method will consider as input a currently solved problem and differentials
with respect to the solution set with the [`ReverseVariablePrimal`](@ref) attribute.
The output problem data differentials can be queried with the
attributes [`ReverseObjectiveFunction`](@ref) and [`ReverseConstraintFunction`](@ref).
"""
function reverse_differentiate! end

"""
    forward_differentiate!(model::Optimizer)

Wrapper method for the forward pass.
This method will consider as input a currently solved problem and
differentials with respect to problem data set with
the [`ForwardObjectiveFunction`](@ref) and  [`ForwardConstraintFunction`](@ref) attributes.
The output solution differentials can be queried with the attribute
[`ForwardVariablePrimal`](@ref).
"""
function forward_differentiate! end

"""
    ForwardObjectiveFunction <: MOI.AbstractModelAttribute

A `MOI.AbstractModelAttribute` to set input data to forward differentiation, that
is, problem input data.
The possible values are any `MOI.AbstractScalarFunction`.
A `MOI.ScalarQuadraticFunction` can only be used in linearly constrained
quadratic models.

For instance, if the objective contains `θ * (x + 2y)`, for the purpose of
computing the derivative with respect to `θ`, the following should be set:
```julia
MOI.set(model, DiffOpt.ForwardObjectiveFunction(), 1.0 * x + 2.0 * y)
```
where `x` and `y` are the relevant `MOI.VariableIndex`.
"""
struct ForwardObjectiveFunction <: MOI.AbstractModelAttribute end

"""
    ForwardConstraintFunction <: MOI.AbstractConstraintAttribute

A `MOI.AbstractConstraintAttribute` to set input data to forward differentiation, that
is, problem input data.

For instance, if the scalar constraint of index `ci` contains `θ * (x + 2y) <= 5θ`,
for the purpose of computing the derivative with respect to `θ`, the following
should be set:
```julia
MOI.set(model, DiffOpt.ForwardConstraintFunction(), ci, 1.0 * x + 2.0 * y - 5.0)
```
Note that we use `-5` as the `ForwardConstraintFunction` sets the tangent of the
ConstraintFunction so we consider the expression `θ * (x + 2y - 5)`.
"""
struct ForwardConstraintFunction <: MOI.AbstractConstraintAttribute end


"""
    ForwardVariablePrimal <: MOI.AbstractVariableAttribute

A `MOI.AbstractVariableAttribute` to get output data from forward
differentiation, that is, problem solution.

For instance, to get the tangent of the variable of index `vi` corresponding to
the tangents given to `ForwardObjectiveFunction` and `ForwardConstraintFunction`, do the
following:
```julia
MOI.get(model, DiffOpt.ForwardVariablePrimal(), vi)
```
"""
struct ForwardVariablePrimal <: MOI.AbstractVariableAttribute end
MOI.is_set_by_optimize(::ForwardVariablePrimal) = true

"""
    ReverseVariablePrimal <: MOI.AbstractVariableAttribute

A `MOI.AbstractVariableAttribute` to set input data to
reverse differentiation, that is, problem solution.

For instance, to set the tangent of the variable of index `vi`, do the
following:
```julia
MOI.set(model, DiffOpt.ReverseVariablePrimal(), x)
```
"""
struct ReverseVariablePrimal <: MOI.AbstractVariableAttribute end

"""
    ReverseObjectiveFunction <: MOI.AbstractModelAttribute

A `MOI.AbstractModelAttribute` to get output data to reverse differentiation,
that is, problem input data.

For instance, to get the tangent of the objective function corresponding to
the tangent given to `ReverseVariablePrimal`, do the
following:
```julia
func = MOI.get(model, DiffOpt.ReverseObjectiveFunction())
```
Then, to get the sensitivity of the linear term with variable `x`, do
```julia
JuMP.coefficient(func, x)
```
To get the sensitivity with respect to the quadratic term with variables `x`
and `y`, do either
```julia
JuMP.coefficient(func, x, y)
```
or
```julia
DiffOpt.quad_sym_half(func, x, y)
```

!!! warning
    These two lines are **not** equivalent in case `x == y`, see
    [`quad_sym_half`](@ref) for the details on the difference between these two
    functions.
"""
struct ReverseObjectiveFunction <: MOI.AbstractModelAttribute end
MOI.is_set_by_optimize(::ReverseObjectiveFunction) = true

"""
    ReverseConstraintFunction

An `MOI.AbstractConstraintAttribute` to get output data to reverse differentiation, that
is, problem input data.

For instance, if the following returns `x + 2y + 5`, it means that the tangent
has coordinate `1` for the coefficient of `x`, coordinate `2` for the
coefficient of `y` and `5` for the function constant.
If the constraint is of the form `func == constant` or `func <= constant`,
the tangent for the constant on the right-hand side is `-5`.
```julia
MOI.get(model, DiffOpt.ReverseConstraintFunction(), ci)
```
"""
struct ReverseConstraintFunction <: MOI.AbstractConstraintAttribute end
MOI.is_set_by_optimize(::ReverseConstraintFunction) = true

"""
    abstract type AbstractModel <: MOI.ModelLike end

Model supporting [`forward_differentiate!`](@ref) and
[`reverse_differentiate!`](@ref).
"""
abstract type AbstractModel <: MOI.ModelLike end

MOI.supports_incremental_interface(::AbstractModel) = true

MOI.is_valid(model::AbstractModel, idx::MOI.Index) = MOI.is_valid(model.model, idx)

function MOI.add_variable(model::AbstractModel)
    return MOI.add_variable(model.model)
end

function MOI.add_variables(model::AbstractModel, n)
    return MOI.add_variables(model.model, n)
end

function MOI.Utilities.pass_nonvariable_constraints(
    dest::AbstractModel,
    src::MOI.ModelLike,
    idxmap::MOIU.IndexMap,
    constraint_types,
)
    MOI.Utilities.pass_nonvariable_constraints(dest.model, src, idxmap, constraint_types)
end

function MOI.Utilities.final_touch(model::AbstractModel, index_map)
    MOI.Utilities.final_touch(model.model, index_map)
end

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

function MOI.add_constraint(model::AbstractModel, func::MOI.AbstractFunction, set::MOI.AbstractSet)
    return MOI.add_constraint(model.model, func, set)
end

function _enlarge_set(vec::Vector, idx, value)
    m = last(idx)
    if length(vec) < m
        n = length(vec)
        resize!(vec, m)
        fill!(view(vec, (n+1):m), NaN)
    end
    vec[idx] = value
    return
end

function MOI.set(model::AbstractModel, ::MOI.VariablePrimalStart, vi::MOI.VariableIndex, value)
    MOI.throw_if_not_valid(model, vi)
    _enlarge_set(model.x, vi.value, value)
end

function MOI.set(model::AbstractModel, ::ForwardObjectiveFunction, objective)
    model.input_cache.objective = objective
    return
end
function MOI.set(model::AbstractModel, ::ReverseVariablePrimal, vi::VI, val)
    model.input_cache.dx[vi] = val
    return
end
function MOI.set(model::AbstractModel,
    ::ForwardConstraintFunction,
    ci::CI{MOI.ScalarAffineFunction{T},S},
    func::MOI.ScalarAffineFunction{T},
) where {T,S}
    model.input_cache.scalar_constraints[ci] = func
    return
end
function MOI.set(model::AbstractModel,
    ::ForwardConstraintFunction,
    ci::CI{MOI.VectorAffineFunction{T},S},
    func::MOI.VectorAffineFunction{T},
) where {T,S}
    model.input_cache.vector_constraints[ci] = func
    return
end

function lazy_combination(op::F, α, a, β, b) where {F<:Function}
    return LazyArrays.ApplyArray(
        op,
        LazyArrays.@~(α .* a),
        LazyArrays.@~(β .* b),
    )
end

function lazy_combination(op::F, α, a, β, b, I::UnitRange) where {F<:Function}
    return lazy_combination(op, α, view(a, I), β, view(b, I))
end
function lazy_combination(op::F, a, b, i::Integer, args::Vararg{Any,N}) where {F<:Function,N}
    return lazy_combination(op, a[i], b, b[i], a, args...)
end
function lazy_combination(op::F, a, b, i::UnitRange, I::UnitRange) where {F<:Function}
    return lazy_combination(op, view(a, i), b, view(b, i), a, I)
end

_lazy_affine(vector, constant::Number) = VectorScalarAffineFunction(vector, constant)
_lazy_affine(matrix, vector) = MatrixVectorAffineFunction(matrix, vector)
function _get_db end
function _get_dA end
function MOI.get(model::AbstractModel, ::ReverseConstraintFunction, ci::CI)
    return _lazy_affine(_get_dA(model, ci), _get_db(model, ci))
end

"""
    π(v::Vector{Float64}, model::MOI.ModelLike, cones::ProductOfSets)

Given a `model`, its `cones`, find the projection of the vectors `v`
of length equal to the number of rows in the conic form onto the cartesian
product of the cones corresponding to these rows.
For more info, refer to https://github.com/matbesancon/MathOptSetDistances.jl
"""
function π(v::Vector{T}, model::MOI.ModelLike, cones::ProductOfSets) where T
    return map_rows(model, cones, Flattened{T}()) do ci, r
        MOSD.projection_on_set(
            MOSD.DefaultDistance(),
            v[r],
            MOI.dual_set(MOI.get(model, MOI.ConstraintSet(), ci))
        )
    end
end


"""
    Dπ(v::Vector{Float64}, model, cones::ProductOfSets)

Given a `model`, its `cones`, find the gradient of the projection of
the vectors `v` of length equal to the number of rows in the conic form onto the
cartesian product of the cones corresponding to these rows.
For more info, refer to https://github.com/matbesancon/MathOptSetDistances.jl
"""
function Dπ(v::Vector{T}, model::MOI.ModelLike, cones::ProductOfSets) where T
    return BlockDiagonals.BlockDiagonal(
        map_rows(model, cones, Nested{Matrix{T}}()) do ci, r
            MOSD.projection_gradient_on_set(
                MOSD.DefaultDistance(),
                v[r],
                MOI.dual_set(MOI.get(model, MOI.ConstraintSet(), ci)),
            )
        end
    )
end

# See the docstring of `map_rows`.
struct Nested{T} end
struct Flattened{T} end

# Store in `x` the values `y` corresponding to the rows `r` and the `k`th
# constraint.
function _assign_mapped!(x, y, r, k, ::Nested)
    x[k] = y
end
function _assign_mapped!(x, y, r, k, ::Flattened)
    x[r] = y
end

# Map the rows corresponding to `F`-in-`S` constraints and store it in `x`.
function _map_rows!(f::Function, x::Vector, model, cones::ProductOfSets, ::Type{F}, ::Type{S}, map_mode, k) where {F, S}
    for ci in MOI.get(model, MOI.ListOfConstraintIndices{F, S}())
        r = MOI.Utilities.rows(cones, ci)
        k += 1
        _assign_mapped!(x, f(ci, r), r, k, map_mode)
    end
    return k
end

# Allocate a vector for storing the output of `map_rows`.
_allocate_rows(cones, ::Nested{T}) where {T} = Vector{T}(undef, length(cones.dimension))
_allocate_rows(cones, ::Flattened{T}) where {T} = Vector{T}(undef, MOI.dimension(cones))

"""
    map_rows(f::Function, model, cones::ProductOfSets, map_mode::Union{Nested{T}, Flattened{T}})

Given a `model`, its `cones` and `map_mode` of type `Nested` (resp.
`Flattened`), return a `Vector{T}` of length equal to the number of cones (resp.
rows) in the conic form where the value for the index (resp. rows) corresponding
to each cone is equal to `f(ci, r)` where `ci` is the corresponding constraint
index in `model` and `r` is a `UnitRange` of the corresponding rows in the conic
form.
"""
function map_rows(f::Function, model, cones::ProductOfSets, map_mode::Union{Nested, Flattened})
    x = _allocate_rows(cones, map_mode)
    k = 0
    for (F, S) in MOI.get(model, MOI.ListOfConstraintTypesPresent())
        # Function barrier for type unstability of `F` and `S`
        # `con_map` is a `MOIU.DoubleDicts.MainIndexDoubleDict`, we index it at `F, S`
        # which returns a `MOIU.DoubleDicts.IndexWithType{F, S}` which is type stable.
        # If we have a small number of different constraint types and many
        # constraint of each type, this mostly removes type unstabilities
        # as most the time is in `_map_rows!` which is type stable.
        k = _map_rows!(f, x, model, cones, F, S, map_mode, k)
    end
    return x
end

function _fill(neg::Function, gradient_cache, input_cache, cones, args...)
    _fill(S -> true, neg, gradient_cache, input_cache, cones, args...)
end
function _fill(filter::Function, neg::Function, gradient_cache, input_cache, cones, args...)
    for (F, S) in keys(input_cache.scalar_constraints.dict)
        filter(S) || continue
        _fill(args..., neg(S), cones, input_cache.scalar_constraints[F,S])
    end
    for (F, S) in keys(input_cache.vector_constraints.dict)
        _fill(args..., neg(S), cones, input_cache.vector_constraints[F,S])
    end
    return
end

function _fill(vector::Vector, neg::Bool, cones, dict)
    for (ci, func) in dict
        r = MOI.Utilities.rows(cones, ci)
        vector[r] = neg ? -MOI.constant(func) : MOI.constant(func)
    end
end
function _fill(I::Vector, J::Vector, V::Vector, neg::Bool, cones, dict)
    for (ci, func) in dict
        r = MOI.Utilities.rows(cones, ci)
        for term in func.terms
            _push_term(I, J, V, neg, r, term)
        end
    end
end
function _push_term(I::Vector, J::Vector, V::Vector, neg::Bool, r::Integer, term::MOI.ScalarAffineTerm)
    push!(I, r)
    push!(J, term.variable.value)
    push!(V, neg ? -term.coefficient : term.coefficient)
end
function _push_term(I::Vector, J::Vector, V::Vector, neg::Bool, r::UnitRange, term::MOI.VectorAffineTerm)
    _push_term(I, J, V, neg, r[term.output_index], term.scalar_term)
end


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

function MOI.set(model::AbstractModel, attr::MOI.AbstractModelAttribute, value)
    MOI.set(model.model, attr, value)
end

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