module __PolishDoughnut

import Gradus:
    CircularOrbits,
    AbstractMetric,
    __BoyerLindquistAD,
    KerrMetric,
    inverse_metric_components,
    metric_components
import ForwardDiff, Roots
import OrdinaryDiffEq: solve, Tsit5, DiscreteCallback, ODEProblem, terminate!
import StaticArrays: SVector

# modified Fuerst & Wu (2004, 2007)
function Î©(m::AbstractMetric, x, râ‚–, n)
    rsinÎ¸ = x[1] * sin(x[2])
    xprime = SVector(rsinÎ¸, Ï€ / 2)
    CircularOrbits.Î©(m, xprime) * (râ‚– / (rsinÎ¸))^n
end

function orbital_energy(m::AbstractMetric, x, râ‚–, n)
    ð�›º = Î©(m, x, râ‚–, n)
    ut_uÏ• = CircularOrbits.ut_uÏ•(ð�›º, inverse_metric_components(m, x))
    CircularOrbits.energy(m, x, ut_uÏ•)
    Î©â‚€ = ð�›º
    g = metric_components(m, x)
    -(g[1] + g[5] * Î©â‚€) / âˆš(-g[1] - 2g[5] * Î©â‚€ - g[4] * Î©â‚€^2)
end

# Younsi et al. (2012): eq. (30)
Î¨â‚�(r, Î¸, M, a, invÎ©, Î£) = M * ((Î£ - 2r^2) / (Î£^2)) * (invÎ© - a * sin(Î¸))^2 + r * sin(Î¸)^2
Î¨â‚�(m, x, ð�›º, Î£) = Î¨â‚�(x[1], x[2], m.M, m.a, ð�›º, Î£)

# Younsi et al. (2012): eq. (31)
Î¨â‚‚(r, Î¸, M, a, invÎ©, Î£, Î”) =
    sin(2Î¸) * ((M * r / (Î£^2)) * (a * invÎ© - (r^2 + a^2))^2 + Î” / 2)
Î¨â‚‚(m, x, ð�›º, Î£, Î”) = Î¨â‚‚(x[1], x[2], m.M, m.a, ð�›º, Î£, Î”)

function isobar_differential(m::KerrMetric, x, invð�›º)
    Î£ = __BoyerLindquistAD.Î£(x[1], m.a, x[2])
    Î” = __BoyerLindquistAD.Î”(x[1], 2 * m.M, m.a)
    Ïˆâ‚� = Î¨â‚�(m, x, invð�›º, Î£)
    Ïˆâ‚‚ = Î¨â‚‚(m, x, invð�›º, Î£, Î”)

    d = inv(âˆš(Î” * Ïˆâ‚�^2 + Ïˆâ‚‚^2) * âˆšinv(Î” / Î£))
    dr = Ïˆâ‚‚ * d
    dÎ¸ = -Ïˆâ‚� * d

    SVector(dr, dÎ¸)
end

function innermost_radius(m, râ‚–, n; init_r = 5.0)
    # we want to find r for which dE/dr == 0
    f = r -> orbital_energy(m, SVector(r, Ï€ / 2), râ‚–, n)
    df = r -> ForwardDiff.derivative(f, r)
    d2f = r -> ForwardDiff.derivative(df, r)
    Roots.find_zero((df, d2f), init_r)
end

function isobar(
    m::AbstractMetric,
    inner_radius,
    râ‚–,
    n;
    Î»_max = 40.0,
    dtmax = 5e-2,
    solver_opts...,
)
    terminate_when_negative =
        DiscreteCallback((u, t, integrator) -> u[1] * cos(u[2]) < 0, terminate!)

    u0 = SVector(inner_radius, Ï€ / 2)

    function drdÎ¸(u, p, Î»)
        ð�›º = Î©(m, u, râ‚–, n)
        isobar_differential(m, u, inv(ð�›º))
    end

    prob = ODEProblem{false}(drdÎ¸, u0, (0.0, Î»_max))
    sol = solve(
        prob,
        Tsit5();
        callback = terminate_when_negative,
        dtmax = dtmax,
        solver_opts...,
    )

    # unpack solution
    r = first.(sol.u)
    z = @. cos(last(sol.u)) * r

    # positivity
    I = @. z > 0
    r[I], z[I]
end

end # module

struct PolishDoughnut{T,F} <: AbstractThickAccretionDisc{T}
    inner_radius::T
    outer_radius::T
    râ‚–::T
    n::T
    f::F
end

function PolishDoughnut(m::AbstractMetric; râ‚– = 12.0, n = 0.21, init_r = 5.0, kwargs...)
    inner_radius = __PolishDoughnut.innermost_radius(m, râ‚–, n; init_r = init_r)
    r, z = __PolishDoughnut.isobar(m, inner_radius, râ‚–, n; kwargs...)
    # interpolate cross section
    f = DataInterpolations.LinearInterpolation(z, r)
    outer_radius = maximum(r)
    PolishDoughnut(inner_radius, outer_radius, râ‚–, n, f)
end

function cross_section(d::PolishDoughnut, Ï�)
    if d.inner_radius â‰¤ Ï� â‰¤ d.outer_radius
        d.f(Ï�)
    else
        zero(typeof(Ï�))
    end
end

export PolishDoughnut
