"""
    ManifoldProjection(
        manifold; nlsolve = missing, save = true, autonomous = nothing,
        manifold_jacobian = nothing, autodiff = nothing,
        resid_prototype = nothing, kwargs...) -> DiscreteCallback

In many cases, you may want to declare a manifold on which a solution lives. Mathematically,
a manifold `M` is defined by a function `g` as the set of points where `g(u) = 0`. An
embedded manifold can be a lower dimensional object which constrains the solution. For
example, `g(u) = E(u) - C` where `E` is the energy of the system in state `u`, meaning that
the energy must be constant (energy preservation). Thus by defining the manifold the
solution should live on, you can retain desired properties of the solution.

`ManifoldProjection` projects the solution of the differential equation to the chosen
manifold `g`, conserving a property while conserving the order. It is a consequence of
convergence proofs both in the deterministic and stochastic cases that post-step projection
to manifolds keep the same convergence rate, thus any algorithm can be easily extended to
conserve properties. If the solution is supposed to live on a specific manifold or conserve
such property, this guarantees the conservation law without modifying the convergence
properties.

# Arguments

  - `manifold`: residual function defining the manifold. For an in-place problem, define
    `manifold(resid, u, p)` or `manifold(resid, u, p, t)`. For an out-of-place problem,
    define `manifold(u, p)` or `manifold(u, p, t)`. The residual is zero on the manifold.

# Keywords

  - `nlsolve = missing`: use the built-in single-factorization projection algorithm. Pass a
    nonlinear solver in the
    [NonlinearSolve.jl format](https://docs.sciml.ai/NonlinearSolve/stable/basics/solve/)
    to select another algorithm. Pass `nothing` to use the NonlinearSolve polyalgorithm.
  - `save::Bool = true`: save immediately after the projection.
  - `autonomous = nothing`: whether `manifold` omits the time argument. Pass `Val(true)`
    or `Val(false)` to avoid runtime branching; `nothing` infers the form during
    initialization.
  - `resid_prototype = nothing`: prototype defining the residual shape for an in-place
    problem. Without one, the residual is assumed to have the same shape as `u`.
  - `autodiff = nothing`: DifferentiationInterface automatic-differentiation backend used
    when `manifold_jacobian` is not supplied.
  - `manifold_jacobian = nothing`: analytic Jacobian of `manifold` with respect to `u`. Use
    the same calling form as `manifold`, with the Jacobian output as the first argument for
    an in-place problem.
  - `kwargs...`: additional keywords passed to the built-in projection algorithm or to
    [NonlinearSolve.jl](https://docs.sciml.ai/NonlinearSolve/stable/basics/solve/) if
    `nlsolve` is not `missing`.

# Returns

  - `DiscreteCallback`: a callback that projects the state after each accepted step. If the
    nonlinear projection does not converge, the callback terminates the integrator with the
    projection solver's unsuccessful return code.

# Throws

  - `ErrorException`: during callback initialization if both `manifold_jacobian` and
    `autodiff` are `nothing`.

# Saveat Warning

Note that the `ManifoldProjection` callback modifies the endpoints of the integration
intervals and thus breaks assumptions of internal interpolations. Because of this, the
values for given by saveat will not be order-matching. However, the interpolation error can
be proportional to the change by the projection, so if the projection is making small
changes then one is still safe. However, if there are large changes from each projection,
you should consider only saving at stopping/projection times. To do this, set `tstops` to
the same values as `saveat`. There is a performance hit by doing so because now the
integrator is forced to stop at every saving point, but this is guaranteed to match the
order of the integrator even with the `ManifoldProjection`.

# References

[1] Ernst Hairer, Christian Lubich, Gerhard Wanner. Geometric Numerical Integration:
Structure-Preserving Algorithms for Ordinary Differential Equations. Berlin ;
New York :Springer, 2002.

# Examples

```julia
using ADTypes, DiffEqCallbacks, OrdinaryDiffEq

function unit_circle(resid, u, p, t)
    resid[1] = sum(abs2, u) - 1
end

function rotation!(du, u, p, t)
    du[1] = -u[2]
    du[2] = u[1]
end

prob = ODEProblem(rotation!, [1.0, 0.0], (0.0, 10.0))
cb = ManifoldProjection(unit_circle; resid_prototype = [0.0], autodiff = AutoFiniteDiff())

sol = solve(prob, Tsit5(); callback = cb)
```
"""
@concrete mutable struct ManifoldProjection
    manifold
    manifold_jacobian
    autodiff
    nlcache::Any
    nlsolve
    kwargs
    autonomous
end

function ManifoldProjection(
        manifold; nlsolve = missing, save = true, autonomous = nothing,
        manifold_jacobian = nothing, autodiff = nothing, kwargs...
    )
    affect! = ManifoldProjection(
        manifold, autodiff, manifold_jacobian, nlsolve, kwargs, autonomous
    )
    return DiscreteCallback(
        Returns(true), affect!; initialize = initialize_manifold_projection,
        save_positions = (false, save)
    )
end

function ManifoldProjection(
        manifold, autodiff, manifold_jacobian, nlsolve, kwargs, autonomous
    )
    wrapped_manifold = wrap_autonomous_function(autonomous, manifold)
    wrapped_manifold_jacobian = wrap_autonomous_function(autonomous, manifold_jacobian)
    return ManifoldProjection(
        wrapped_manifold, wrapped_manifold_jacobian,
        autodiff, nothing, nlsolve, kwargs, autonomous
    )
end

function (proj::ManifoldProjection)(integrator)
    # update current time if residual function is time-dependent
    proj.manifold.t = integrator.t
    proj.manifold_jacobian !== nothing && (proj.manifold_jacobian.t = integrator.t)

    SciMLBase.reinit!(proj.nlcache, integrator.u; integrator.p)
    _, u, retcode = SciMLBase.solve!(proj.nlcache)

    if !SciMLBase.successful_retcode(retcode)
        SciMLBase.terminate!(integrator, retcode)
        return
    end

    return copyto!(integrator.u, u)
end

function initialize_manifold_projection(cb, u, t, integrator)
    return initialize_manifold_projection(cb.affect!, u, t, integrator)
end
function initialize_manifold_projection(affect!::ManifoldProjection, u, t, integrator)
    if affect!.autonomous === nothing
        autonomous = maximum(SciMLBase.numargs(affect!.manifold.f)) ==
            2 + SciMLBase.isinplace(integrator.f)
        affect!.manifold.autonomous = autonomous
        affect!.manifold_jacobian !== nothing &&
            (affect!.manifold_jacobian.autonomous = autonomous)
    end

    affect!.manifold.t = t
    affect!.manifold_jacobian !== nothing && (affect!.manifold_jacobian.t = t)

    if affect!.nlsolve === missing
        cache = init_manifold_projection(
            Val(SciMLBase.isinplace(integrator.f)), affect!.manifold, affect!.autodiff,
            affect!.manifold_jacobian, u, integrator.p; affect!.kwargs...
        )
    else
        cache = init_manifold_projection_nonlinear_problem(
            Val(SciMLBase.isinplace(integrator.f)), affect!.manifold, affect!.autodiff,
            affect!.manifold_jacobian, u, integrator.p, affect!.nlsolve; affect!.kwargs...
        )
    end
    affect!.nlcache = cache
    return derivative_discontinuity!(integrator, false)
end

export ManifoldProjection

# wrapper for non-autonomous functions
function wrap_autonomous_function(autonomous::Union{Val{true}, Val{false}}, g)
    g === nothing && return nothing
    return TypedNonAutonomousFunction{_unwrap_val(autonomous)}(g, nothing)
end
function wrap_autonomous_function(autonomous::Union{Bool, Nothing}, g)
    g === nothing && return nothing
    autonomous = autonomous === nothing ? false : autonomous
    return UntypedNonAutonomousFunction(autonomous, g, nothing)
end

abstract type AbstractNonAutonomousFunction end

@concrete mutable struct TypedNonAutonomousFunction{autonomous} <:
    AbstractNonAutonomousFunction
    f
    t::Any
end

(f::TypedNonAutonomousFunction{false})(res, u, p) = f.f(res, u, p, f.t)
(f::TypedNonAutonomousFunction{true})(res, u, p) = f.f(res, u, p)

(f::TypedNonAutonomousFunction{false})(u, p) = f.f(u, p, f.t)
(f::TypedNonAutonomousFunction{true})(u, p) = f.f(u, p)

@concrete mutable struct UntypedNonAutonomousFunction <: AbstractNonAutonomousFunction
    autonomous::Bool
    f
    t::Any
end

function (f::UntypedNonAutonomousFunction)(res, u, p)
    return f.autonomous ? f.f(res, u, p) : f.f(res, u, p, f.t)
end
(f::UntypedNonAutonomousFunction)(u, p) = f.autonomous ? f.f(u, p) : f.f(u, p, f.t)

# This is solving the langrange multiplier formulation. This is more accurate but at the
# same time significantly more expensive.
@concrete mutable struct NonlinearSolveManifoldProjectionCache{iip}
    manifold
    p
    λ
    z
    ũ
    gu_cache
    nlcache

    first_call::Bool
    J
    manifold_jacobian
    autodiff
    di_extras
end

function SciMLBase.reinit!(
        cache::NonlinearSolveManifoldProjectionCache{iip}, u; p = cache.p
    ) where {iip}
    if !cache.first_call || (cache.ũ !== u || cache.p !== p)
        compute_manifold_jacobian!(
            cache.J, cache.manifold_jacobian, cache.autodiff,
            Val(iip), cache.manifold, cache.gu_cache, u, p, cache.di_extras
        )
    end
    cache.first_call = false
    cache.ũ = u
    cache.p = p

    cache.z[1:length(cache.λ)] .= false
    cache.z[(length(cache.λ) + 1):end] .= vec(u)
    return SciMLBase.reinit!(cache.nlcache, cache.z; p = (u, cache.J, p))
end

function init_manifold_projection_nonlinear_problem(
        IIP::Val{iip}, manifold, autodiff, manifold_jacobian, ũ, p, alg;
        resid_prototype = nothing, kwargs...
    ) where {iip}
    if iip
        if resid_prototype !== nothing
            gu = similar(resid_prototype)
            λ = similar(resid_prototype)
        else
            @warn "`resid_prototype` not provided for in-place problem. Assuming size of \
                   output is the same as input. This might be incorrect." maxlog = 1
            gu = similar(ũ)
            λ = similar(ũ)
        end
    else
        gu = nothing
        λ = manifold(ũ, p)
    end

    J,
        di_extras = setup_manifold_jacobian(
        manifold_jacobian, autodiff, IIP, manifold,
        gu, ũ, p
    )
    z = vcat(vec(λ), vec(ũ))

    nlfunc = if iip
        let λlen = length(λ), λsz = size(λ), zsz = size(ũ)
            @views (
                resid, u,
                ps,
            ) -> begin
                ũ2, J2, p2 = ps
                λ2, z2 = u[1:λlen], u[(λlen + 1):end]
                manifold(reshape(resid[1:λlen], λsz), reshape(z2, zsz), p2)
                resid[(λlen + 1):end] .= z2 .- vec(ũ2) .+ vec(vec(J2' * λ2))
            end
        end
    else
        let λlen = length(λ), zsz = size(ũ)
            @views (u, ps) -> begin
                ũ2, J2, p2 = ps
                λ2, z2 = u[1:λlen], u[(λlen + 1):end]
                gz = vec(manifold(reshape(z2, zsz), p2))
                resid = z2 .- vec(ũ2) .+ vec(J2' * λ2)
                return vcat(gz, resid)
            end
        end
    end

    nlprob = NonlinearProblem(NonlinearFunction{iip}(nlfunc), z, (ũ, J, p))
    nlcache = SciMLBase.init(nlprob, alg; kwargs...)

    return NonlinearSolveManifoldProjectionCache{iip}(
        manifold, p, λ, z, ũ, gu, nlcache, true, J, manifold_jacobian, autodiff, di_extras
    )
end

@views function SciMLBase.solve!(cache::NonlinearSolveManifoldProjectionCache{iip}) where {iip}
    sol = SciMLBase.solve!(cache.nlcache)
    (; u, retcode) = sol
    λ = reshape(u[1:length(cache.λ)], size(cache.λ))
    ũ = reshape(u[(length(cache.λ) + 1):end], size(cache.ũ))
    return λ, ũ, retcode
end

# This is the algorithm described in Hairer III.
@concrete mutable struct SingleFactorizeManifoldProjectionCache{iip}
    manifold
    p
    abstol
    maxiters::Int

    ũ
    JJᵀfact::Any  # LU might fail and we might end up doing QR
    u_cache
    λ_cache
    gu_cache

    first_call::Bool
    J
    JJᵀ
    manifold_jacobian
    autodiff
    di_extras
end

function SciMLBase.reinit!(
        cache::SingleFactorizeManifoldProjectionCache{iip}, u; p = cache.p
    ) where {iip}
    if !cache.first_call || (cache.ũ !== u || cache.p !== p)
        compute_manifold_jacobian!(
            cache.J, cache.manifold_jacobian, cache.autodiff,
            Val(iip), cache.manifold, cache.gu_cache, u, p, cache.di_extras
        )
        mul!(cache.JJᵀ, cache.J, cache.J')
        cache.JJᵀfact = safe_factorize!(cache.JJᵀ)
    end
    cache.first_call = false
    cache.ũ = u
    return cache.p = p
end

default_abstol(::Type{T}) where {T} = real(oneunit(T)) * (eps(real(one(T))))^(4 // 5)

function init_manifold_projection(
        IIP::Val{iip}, manifold, autodiff, manifold_jacobian, ũ,
        p; abstol = default_abstol(eltype(ũ)), maxiters = 1000,
        resid_prototype = nothing, kwargs...
    ) where {iip}
    if iip
        if resid_prototype !== nothing
            gu = similar(resid_prototype)
            λ = similar(resid_prototype)
        else
            @warn "`resid_prototype` not provided for in-place problem. Assuming size of \
                   output is the same as input. This might be incorrect." maxlog = 1
            gu = similar(ũ)
            λ = similar(ũ)
        end
    else
        gu = nothing
        λ = manifold(ũ, p)
    end

    J,
        di_extras = setup_manifold_jacobian(
        manifold_jacobian, autodiff, IIP, manifold,
        gu, ũ, p
    )
    JJᵀ = J * J'
    JJᵀfact = safe_factorize!(JJᵀ)

    return SingleFactorizeManifoldProjectionCache{iip}(
        manifold, p, abstol, maxiters, ũ, JJᵀfact, similar(ũ), λ, gu,
        true, J, JJᵀ, manifold_jacobian, autodiff, di_extras
    )
end

function SciMLBase.solve!(cache::SingleFactorizeManifoldProjectionCache{iip}) where {iip}
    fill!(cache.λ_cache, false)
    ũ = cache.ũ
    gu = cache.gu_cache

    copyto!(cache.u_cache, ũ)
    internal_solve_failed = true

    if cache.gu_cache !== nothing
        cache.manifold(gu, ũ, cache.p)
    else
        gu = cache.manifold(ũ, cache.p)
    end

    for _ in 1:(cache.maxiters)
        if maximum(abs, gu) < cache.abstol
            internal_solve_failed = false
            break
        end

        δλ = cache.JJᵀfact \ gu
        @. cache.λ_cache -= δλ

        mul!(vec(cache.u_cache), cache.J', vec(cache.λ_cache))
        cache.u_cache += ũ
        if cache.gu_cache !== nothing
            cache.manifold(gu, cache.u_cache, cache.p)
        else
            gu = cache.manifold(cache.u_cache, cache.p)
        end
    end

    return (
        cache.λ_cache, cache.u_cache,
        ifelse(internal_solve_failed, ReturnCode.ConvergenceFailure, ReturnCode.Success),
    )
end

function setup_manifold_jacobian(
        manifold_jacobian::M, autodiff, ::Val{iip}, manifold, gu, ũ, p
    ) where {M, iip}
    if iip
        J = similar(ũ, promote_type(eltype(gu), eltype(ũ)), (length(gu), length(ũ)))
        manifold_jacobian(J, ũ, p)
    else
        J = manifold_jacobian(ũ, p)
    end
    return J, nothing
end

function setup_manifold_jacobian(
        ::Nothing, autodiff, ::Val{iip}, manifold, gu, ũ, p
    ) where {iip}
    if iip
        di_extras = DI.prepare_jacobian(manifold, gu, autodiff, ũ, Constant(p))
        J = DI.jacobian(manifold, gu, di_extras, autodiff, ũ, Constant(p))
    else
        di_extras = DI.prepare_jacobian(manifold, autodiff, ũ, Constant(p))
        J = DI.jacobian(manifold, di_extras, autodiff, ũ, Constant(p))
    end
    return J, di_extras
end

function setup_manifold_jacobian(
        ::Nothing, ::Nothing, ::Val{iip}, manifold, gu, ũ, p
    ) where {iip}
    error("`autodiff` is set to `nothing` and analytic manifold jacobian is not provided.")
end

function compute_manifold_jacobian!(
        J, manifold_jacobian, autodiff, ::Val{iip},
        manifold, gu, ũ, p, di_extras
    ) where {iip}
    if iip
        manifold_jacobian(J, ũ, p)
    else
        J = manifold_jacobian(ũ, p)
    end
    return J
end

function compute_manifold_jacobian!(
        J, ::Nothing, autodiff, ::Val{iip}, manifold, gu,
        ũ, p, di_extras
    ) where {iip}
    if iip
        DI.jacobian!(manifold, gu, J, di_extras, autodiff, ũ, Constant(p))
    else
        DI.jacobian!(manifold, J, di_extras, autodiff, ũ, Constant(p))
    end
    return J
end

_issquare_matrix(A::AbstractMatrix) = size(A, 1) == size(A, 2)

function safe_factorize!(A::AbstractMatrix)
    if _issquare_matrix(A)
        fact = LinearAlgebra.cholesky(A; check = false)
        fact_successful(fact) && return fact
    elseif size(A, 1) > size(A, 2)
        fact = LinearAlgebra.qr(A)
        fact_successful(fact) && return fact
    end
    return LinearAlgebra.qr!(A, LinearAlgebra.ColumnNorm())
end

function fact_successful(F::LinearAlgebra.QRCompactWY)
    m, n = size(F)
    U = view(F.factors, 1:min(m, n), 1:n)
    return all(!iszero, Iterators.reverse(@view U[diagind(U)]))
end
function fact_successful(F::FT) where {FT}
    return hasmethod(LinearAlgebra.issuccess, (FT,)) ? LinearAlgebra.issuccess(F) : true
end
