"""
Exact completion probabilities for a chain of sequential exponential stages (a phase-type distribution),
optionally with a constant hazard h of the biosphere dying before the chain completes.

Method: uniformisation with non-negative matrices and repeated squaring (no cancellation), tabulated per
sample on a doubling time grid. Times are in reference windows (5.6 Gyr). Adapted from the round-2 code
auditor's exactchain.py, with the killing hazard added and re-tested against closed forms (run this file).
"""
import math
import numpy as np

K = 16
TK = 2.0 ** np.arange(-8, 9)                   # ascending times in reference windows; 1.0 (the Sun-like window) is a grid point
TMAX = TK[-1]
LAMBDA = 60.0                                   # uniformisation rate; stage rates are clipped at LAMBDA - 10.5 and the hazard at 10 (faster stages are effectively instant)
R_FLOOR = 1e-10


def logcdf_grid(R, h=None):
    """R: (n, k) stage rates per reference window; h: (n,) killing hazard per window or None.
    Returns log10 P(all k stages done by TK[j], not killed before completion) as (n, K+1)."""
    n, k = R.shape
    hz = np.zeros(n) if h is None else np.minimum(h, 10.0)     # the model passes at most -ln(1e-3) = 6.9
    Rc = np.clip(R, R_FLOOR, LAMBDA - 10.5)
    logscale = np.where(R < R_FLOOR, np.log10(np.maximum(R, 1e-300)) - np.log10(R_FLOOR), 0.0).sum(1)
    A = np.zeros((n, k + 1, k + 1))
    for i in range(k):
        A[:, i, i] = LAMBDA - Rc[:, i] - hz
        A[:, i, i + 1] = Rc[:, i]
    A[:, k, k] = LAMBDA                         # completion is absorbing and no longer exposed to the hazard
    t0 = TK[0]
    term = np.broadcast_to(np.eye(k + 1), (n, k + 1, k + 1)).copy()
    Pm = term.copy()
    for j in range(1, 16):
        term = term @ A * (t0 / j)
        Pm = Pm + term
    Pm = Pm * math.exp(-LAMBDA * t0)
    out = np.empty((n, K + 1))
    for j in range(K + 1):
        out[:, j] = np.log10(np.maximum(Pm[:, 0, k], 1e-300)) + logscale
        if j < K:
            Pm = Pm @ Pm
    return out


def _cll(logP):
    """complementary log-log transform, log10(-ln(1 - P)), which is exactly linear in log t for a single stage and for hard chains."""
    P = np.minimum(10 ** np.maximum(logP, -300.0), 1 - 1e-15)
    return np.where(logP < -8, logP, np.log10(np.maximum(-np.log1p(-P), 1e-300)))


def interp_log(logP, s):
    """logP: (n, K+1) on TK; s: (n,) times -> log10 P at s. Interpolation is linear in log t on the complementary
    log-log scale (power-law extrapolation below TK[0])."""
    ls = np.log(np.maximum(s, 1e-12))
    lt = np.log(TK)
    j = np.clip(np.searchsorted(lt, ls), 1, K)
    w = (ls - lt[j - 1]) / (lt[j] - lt[j - 1])
    idx = np.arange(len(s))
    y = (1 - w) * _cll(logP[idx, j - 1]) + w * _cll(logP[idx, j])
    return np.where(y < -8, y, np.log10(np.maximum(-np.expm1(-10 ** np.minimum(y, 2.0)), 1e-300)))


if __name__ == "__main__":
    np.set_printoptions(precision=5, linewidth=160)
    R = np.array([[0.2], [3.0], [1e-30]])
    lp = logcdf_grid(R)
    for s in (0.5, 1.0, 7.3):
        print("1 stage, no hazard, s =", s, 10 ** interp_log(lp, np.full(3, s)), "exact", 1 - np.exp(-R[:, 0] * s))
    R = np.array([[0.5, 0.5001], [1e-3, 2e-4], [0.105, 1e-6], [20.0, 0.01]])
    lp = logcdf_grid(R)
    for s in (1.0, 5.36):
        r1, r2 = R[:, 0], R[:, 1]
        ex = 1 - (r2 * np.exp(-r1 * s) - r1 * np.exp(-r2 * s)) / (r2 - r1)
        print("2 stages, no hazard, s =", s, 10 ** interp_log(lp, np.full(4, s)), "exact", ex)
    # with a hazard: 1 stage -> r/(r+h) (1 - exp(-(r+h)s)); 2 hard stages -> r1 r2 * int_0^s t e^{-ht} dt
    h = np.array([0.7, 2.5, 3.4])
    R = np.array([[0.2], [3.0], [1e-30]])
    lp = logcdf_grid(R, h)
    for s in (0.5, 1.0, 7.3):
        r = R[:, 0]
        print("1 stage, hazard, s =", s, 10 ** interp_log(lp, np.full(3, s)), "exact", r / (r + h) * (1 - np.exp(-(r + h) * s)))
    R = np.array([[1e-4, 3e-5], [1e-12, 1e-15], [2e-3, 1e-3]])
    lp = logcdf_grid(R, h)
    for s in (1.0, 4.0):
        ex = R[:, 0] * R[:, 1] * (1 - np.exp(-h * s) * (1 + h * s)) / h ** 2
        print("2 hard stages, hazard, s =", s, 10 ** interp_log(lp, np.full(3, s)), "exact (small-rate limit)", ex)
    R = np.full((1, 5), 1e-3)
    print("5 hard stages at s = 1:", 10 ** interp_log(logcdf_grid(R), np.ones(1)), " (r s)^5/5! =", 1e-15 / 120)