zephyrus.atomic_data

def badnell_alpha_rr(T: float, coeffs: tuple = BADNELL_N) -> float:
    """Radiative recombination coefficient, cm^3 s^-1 (Badnell 2006, Eqs. 1-2)."""
    T0, T1, T2, A, B, C = coeffs
    bp = B + C * math.exp(-T2 / T)
    s0 = math.sqrt(T / T0)
    s1 = math.sqrt(T / T1)
    return A / (s0 * (1.0 + s0) ** (1.0 - bp) * (1.0 + s1) ** (1.0 + bp))

def alpha_case_b(element: str, T: float) -> float:
    """Case B recombination coefficient, cm^3 s^-1, with temperature scaling.

    Applies the Murray-Clay et al. (2009) hydrogen temperature dependence
    ``(T / 1e4 K)^-0.9`` to every element: the heavies' archival values
    exist only at 1e4 K, so extending hydrogen's exponent to them is a
    documented approximation that matters once the thermostat moves the
    wind temperature away from 1e4 K. Elements without a tabulated value
    take the atomic-O coefficient (the nearest heavy; callers record the
    coefficient provenance upstream).
    """
    a0 = CASE_B_1E4K.get(element)
    if a0 is None:
        a0 = CASE_B_1E4K['O']
    return a0 * (T / 1.0e4) ** -0.9

def co2_band_cooling(n_co2: float, colliders: dict, T: float, col_co2: float = 0.0) -> float:
    """CO2 15 micron band volumetric cooling, erg s^-1 cm^-3.

    Johnstone et al. (2018) Eqs. (34)-(38) with no stellar infrared pumping
    (their S_IR term is zero in the escaping-region application).
    ``n_co2`` [cm^-3]; ``colliders`` maps species to number densities
    [cm^-3]; ``col_co2`` is the overlying CO2 column [cm^-2] for the escape
    probability, whose zero-column limit is the non-LTE ceiling 0.5.
    Excitation rates follow from detailed balance,
    ``k_e = 2 k_d exp(-h nu / k_B T)``.

    The abundances are molecular because this is a base-region coolant, on
    deexcitation fits measured over roughly 150 to 500 K. Above CO2
    dissociation the band should be absent rather than small, which this
    function does not enforce.
    """
    sn = SIGMA_CO2_15UM * col_co2
    if sn > 2.0:
        eps = 0.7202 * sn**-0.613
    elif sn > 0.0:
        # The fit rises through 0.5 below sn = 3.4e-4 and diverges; half the
        # photons escaping is the non-LTE ceiling of this two-level band,
        # and capping there keeps the zero-column case continuous.
        eps = min(0.4732 * sn**-0.0069, 0.5)
    else:
        eps = 0.5
    ke_sum = kd_sum = 0.0
    for sp, n in colliders.items():
        if sp not in CO2_KD or n <= 0.0:
            continue
        a, b = CO2_KD[sp]
        kd = a * T**b
        ke = 2.0 * kd * math.exp(-T_15UM / T)
        kd_sum += kd * n
        ke_sum += ke * n
    if ke_sum == 0.0:
        return 0.0
    n_star = ke_sum * n_co2 / (ke_sum + kd_sum + A10_CO2 * eps)
    return HNU_15UM * A10_CO2 * eps * n_star

def o_finestructure_cooling(n_o: float, T: float) -> float:
    """Atomic O fine-structure cooling at 63 and 147 micron, erg s^-1 cm^-3.

    The closed LTE form of Johnstone et al. (2018) Eqs. (41)-(43), tracing
    to Bates (1951); ``n_o`` is the atomic oxygen density [cm^-3].
    """
    den = 1.0 + 0.6 * math.exp(-228.0 / T) + 0.2 * math.exp(-326.0 / T)
    q63 = 1.67e-18 * math.exp(-228.0 / T) * n_o / den
    q147 = 4.59e-20 * math.exp(-326.0 / T) * n_o / den
    return q63 + q147