"""Core numerical routines for the Earth-timeline (timing) study.

All times are in Gyr, all rates in Gyr^-1.

Conventions
-----------
t0      : age (Ga) at which Earth became habitable (time zero of every model)
t_obs   : elapsed time from t0 to the arrival of observers (= t0, since "now" = 0 Ga)
T       : total habitable window measured from t0 (= t0 + future habitable lifetime)
t_Lp    : upper bound on the abiogenesis delay (= t0 - age of earliest evidence of life)

Selection-conditioned likelihood (used everywhere):
    L(rates) = p(observed timeline | rates) / P(all steps complete before T | rates)
because an observer can only be looking at a timeline on which every required
step succeeded inside the window.
"""
from __future__ import annotations

import numpy as np
from scipy import special

# ----------------------------------------------------------------------------
# Hypoexponential CDF, numerically stable over ~40 decades of rates
# ----------------------------------------------------------------------------

CHUNK = 20_000
RATE_CAP = 1.0e5  # lambda*T above this is treated as lambda*T = RATE_CAP (step is "instant");
                  # relative error on the CDF is < n / RATE_CAP.


def log_hypoexp_cdf(lam, T, taylor_terms: int = 22):
    """log P(E_1 + ... + E_n < T) for independent E_i ~ Exp(lam_i).

    Parameters
    ----------
    lam : array (M, n) of rates (any positive values, may span many decades)
    T   : scalar or array (M,)

    Method: the CDF equals prod(lam_i T) * [exp(Z)]_{0,n} where Z is the
    (n+1)x(n+1) upper-bidiagonal matrix with diagonal (-lam_1 T, ..., -lam_n T, 0)
    and unit super-diagonal (Opitz formula for divided differences of exp).
    Z is a Metzler matrix, so exp(Z) is entrywise non-negative and the
    scaling-and-squaring recursion involves no cancellation after the Taylor
    stage.  Working with a unit super-diagonal keeps the result O(1/n!) even
    when the rates are 1e-100, so nothing underflows.
    """
    lam = np.atleast_2d(np.asarray(lam, dtype=float))
    M, n = lam.shape
    T = np.broadcast_to(np.asarray(T, dtype=float), (M,))
    if M > CHUNK:                             # keep the work arrays in cache
        return np.concatenate([log_hypoexp_cdf(lam[i:i + CHUNK], T[i:i + CHUNK], taylor_terms)
                               for i in range(0, M, CHUNK)])
    x = lam * T[:, None]                      # dimensionless rates
    xc = np.minimum(x, RATE_CAP)
    Z = np.zeros((M, n + 1, n + 1))
    idx = np.arange(n)
    Z[:, idx, idx] = -xc
    Z[:, idx, idx + 1] = 1.0
    norm = max(float(xc.max(initial=0.0)), 1.0)
    s = int(np.ceil(np.log2(norm / 0.25))) if norm > 0.25 else 0
    A = Z / (2.0 ** s)
    # Taylor series for exp(A)
    E = np.broadcast_to(np.eye(n + 1), A.shape).copy()
    term = E.copy()
    for k in range(1, taylor_terms + 1):
        term = np.matmul(term, A) / k
        E += term
    # The super-diagonal was scaled by 2^-s as well; undo that analytically at the end:
    # exp(Z)[0,n] is obtained by squaring exp(A) s times.
    for _ in range(s):
        E = np.matmul(E, E)
    G = E[:, 0, n]
    G = np.maximum(G, 1e-300)
    return np.sum(np.log(xc), axis=1) + np.log(G)


def hypoexp_cdf(lam, T):
    return np.exp(log_hypoexp_cdf(lam, T))


def log_hypoexp_cdf2(l1, l2, T):
    """Exact, fast and stable log P(E1 + E2 < T) for two exponential waits.

    Uses F = a b [g(b) - g(a)] / (a - b) with a = l1 T, b = l2 T and
    g(x) = (1 - e^-x)/x, a power series when both are small, and the
    derivative form when a ~ b.  Vectorised; no matrix exponentials.
    """
    a, b = np.broadcast_arrays(np.asarray(l1, float) * T, np.asarray(l2, float) * T)
    hi = np.maximum(a, b)
    lo = np.minimum(a, b)
    out = np.empty(hi.shape)
    small = hi < 1e-3
    close = (~small) & ((hi - lo) < 1e-6 * hi)
    gen = ~(small | close)
    # series: F = ab/2 * (1 - (a+b)/3 + (a^2+ab+b^2)/12 - ...)
    x, y = hi[small], lo[small]
    out[small] = np.log(x) + np.log(y) - np.log(2.0) + np.log1p(-(x + y) / 3 + (x * x + x * y + y * y) / 12)
    x = hi[close]
    out[close] = np.log(-np.expm1(-x) - x * np.exp(-x))
    x, y = hi[gen], lo[gen]
    g = lambda z: special.exprel(-z)
    out[gen] = np.log(x) + np.log(y) + np.log(g(y) - g(x)) - np.log(x - y)
    return out


# ----------------------------------------------------------------------------
# Small helpers
# ----------------------------------------------------------------------------

def log_exprel(y):
    """log((e^y - 1)/y), stable for |y| small and for y large (no overflow)."""
    y = np.asarray(y, dtype=float)
    out = np.empty_like(y)
    big = y > 30.0
    out[~big] = np.log(special.exprel(y[~big]))
    out[big] = y[big] - np.log(y[big]) + np.log1p(-np.exp(-y[big]))
    return out


def weighted_quantile(x, w, q):
    x = np.asarray(x, float)
    w = np.asarray(w, float)
    o = np.argsort(x)
    x = x[o]
    w = w[o]
    c = np.cumsum(w)
    c = (c - 0.5 * w) / c[-1]
    return np.interp(q, c, x)


def ess(w):
    w = np.asarray(w, float)
    return w.sum() ** 2 / np.sum(w ** 2)


# ----------------------------------------------------------------------------
# Likelihoods
# ----------------------------------------------------------------------------

def loglike_steps(lam, t_Lp, s_later, T, first_is_bound=True, return_logF=False):
    """Selection-conditioned log-likelihood for a chain of n sequential Poisson steps.

    lam      : (M, n) rates; step 1 is abiogenesis
    t_Lp     : (M,) or scalar; abiogenesis happened at some unknown time in [0, t_Lp]
               (first_is_bound=True) or exactly at t_Lp (first_is_bound=False)
    s_later  : (M, n-1) or (n-1,) absolute completion times (since t0) of steps 2..n,
               the last one being t_obs
    T        : (M,) or scalar habitable window

    Returns log L = log p(timeline | lam) - log P(sum of n waits < T | lam).
    The numerator is a density in (s_2..s_n) and a probability in the first step
    when first_is_bound=True.
    """
    lam = np.atleast_2d(np.asarray(lam, float))
    M, n = lam.shape
    t_Lp = np.broadcast_to(np.asarray(t_Lp, float), (M,))
    T = np.broadcast_to(np.asarray(T, float), (M,))
    s_later = np.broadcast_to(np.asarray(s_later, float), (M, n - 1)) if n > 1 else np.zeros((M, 0))
    l1 = lam[:, 0]
    if n == 1:
        if first_is_bound:
            num = np.log(-np.expm1(-l1 * t_Lp))
        else:
            num = np.log(l1) - l1 * t_Lp
    else:
        l2 = lam[:, 1]
        s2 = s_later[:, 0]
        if first_is_bound:
            # int_0^a l1 e^{-l1 t} l2 e^{-l2 (s2 - t)} dt = l1 l2 e^{-l2 s2} a exprel((l2-l1) a)
            num = (np.log(l1) + np.log(l2) - l2 * s2 + np.log(t_Lp)
                   + log_exprel((l2 - l1) * t_Lp))
        else:
            num = np.log(l1) - l1 * t_Lp + np.log(l2) - l2 * (s2 - t_Lp)
        for i in range(2, n):
            num = num + np.log(lam[:, i]) - lam[:, i] * (s_later[:, i - 1] - s_later[:, i - 2])
    if n == 1:
        logF = np.log(-np.expm1(-l1 * T))
    elif n == 2:
        logF = log_hypoexp_cdf2(lam[:, 0] * T, lam[:, 1] * T, 1.0)
    else:
        logF = log_hypoexp_cdf(lam, T)
    if return_logF:
        return num - logF, logF
    return num - logF


def kipping_like(lL, lI, tLp, tIp, T, nquad: int = 200):
    """Kipping (2020) likelihood, eq. 18: P(t_L < t_L', t_I' < t_I < t_I' + t_L' | rates),
    normalised by P(t_L + t_I < T).  Evaluated by Gauss-Legendre quadrature over t_L so
    that both branches (T > or < 2 t_L' + t_I') are handled by the same code.
    lL, lI broadcastable arrays.
    """
    lL, lI = np.broadcast_arrays(np.asarray(lL, float), np.asarray(lI, float))
    shp = lL.shape
    lL = lL.ravel()
    lI = lI.ravel()
    x, w = np.polynomial.legendre.leggauss(nquad)
    # substitute t = tLp * u^p to resolve very fast abiogenesis rates near t = 0
    # (use piecewise: analytic when the clipping by T is inactive)
    out = np.empty_like(lL)
    if T >= 2 * tLp + tIp:
        num = (-np.expm1(-lL * tLp)) * np.exp(-lI * tIp) * (-np.expm1(-lI * tLp))
    else:
        # region: 0<tL<tLp, tIp < tI < min(tIp+tLp, T - tL)
        # split at tL* = T - tIp - tLp (>= 0?) below which the upper limit is tIp+tLp
        ts = max(T - tIp - tLp, 0.0)
        part1 = (-np.expm1(-lL * ts)) * np.exp(-lI * tIp) * (-np.expm1(-lI * tLp))
        # part 2: tL in [ts, min(tLp, T - tIp)], upper limit T - tL
        hi = min(tLp, T - tIp)
        t = 0.5 * (hi - ts) * (x + 1) + ts
        wt = 0.5 * (hi - ts) * w
        integrand = (lL[:, None] * np.exp(-lL[:, None] * t[None, :])
                     * (np.exp(-lI[:, None] * tIp) - np.exp(-lI[:, None] * (T - t[None, :]))))
        part2 = integrand @ wt
        num = part1 + part2
    logF = log_hypoexp_cdf2(lL, lI, T)
    out = num / np.exp(logF)
    return out.reshape(shp)


# ----------------------------------------------------------------------------
# Closed forms for the hard-step (all later rates -> 0) limits
# ----------------------------------------------------------------------------

def bmax_hardsteps(m, tLp, tobs, T, delay=0.0):
    """Limiting Bayes factor  L(lambda_L -> inf) / L(lambda_L -> 0)  when abiogenesis is
    followed by m hard steps (rates -> 0) plus a fixed cumulative delay `delay`,
    data = {abiogenesis before tLp, observers at tobs}.

        B = (T - D) * m * (tobs - D)^(m-1) / ((m+1) * [(tobs - D)^m - (tobs - D - tLp)^m])

    m = 1, D = 0 gives Kipping's T / (2 tLp).  Valid for tLp <= tobs - D and tobs <= T.
    """
    m = np.asarray(m, float)
    a = tobs - delay
    return (T - delay) * m * a ** (m - 1) / ((m + 1) * (a ** m - (a - tLp) ** m))


def plateau_ratio_general(later_rates, tLp, tobs, T, ngrid=2001, h=2e-3):
    """Ratio of the two likelihood plateaus, L(lambda_L -> inf) / L(lambda_L -> 0), when
    abiogenesis is followed by ANY chain of Poisson steps with the given rates (not only hard
    ones).  Data: abiogenesis before tLp, observers exactly at tobs; likelihood conditioned on
    finishing before T.  With F and g the CDF and density of the post-abiogenesis duration,

        L(lambda_L) = int_0^tLp lambda_L e^(-lambda_L t) g(tobs - t) dt
                      / int_0^T lambda_L e^(-lambda_L t) F(T - t) dt
        B = g(tobs) * int_0^T F(s) ds / ( F(T) * [F(tobs) - F(tobs - tLp)] ).

    F proportional to s^m (m hard steps) recovers bmax_hardsteps.  F is evaluated with the
    hypoexponential routine; g by a Richardson-extrapolated central difference.
    """
    lam = np.asarray(later_rates, float)[None, :]
    s = np.linspace(0.0, T, ngrid)[1:]
    F = np.exp(log_hypoexp_cdf(np.repeat(lam, len(s), axis=0), s))
    F = np.concatenate([[0.0], F])
    w = np.ones(ngrid)                       # Simpson weights (ngrid odd)
    w[1:-1:2], w[2:-1:2] = 4.0, 2.0
    intF = float(np.sum(w * F) * (T / (ngrid - 1)) / 3.0)
    f = lambda t: float(np.exp(log_hypoexp_cdf(lam, t))[0]) if t > 0 else 0.0
    d1 = (f(tobs + h) - f(tobs - h)) / (2 * h)
    d2 = (f(tobs + 2 * h) - f(tobs - 2 * h)) / (4 * h)
    g = (4 * d1 - d2) / 3.0
    dF = f(tobs) - f(tobs - tLp)
    if not (dF > 1e-12 and g > 0):          # every later step so fast that F ~ 1 long before tobs:
        return 0.0                          # the ratio tends to 0 there (see the one-step closed form)
    return g * intF / (f(T) * dF)


def plateau_ratio_one_later_step(r, tLp, T):
    """Closed forms of the plateau ratio for ONE later step of rate r (any value).
    Returns (exact-arrival-time data, Kipping's interval data).
        exact:    B = r (T - E/r) / (E (e^(r tLp) - 1)),   E = 1 - e^(-r T)
        interval: B = (T - E/r) / (tLp E)                  [valid for T >= 2 tLp + tIp]
    Both tend to T / (2 tLp) as r -> 0.  The first peaks and then falls to 0; the second rises
    monotonically to T / tLp."""
    r = np.asarray(r, float)
    E = -np.expm1(-r * T)
    return r * (T - E / r) / (E * np.expm1(r * tLp)), (T - E / r) / (tLp * E)


def kipping_bound(tLp, tIp, T):
    """Kipping (2020) eq. 15 / Kipping (2025) eq. 1 (both branches)."""
    if T >= 2 * tLp + tIp:
        return T / (2 * tLp)
    return T * tLp / (4 * (T - tIp) * tLp - 2 * tLp ** 2 - (T - tIp) ** 2)


def hypoexp_logpdf(lam, t):
    """log density of E_1 + ... + E_k at t:  pdf = lam_k * [F_{k-1}(t) - F_k(t)].
    Only used for validation with small rates (the difference cancels when lam_k t >> 1)."""
    lam = np.atleast_2d(np.asarray(lam, float))
    k = lam.shape[1]
    Fk = np.exp(log_hypoexp_cdf(lam, t))
    Fkm1 = np.exp(log_hypoexp_cdf(lam[:, :-1], t)) if k > 1 else 1.0
    return np.log(lam[:, -1]) + np.log(Fkm1 - Fk)