"""SNAPSHOT (copied 2026-10-12) of studies/mdwarf_early/code/earliness.py, the earliness / birth-rank model
of the first-wave study `mdwarf_early`.  Copied so that this study does not change if that file is
revised.  The only edit is that the constants imported from mdwarf_early/code/common.py are in-lined.
Original docstring follows.

Part 3: humanity's rank in the distribution of civilisation birth dates over all cosmic time.

Model.  Stars of mass M form at cosmic time b at a rate psi(b) xi(M).  Each hosts (with weight
w(M)) one habitable planet whose habitable window is L(M) = min(tau_MS(M), Lmax).  With n hard
steps the probability that a civilisation has appeared by planet age a is proportional to a^n for
a <= L (and stays at L^n afterwards).  The number of civilisations born before cosmic time t is

    N(<t) = int dM xi(M) w(M) int_0^t db psi(b) min(t - b, L(M))^n ,

and humanity's rank is N(<t0) / N(<infinity) with t0 = 13.79 Gyr.

Two readings of the rank (see report.md, Method):
  * for fixed (n, Lmax, w) the rank is the likelihood ratio of this no-deadline model to the same
    model with a deadline just after today, for the datum "we appeared at t0";
  * the law a^n is NOT saturated: it is extrapolated to planet ages of 1e3-1e4 Gyr, which is an
    assumption equivalent to choosing a very large Lmax.
"""
from functools import lru_cache

import numpy as np
from scipy.special import gammaln

# --- constants in-lined from studies/mdwarf_early/code/common.py ---
T0_GYR = 13.79                # age of the Universe today (Planck 2018)
M_DWARF_MAX_MASS = 0.60       # Msun
A_EARTH_GYR = 4.54
H0 = 67.4
OMEGA_M = 0.315
OMEGA_L = 1.0 - OMEGA_M
H0_PER_GYR = H0 * 1.0227e-3

# ----------------------------------------------------------------------------------------------
# ingredients
# ----------------------------------------------------------------------------------------------


def scale_factor(t_gyr):
    """a(t) for flat LCDM (matter + Lambda), a(t0) = 1."""
    t = np.asarray(t_gyr, dtype=float)
    return (OMEGA_M / OMEGA_L) ** (1 / 3) * np.sinh(1.5 * np.sqrt(OMEGA_L) * H0_PER_GYR * t) ** (2 / 3)


def age_today():
    return 2.0 / (3.0 * np.sqrt(OMEGA_L) * H0_PER_GYR) * np.arcsinh(np.sqrt(OMEGA_L / OMEGA_M))


def md14(one_plus_z):
    """Madau & Dickinson (2014) eq. 15, Msun/yr/Mpc^3."""
    return 0.015 * one_plus_z ** 2.7 / (1.0 + (one_plus_z / 2.9) ** 5.6)


def sfr(t_gyr, model="MD14"):
    """Cosmic star-formation-rate density versus cosmic time (arbitrary but consistent units).

    MD14        Madau & Dickinson fit, analytically continued to the future through 1+z = 1/a(t)
    MD14_exp    MD14 to today, then exponential decline with their quoted 3.9 Gyr e-folding time
    Sobral13    log10 psi = -2.1/(1+z) (Sobral et al. 2013) for z < 2.23 and into the future,
                MD14 shape (rescaled to match) at z > 2.23
    BP15_long   MD14 to today, then a slow exponential decline whose integral is 10x the stellar
                mass formed so far (upper bound of Behroozi & Peeples 2015: all gas now in haloes forms stars)
    Hanson      t^3 exp(-t/4 Gyr), the 'habitable SFR' of Hanson et al. 2021 (peak 12 Gyr)
    """
    t = np.asarray(t_gyr, dtype=float)
    # far-future times overflow sinh -> a = inf -> 1+z = 0 -> psi = 0, which is the intended limit
    with np.errstate(over="ignore", divide="ignore", invalid="ignore"):
        return _sfr(t, model)


def _sfr(t, model):
    opz = 1.0 / scale_factor(np.maximum(t, 1e-6))
    t0 = age_today()
    if model == "MD14":
        return md14(opz)
    if model == "MD14_exp":
        return np.where(t <= t0, md14(opz), md14(1.0) * np.exp(-(t - t0) / 3.9))
    if model == "Sobral13":
        z_join = 2.23
        sob = 10.0 ** (-2.1 / opz)
        scale = 10.0 ** (-2.1 / (1 + z_join)) / md14(1 + z_join)
        return np.where(opz <= 1 + z_join, sob, md14(opz) * scale)
    if model == "BP15_long":
        tt = np.linspace(1e-3, t0, 4000)
        past = np.trapezoid(md14(1.0 / scale_factor(tt)), tt)
        tau = 10.0 * past / md14(1.0)
        return np.where(t <= t0, md14(opz), md14(1.0) * np.exp(-(t - t0) / tau))
    if model == "Hanson":
        return t ** 3 * np.exp(-t / 4.0)
    raise ValueError(model)


def ms_lifetime_gyr(m):
    """Main-sequence lifetime, piecewise power law of Adams & Laughlin (1997) / Loeb et al. (2016)."""
    m = np.asarray(m, dtype=float)
    return np.where(m >= 0.75, 10.0 * m ** -2.5,
                    np.where(m >= 0.25, 7.6 * m ** -3.5, 53.0 * m ** -2.1))


def imf_number(m, kind="Kroupa"):
    """dN/dM (unnormalised)."""
    m = np.asarray(m, dtype=float)
    if kind == "Kroupa":      # Kroupa 2001
        return np.where(m < 0.5, m ** -1.3, 0.5 * m ** -2.3)
    if kind == "Chabrier":    # Chabrier 2003 single-star disk IMF
        lognorm = 0.158 * np.exp(-(np.log10(m) - np.log10(0.079)) ** 2 / (2 * 0.69 ** 2)) / m
        pl = 0.158 * np.exp(-(0 - np.log10(0.079)) ** 2 / (2 * 0.69 ** 2)) * m ** -2.3
        return np.where(m < 1.0, lognorm, pl)
    raise ValueError(kind)


# Mass cells: 320 logarithmic cells between 0.08 and 3 Msun; MASS_GRID holds the cell mid-points
# (mid-point rule).  The first-pass version put grid POINTS on 0.08 and 3.0 with full weight, which
# over-counted half a cell at each end and placed host-mass cuts up to one cell (1.1% in mass) off;
# for steep integrands (L^n with n >= 4) that biased ranks by 3-8%.  Host-mass cuts now weight the
# cell that contains the cut by the fraction of it above the cut (host_weight).
N_MASS = 320
MASS_EDGES = np.exp(np.linspace(np.log(0.08), np.log(3.0), N_MASS + 1))
MASS_GRID = np.sqrt(MASS_EDGES[1:] * MASS_EDGES[:-1])
DLNM = np.log(MASS_EDGES[1] / MASS_EDGES[0])


def host_weight(m_cut):
    """Weight 1 for mass cells entirely above m_cut, 0 for cells below, and the fraction (in ln M)
    above the cut for the cell that contains it."""
    return np.clip((np.log(MASS_EDGES[1:]) - np.log(m_cut)) / DLNM, 0.0, 1.0)


@lru_cache(maxsize=None)
def _time_grids(model, t_end=3.0e4, n_past=3000):
    """Birth-time grid to today and the total integral of psi to the far future."""
    t0 = age_today()
    b = np.linspace(1e-3, t0, n_past)
    psi = sfr(b, model)
    tf = np.concatenate([np.linspace(t0, 200.0, 20000)[1:], np.exp(np.linspace(np.log(200.0), np.log(t_end), 4000))[1:]])
    tf = np.concatenate([[t0], tf])
    fut = np.trapezoid(sfr(tf, model), tf)
    past = np.trapezoid(psi, b)
    return b, psi, past, past + fut


def birth_integrals(n, lmax, model="MD14", imf="Kroupa", f_window=1.0):
    """Per-mass numerator and denominator of the rank:
       num[M] = xi(M) dM int_0^t0 psi(b) min(t0-b, L)^n db ;  den[M] = xi(M) dM L^n Psi_total."""
    t0 = age_today()
    b, psi, past, total = _time_grids(model)
    L = np.minimum(f_window * ms_lifetime_gyr(MASS_GRID), lmax)
    age = t0 - b
    amin = np.minimum(age[None, :], L[:, None])
    inner = np.trapezoid(psi[None, :] * amin ** n, b, axis=1)
    xi = imf_number(MASS_GRID, imf) * MASS_GRID * DLNM
    return xi * inner, xi * L ** n * total


def rank_from_weights(num, den, w):
    """w: (..., n_mass) habitability weights."""
    return (w * num).sum(axis=-1) / (w * den).sum(axis=-1)


def step_weights(p_m, m_cut=M_DWARF_MAX_MASS):
    """Weight 1 above m_cut and p_m below it."""
    hw = host_weight(m_cut)
    return hw + (1.0 - hw) * p_m


def rank_step(n, lmax, p_m, model="MD14", imf="Kroupa", m_cut=M_DWARF_MAX_MASS, m_min_host=None, f_window=1.0):
    num, den = birth_integrals(n, lmax, model, imf, f_window)
    w = step_weights(p_m, m_cut)
    if m_min_host is not None:
        w = w * host_weight(m_min_host)
    return float(rank_from_weights(num, den, w))


def critical_p_m(n, lmax, target, model="MD14", imf="Kroupa", m_cut=M_DWARF_MAX_MASS, f_window=1.0):
    """Largest weight p_M of hosts below m_cut for which rank >= target.  Returns (p_crit, rank_without_them)."""
    num, den = birth_integrals(n, lmax, model, imf, f_window)
    hw = host_weight(m_cut)
    A, B = (num * hw).sum(), (num * (1 - hw)).sum()
    C, D = (den * hw).sum(), (den * (1 - hw)).sum()
    r0 = A / C
    r1 = (A + B) / (C + D)
    if r1 >= target:
        return 1.0, r0            # not early even with all M dwarfs
    if r0 < target:
        return 0.0, r0            # early even with no M dwarfs
    return float((A - target * C) / (target * D - B)), r0


def self_datum_likelihood(DEN, W, n_values, L_grid, model="MD14", f_window=1.0, a_earth=A_EARTH_GYR):
    """Likelihood of our own datum under the no-deadline model, on a (n, Lmax) grid.

    Datum D = (birth date t0, host mass 1 Msun, planet age a_earth).  The density of civilisations
    over (t, M, a) is psi(t-a) xi(M) w(M) n a^(n-1) 1[a < L(M)], and the total number is
    Psi_tot int xi w L^n dM, so

        p(D | n, Lmax, w) = psi(t0-a) xi(1) w(1) n a^(n-1) 1[a < L(1)] / (Psi_tot int xi w L^n dM).

    psi(t0-a) xi(1) / Psi_tot does not depend on (n, Lmax, w) and w(1 Msun) = 1 in every scenario,
    so we return   like[s, n, L] = n a^(n-1) 1[a < L(1)] / <w L^n>   [Gyr^-1],
    where <.> is the IMF number average over 0.08-3 Msun.

    DEN: (n, L, mass) array from birth_integrals (= xi dM L^n Psi_tot);  W: (Ns, mass) weights.
    """
    _, _, _, total = _time_grids(model)
    xi = imf_number(MASS_GRID, "Kroupa") * MASS_GRID * DLNM
    mean_wLn = np.einsum("sm,nlm->snl", W, DEN) / (total * xi.sum())
    n = np.asarray(n_values, dtype=float)[None, :, None]
    L_sun = np.minimum(f_window * float(ms_lifetime_gyr(1.0)), np.asarray(L_grid, dtype=float))[None, None, :]
    return np.where(a_earth < L_sun, n * a_earth ** (n - 1.0) / mean_wLn, 0.0)


def rank_host_cut(n, lmax, m_cut, model="MD14", imf="Kroupa", f_window=1.0):
    """Rank when only hosts with M >= m_cut are habitable (weight 1), all others weight 0."""
    num, den = birth_integrals(n, lmax, model, imf, f_window)
    w = host_weight(m_cut)
    return float((w * num).sum() / (w * den).sum())


# ----------------------------------------------------------------------------------------------
# replication of Hanson et al. (2021) equation 1 / figure 2
# ----------------------------------------------------------------------------------------------
def hanson_rank(n, lbar, phi=4.0, chi=12.0, kappa=0.0, l_min=1.0):
    """Percentile rank of t0 = 13.8 Gyr under Hanson et al. eq. 1:
       alpha(t) = q int (t-b)^(n-1) rho(b) [H(Lbar) - H(t-b)] db,  rho = b^lam exp(-b/phi), lam = chi/phi,
       H(L) = L^((3-2 kappa)/6) (number of planets with lifetime below L)."""
    lam = chi / phi
    h = (3.0 - 2.0 * kappa) / 6.0
    t0 = 13.8

    def surv(a):  # number of planets whose lifetime exceeds age a (unnormalised)
        a = np.maximum(a, l_min) if h < 0 else a
        return np.sign(h) * (lbar ** h - a ** h)

    # cumulative inner integral G(x) = int_0^x a^(n-1) surv(a) da on a log grid
    a = np.exp(np.linspace(np.log(1e-4), np.log(lbar), 6000))
    f = a ** (n - 1) * surv(a)
    G = np.concatenate([[0.0], np.cumsum(0.5 * (f[1:] + f[:-1]) * np.diff(a))])
    b = np.linspace(1e-4, t0, 4000)
    rho = b ** lam * np.exp(-b / phi)
    num = np.trapezoid(rho * np.interp(np.minimum(t0 - b, lbar), a, G), b)
    rho_tot = np.exp(gammaln(lam + 1) + (lam + 1) * np.log(phi))
    return num / (rho_tot * G[-1])