abstract type AbstractCollocation end

struct DynamicOptSolution
    model::Any
    sol::ODESolution
    input_sol::Union{Nothing, ODESolution}
end

function Base.show(io::IO, sol::DynamicOptSolution)
    println("retcode: ", sol.sol.retcode, "\n")

    println("Optimal control solution for following model:\n")
    show(sol.model)

    return print("\n\nPlease query the model using sol.model, the solution trajectory for the system using sol.sol, or the solution trajectory for the controllers using sol.input_sol.")
end

"""
    JuMPDynamicOptProblem(sys::System, op, tspan; dt, steps, guesses, kwargs...)

Convert a System representing an optimal control system into a JuMP model
for solving using optimization. Must provide either `dt`, the timestep between collocation
points (which, along with the timespan, determines the number of points), or directly
provide the number of points as `steps`.

To construct the problem, please load InfiniteOpt along with ModelingToolkitBase.
"""
function JuMPDynamicOptProblem end
"""
    InfiniteOptDynamicOptProblem(sys::System, op, tspan; dt)

Convert a System representing an optimal control system into a InfiniteOpt model
for solving using optimization. Must provide `dt` for determining the length
of the interpolation arrays.

Related to `JuMPDynamicOptProblem`, but directly adds the differential equations
of the system as derivative constraints, rather than using a solver tableau.

To construct the problem, please load InfiniteOpt along with ModelingToolkitBase.
"""
function InfiniteOptDynamicOptProblem end
"""
    CasADiDynamicOptProblem(sys::System, op, tspan; dt, steps, guesses, kwargs...)

Convert a System representing an optimal control system into a CasADi model
for solving using optimization. Must provide either `dt`, the timestep between collocation
points (which, along with the timespan, determines the number of points), or directly
provide the number of points as `steps`.

To construct the problem, please load CasADi along with ModelingToolkitBase.
"""
function CasADiDynamicOptProblem end
"""
    PyomoDynamicOptProblem(sys::System, op, tspan; dt, steps)

Convert a System representing an optimal control system into a Pyomo model
for solving using optimization. Must provide either `dt`, the timestep between collocation
points (which, along with the timespan, determines the number of points), or directly
provide the number of points as `steps`.

To construct the problem, please load Pyomo along with ModelingToolkitBase.
"""
function PyomoDynamicOptProblem end

### Collocations
"""
JuMP Collocation solver. Takes two arguments:
- `solver`: a optimization solver such as Ipopt
- `tableau`: An ODE RK tableau. Load a tableau by calling a function like `constructRK4` and may be found at https://docs.sciml.ai/DiffEqDevDocs/stable/internals/tableaus/. If this argument is not passed in, the solver will default to Radau second order.
"""
function JuMPCollocation end
"""
InfiniteOpt Collocation solver.
- `solver`: an optimization solver such as Ipopt
- `derivative_method`: the method used by InfiniteOpt to compute derivatives. The list of possible options can be found at https://infiniteopt.github.io/InfiniteOpt.jl/stable/guide/derivative/. Defaults to FiniteDifference(Backward()).
"""
function InfiniteOptCollocation end
"""
CasADi Collocation solver.
- `solver`: an optimization solver such as Ipopt. Should be given as a string or symbol in all lowercase, e.g. "ipopt"
- `tableau`: An ODE RK tableau. Load a tableau by calling a function like `constructRK4` and may be found at https://docs.sciml.ai/DiffEqDevDocs/stable/internals/tableaus/. If this argument is not passed in, the solver will default to Radau second order.
"""
function CasADiCollocation end
"""
Pyomo Collocation solver.
- `solver`: an optimization solver such as Ipopt. Should be given as a string or symbol in all lowercase, e.g. "ipopt"
- `derivative_method`: a derivative method from Pyomo. The choices here are ForwardEuler, BackwardEuler, MidpointEuler, LagrangeRadau, or LagrangeLegendre. The last two should additionally have a number indicating the number of collocation points per timestep, e.g. PyomoCollocation("ipopt", LagrangeRadau(3)). Defaults to LagrangeRadau(5).
"""
function PyomoCollocation end

function warn_overdetermined(sys, op)
    cstrs = constraints(sys)
    init_conds = filter(x -> value(x) ∈ Set(unknowns(sys)), [k for (k, v) in op])
    return if !isempty(cstrs)
        (length(cstrs) + length(init_conds) > length(unknowns(sys))) &&
            @warn "The control problem is overdetermined. The total number of conditions (# constraints + # fixed initial values given by op) exceeds the total number of states. The solvers will default to doing a nonlinear least-squares optimization."
    end
end

"""
Default ODE Tableau: RadauIIA5
"""
function constructDefault(T::Type = Float64)
    sq6 = sqrt(6)
    A = [
        11 // 45 - 7sq6 / 360 37 // 225 - 169sq6 / 1800 -2 // 225 + sq6 / 75
        37 // 225 + 169sq6 / 1800 11 // 45 + 7sq6 / 360 -2 // 225 - sq6 / 75
        4 // 9 - sq6 / 36 4 // 9 + sq6 / 36 1 // 9
    ]
    c = [2 // 5 - sq6 / 10; 2 / 5 + sq6 / 10; 1]
    α = [4 // 9 - sq6 / 36; 4 // 9 + sq6 / 36; 1 // 9]
    A = map(T, A)
    α = map(T, α)
    c = map(T, c)

    return DiffEqBase.ImplicitRKTableau(A, c, α, 5)
end

is_explicit(tableau) = tableau isa DiffEqBase.ExplicitRKTableau

@fallback_iip_specialize function SciMLBase.ODEInputFunction{iip, specialize}(
        sys::System;
        inputs = inputs(sys),
        disturbance_inputs = disturbances(sys),
        u0 = nothing, tgrad = false,
        jac = false, controljac = false,
        p = nothing, t = nothing,
        eval_expression = false,
        sparse = false, simplify = false,
        eval_module = @__MODULE__,
        steady_state = false,
        checkbounds = false,
        sparsity = false,
        analytic = nothing,
        initialization_data = nothing,
        cse = true,
        kwargs...
    ) where {iip, specialize}
    f, _,
        _ = generate_control_function(
        sys, inputs, disturbance_inputs; eval_module, cse, kwargs...
    )
    f = f[1]

    if tgrad
        _tgrad = generate_tgrad(
            sys;
            simplify = simplify,
            expression = Val{true},
            wrap_gfw = Val{true},
            expression_module = eval_module, cse,
            checkbounds = checkbounds, kwargs...
        )
    else
        _tgrad = nothing
    end

    if jac
        _jac = generate_jacobian(
            sys;
            simplify = simplify, sparse = sparse,
            expression = Val{true},
            wrap_gfw = Val{true},
            expression_module = eval_module, cse,
            checkbounds = checkbounds, kwargs...
        )
    else
        _jac = nothing
    end

    if controljac
        _cjac = generate_control_jacobian(
            sys;
            simplify = simplify, sparse = sparse,
            expression = Val{true}, wrap_gfw = Val{true},
            expression_module = eval_module, cse,
            checkbounds = checkbounds, kwargs...
        )
    else
        _cjac = nothing
    end

    M = calculate_massmatrix(sys)
    _M = concrete_massmatrix(M; sparse, u0)

    observedfun = ObservedFunctionCache(
        sys; steady_state, eval_expression, eval_module, checkbounds, cse
    )

    _W_sparsity = W_sparsity(sys)
    W_prototype = calculate_W_prototype(_W_sparsity; u0, sparse)
    if sparse
        uElType = u0 === nothing ? Float64 : eltype(u0)
        controljac_prototype = similar(calculate_control_jacobian(sys), uElType)
    else
        controljac_prototype = nothing
    end

    ODEInputFunction{iip, specialize}(
        f;
        sys = sys,
        jac = _jac === nothing ? nothing : _jac,
        controljac = _cjac === nothing ? nothing : _cjac,
        tgrad = _tgrad === nothing ? nothing : _tgrad,
        mass_matrix = _M,
        jac_prototype = W_prototype,
        controljac_prototype = controljac_prototype,
        observed = observedfun,
        sparsity = sparsity ? _W_sparsity : nothing,
        analytic = analytic,
        initialization_data
    )
end

# returns the JuMP timespan, the number of steps, and whether it is a free time problem.
function process_tspan(tspan, dt, steps)
    is_free_time = false
    symbolic_type(tspan[1]) !== NotSymbolic() &&
        error("Free initial time problems are not currently supported by the collocation solvers.")

    if isnothing(dt) && isnothing(steps)
        error("Must provide either the dt or the number of intervals to the collocation solvers (JuMP, InfiniteOpt, CasADi).")
    elseif symbolic_type(tspan[1]) === ScalarSymbolic() ||
            symbolic_type(tspan[2]) === ScalarSymbolic()
        isnothing(steps) &&
            error("Free final time problems require specifying the number of steps using the keyword arg `steps`, rather than dt.")
        isnothing(dt) ||
            @warn "Specified dt for free final time problem. This will be ignored; dt will be determined by the number of timesteps."

        return (0, 1), steps, true
    else
        isnothing(steps) ||
            @warn "Specified number of steps for problem with concrete tspan. This will be ignored; number of steps will be determined by dt."

        return tspan, length(tspan[1]:dt:tspan[2]), false
    end
end

function get_discrete_time_evaluations(expr)
    vars = Symbolics.get_variables(expr)

    # Group by base variable
    result = Dict()

    for v in vars
        if iscall(v)
            args = arguments(v)
            if length(args) == 1 && args[1] isa Number
                base_var = operation(v)
                time_point = args[1]

                if !haskey(result, base_var)
                    result[base_var] = Float64[]
                end
                push!(result[base_var], time_point)
            end
        end
    end

    # Sort and unique the time points
    for (var, times) in result
        result[var] = sort!(unique!(times))
    end

    return result
end

function check_collocation_time_mismatch(sys, expected_tsteps, tsteps)
    return if collect(expected_tsteps) ≠ tsteps
        eval_times = get_discrete_time_evaluations(cost(sys))
        for (var, ts) in eval_times
            tdiff = setdiff(ts, expected_tsteps)
            @info tdiff
            if !isempty(tdiff)
                error(
                    "$var is evaluated inside the cost function at $(length(tdiff)) points " *
                        "that are not in the $(length(expected_tsteps)) collocation points. " *
                        "Cost evaluation points must align with the collocation grid. " *
                        "Adjust the dt to match the time points used in the cost function."
                )
            end
        end
        if length(expected_tsteps) ≠ length(tsteps)
            error("Found extra $(abs(length(expected_tsteps) - length(tsteps))) collocation points.")
        elseif maximum(abs.(expected_tsteps .- tsteps)) > 1.0e-10
            error("The collocation points differ from the expected ones by more than 1e-10.")
        end
    end
end

##########################
### MODEL CONSTRUCTION ###
##########################
function process_DynamicOptProblem(
        prob_type::Type{<:SciMLBase.AbstractDynamicOptProblem}, model_type, sys::System, op, tspan;
        dt = nothing,
        steps = nothing,
        tune_parameters = false,
        guesses = Dict(),
        bounds = Dict(), kwargs...
    )
    warn_overdetermined(sys, op)
    ctrls = inputs(sys)
    states = unknowns(sys)
    tunable_params = tune_parameters ? tunable_parameters(sys) : []

    stidxmap = Dict([v => i for (i, v) in enumerate(states)])
    op = Dict([default_toterm(value(k)) => v for (k, v) in op])
    bounds = Dict([default_toterm(value(k)) => v for (k, v) in bounds])
    u0_idxs = has_alg_eqs(sys) ? collect(1:length(states)) :
        [stidxmap[default_toterm(k)] for (k, v) in op if haskey(stidxmap, k)]

    _op = has_alg_eqs(sys) ? op : merge(Dict(op), Dict(guesses))
    f, u0,
        p = process_SciMLProblem(
        ODEInputFunction, sys, _op;
        t = tspan !== nothing ? tspan[1] : tspan, kwargs...
    )
    model_tspan, steps, is_free_t = process_tspan(tspan, dt, steps)
    warn_overdetermined(sys, op)

    # Build pmap for symbolic substitution in costs/constraints/bounds
    all_parameters = default_toterm.(parameters(sys))
    # Extract all parameter values from processed p (which has defaults filled in)
    getter = SymbolicIndexingInterface.getp(sys, all_parameters)
    pmap = AnyDict(all_parameters .=> getter(p))

    # Filter out tunable parameters - they should remain symbolic
    tunable_set = Set(default_toterm.(tunable_params))
    pmap = filter(kvp -> first(kvp) ∉ tunable_set, pmap)

    # Resolve parameter bindings so observed equations and constraints
    # referencing pre-binding names can be fully substituted.
    for (k, v) in bindings(sys)
        v === COMMON_MISSING && continue
        haskey(pmap, v) && !haskey(pmap, k) && (pmap[k] = pmap[v])
    end

    c0 = value.([pmap[c] for c in ctrls])
    p0, _ = SciMLStructures.canonicalize(SciMLStructures.Tunable(), p)

    tsteps = LinRange(model_tspan[1], model_tspan[2], steps)
    model = generate_internal_model(model_type)
    generate_time_variable!(model, model_tspan, tsteps)
    U = generate_state_variable!(model, u0, length(states), tsteps)
    V = generate_input_variable!(model, c0, length(ctrls), tsteps)
    P = generate_tunable_params!(model, p0, length(tunable_params))
    # Add the symbolic representation of the tunable parameters to the map
    # The order of the Tunable section is given by the tunable_parameters(sys)
    # Some backends need symbolic accessors instead of raw variables
    P_syms = [get_param_for_pmap(model, P, i) for i in eachindex(tunable_params)]
    P_backend = needs_individual_tunables(model) ? P_syms : P

    final_t_ic = get(pmap, tspan[2], tspan[2])
    @assert symbolic_type(final_t_ic) isa NotSymbolic "The initial condition for the final time could not be determined from $final_t_ic."
    tₛ = generate_timescale!(model, final_t_ic, is_free_t)
    fullmodel = model_type(model, U, V, P_backend, tₛ, is_free_t, tsteps)

    merge!(pmap, Dict(tunable_params .=> P_syms))

    set_variable_bounds!(fullmodel, sys, pmap, tspan[2], tunable_params, bounds)
    add_cost_function!(fullmodel, sys, tspan, pmap)
    add_user_constraints!(fullmodel, sys, tspan, pmap)
    add_initial_constraints!(fullmodel, u0, u0_idxs, model_tspan[1])

    return prob_type(f, u0, tspan, p, fullmodel, kwargs...), pmap
end

function generate_time_variable! end
function generate_internal_model end
function generate_state_variable! end
function generate_input_variable! end
function generate_tunable_params! end
function generate_timescale! end
function add_initial_constraints! end
function add_constraint! end
# Default: return P[i] directly. Symbolic backends (like Pyomo) can override this.
get_param_for_pmap(model, P, i) = P isa AbstractArray ? P[i] : P
# Some backends need symbolic accessors instead of raw variables (CasADi in particular)
needs_individual_tunables(model) = false

function f_wrapper(f, Uₙ, Vₙ, p, P, t)
    if isempty(P)
        # no tunable parameters
        return f(Uₙ, Vₙ, p, t)
    end
    return if SciMLStructures.isscimlstructure(p)
        _, repack, _ = SciMLStructures.canonicalize(SciMLStructures.Tunable(), p)
        p′ = repack(P)
        f(Uₙ, Vₙ, p′, t)
    else
        f(Uₙ, Vₙ, P, t)
    end
end

"""
    extract_variable_bounds(sys, pmap, tf, tunable_params)

Extract and parameter-substitute variable bounds from the system.
Returns `(; state_bounds, input_bounds, param_bounds, tf_bounds)` where each
`*_bounds` is a `Dict{Int, Tuple{Any, Any}}` mapping variable index to `(lo, hi)`,
and `tf_bounds` is either `nothing` or a `(lo, hi)` tuple.
"""
function extract_variable_bounds(sys, pmap, tf, tunable_params, user_bounds = Dict())
    state_bounds = _extract_bounds(unknowns(sys), pmap)
    input_bounds = _extract_bounds(inputs(sys), pmap)
    param_bounds = _extract_bounds(tunable_params, pmap)
    # Merge user-provided bounds (override metadata bounds)
    dvs = unknowns(sys)
    ctrls = inputs(sys)
    for (var, (lo, hi)) in user_bounds
        idx = findfirst(v -> isequal(v, var), dvs)
        if !isnothing(idx)
            state_bounds[idx] = (lo, hi)
            continue
        end
        idx = findfirst(v -> isequal(v, var), ctrls)
        if !isnothing(idx)
            input_bounds[idx] = (lo, hi)
        end
    end
    tf_bounds = if symbolic_type(tf) === ScalarSymbolic() && hasbounds(tf)
        lo, hi = getbounds(tf)
        (SymbolicUtils.unwrap_const(unwrap(Symbolics.fixpoint_sub(lo, pmap))),
            SymbolicUtils.unwrap_const(unwrap(Symbolics.fixpoint_sub(hi, pmap))))
    else
        nothing
    end
    return (; state_bounds, input_bounds, param_bounds, tf_bounds)
end

function _extract_bounds(vars, pmap)
    bounds = Dict{Int, Tuple{Any, Any}}()
    for (i, v) in enumerate(vars)
        if hasbounds(v)
            lo, hi = getbounds(v)
            lo = SymbolicUtils.unwrap_const(unwrap(Symbolics.fixpoint_sub(lo, pmap)))
            hi = SymbolicUtils.unwrap_const(unwrap(Symbolics.fixpoint_sub(hi, pmap)))
            bounds[i] = (lo, hi)
        end
    end
    return bounds
end

function set_variable_bounds! end

is_free_final(model) = model.is_free_final

function add_cost_function!(model, sys, tspan, pmap)
    jcosts = cost(sys)
    if Symbolics._iszero(jcosts)
        set_objective!(model, 0)
        return
    end

    # First resolve observed variables so that EvalAt evaluations like
    # obs_var(1.0) or obs_var(tf) are expanded before model-var rules.
    # Parameter rules are intentionally excluded: they would fold initial
    # guesses (e.g. u(t) => 0.0) into the cost instead of letting
    # model-var rules map them to solver variables (e.g. V[1](t)).
    obs_rules = Dict{Any, Any}()
    get_observed_substitution_rules!(obs_rules, sys; tspan)
    jcosts = fixpoint_sub(jcosts, obs_rules; fold = Val(true), filterer = Returns(true))

    rules = Dict{Any, Any}()
    get_model_vars_substitution_rules!(rules, model, sys, tspan)
    get_param_substitution_rules!(rules, pmap)
    get_integral_substitution_rules!(rules, model, jcosts, tspan)
    jcosts = substitute(jcosts, rules; fold = Val(true), filterer = Returns(true))
    return set_objective!(model, value(jcosts))
end

function get_integral_substitution_rules!(rules::Dict{Any, Any}, model, expr, tspan)
    for int in collect_applied_operators(expr, Symbolics.Integral)
        op = operation(int)
        arg = only(arguments(value(int)))
        lo, hi = value.((op.domain.domain.left, op.domain.domain.right))
        lo, hi = process_integral_bounds(model, (lo, hi), tspan)
        arg = substitute(arg, rules; fold = Val(true), filterer = Returns(true))
        rules[int] = lowered_integral(model, arg, lo, hi)
    end
    return
end

function process_integral_bounds(model, integral_span, tspan)
    return if is_free_final(model) && isequal(integral_span, tspan)
        integral_span = (0, 1)
    elseif is_free_final(model)
        error("Free final time problems cannot handle partial timespans.")
    else
        (lo, hi) = integral_span
        (lo < tspan[1] || hi > tspan[2]) &&
            error("Integral bounds are beyond the timespan.")
        integral_span
    end
end

function get_model_vars_substitution_rules!(rules::Dict{Any, Any}, model, sys, tspan)
    x_ops = [operation(unwrap(st)) for st in unknowns(sys)]
    c_ops = [operation(unwrap(ct)) for ct in inputs(sys)]
    t = get_iv(sys)
    merge!(rules, whole_t_map(model, t, x_ops, c_ops))
    (ti, tf) = tspan
    if symbolic_type(tf) === ScalarSymbolic()
        _tf = model.tₛ + ti
        merge!(rules, free_t_map(model, tf, x_ops, c_ops))
        rules[tf] = _tf
    end
    merge!(rules, fixed_t_map(model, x_ops, c_ops))
    return nothing
end

function get_observed_substitution_rules!(rules::Dict{Any, Any}, sys; tspan = nothing)
    # add the substitution rules for the observed variables
    subs = get_substitutions(sys)
    merge!(rules, subs)
    # Also add operator-level rules so that concrete time evaluations
    # like obs_var(1.0) from EvalAt are substituted correctly.
    # The operator lambda handles numeric time points (combine_fold can call it
    # when all args are Const). For symbolic time points like tf, we add
    # explicit rules since combine_fold won't call the lambda with symbolic args.
    iv = get_iv(sys)
    for (lhs, rhs) in subs
        op = operation(unwrap(lhs))
        haskey(rules, op) && continue
        # filterer = Returns(true) is needed to recurse inside the (t) arguments
        # of the observed expression;
        # e.g. for obs_val = a*x(t) and a cost that has obs_val(1.0), we want to get a*x(1.0)
        rules[op] = t -> substitute(rhs, Dict(iv => t); filterer = Returns(true))
        # Add explicit rules for symbolic tspan endpoints (e.g. obs_val(tf))
        if tspan !== nothing
            for ti in tspan
                symbolic_type(ti) === ScalarSymbolic() || continue
                evaluated = substitute(rhs, Dict(iv => value(ti)); filterer = Returns(true))
                rules[op(value(ti))] = evaluated
            end
        end
    end
    return nothing
end

"""Mappings for variables that depend on the final time parameter, x(tf)."""
function free_t_map(m, tf, x_ops, c_ops)
    return Dict(
        [
            [x(tf) => lowered_var(m, :U, i, 1) for (i, x) in enumerate(x_ops)];
            [c(tf) => lowered_var(m, :V, i, 1) for (i, c) in enumerate(c_ops)]
        ]
    )
end

"""Mappings for variables that cover the whole timespan, x(t)."""
function whole_t_map(m, t, x_ops, c_ops)
    return Dict(
        [
            [v(t) => lowered_var(m, :U, i, t) for (i, v) in enumerate(x_ops)];
            [v(t) => lowered_var(m, :V, i, t) for (i, v) in enumerate(c_ops)]
        ]
    )
end

"""Mappings for variables that cover the whole timespan, x(t)."""
function fixed_t_map(m, x_ops, c_ops)
    return Dict(
        [
            [v => (t -> lowered_var(m, :U, i, t)) for (i, v) in enumerate(x_ops)];
            [v => (t -> lowered_var(m, :V, i, t)) for (i, v) in enumerate(c_ops)]
        ]
    )
end

function process_integral_bounds end
function lowered_integral end
function lowered_derivative end
function lowered_var end
function fixed_t_map end

function add_user_constraints!(model, sys, tspan, pmap)
    jconstraints = get_constraints(sys)
    (isnothing(jconstraints) || isempty(jconstraints)) && return nothing
    cons_dvs,
        cons_ps = process_constraint_system(
        jconstraints, Set(unknowns(sys)), parameters(sys), get_iv(sys); validate = false
    )

    is_free_final(model) && check_constraint_vars(cons_dvs)

    # First resolve observed variables so that EvalAt evaluations
    # are expanded before model-var rules are applied.
    obs_rules = Dict{Any, Any}()
    get_observed_substitution_rules!(obs_rules, sys; tspan)
    jconstraints = fixpoint_sub(jconstraints, obs_rules; fold = Val(true), filterer = Returns(true))

    rules = Dict{Any, Any}()
    get_toterm_substitution_rules!(rules, cons_dvs)
    get_model_vars_substitution_rules!(rules, model, sys, tspan)
    get_param_substitution_rules!(rules, pmap)
    # `fixpoint_sub` to recursively substitute into `toterm` rules
    jconstraints = fixpoint_sub(jconstraints, rules; fold = Val(true), filterer = Returns(true))

    for c in jconstraints
        add_constraint!(model, c)
    end
    return
end

function add_equational_constraints!(model, sys, pmap, tspan)
    rules = Dict{Any, Any}()
    get_observed_substitution_rules!(rules, sys)
    get_model_vars_substitution_rules!(rules, model, sys, tspan)
    get_param_substitution_rules!(rules, pmap)
    get_differential_substitution_rules!(rules, model, sys)
    diff_eqs = fixpoint_sub(diff_equations(sys), rules; fold = Val(true), filterer = Returns(true))
    for eq in diff_eqs
        add_constraint!(model, eq.lhs ~ unwrap_const(eq.rhs) * model.tₛ)
    end

    alg_eqs = fixpoint_sub(alg_equations(sys), rules; fold = Val(true), filterer = Returns(true))
    for eq in alg_eqs
        add_constraint!(model, eq.lhs ~ eq.rhs)
    end
    return
end

function set_objective! end
objective_value(sol::DynamicOptSolution) = objective_value(sol.model)

function get_differential_substitution_rules!(rules::Dict{Any, Any}, model, sys)
    t = get_iv(sys)
    D = Differential(t)
    diffsubmap = Dict(
        [
            D(unk) => lowered_derivative(model, i)
                for (i, unk) in enumerate(unknowns(sys))
        ]
    )
    merge!(rules, diffsubmap)
    return nothing
end

function get_toterm_substitution_rules!(rules::Dict{Any, Any}, vars)
    for u in vars
        ttu = default_toterm(unwrap(u))
        isequal(u, ttu) || (rules[u] = ttu)
    end
    return
end

function get_param_substitution_rules!(rules::Dict{Any, Any}, pmap)
    return left_merge!(rules, pmap)
end

function check_constraint_vars(vars)
    for u in vars
        x = operation(u)
        t = only(arguments(u))
        if (symbolic_type(t) === NotSymbolic())
            error("Provided specific time constraint in a free final time problem. This is not supported by the collocation solvers at the moment. The offending variable is $u. Specific-time user constraints can only be specified at the end of the timespan.")
        end
    end
    return
end

########################
### SOLVER UTILITIES ###
########################
"""
Add the solve constraints, set the solver (Ipopt, e.g.) and solver options, optimize the model.
"""
function prepare_and_optimize! end
function get_t_values end
function get_U_values end
function get_V_values end
function get_P_values end
function successful_solve end

"""
    solve(prob::AbstractDynamicOptProblem, solver::AbstractCollocation; verbose = false, kwargs...)

- kwargs are used for other options. For example, the `plugin_options` and `solver_options` will propagated to the Opti object in CasADi.
"""
function CommonSolve.solve(
        prob::SciMLBase.AbstractDynamicOptProblem,
        solver::AbstractCollocation; verbose = false, kwargs...
    )
    solved_model = prepare_and_optimize!(prob, solver; verbose, kwargs...)

    ts = get_t_values(solved_model)
    Us = get_U_values(solved_model)
    Vs = get_V_values(solved_model)
    Ps = get_P_values(solved_model)
    is_free_final(prob.wrapped_model) && (ts .+ prob.tspan[1])

    # update the parameters with the ones in the solved_model
    if !isempty(Ps)
        new_p = SciMLStructures.replace(SciMLStructures.Tunable(), prob.p, Ps)
        new_prob = remake(prob, p = new_p)
    else
        new_prob = prob
    end
    ode_sol = SciMLBase.build_solution(new_prob, solver, ts, Us)
    input_sol = isnothing(Vs) ? nothing : SciMLBase.build_solution(new_prob, solver, ts, Vs)

    if !successful_solve(solved_model)
        ode_sol = SciMLBase.solution_new_retcode(
            ode_sol, SciMLBase.ReturnCode.ConvergenceFailure
        )
        !isnothing(input_sol) && (
            input_sol = SciMLBase.solution_new_retcode(
                input_sol, SciMLBase.ReturnCode.ConvergenceFailure
            )
        )
    end
    return DynamicOptSolution(solved_model, ode_sol, input_sol)
end
