_promote_disc_for_transfer_functions(d::AbstractAccretionDisc) = (d, d)
function _promote_disc_for_transfer_functions(::ThinDisc{T}) where {T}
    plane = DatumPlane(zero(T))
    plane, plane
end

struct _TransferFunctionSetup{T,A}
    h::T
    Œ∏_offset::T
    "Tolerance for root finding"
    zero_atol::T
    "Unstable radius with respect to when finer tolerances are needed for the Jacobian evaluation"
    unstable_radius::T
    Œ±‚ÇÄ::T
    Œ≤‚ÇÄ::T
    N::Int
    N_extrema::Int
    root_solver::A
end

function _TransferFunctionSetup(
    m::AbstractMetric{T},
    d::AbstractAccretionGeometry;
    Œ∏_offset = T(0.6),
    zero_atol = T(1e-7),
    N = 80,
    N_extrema = 17,
    Œ±‚ÇÄ = 0,
    Œ≤‚ÇÄ = 0,
    h = T(1e-6),
    root_solver = nothing,
    kwargs...,
) where {T}
    # specialize the algorithm depending on whether we are calculating for a thin or thick disc
    _alg = if isnothing(root_solver)
        if d isa AbstractThickAccretionDisc
            RootsAlg()
        else
            NonLinearAlg()
        end
    else
        root_solver
    end
    setup = _TransferFunctionSetup{T,typeof(_alg)}(
        h,
        Œ∏_offset,
        zero_atol,
        Gradus.isco(m) + 1,
        convert(T, Œ±‚ÇÄ),
        convert(T, Œ≤‚ÇÄ),
        N,
        N_extrema,
        _alg,
    )
    kwargs, setup
end

_is_unstable(setup::_TransferFunctionSetup, d::AbstractAccretionDisc, r) =
    (d isa AbstractThickAccretionDisc) && (r < setup.unstable_radius)

_calculate_transfer_function(r‚Çë, g, g‚ú∂, J) = @. (1 / (œÄ * r‚Çë)) * g * ‚àö(g‚ú∂ * (1 - g‚ú∂)) * J

function _adjust_extrema!(arr::AbstractArray{T}, min = zero(T), max = one(T)) where {T}
    arr[1] = min
    arr[end] = max
end

function _make_sorted_with_adjustments!(g1, f1, t1, g2, f2, t2; h = 1e-6)
    I1 = sortperm(g1)
    I2 = sortperm(g2)

    Base.permute!(g1, I1)
    Base.permute!(f1, I1)
    Base.permute!(t1, I1)

    Base.permute!(g2, I2)
    Base.permute!(f2, I2)
    Base.permute!(t2, I2)

    # at low inclination angles, the range of g is so tight that the edges are
    # very problematic due to the g‚ú∂ * (1 - g‚ú∂) terms, sending the branches to zero
    # so we use an offset to avoid these that is small enough to not impact results at
    # high inclination angles
    J1 = @. (g1 < 1 - h) & (g1 > h)
    J2 = @. (g2 < 1 - h) & (g2 > h)
    g1 = g1[J1]
    g2 = g2[J2]
    f1 = f1[J1]
    f2 = f2[J2]

    # save endpoints in time, since these will be pretty accurate
    t_lo = (t1[1] + t2[1]) / 2
    t_hi = (t1[end] + t2[end]) / 2
    # mask
    t1 = t1[J1]
    t2 = t2[J2]
    # restore endpoints
    _adjust_extrema!(t1, t_lo, t_hi)
    _adjust_extrema!(t2, t_lo, t_hi)
    _adjust_extrema!(g1)
    _adjust_extrema!(g2)
    g1, f1, t1, g2, f2, t2
end

function splitbranches(ctf::CunninghamTransferData{T}) where {T}
    # first things first we want to pull out the extremal
    # as these values are mutual to both branches, and cap off the
    # extrema. this way we avoid any accidental linear extrapolation
    # which might occur if N is small

    _, imin = findmin(ctf.g‚ú∂)
    _, imax = findmax(ctf.g‚ú∂)
    i1, i2 = imax > imin ? (imin, imax) : (imax, imin)

    if (i1 == i2)
        error("Resolved same min/max for r‚Çë = $(ctf.r‚Çë)")
    end

    # branch sizes, with duplicate extrema
    N1 = i2 - i1 + 1
    N2 = length(ctf.f) - N1 + 2
    # allocate branches
    #¬†todo: do this all at once with a matrix
    branch1_f = zeros(T, N1)
    branch2_f = zeros(T, N2)
    branch1_t = zeros(T, N1)
    branch2_t = zeros(T, N2)
    branch1_g‚ú∂ = zeros(T, N1)
    branch2_g‚ú∂ = zeros(T, N2)

    for (i, j) in enumerate(i1:i2)
        branch1_f[i] = ctf.f[j]
        branch1_g‚ú∂[i] = ctf.g‚ú∂[j]
        branch1_t[i] = ctf.t[j]
    end
    for (i, j) in enumerate(Iterators.flatten((1:i1, i2:length(ctf.f))))
        branch2_f[i] = ctf.f[j]
        branch2_g‚ú∂[i] = ctf.g‚ú∂[j]
        branch2_t[i] = ctf.t[j]
    end

    #¬†determine which is the upper branch
    # return:  (lower_g‚ú∂, lower_f, upper_g‚ú∂, upper_f)
    if branch1_f[2] > branch1_f[1]
        branch2_g‚ú∂, branch2_f, branch2_t, branch1_g‚ú∂, branch1_f, branch1_t
    else
        branch1_g‚ú∂, branch1_f, branch1_t, branch2_g‚ú∂, branch2_f, branch2_t
    end
end

function interpolate_branches(ctf::CunninghamTransferData{T}; h = 1e-6) where {T}
    (lower_g‚ú∂, lower_f, lower_t, upper_g‚ú∂, upper_f, upper_t) = splitbranches(ctf)
    (lower_g‚ú∂, lower_f, lower_t, upper_g‚ú∂, upper_f, upper_t) =
        _make_sorted_with_adjustments!(
            lower_g‚ú∂,
            lower_f,
            lower_t,
            upper_g‚ú∂,
            upper_f,
            upper_t;
            h = h,
        )
    lower_branch = _make_interpolation(lower_g‚ú∂, lower_f)
    lower_time_branch = _make_interpolation(lower_g‚ú∂, lower_t)
    upper_branch = _make_interpolation(upper_g‚ú∂, upper_f)
    upper_time_branch = _make_interpolation(upper_g‚ú∂, upper_t)
    TransferBranches{false}(
        upper_branch,
        lower_branch,
        upper_time_branch,
        lower_time_branch,
        ctf.gmin,
        ctf.gmax,
        ctf.r‚Çë,
    )
end

function _setup_workhorse_jacobian_with_kwargs(
    setup::_TransferFunctionSetup,
    m::AbstractMetric,
    x,
    d::AbstractAccretionDisc,
    r‚Çë;
    max_time = 2 * x[2],
    offset_max = 0.4r‚Çë + 10,
    redshift_pf = ConstPointFunctions.redshift(m, x),
    jacobian_disc = d,
    tracer_kwargs...,
)
    # underscores to avoid boxing variables
    function _jacobian_function(Œ±_, Œ≤_)
        jacobian_‚àÇŒ±Œ≤_‚àÇgr(
            m,
            x,
            jacobian_disc,
            Œ±_,
            Œ≤_,
            max_time;
            redshift_pf = redshift_pf,
            tracer_kwargs...,
        )
    end
    function _workhorse(Œ∏)
        r, gp = find_offset_for_radius(
            m,
            x,
            d,
            r‚Çë,
            Œ∏;
            zero_atol = setup.zero_atol,
            offset_max = offset_max,
            max_time = max_time,
            root_solver = setup.root_solver,
            Œ≤‚ÇÄ = setup.Œ≤‚ÇÄ,
            Œ±‚ÇÄ = setup.Œ±‚ÇÄ,
            tracer_kwargs...,
        )
        if isnan(r)
            error(
                "Transfer function integration failed (r‚Çë=$r‚Çë, Œ∏=$Œ∏, offset_max = $offset_max).",
            )
        end
        g = redshift_pf(m, gp, max_time)
        (g, gp, r)
    end
    (_workhorse, _jacobian_function, tracer_kwargs)
end
function _rear_workhorse(
    setup::_TransferFunctionSetup,
    m::AbstractMetric,
    x,
    d::AbstractAccretionDisc,
    r‚Çë;
    kwargs...,
)
    disc, jacobian_disc = _promote_disc_for_transfer_functions(d)
    workhorse, jacobian_function, _ = _setup_workhorse_jacobian_with_kwargs(
        setup,
        m,
        x,
        disc,
        r‚Çë;
        jacobian_disc = jacobian_disc,
        kwargs...,
    )
    function _disc_workhorse(Œ∏::T)::NTuple{3,T} where {T}
        g, gp, r = workhorse(Œ∏)
        Œ±, Œ≤ = _rŒ∏_to_Œ±Œ≤(r, Œ∏; Œ±‚ÇÄ = setup.Œ±‚ÇÄ, Œ≤‚ÇÄ = setup.Œ≤‚ÇÄ)
        g, jacobian_function(Œ±, Œ≤), gp.x[1]
    end
end

function _rear_workhorse(
    setup::_TransferFunctionSetup,
    m::AbstractMetric,
    x,
    d::AbstractThickAccretionDisc,
    r‚Çë;
    max_time = 2 * x[2],
    offset_max = 0.4r‚Çë + 10,
    kwargs...,
)
    plane = datumplane(d, r‚Çë)
    datum_workhorse, jacobian_function, tracer_kwargs =
        _setup_workhorse_jacobian_with_kwargs(
            setup,
            m,
            x,
            plane,
            r‚Çë;
            max_time = max_time,
            offset_max = offset_max,
            jacobian_disc = d,
            kwargs...,
        )
    function _thick_workhorse(Œ∏::T)::NTuple{4,T} where {T}
        g, gp, r = datum_workhorse(Œ∏)
        r‚Çä = try
            r_thick, _ = _find_offset_for_radius(
                m,
                x,
                d,
                r‚Çë,
                Œ∏;
                initial_r = r,
                zero_atol = setup.zero_atol,
                root_solver = setup.root_solver,
                offset_max = offset_max,
                max_time = max_time,
                Œ≤‚ÇÄ = setup.Œ≤‚ÇÄ,
                Œ±‚ÇÄ = setup.Œ±‚ÇÄ,
                tracer_kwargs...,
                # don't echo warnings
                warn = false,
            )
            r_thick
        catch
            # if we fail, for whatever reason, to root solve on the thick discs, 
            # we don't care, we just need a NaN value and then set that point to 
            # "not visible"
            NaN
        end
        is_visible, J = if !isnan(r‚Çä) && isapprox(r, r‚Çä, atol = 1e-3)
            # trace jacobian on updated impact parameters
            Œ±, Œ≤ = _rŒ∏_to_Œ±Œ≤(r‚Çä, Œ∏; Œ±‚ÇÄ = setup.Œ±‚ÇÄ, Œ≤‚ÇÄ = setup.Œ≤‚ÇÄ)
            _J = jacobian_function(Œ±, Œ≤)
            isfinite(_J), _J
        else
            false, NaN
        end
        (g, J, gp.x[1], is_visible)
    end
end

function _cunningham_transfer_function!(
    data::_TransferDataAccumulator,
    setup::_TransferFunctionSetup,
    workhorse,
    Œ∏iterator,
    r‚Çë,
)
    Œ∏_offset = setup.Œ∏_offset
    for (i, Œ∏) in enumerate(Œ∏iterator)
        Œ∏_corrected = Œ∏ + 1e-4
        insert_data!(data, i, Œ∏_corrected, workhorse(Œ∏))
    end

    gmin_candidate, gmax_candidate = _search_extremal!(data, workhorse, Œ∏_offset)
    gmin, gmax = _check_gmin_gmax(gmin_candidate, gmax_candidate, r‚Çë, data.gs)
    # we might not have used all of the memory we allocated so let's clean up
    remove_unused_elements!(data)
    # sort everything by angle
    sort!(data)

    # convert from ‚àÇg to ‚àÇg‚ú∂
    @. data.Js = (gmax - gmin) * data.Js

    @inbounds for i in eachindex(data)
        g‚ú∂ = g_to_g‚ú∂(data.gs[i], gmin, gmax)
        # Js is now storing f
        data.Js[i] = _calculate_transfer_function(r‚Çë, data.gs[i], g‚ú∂, data.Js[i])
        # gs is now storing g‚ú∂
        data.gs[i] = g‚ú∂
    end

    gmin, gmax
end

function cunningham_transfer_function(
    m::AbstractMetric{Q},
    x,
    d::AbstractAccretionDisc,
    r‚Çë::T;
    kwargs...,
) where {Q,T}
    solver_kwargs, setup = _TransferFunctionSetup(m, d; kwargs...)
    cunningham_transfer_function(setup, m, x, d, r‚Çë; solver_kwargs...)
end

function cunningham_transfer_function(
    setup::_TransferFunctionSetup,
    m::AbstractMetric,
    x,
    d::AbstractAccretionDisc,
    r‚Çë::T,
    ;
    chart = chart_for_metric(m, 10 * x[2]),
    max_time = 10 * x[2],
    solver_kwargs...,
) where {T}
    M = setup.N + 2 * setup.N_extrema
    K = setup.N √∑ 5
    data = _TransferDataAccumulator(T, M, setup.N)
    # sample so that the expected minima and maxima (0 and œÄ)
    Œ∏iterator = Iterators.flatten((
        range(0 - 2setup.Œ∏_offset, 0 + 2setup.Œ∏_offset, K),
        range(-œÄ / 2, 3œÄ / 2, setup.N - 2 * K),
        range(œÄ - 2setup.Œ∏_offset, œÄ + 2setup.Œ∏_offset, K),
    ))
    workhorse = _rear_workhorse(
        setup,
        m,
        x,
        d,
        r‚Çë;
        max_time = max_time,
        chart = chart,
        solver_kwargs...,
    )
    gmin, gmax = _cunningham_transfer_function!(data, setup, workhorse, Œ∏iterator, r‚Çë)
    CunninghamTransferData(
        data.data[2, :],
        data.data[3, :],
        data.data[4, :],
        gmin,
        gmax,
        r‚Çë,
    )
end

# find gmin and gmax via GoldenSection
# storing all attempts in gs and Js
function _search_extremal!(data::_TransferDataAccumulator, workhorse, offset)
    # need to specify the type to avoid boxing
    i::Int = data.cutoff
    N = (lastindex(data) - data.cutoff) √∑ 2 - 1

    function _gmin_finder(Œ∏)
        if i >= lastindex(data)
            error("i >= lastindex(data): $i >= $(lastindex(data))")
        end
        i += 1
        # avoid poles
        if abs(Œ∏) < 1e-4 || abs(abs(Œ∏) - œÄ) < 1e-4
            Œ∏ += 1e-4
        end
        insert_data!(data, i, Œ∏, workhorse(Œ∏))
        data.gs[i]
    end
    function _gmax_finder(Œ∏)
        -_gmin_finder(Œ∏)
    end

    # stride either side of our best guess so far
    res_min = Optim.optimize(
        _gmin_finder,
        0 - offset,
        0 + offset,
        GoldenSection(),
        iterations = N,
    )
    res_max = Optim.optimize(
        _gmax_finder,
        œÄ - offset,
        œÄ + offset,
        GoldenSection(),
        iterations = N,
    )

    # unpack result, remembering that maximum is inverted
    Optim.minimum(res_min), -Optim.minimum(res_max)
end

function interpolated_transfer_branches(
    m::AbstractMetric,
    x,
    d::AbstractAccretionDisc,
    radii;
    kwargs...,
)
    solver_kwargs, setup = _TransferFunctionSetup(m, d; kwargs...)
    interpolated_transfer_branches(setup, m, x, d, radii; solver_kwargs...)
end

function interpolated_transfer_branches(
    setup::_TransferFunctionSetup,
    m::AbstractMetric{T},
    x,
    d,
    radii;
    verbose = false,
    solver_kwargs...,
) where {T}
    progress_bar = init_progress_bar("Transfer functions:", length(radii), verbose)
    # IILF for calculating the interpolated branches
    ùîâ =
        r‚Çë -> begin
            ctf = cunningham_transfer_function(
                setup,
                m,
                x,
                d,
                r‚Çë,
                ;
                verbose = verbose,
                solver_kwargs...,
            )
            itp = interpolate_branches(ctf; h = setup.h)
            ProgressMeter.next!(progress_bar)
            itp
        end
    # calculate interpolated transfer functions for each emission radius
    InterpolatingTransferBranches(_threaded_map(ùîâ, radii))
end

function transfer_function_grid(
    m::AbstractMetric,
    x,
    d::AbstractAccretionGeometry,
    radii;
    Ng = 20,
    kwargs...,
)
    itfs = interpolated_transfer_branches(m, x, d, radii; kwargs...)
    transfer_function_grid(itfs, Ng)
end

function transfer_function_grid(itfs::InterpolatingTransferBranches, Ng::Int)
    g‚ú∂_grid = collect(range(0, 1, Ng))
    r_grid = itfs.radii

    lower_f = [branch.lower_f(g) for g in g‚ú∂_grid, branch in itfs.branches]
    upper_f = [branch.upper_f(g) for g in g‚ú∂_grid, branch in itfs.branches]
    lower_t = [branch.lower_t(g) for g in g‚ú∂_grid, branch in itfs.branches]
    upper_t = [branch.upper_t(g) for g in g‚ú∂_grid, branch in itfs.branches]

    CunninghamTransferGrid(
        r_grid,
        g‚ú∂_grid,
        itfs.gmin,
        itfs.gmax,
        lower_f,
        upper_f,
        lower_t,
        upper_t,
    )
end

function make_transfer_function_table(
    M::Type{<:KerrMetric},
    d::AbstractAccretionDisc,
    a_range::AbstractVector,
    Œ∏_range::AbstractVector;
    verbose = true,
    r_max = 500.0,
    n_radii = 150,
    kwargs...,
)
    function _mapper(a, Œ∏)
        if verbose
            @info "a, Œ∏ = $((a, Œ∏))"
        end
        m = M(1.0, a)
        x = SVector(0.0, 10000.0, deg2rad(Œ∏), 0.0)

        radii = Gradus.Grids._inverse_grid(Gradus.isco(m) + 1e-2, r_max, n_radii)
        Gradus.transfer_function_grid(m, x, d, radii; verbose = verbose, kwargs...)
    end
    grids = [_mapper(a, Œ∏) for a in a_range, Œ∏ in Œ∏_range]
    Gradus.CunninghamTransferTable((a_range, Œ∏_range), grids)
end

export CunninghamTransferData,
    TransferBranches,
    InterpolatingTransferBranches,
    splitbranches,
    interpolate_branches,
    cunningham_transfer_function,
    make_transfer_function_table
