MTKInfiniteOptExt.jl

module MTKInfiniteOptExt
using ModelingToolkitBase
using OrderedCollections: OrderedSet
import Symbolics: SymbolicT
using InfiniteOpt
using DiffEqBase
using SciMLStructures
using LinearAlgebra
using StaticArrays
using UnPack
using Symbolics: unwrap
import SymbolicUtils
import NaNMath
const MTK = ModelingToolkitBase

function __init__()
    # Workaround for Julia 1.10 compiler bug (Issue #4211)
    # On Julia 1.10, loading this extension can cause collect_vars! to fail
    # to properly call collect_var! until collect_var! has been called directly.
    # This forces correct method compilation before user code runs.
    if VERSION < v"1.11"
        _force_collect_var_compilation()
    end
    return nothing
end

# Helper function to force collect_var! compilation on Julia 1.10.
# This works around a method invalidation bug where the mutual recursion between
# collect_vars! and collect_var! breaks after this extension loads.
function _force_collect_var_compilation()
    MTK.@variables _dummy_t
    MTK.@parameters _dummy_p
    MTK.@variables _dummy_x(_dummy_t) = _dummy_p
    _us = OrderedSet{SymbolicT}()
    _ps = OrderedSet{SymbolicT}()
    MTK.collect_var!(_us, _ps, unwrap(_dummy_x), unwrap(_dummy_t); depth = 0)
    return nothing
end

struct InfiniteOptModel
    model::InfiniteModel
    U::Vector{<:AbstractVariableRef}
    V::Vector{<:AbstractVariableRef}
    P::Vector{<:AbstractVariableRef}
    tₛ::AbstractVariableRef
    is_free_final::Bool
    tsteps::LinRange{Float64, Int}
end

struct JuMPDynamicOptProblem{uType, tType, isinplace, P, F, K} <:
    SciMLBase.AbstractDynamicOptProblem{uType, tType, isinplace}
    f::F
    u0::uType
    tspan::tType
    p::P
    wrapped_model::InfiniteOptModel
    kwargs::K

    function JuMPDynamicOptProblem(f, u0, tspan, p, model, kwargs)
        return new{
            typeof(u0), typeof(tspan), SciMLBase.isinplace(f, 5),
            typeof(p), typeof(f), typeof(kwargs),
        }(f, u0, tspan, p, model, kwargs)
    end
    function JuMPDynamicOptProblem(f, u0, tspan, p, model; kwargs...)
        return JuMPDynamicOptProblem(f, u0, tspan, p, model, kwargs)
    end
end

struct InfiniteOptDynamicOptProblem{uType, tType, isinplace, P, F, K} <:
    SciMLBase.AbstractDynamicOptProblem{uType, tType, isinplace}
    f::F
    u0::uType
    tspan::tType
    p::P
    wrapped_model::InfiniteOptModel
    kwargs::K

    function InfiniteOptDynamicOptProblem(f, u0, tspan, p, model, kwargs)
        return new{
            typeof(u0), typeof(tspan), SciMLBase.isinplace(f),
            typeof(p), typeof(f), typeof(kwargs),
        }(f, u0, tspan, p, model, kwargs)
    end
    function InfiniteOptDynamicOptProblem(f, u0, tspan, p, model; kwargs...)
        return InfiniteOptDynamicOptProblem(f, u0, tspan, p, model, kwargs)
    end
end

MTK.generate_internal_model(m::Type{InfiniteOptModel}) = InfiniteModel()
function MTK.generate_time_variable!(m::InfiniteModel, tspan, tsteps)
    return @infinite_parameter(m, t in [tspan[1], tspan[2]], num_supports = length(tsteps))
end
function MTK.generate_state_variable!(m::InfiniteModel, u0::Vector, ns, ts)
    return @variable(m, U[i = 1:ns], Infinite(m[:t]), start = u0[i])
end
function MTK.generate_input_variable!(m::InfiniteModel, c0, nc, ts)
    return @variable(m, V[i = 1:nc], Infinite(m[:t]), start = c0[i])
end
function MTK.generate_tunable_params!(m::InfiniteModel, p0, np)
    return @variable(m, P[i = 1:np], start = p0[i])
end

function MTK.generate_timescale!(m::InfiniteModel, guess, is_free_t)
    @variable(m, tₛ ≥ 0, start = guess)
    if !is_free_t
        fix(tₛ, 1, force = true)
        set_start_value(tₛ, 1)
    end
    return tₛ
end

function MTK.add_constraint!(m::InfiniteOptModel, expr::Union{Equation, Inequality})
    return if expr isa Equation
        @constraint(m.model, SymbolicUtils.unwrap_const(expr.lhs) - SymbolicUtils.unwrap_const(expr.rhs) == 0)
    elseif expr.relational_op === Symbolics.geq
        @constraint(m.model, SymbolicUtils.unwrap_const(expr.lhs) - SymbolicUtils.unwrap_const(expr.rhs) ≥ 0)
    else
        @constraint(m.model, SymbolicUtils.unwrap_const(expr.lhs) - SymbolicUtils.unwrap_const(expr.rhs) ≤ 0)
    end
end
MTK.set_objective!(m::InfiniteOptModel, expr) = @objective(m.model, Min, SymbolicUtils.unwrap_const(expr))

function MTK.set_variable_bounds!(m::InfiniteOptModel, sys, pmap, tf, tunable_params, user_bounds = Dict())
    (; state_bounds, input_bounds, param_bounds, tf_bounds) = MTK.extract_variable_bounds(sys, pmap, tf, tunable_params, user_bounds)
    for (i, (lo, hi)) in state_bounds
        set_lower_bound(m.U[i], lo)
        set_upper_bound(m.U[i], hi)
    end
    for (i, (lo, hi)) in input_bounds
        set_lower_bound(m.V[i], lo)
        set_upper_bound(m.V[i], hi)
    end
    for (i, (lo, hi)) in param_bounds
        set_lower_bound(m.P[i], lo)
        set_upper_bound(m.P[i], hi)
    end
    return if !isnothing(tf_bounds)
        set_lower_bound(m.tₛ, tf_bounds[1])
        set_upper_bound(m.tₛ, tf_bounds[2])
    end
end

function MTK.JuMPDynamicOptProblem(
        sys::System, op, tspan;
        dt = nothing,
        steps = nothing,
        tune_parameters = false,
        guesses = Dict(),
        bounds = Dict(), kwargs...
    )
    prob,
        _ = MTK.process_DynamicOptProblem(
        JuMPDynamicOptProblem, InfiniteOptModel, sys,
        op, tspan; dt, steps, tune_parameters, guesses, bounds, kwargs...
    )
    return prob
end

function MTK.InfiniteOptDynamicOptProblem(
        sys::System, op, tspan;
        dt = nothing,
        steps = nothing,
        tune_parameters = false,
        guesses = Dict(),
        bounds = Dict(), kwargs...
    )
    prob,
        pmap = MTK.process_DynamicOptProblem(
        InfiniteOptDynamicOptProblem, InfiniteOptModel,
        sys, op, tspan; dt, steps, tune_parameters, guesses, bounds, kwargs...
    )
    MTK.add_equational_constraints!(prob.wrapped_model, sys, pmap, tspan)
    return prob
end

function MTK.lowered_integral(model::InfiniteOptModel, expr, lo, hi)
    return model.tₛ * InfiniteOpt.∫(SymbolicUtils.unwrap_const(expr), model.model[:t], lo, hi)
end
MTK.lowered_derivative(model::InfiniteOptModel, i) = ∂(model.U[i], model.model[:t])

function MTK.process_integral_bounds(model::InfiniteOptModel, integral_span, tspan)
    return if MTK.is_free_final(model) && isequal(integral_span, tspan)
        integral_span = (0, 1)
    elseif MTK.is_free_final(model)
        error("Free final time problems cannot handle partial timespans.")
    else
        integral_span
    end
end

function MTK.add_initial_constraints!(m::InfiniteOptModel, u0, u0_idxs, ts)
    for i in u0_idxs
        fix(m.U[i](0), u0[i], force = true)
    end
    return
end

function MTK.lowered_var(m::InfiniteOptModel, uv, i, t)
    X = getfield(m, uv)
    return t isa Union{Num, SymbolicT} ? X[i] : X[i](t)
end

function add_solve_constraints!(prob::JuMPDynamicOptProblem, tableau)
    @unpack A, α, c = tableau
    @unpack wrapped_model, f, p = prob
    @unpack tₛ, U, V, P, model = wrapped_model
    t = model[:t]
    tsteps = supports(t)

    # InfiniteOpt can introduce additional collocation points
    # Make sure that the collocation points are correct.
    MTK.check_collocation_time_mismatch(prob.f.sys, wrapped_model.tsteps, tsteps)

    dt = (tsteps[end] - tsteps[1]) / (length(tsteps) - 1)

    nᵤ = length(U)
    nᵥ = length(V)
    return if MTK.is_explicit(tableau)
        K = Any[]
        for τ in tsteps[1:(end - 1)]
            for (i, h) in enumerate(c)
                ΔU = sum([A[i, j] * K[j] for j in 1:(i - 1)], init = zeros(nᵤ))
                Uₙ = [U[i](τ) + ΔU[i] * dt for i in 1:nᵤ]
                Vₙ = [V[i](τ) for i in 1:nᵥ]
                Kₙ = tₛ * MTK.f_wrapper(f, Uₙ, Vₙ, p, P, τ + h * dt)
                push!(K, Kₙ)
            end
            ΔU = dt * sum([α[i] * K[i] for i in 1:length(α)])
            @constraint(
                model, [n = 1:nᵤ], U[n](τ) + ΔU[n] == U[n](τ + dt),
                base_name = "solve_time_$τ"
            )
            empty!(K)
        end
    else
        K = @variable(model, K[1:length(α), 1:nᵤ], Infinite(model[:t]))
        ΔUs = A * K
        ΔU_tot = dt * (K' * α)
        for τ in tsteps[1:(end - 1)]
            for (i, h) in enumerate(c)
                ΔU = @view ΔUs[i, :]
                Uₙ = U + ΔU * dt
                @constraint(
                    model, [j = 1:nᵤ], K[i, j] == (tₛ * MTK.f_wrapper(f, Uₙ, V, p, P, τ + h * dt)[j]),
                    DomainRestriction(==(τ), t), base_name = "solve_K$i($τ)"
                )
            end
            @constraint(
                model,
                [n = 1:nᵤ], U[n](τ) + ΔU_tot[n] == U[n](min(τ + dt, tsteps[end])),
                DomainRestriction(==(τ), t), base_name = "solve_U($τ)"
            )
        end
    end
end

struct JuMPCollocation <: AbstractCollocation
    solver::Any
    tableau::DiffEqBase.ODERKTableau
end
function MTK.JuMPCollocation(solver, tableau = MTK.constructDefault())
    return JuMPCollocation(solver, tableau)
end

struct InfiniteOptCollocation <: AbstractCollocation
    solver::Any
    derivative_method::InfiniteOpt.AbstractDerivativeMethod
end
function MTK.InfiniteOptCollocation(
        solver, derivative_method = InfiniteOpt.FiniteDifference(InfiniteOpt.Backward())
    )
    return InfiniteOptCollocation(solver, derivative_method)
end

function MTK.prepare_and_optimize!(
        prob::JuMPDynamicOptProblem, solver::JuMPCollocation; verbose = false, kwargs...
    )
    model = prob.wrapped_model.model
    verbose || set_silent(model)
    # Unregister current solver constraints
    for con in all_constraints(model)
        if occursin("solve", JuMP.name(con))
            unregister(model, Symbol(JuMP.name(con)))
            delete(model, con)
        end
    end
    unregister(model, :K)
    for var in all_variables(model)
        if occursin("K", JuMP.name(var))
            unregister(model, Symbol(JuMP.name(var)))
            delete(model, var)
        end
    end
    add_solve_constraints!(prob, solver.tableau)
    set_optimizer(model, solver.solver)
    optimize!(model)
    return model
end

function MTK.prepare_and_optimize!(
        prob::InfiniteOptDynamicOptProblem,
        solver::InfiniteOptCollocation; verbose = false, kwargs...
    )
    model = prob.wrapped_model.model
    verbose || set_silent(model)
    set_derivative_method(model[:t], solver.derivative_method)
    set_optimizer(model, solver.solver)
    optimize!(model)
    return model
end

function MTK.get_V_values(m::InfiniteModel)
    nt = length(supports(m[:t]))
    return if !isempty(m[:V])
        V_vals = value.(m[:V])
        V_vals = [[V_vals[i][j] for i in 1:length(V_vals)] for j in 1:nt]
    else
        nothing
    end
end
function MTK.get_U_values(m::InfiniteModel)
    nt = length(supports(m[:t]))
    U_vals = value.(m[:U])
    return U_vals = [[U_vals[i][j] for i in 1:length(U_vals)] for j in 1:nt]
end
MTK.get_P_values(m::InfiniteModel) = value(m[:P])
MTK.get_t_values(m::InfiniteModel) = value(m[:tₛ]) * supports(m[:t])
MTK.objective_value(m::InfiniteModel) = InfiniteOpt.objective_value(m)

function MTK.successful_solve(model::InfiniteModel)
    tstatus = termination_status(model)
    pstatus = primal_status(model)
    !has_values(model) &&
        error("Model not solvable; please report this to github.com/SciML/ModelingToolkitBase.jl with a MWE.")

    return pstatus === FEASIBLE_POINT &&
        (
        tstatus === OPTIMAL || tstatus === LOCALLY_SOLVED || tstatus === ALMOST_OPTIMAL ||
            tstatus === ALMOST_LOCALLY_SOLVED
    )
end

import InfiniteOpt: JuMP, GeneralVariableRef

for ff in [acos, log1p, acosh, log2, asin, tan, atanh, cos, log, sin, log10, sqrt]
    f = nameof(ff)
    # These need to be defined so that JuMP can trace through functions built by Symbolics
    @eval NaNMath.$f(x::GeneralVariableRef) = Base.$f(x)
end

# JuMP variables and Symbolics variables never compare equal. When tracing through dynamics, a function argument can be either a JuMP variable or A Symbolics variable, it can never be both.
function Base.isequal(
        ::SymbolicT,
        ::Union{JuMP.GenericAffExpr, JuMP.GenericQuadExpr, JuMP.GenericNonlinearExpr}
    )
    return false
end
function Base.isequal(
        ::Union{JuMP.GenericAffExpr, JuMP.GenericQuadExpr, JuMP.GenericNonlinearExpr},
        ::SymbolicT
    )
    return false
end
end