zephyrus.hydrostatic

def volkov_flat_factor(lam: float) -> float:
    """Kinetic enhancement C(lambda) on the Jeans flux, dimensionless.

    Direct simulation Monte Carlo runs exceed the Jeans flux by a factor
    1.7 near lambda = 6, falling to 1.4 by lambda = 15 (Volkov et al.
    2011); linear between, held at the endpoint values outside, where the
    caller flags the extrapolation on either side. Below lambda = 6 the
    held value likely understates the flux, and trace light species on a
    heavy background reach that side.
    """
    if lam <= 6.0:
        return 1.7
    if lam >= 15.0:
        return 1.4
    return 1.7 + (1.4 - 1.7) * (lam - 6.0) / (15.0 - 6.0)

def jeans_effusion_velocity(T: float, m: float, lam: float) -> float:
    """Jeans effusion velocity sqrt(kT/(2 pi m)) (1 + lambda) exp(-lambda), m/s."""
    return math.sqrt(kb * T / (2.0 * math.pi * m)) * (1.0 + lam) * math.exp(-lam)

def bates_extension(
    profile,
    M_p: float,
    T_exo: float,
    gamma: float = 0.75,
    zeta_max: float = 40.0,
    n_levels: int = 400,
) -> dict:
    """Bates upper structure above the topmost profile level.

    The hydrostatic relation on the Bates temperature profile in
    ``zeta = ln(p_top / p)``, with composition and mean mass frozen at the
    anchor, integrates in closed form: ``d(1/r)/dzeta = -kB T / (G M mu)``
    and the Bates profile integrates analytically, so the structure carries
    no step error at any resolution. It ends, flagged ``unbound``, where the
    local Jeans parameter falls to 2: an isothermal-tail structure is
    unbound beyond that point and the geometry belongs to the overflow and
    boil-off machinery, not to a hydrostatic exosphere.

    Returns a dict of arrays on ``n_levels`` points spanning the extension
    (``zeta``, ``p`` [Pa], ``r`` [m], ``T`` [K], ``n`` [m^-3]), the end
    point ``zeta_end``, the frozen ``mu`` [kg], the normalized species
    ``vmr``, the anchor values ``p0`` and ``r0``, the Bates parameters
    (``T_exo``, ``T_top``, ``gamma``, ``M_p``), and ``unbound``.
    ``bates_state`` evaluates the same structure at any ``zeta``.
    """
    p0 = float(profile.p[-1])
    r0 = float(profile.r[-1])
    t_top = float(profile.T[-1])
    mu = float(profile.mmw[-1])
    # An exobase colder than the anchor would bind the exobase more strongly
    # than its anchor, inverting the construction. Floored and flagged, not
    # raised, since a coupled run's profile top can warm past the prescription.
    floored = T_exo < t_top
    if floored:
        T_exo = t_top
    # Every species at the anchor is carried, however thin: a trace light
    # species can dominate the exospheric loss.
    vmr = {
        sp: float(np.asarray(v)[-1])
        for sp, v in profile.vmr.items()
        if float(np.asarray(v)[-1]) > 0.0
    }
    tot = sum(vmr.values())
    if tot <= 0.0:
        raise ValueError('no species present at the profile top')
    vmr = {sp: x / tot for sp, x in vmr.items()}
    ext = dict(mu=mu, vmr=vmr, p0=p0, r0=r0, T_exo=T_exo, T_top=t_top, gamma=gamma, M_p=M_p)

    def excess_lambda(z: float) -> float:
        _r, t, _n, inv_r = _bates_point(ext, z)
        return G * M_p * mu * inv_r / (kb * t) - 2.0

    if excess_lambda(zeta_max) >= 0.0:
        zeta_end, unbound = zeta_max, False
    elif excess_lambda(0.0) <= 0.0:
        zeta_end, unbound = 0.0, True
    else:
        # The Jeans parameter falls monotonically with zeta (1/r falls and
        # T does not), so the crossing is unique.
        zeta_end = brentq(excess_lambda, 0.0, zeta_max, xtol=1e-12, rtol=1e-12)
        unbound = True
    zeta = np.linspace(0.0, zeta_end, n_levels) if zeta_end > 0.0 else np.zeros(1)
    r, T, n, _inv = bates_state(ext, zeta)
    ext.update(
        zeta=zeta,
        p=p0 * np.exp(-zeta),
        r=r,
        T=T,
        n=n,
        zeta_end=float(zeta_end),
        unbound=unbound,
        t_exo_floored=floored,
    )
    return ext

def bates_state(ext: dict, zeta):
    """Radius [m], temperature [K], number density [m^-3], and 1/r at ``zeta``.

    ``T = T_exo - (T_exo - T_top) exp(-gamma zeta)`` integrates to
    ``T_exo zeta - (T_exo - T_top)(1 - exp(-gamma zeta)) / gamma``, which
    gives ``1/r`` exactly. Beyond the point where ``1/r`` reaches zero the
    radius is infinite, and ``bates_extension`` stops well before it.
    """
    zeta = np.asarray(zeta, dtype=float)
    t_exo, t_top, gamma = ext['T_exo'], ext['T_top'], ext['gamma']
    gz = gamma * zeta
    T = t_exo - (t_exo - t_top) * np.exp(-gz)
    # (1 - e^-gz) / gamma without its cancellation at small gamma zeta.
    with np.errstate(divide='ignore', invalid='ignore'):
        rise = np.where(np.abs(gz) > 1e-8, -np.expm1(-gz) / gamma, zeta * (1.0 - 0.5 * gz))
    t_int = t_exo * zeta - (t_exo - t_top) * rise
    inv_r = 1.0 / ext['r0'] - kb * t_int / (G * ext['M_p'] * ext['mu'])
    with np.errstate(divide='ignore'):
        r = np.where(inv_r > 0.0, 1.0 / np.where(inv_r > 0.0, inv_r, 1.0), np.inf)
    n = ext['p0'] * np.exp(-zeta) / (kb * T)
    return r, T, n, inv_r

def find_exobase(ext: dict, M_p: float, n_scan: int = 201) -> tuple[float, dict]:
    """Exobase on the extension, as ``(zeta_exo, flags)``.

    The level where the Maxwell mean free path ``1/(sqrt(2) sigma n)`` with
    the mixture cross section equals the local scale height
    ``kB T r^2 / (G M mu)``, located by a bracketing root find on the
    closed-form structure rather than snapped to a grid level, so the
    exobase state does not depend on any quadrature resolution. ``n_scan``
    points bracket the first crossing. When the extension never reaches
    that point its end is used, flagged ``exobase_not_reached``. An exobase
    at the anchor itself returns ``zeta_exo = 0``, flagged
    ``exobase_at_anchor``.
    """
    flags: dict = {}
    zeta_end = float(ext['zeta_end'])

    def excess(z: float) -> float:
        r, t, n, _inv = _bates_point(ext, z)
        sigma, _prov = sigma_mixture(ext['vmr'], t)
        mfp = 1.0 / (math.sqrt(2.0) * sigma * n)
        return mfp / (kb * t * r**2 / (G * M_p * ext['mu'])) - 1.0

    if excess(0.0) >= 0.0:
        flags['exobase_at_anchor'] = True
        return 0.0, flags
    if zeta_end <= 0.0:
        flags['exobase_not_reached'] = True
        return 0.0, flags
    grid = np.linspace(0.0, zeta_end, max(2, int(n_scan)))
    lo = 0.0
    for z in grid[1:]:
        if excess(float(z)) >= 0.0:
            return float(brentq(excess, lo, float(z), xtol=1e-12, rtol=1e-12)), flags
        lo = float(z)
    flags['exobase_not_reached'] = True
    return zeta_end, flags

def hydrostatic_rates(
    profile,
    M_p: float,
    T_exo: float,
    gamma_bates: float = 0.75,
    kzz_default: float = 3.0e2,
    n_levels: int = 400,
) -> tuple[dict, dict]:
    """Per-species hydrostatic escape mapped onto per-element rates [kg/s].

    Builds the Bates extension at the prescribed exobase temperature,
    locates the exobase, and combines the kinetic-corrected Jeans flux with
    the diffusion-limited supply by the harmonic mean, species by species
    (the module notes give the construction and its provenance). The
    species are the ones the supplied profile carries at its top level (the
    profile chemistry decides how atomized the exobase gas is); their rates
    map onto elements stoichiometrically at output.

    The eddy diffusion coefficient comes from the profile's top level when
    a ``kzz`` column is present, else ``kzz_default`` [m^2/s]. Species
    whose supply-free Jeans rate already falls below one proton mass per
    year skip the supply integrals (``pruned`` in the per-species detail)
    and return that rate at their profile-top mole fraction. It is an
    estimate, not a bound, since diffusive enrichment can raise a light
    species' full rate above it (2 to 5 times on the trace hydrogen and
    helium profiles tested), enough to carry a rate just under the floor
    above it. The shortcut exists because the integrals dominate
    the cost of the branch on many-species profiles.

    ``n_levels`` sets the quadrature resolution of the supply integrals,
    which are first-order accurate in the log-pressure step;
    ``hydrostatic_rates_refined`` refines it to a target. The exobase does
    not depend on it.

    Returns ``(per_element, detail)``; the detail dict carries the exobase
    state, both escape temperatures, the ``dominant`` species that supplies
    itself without a diffusion cap, the per-species terms, coefficient
    provenance, and flags (including ``hydrostatic_lower_limit``, which is
    always on: nonthermal loss channels are absent).
    """
    structure = _exobase_structure(profile, M_p, T_exo, gamma_bates)
    return _rates_on_structure(profile, M_p, structure, kzz_default, n_levels)

def hydrostatic_rates_refined(
    profile,
    M_p: float,
    T_exo: float,
    gamma_bates: float = 0.75,
    kzz_default: float = 3.0e2,
    n_levels_min: int = 200,
    n_levels_max: int = 3200,
    rtol: float = 1.0e-2,
) -> tuple[dict, dict]:
    """Hydrostatic rates refined until the quadrature stops moving them.

    The exobase is located once, on the closed-form structure, so only the
    supply integrals depend on the grid. They are first-order accurate in
    the log-pressure step, so the change between a grid and its refinement
    estimates the remaining error. The resolution doubles from
    ``n_levels_min`` until two successive relative changes fall below
    ``rtol``, in the bulk rate and in every element whose rate clears the
    rate floor, or until ``n_levels_max`` is reached. Two changes, so a
    single chance agreement does not end the refinement, and every element,
    since a trace whose supply integral is still moving can sit under a bulk
    the dominant species sets alone. Meeting the target needs
    ``n_levels_max`` of at least four times ``n_levels_min``.
    The finest grid's rates are returned, never an extrapolation.

    The detail dict gains a ``convergence`` entry recording the levels used,
    the last relative change in the bulk rate and the worst one across the
    elements that clear the rate floor, the grid history, and whether the target was met. A call that
    exhausts ``n_levels_max`` without meeting it reports ``converged`` False
    rather than raising: the rate is still the best available and the
    caller can see how far from the target it is.
    """
    structure = _exobase_structure(profile, M_p, T_exo, gamma_bates)
    n = max(2, int(n_levels_min))
    n_max = max(n, int(n_levels_max))
    per_element, detail = _rates_on_structure(profile, M_p, structure, kzz_default, n)
    history = [(n, sum(per_element.values()))]
    rel_bulk = math.inf
    rel_species = math.inf
    below = 0
    converged = False
    while n < n_max:
        n_fine = min(2 * n, n_max)
        per_fine, detail_fine = _rates_on_structure(
            profile, M_p, structure, kzz_default, n_fine
        )
        rel_bulk = _relative_change(sum(per_element.values()), sum(per_fine.values()))
        rel_species = max(
            (
                _relative_change(per_element.get(el, 0.0), per_fine.get(el, 0.0))
                for el in set(per_element) | set(per_fine)
                if max(per_element.get(el, 0.0), per_fine.get(el, 0.0)) >= RATE_FLOOR_KG_S
            ),
            default=0.0,
        )
        per_element, detail, n = per_fine, detail_fine, n_fine
        history.append((n, sum(per_element.values())))
        below = below + 1 if max(rel_bulk, rel_species) <= rtol else 0
        if below >= 2:
            converged = True
            break
    detail['convergence'] = dict(
        n_levels=n,
        n_levels_min=int(n_levels_min),
        n_levels_max=n_max,
        rtol=float(rtol),
        rel_change_bulk=float(rel_bulk),
        rel_change_worst_species=float(rel_species),
        converged=bool(converged),
        history=history,
    )
    return per_element, detail

def gate_unstable(
    T_exo: float, detail: dict, gate: str, f_plus_exo: float
) -> tuple[bool, bool]:
    """Escape-temperature gate on the hydrostatic branch.

    The hydrostatic equilibrium is unstable when the exobase temperature
    exceeds half the gating escape temperature (the stability criterion of
    Chatterjee & Pierrehumbert 2026, their Figure 10). ``gate`` selects the
    neutral or the plasma escape temperature; both are always computed, and
    the returned ``contested`` marks points where the two conventions
    disagree, which is where the (unmodeled) ion physics decides the branch
    assignment and both branch rates belong in the diagnostics.
    ``f_plus_exo`` is carried for the caller's record and does not gate.
    """
    un_neutral = T_exo > 0.5 * detail['T_esc_neutral']
    un_plasma = T_exo > 0.5 * detail['T_esc_plasma']
    unstable = un_plasma if gate == 'plasma' else un_neutral
    return unstable, (un_neutral != un_plasma)