module CircularOrbits
import ..StaticArrays: SVector, @SVector
import ..Gradus:
    AbstractStaticAxisSymmetric,
    metric_components,
    metric_jacobian,
    inverse_metric_components,
    MuladdMacro

MuladdMacro.@muladd begin
    @inline function _Î©_analytic(âˆ‚rg, contra_rotating)
        Î” = âˆš(âˆ‚rg[5]^2 - âˆ‚rg[1] * âˆ‚rg[4])
        if contra_rotating
            -(âˆ‚rg[5] + Î”) / âˆ‚rg[4]
        else
            -(âˆ‚rg[5] - Î”) / âˆ‚rg[4]
        end
    end

    function Î©(m::AbstractStaticAxisSymmetric, rÎ¸; contra_rotating = false)
        _, jacs = metric_jacobian(m, rÎ¸)
        âˆ‚rg = jacs[:, 1]
        _Î©_analytic(âˆ‚rg, contra_rotating)
    end

    function ut_uÏ•(ğ�›º::Number, ginv)
        ğ�’œ = -(ğ�›º * ginv[1] - ginv[5])
        â„¬ = (ğ�›º * ginv[5] - ginv[4])

        denom = â„¬^2 * ginv[1] + 2 * ğ�’œ * â„¬ * ginv[5] + ğ�’œ^2 * ginv[4]
        d = -sign(denom) * sqrt(inv(abs(denom)))
        ut = â„¬ * d
        uÏ• = ğ�’œ * d

        SVector(ut, uÏ•)
    end

    function ut_uÏ•(
        m::AbstractStaticAxisSymmetric,
        rÎ¸,
        ginv = inverse_metric_components(m, rÎ¸);
        kwargs...,
    )
        ut_uÏ•(Î©(m, rÎ¸; kwargs...), ginv)
    end

    # these 4 functions can be overwritten for a specific
    # metric, e.g. Kerr-Newman
    function energy(::AbstractStaticAxisSymmetric, rÎ¸, utuÏ•)
        -utuÏ•[1]
    end
    function angmom(::AbstractStaticAxisSymmetric, rÎ¸, utuÏ•)
        utuÏ•[2]
    end
    # this component doesn't actually seem to correctly constrain the geodesic
    # to being light- / null-like, or timelike. maybe revist? else call constrain before returning
    vt(::AbstractStaticAxisSymmetric, rÎ¸, ginv, utuÏ•) =
        ginv[1] * utuÏ•[1] + ginv[5] * utuÏ•[2]
    vÏ•(::AbstractStaticAxisSymmetric, rÎ¸, ginv, utuÏ•) =
        ginv[5] * utuÏ•[1] + ginv[4] * utuÏ•[2]

    # dispatch helpers
    energy(
        m::AbstractStaticAxisSymmetric,
        rÎ¸::SVector;
        contra_rotating = false,
        kwargs...,
    ) = energy(m, rÎ¸, ut_uÏ•(m, rÎ¸; contra_rotating = contra_rotating, kwargs...); kwargs...)
    angmom(
        m::AbstractStaticAxisSymmetric,
        rÎ¸::SVector;
        contra_rotating = false,
        kwargs...,
    ) = angmom(m, rÎ¸, ut_uÏ•(m, rÎ¸; contra_rotating = contra_rotating, kwargs...); kwargs...)
    energy(m::AbstractStaticAxisSymmetric, r::Number; kwargs...) =
        energy(m, SVector(r, Ï€ / 2); kwargs...)
    angmom(m::AbstractStaticAxisSymmetric, r::Number; kwargs...) =
        angmom(m, SVector(r, Ï€ / 2); kwargs...)

    function energy_angmom(m::AbstractStaticAxisSymmetric, rÎ¸::SVector; kwargs...)
        utuÏ• = ut_uÏ•(m, rÎ¸; kwargs...)
        energy(m, rÎ¸, utuÏ•; kwargs...), angmom(m, rÎ¸, utuÏ•; kwargs...)
    end
    energy_angmom(m::AbstractStaticAxisSymmetric, r::Number; kwargs...) =
        energy_angmom(m, SVector(r, Ï€ / 2); kwargs...)

    function vt(
        m::AbstractStaticAxisSymmetric,
        rÎ¸::SVector;
        contra_rotating = false,
        kwargs...,
    )
        ginv = inverse_metric_components(m, rÎ¸)
        utuÏ• = ut_uÏ•(m, rÎ¸, ginv; contra_rotating = contra_rotating, kwargs...)
        vt(m, ginv, rÎ¸, utuÏ•)
    end
    vt(m::AbstractStaticAxisSymmetric, r::Number; kwargs...) =
        vt(m, SVector(r, Ï€ / 2); kwargs...)

    function vÏ•(
        m::AbstractStaticAxisSymmetric,
        rÎ¸::SVector;
        contra_rotating = false,
        kwargs...,
    )
        ginv = inverse_metric_components(m, rÎ¸)
        utuÏ• = ut_uÏ•(m, rÎ¸, ginv; contra_rotating = contra_rotating, kwargs...)
        vÏ•(m, rÎ¸, ginv, utuÏ•)
    end
    vÏ•(m::AbstractStaticAxisSymmetric, r::Number; kwargs...) =
        vÏ•(m, SVector(r, Ï€ / 2); kwargs...)

    function fourvelocity(m::AbstractStaticAxisSymmetric, rÎ¸; kwargs...)
        ginv = inverse_metric_components(m, rÎ¸)
        utuÏ• = ut_uÏ•(m, rÎ¸, ginv; kwargs...)

        SVector(vt(m, rÎ¸, ginv, utuÏ•), 0, 0, vÏ•(m, rÎ¸, ginv, utuÏ•))
    end
    fourvelocity(m::AbstractStaticAxisSymmetric, r::Number; kwargs...) =
        fourvelocity(m, SVector(r, Ï€ / 2); kwargs...)
    fourvelocity(m::AbstractStaticAxisSymmetric, x::SVector{4}; kwargs...) =
        fourvelocity(m, SVector(x[2], x[3]); kwargs...)

    """
    NB: this function is only valid **at the ISCO**.
    """
    function plunging_fourvelocity(
        m::AbstractStaticAxisSymmetric,
        rÎ¸;
        contra_rotating = false,
        kwargs...,
    )
        g = metric_components(m, rÎ¸)
        ginv = inverse_metric_components(g)
        utuÏ• = ut_uÏ•(m, rÎ¸, ginv; contra_rotating = contra_rotating, kwargs...)
        E = energy(m, rÎ¸, utuÏ•; kwargs...)
        L = angmom(m, rÎ¸, utuÏ•; kwargs...)
        _vt = vt(m, rÎ¸, ginv, utuÏ•; kwargs...)
        _vÏ• = vÏ•(m, rÎ¸, ginv, utuÏ•; kwargs...)

        nom = ginv[1] * E^2 - 2ginv[5] * E * L + ginv[4] * L^2 + 1
        denom = -g[2]

        @SVector[_vt, -sqrt(abs(nom / denom)), 0.0, _vÏ•]
    end
    plunging_fourvelocity(m::AbstractStaticAxisSymmetric, r::Number; kwargs...) =
        plunging_fourvelocity(m, @SVector([r, Ï€ / 2]); kwargs...)
    plunging_fourvelocity(m::AbstractStaticAxisSymmetric, x::SVector{4}; kwargs...) =
        fourvelocity(m, SVector(x[2], x[3]); kwargs...)

end # mulladd macro
end # module

_keplerian_velocity_projector(m::AbstractMetric, ::AbstractAccretionGeometry) =
    _keplerian_velocity_projector(m)
function _keplerian_velocity_projector(m::AbstractMetric)
    r_isco = isco(m)
    interp = interpolate_plunging_velocities(m)

    function _keplerian_project(x::SVector{4})
        r = _equatorial_project(x)
        if r < r_isco
            vtemp = interp(r)
            SVector(vtemp[1], -vtemp[2], vtemp[3], vtemp[4])
        else
            CircularOrbits.fourvelocity(m, r)
        end
    end
end

export CircularOrbits
