module ManifoldsOrdinaryDiffEqDiffEqCallbacksExt

using Manifolds
using Manifolds:
    IntegratorTerminatorNearChartBoundary,
    affine_connection,
    get_chart_index,
    transition_map!,
    transition_map_diff!
import Manifolds: solve_chart_exp_ode, solve_chart_parallel_transport_ode
using ManifoldsBase

using DiffEqCallbacks
using OrdinaryDiffEq: OrdinaryDiffEq, SciMLBase, Rodas5, AutoVern9, ODEProblem, solve

using RecursiveArrayTools: ArrayPartition

"""
    (int_term::IntegratorTerminatorNearChartBoundary)(u, t, integrator)

Terminate integration when integrator goes too closely to chart boundary.
Closeness is determined by `ϵ` value of [`IntegratorTerminatorNearChartBoundary`](@ref)
`int_term`.

# Arguments:

- `int_term`: object containing keyword arguments for `check_chart_switch`, such as
  the desired maximum distance to boundary,
- `u`: parameters of a point at which the integrator is solving a differential equation.
- `t`: time parameter of the integrator
- `integrator`: state of the integrator. Internal parameters are expected to contained
  the manifold on which the equation is solved, the atlas and the current chart index.
"""
function (int_term::IntegratorTerminatorNearChartBoundary)(u, t, integrator)
    (M, A, i) = integrator.p
    if check_chart_switch(M, A, i, u.x[1]; int_term.check_chart_switch_kwargs...)
        # switch charts
        OrdinaryDiffEq.terminate!(integrator)
    end
    return u
end

@doc raw"""
    StitchedChartSolution{Prob,TM<:AbstractManifold,TA<:AbstractAtlas,TChart}

Solution of an ODE on a manifold `M` in charts of an [`AbstractAtlas`](@ref) `A`.

When `StitchedChartSolution{:Exp}` is used as a function with a number `t` as an argument, a
pair `(p, X)` is returned such that $p\in \mathcal{M}$ is the point at time `t` of the
geodesic and $X \in T_p \mathcal{M}$ is the velocity of the geodesic at that point.
Similarly, `StitchedChartSolution{:PT}` called with number `t` returns a triple `(p, X, Y)`
where `(p, X)` corresponds to the geodesic along which the vector is transported
and $Y\in T_p\mathcal{M}$ is the vector transported to `p`.
"""
struct StitchedChartSolution{Prob, TM <: AbstractManifold, TA <: AbstractAtlas, TChart}
    M::TM
    A::TA
    sols::Vector{Tuple{SciMLBase.AbstractODESolution, TChart}}
end

function StitchedChartSolution(
        M::AbstractManifold, A::AbstractAtlas, problem::Symbol, TChart
    )
    return StitchedChartSolution{problem, typeof(M), typeof(A), TChart}(M, A, [])
end

function (scs::StitchedChartSolution{:Exp})(t::Real)
    if t < scs.sols[1][1].t[1]
        throw(DomainError("Time $t is outside of the solution."))
    end
    for (sol, i) in scs.sols
        if t <= sol.t[end]
            B = induced_basis(scs.M, scs.A, i)
            solt = sol(t)
            p = get_point(scs.M, scs.A, i, solt.x[1])
            X = get_vector(scs.M, p, solt.x[2], B)
            return (p, X)
        end
    end
    throw(
        DomainError(
            "Time $t is outside of the solution (solution time range is [$(scs.sols[1][1].t[1]), $(scs.sols[end][1].t[end])]).",
        ),
    )
end
function (scs::StitchedChartSolution{:PT})(t::Real)
    if t < scs.sols[1][1].t[1]
        throw(DomainError("Time $t is outside of the solution."))
    end
    for (sol, i) in scs.sols
        if t <= sol.t[end]
            B = induced_basis(scs.M, scs.A, i)
            solt = sol(t)
            p = get_point(scs.M, scs.A, i, solt.x[1])
            X = get_vector(scs.M, p, solt.x[2], B)
            Y = get_vector(scs.M, p, solt.x[3], B)
            return (p, X, Y)
        end
    end
    throw(
        DomainError(
            "Time $t is outside of the solution (solution time range is [$(scs.sols[1][1].t[1]), $(scs.sols[end][1].t[end])]).",
        ),
    )
end

function (scs::StitchedChartSolution)(t::AbstractArray)
    return map(scs, t)
end

function chart_exp_problem(u, params, t)
    M, A, i = params
    a = u.x[1]
    dx = u.x[2]
    ddx = -affine_connection(M, A, i, a, dx, dx)
    return ArrayPartition(dx, ddx)
end

"""
    solve_chart_exp_ode(
        M::AbstractManifold, a, Xc, A::AbstractAtlas, i0;
        solver=AutoVern9(Rodas5()),
        final_time::Real=1.0,
        check_chart_switch_kwargs=NamedTuple(),
        kwargs...,
    )

Solve geodesic ODE on a manifold `M` from point of coordinates `a` in chart `i0` from an
[`AbstractAtlas`](@ref) `A` in direction of coordinates `Xc` in the induced basis.
The geodesic is solved up to time `final_time` (by default equal to 1).

## Chart switching

If the solution exceeds the domain of chart `i0` (which is detected using the
`check_chart_switch` function with additional keyword arguments `check_chart_switch_kwargs`),
a new chart is selected using `get_chart_index` on the final point in the old chart.

## Returned value

The function returns an object of type `StitchedChartSolution{:Exp}` to represent the
geodesic.
"""
function solve_chart_exp_ode(
        M::AbstractManifold, a, Xc, A::AbstractAtlas, i0;
        solver = AutoVern9(Rodas5()),
        final_time::Real = 1.0,
        check_chart_switch_kwargs = NamedTuple(),
        kwargs...,
    )
    u0 = ArrayPartition(copy(a), copy(Xc))
    cur_i = i0
    # callback stops solver when we get too close to chart boundary
    cb = FunctionCallingCallback(
        IntegratorTerminatorNearChartBoundary(check_chart_switch_kwargs); func_start = false
    )
    retcode = SciMLBase.ReturnCode.Terminated
    init_time = zero(final_time)
    sols = StitchedChartSolution(M, A, :Exp, typeof(i0))
    while retcode === SciMLBase.ReturnCode.Terminated && init_time < final_time
        params = (M, A, cur_i)
        prob =
            ODEProblem(chart_exp_problem, u0, (init_time, final_time), params; callback = cb)
        sol = solve(prob, solver; kwargs...)
        retcode = sol.retcode
        init_time = sol.t[end]::typeof(final_time)
        push!(sols.sols, (sol, cur_i))
        # here we switch charts
        a_final = sol.u[end].x[1]::typeof(a)
        new_i = get_chart_index(M, A, cur_i, a_final)
        if new_i !== cur_i
            transition_map!(M, u0.x[1], A, cur_i, new_i, a_final)
            transition_map_diff!(
                M, u0.x[2], A, cur_i, a_final, sol.u[end].x[2]::typeof(Xc), new_i
            )
            cur_i = new_i
        end
    end
    return sols
end

function chart_pt_problem(u, params, t)
    M, A, i = params
    a = u.x[1]
    dx = u.x[2]
    dY = u.x[3]

    ddx = -affine_connection(M, A, i, a, dx, dx)
    ddY = -affine_connection(M, A, i, a, dx, dY)
    return ArrayPartition(dx, ddx, ddY)
end

"""
    solve_chart_parallel_transport_ode(
        M::AbstractManifold, a, Xc, A::AbstractAtlas, i0, Yc;
        solver=AutoVern9(Rodas5()), check_chart_switch_kwargs=NamedTuple(), final_time=1.0,
        kwargs...
    )

Parallel transport vector with coordinates `Yc` along geodesic on a manifold `M` from point of
coordinates `a` in a chart `i0` from an [`AbstractAtlas`](@ref) `A` in direction of
coordinates `Xc` in the induced basis.
"""
function solve_chart_parallel_transport_ode(
        M::AbstractManifold, a, Xc, A::AbstractAtlas, i0, Yc;
        solver = AutoVern9(Rodas5()), final_time = 1.0, check_chart_switch_kwargs = NamedTuple(),
        kwargs...
    )
    u0 = ArrayPartition(copy(a), copy(Xc), copy(Yc))
    cur_i = i0
    # callback stops solver when we get too close to chart boundary
    cb = FunctionCallingCallback(
        IntegratorTerminatorNearChartBoundary(check_chart_switch_kwargs); func_start = false
    )
    retcode = SciMLBase.ReturnCode.Terminated
    init_time = zero(final_time)
    sols = StitchedChartSolution(M, A, :PT, typeof(i0))
    while retcode === SciMLBase.ReturnCode.Terminated && init_time < final_time
        params = (M, A, cur_i)
        prob =
            ODEProblem(chart_pt_problem, u0, (init_time, final_time), params; callback = cb)
        sol = solve(prob, solver; kwargs...)
        retcode = sol.retcode
        init_time = sol.t[end]::typeof(final_time)
        push!(sols.sols, (sol, cur_i))
        # here we switch charts
        a_final = sol.u[end].x[1]::typeof(a)
        new_i = get_chart_index(M, A, cur_i, a_final)
        transition_map!(M, u0.x[1], A, cur_i, new_i, a_final)
        transition_map_diff!(
            M, u0.x[2], A, cur_i, a_final, sol.u[end].x[2]::typeof(Xc), new_i
        )
        transition_map_diff!(
            M, u0.x[3], A, cur_i, a_final, sol.u[end].x[3]::typeof(Yc), new_i
        )
        cur_i = new_i
    end
    return sols
end

end
