"""Likelihood of a dated sequence of transitions under a chain of steps, each of which is either

  'E'  easy and gated : waits for its gate g_i and for the previous step, then completes after an
                        exponential lag with rate lam_i (the gating picture of Mills et al. 2025);
  'G'  hard and gated : cannot start before its gate or the previous step, then is so improbable per unit
                        time that, given completion, its time is uniform over what is allowed;
  'H'  hard, ungated  : as 'G' but ignores the gate (Carter's original model).

Conditional on all n steps completing before the end of the window T, the joint density of the completion
times is  prod_E lam e^{-lam (t_i - s_i)} 1[t_i >= s_i]  x  prod_hard 1[t_i >= s_i]  / Z ,
with s_i = max(t_{i-1}, g_i) ('E', 'G') or t_{i-1} ('H').  For hard steps this is the lam -> 0 limit of the
exponential kernel after conditioning, which is why no rate appears.  Z is the integral of the numerator over
all admissible times and is computed here by a forward pass on a time grid, vectorised over Monte Carlo draws.

All 'H'   : density n! / (T - t0)^n, the uniform order statistics of the hard-steps model.
All 'E'   : product of exponential lag densities (Z ~ 1 when the gates open long before T).
"""
import numpy as np
from scipy.special import exp1


def numerator(kinds, t, gates, lam, t0):
    """t, gates, lam: (D, n); t0: (D,).  Returns (D,) unnormalised density (hard kernels have value 1/Gyr)."""
    D, n = t.shape
    out = np.ones(D)
    prev = t0.copy()
    for i, k in enumerate(kinds):
        s = prev if k == "H" else np.maximum(prev, gates[:, i])
        ok = t[:, i] >= s
        if k == "E":
            lag = np.where(ok, t[:, i] - s, 0.0)
            out *= np.where(ok, lam[:, i] * np.exp(-lam[:, i] * lag), 0.0)
        else:
            out *= ok
        prev = t[:, i]
    return out


def _deposit(arr, pos, mass, dx):
    """Add `mass` (D,) at continuous position `pos` (D,), split linearly between the two nearest bin centres."""
    nbin, D = arr.shape
    cols = np.arange(D)
    f = pos / dx - 0.5
    j = np.floor(f).astype(int)
    w = f - j
    j1 = np.clip(j, 0, nbin - 1)
    j2 = np.clip(j + 1, 0, nbin - 1)
    np.add.at(arr, (j1, cols), mass * (1.0 - w))
    np.add.at(arr, (j2, cols), mass * w)


def normaliser(kinds, gates, lam, t0, T, dx=0.02):
    """Forward pass on a grid of bin width dx (Gyr), arrays laid out (bin, draw). Returns Z (D,).

    Mass is carried on bin centres.  The start time and the gate-release times are deposited by linear
    interpolation between neighbouring centres, and easy steps use the exact exponential survival over each
    bin, so the result is second-order accurate in dx and insensitive to dx even when the mean lag is shorter
    than a bin."""
    D, n = gates.shape
    nbin = int(np.ceil((T.max() + 3 * dx) / dx))
    centre = (np.arange(nbin) + 0.5) * dx
    arr = np.zeros((nbin, D))
    _deposit(arr, t0, np.ones(D), dx)
    for i, k in enumerate(kinds):
        if k in ("E", "G"):
            g = gates[:, i]
            before = centre[:, None] < g[None, :]
            held = np.where(before, arr, 0.0).sum(axis=0)
            m_in = np.where(before, 0.0, arr)
            has = held > 0
            if has.any():
                tmp = np.zeros_like(arr[:, has])
                _deposit(tmp, np.clip(g[has], 0.5 * dx, None), held[has], dx)
                m_in[:, has] += tmp
        else:
            m_in = arr
        if k == "E":
            a = np.exp(-lam[:, i] * dx)
            ah = np.exp(-lam[:, i] * dx / 2.0)
            comp = np.empty_like(m_in)
            w = np.zeros(D)
            for b in range(nbin):
                mb = m_in[b]
                comp[b] = w * (1.0 - a) + mb * (1.0 - ah)
                w = w * a + mb * ah
        else:
            W = np.cumsum(m_in, axis=0)
            comp = dx * (W - 0.5 * m_in)
        arr = comp
    left = np.arange(nbin)[:, None] * dx
    wgt = np.clip((T[None, :] - left) / dx, 0.0, 1.0)
    return (arr * wgt).sum(axis=0)


def lag_density_marginal(lag, prior):
    """E over the prior on the mean lag tau of (1/tau) exp(-lag/tau), for lag >= 0.
    prior = ('loguniform', lo, hi) or ('uniform', 0, hi), in Gyr."""
    kind, lo, hi = prior
    lag = np.maximum(lag, 1e-9)
    if kind == "loguniform":
        return (np.exp(-lag / hi) - np.exp(-lag / lo)) / (lag * np.log(hi / lo))
    return exp1(lag / hi) / hi


def numerator_marginal(kinds, t, gates, t0, prior):
    """As numerator(), with the mean lag of every easy step integrated analytically over its prior."""
    D, n = t.shape
    out = np.ones(D)
    prev = t0.copy()
    for i, k in enumerate(kinds):
        s = prev if k == "H" else np.maximum(prev, gates[:, i])
        ok = t[:, i] >= s
        if k == "E":
            out *= np.where(ok, lag_density_marginal(np.where(ok, t[:, i] - s, 1.0), prior), 0.0)
        else:
            out *= ok
        prev = t[:, i]
    return out


def likelihood(kinds, t, gates, lam, t0, T, prior=None, dx=0.02):
    """Per-draw estimate of the chain density, averaged over the prior on the mean lags.

    Without `prior`: N(tau) / Z(tau) at the sampled lags.  With `prior`: the numerator is integrated over the
    lag prior analytically and the tau-dependence of the normaliser is kept as a correction,
        N_bar / Z_ref + N(tau) (1/Z(tau) - 1/Z_ref),
    which has the same expectation over tau but far smaller Monte Carlo variance (Z_ref uses the geometric
    mean of the prior for every easy step)."""
    num = numerator(kinds, t, gates, lam, t0)
    if "E" not in kinds:
        return num / normaliser(kinds, gates, lam, t0, T, dx=dx)
    Z = normaliser(kinds, gates, lam, t0, T, dx=dx)
    if prior is None:
        return num / Z
    kind, lo, hi = prior
    tau_ref = np.sqrt(lo * hi) if kind == "loguniform" else 0.5 * hi
    Zref = normaliser(kinds, gates, np.full_like(lam, 1.0 / tau_ref), t0, T, dx=dx)
    nbar = numerator_marginal(kinds, t, gates, t0, prior)
    return nbar / Zref + num * (1.0 / Z - 1.0 / Zref)


def selftest():
    """Compare the grid normaliser with closed forms."""
    rng = np.random.default_rng(1)
    D = 6
    t0 = rng.uniform(0, 0.4, D)
    T = rng.uniform(5.0, 6.0, D)
    out = {}
    # all hard, ungated: Z = (T - t0)^n / n!
    for n in (1, 3, 5):
        g = np.zeros((D, n)); lam = np.ones((D, n))
        Z = normaliser(["H"] * n, g, lam, t0, T)
        exact = (T - t0) ** n / np.prod(np.arange(1, n + 1))
        out["allH_n%d" % n] = float(np.max(np.abs(Z / exact - 1)))
    # one easy step: Z = 1 - exp(-lam (T - max(t0, g)))
    g = rng.uniform(0.0, 3.0, (D, 1)); lam = rng.uniform(2, 100, (D, 1))
    Z = normaliser(["E"], g, lam, t0, T)
    exact = 1 - np.exp(-lam[:, 0] * (T - np.maximum(t0, g[:, 0])))
    out["oneE"] = float(np.max(np.abs(Z / exact - 1)))
    # two gated hard steps: volume of {s1 < t1 < t2 < T, t2 > g2}, s1 = max(t0, g1), g2 > s1
    g = np.column_stack([rng.uniform(0.5, 1.0, D), rng.uniform(2.0, 3.0, D)]); lam = np.ones((D, 2))
    Z = normaliser(["G", "G"], g, lam, t0, T)
    s1 = np.maximum(t0, g[:, 0]); g2 = g[:, 1]
    exact = (g2 - s1) * (T - g2) + 0.5 * (T - g2) ** 2
    out["twoG"] = float(np.max(np.abs(Z / exact - 1)))
    # easy then hard-gated, Monte Carlo check
    g = np.column_stack([np.full(D, 1.0), np.full(D, 2.0)]); lam = np.full((D, 2), 5.0)
    t0b = np.zeros(D); Tb = np.full(D, 5.0)
    Z = normaliser(["E", "G"], g, lam, t0b, Tb)
    r = np.random.default_rng(2)
    t1 = 1.0 + r.exponential(1 / 5.0, 2_000_000)
    exact = np.mean(np.clip(5.0 - np.maximum(t1, 2.0), 0, None))
    out["E_then_G_mc"] = float(np.max(np.abs(Z / exact - 1)))
    return out


if __name__ == "__main__":
    for k, v in selftest().items():
        print("%-14s max relative error %.2e" % (k, v))