"""Secular circular orbital balances, separate from frozen-potential stream paths.

The independent numerical state is (orbital angular momentum, escaped angular
momentum). Specific angular momenta follow the actual orbital lever arms.
Orbital energy is a diagnostic, not a conserved total stellar energy.
"""
from __future__ import annotations

import numpy as np
from scipy.integrate import solve_ivp


def orbital_angular_momentum(donor, accretor, separation):
    return donor*accretor*np.sqrt(separation/(donor+accretor))


def integrate_transfer(donor, accretor, separation, removed, retained, channel, gamma=1,
                       wind_fraction=0, reemission_fraction=0, ring_fraction=0, ring_gamma=2):
    initial_angular = orbital_angular_momentum(donor, accretor, separation)
    def right_hand_side(stripped, state):
        md = donor-stripped
        ma = accretor+retained*stripped
        total = md+ma
        if channel == "constant-gamma":
            relative_loss = (1-retained)*gamma/total
        elif channel == "donor-wind":
            relative_loss = (1-retained)*ma/(md*total)
        elif channel == "reemission":
            relative_loss = (1-retained)*md/(ma*total)
        elif channel == "mesa-mixture":
            relative_loss = (wind_fraction*ma/(md*total)
                             + reemission_fraction*md/(ma*total)
                             + ring_fraction*ring_gamma*total/(md*ma))
        else:
            raise ValueError(channel)
        loss = state[0]*relative_loss
        return [-loss, loss]
    output = solve_ivp(right_hand_side, [0, removed], [initial_angular, 0],
                       method="DOP853", rtol=2e-12, atol=2e-14, dense_output=True)
    if not output.success:
        raise RuntimeError(output.message)
    stripped = np.linspace(0, removed, 101)
    md = donor-stripped
    ma = accretor+retained*stripped
    angular, escaped_angular = output.sol(stripped)
    distances = angular**2*(md+ma)/(md*ma)**2
    return {"stripped": stripped, "donor": md, "accretor": ma,
            "separation": distances, "angular": angular,
            "escaped_angular": escaped_angular, "initial_angular": initial_angular,
            "steps": len(output.t)}


def constant_gamma_closed(donor, accretor, separation, stripped, retained, gamma):
    md, ma = donor-stripped, accretor+retained*stripped
    return separation*((md+ma)/(donor+accretor))**(2*gamma+1)*(donor*accretor/(md*ma))**2


def native_mesa_closed(donor, accretor, separation, stripped, wind, reemission, ring, ring_gamma):
    """Published mixed-loss algebra used by MESA25.12.1 jdot_ml_check.

    This evaluates its mathematical reference, not the MESA stellar solver.
    MESA alpha/beta/delta are loss fractions; our retained fraction is epsilon.
    """
    retained = 1-wind-reemission-ring
    if not 0 < retained < 1:
        raise ValueError("Native closed form is singular at retained 0 or 1; test limits separately")
    q0 = donor/accretor
    q = (donor-stripped)/(accretor+retained*stripped)
    logarithm = (2*(wind+ring_gamma*ring-1)*np.log(q/q0)
                 + (-wind-reemission+ring)/(1-retained)*np.log((q+1)/(q0+1))
                 + (3+2*(wind*retained**2+reemission+ring_gamma*ring*(1-retained)**2)
                    /(retained*(1-retained)))*np.log((retained*q+1)/(retained*q0+1)))
    return separation*np.exp(logarithm)