"""
npsurv.py -- NONPARAMETRIC BAYESIAN SURVIVAL: piecewise-exponential hazard (from scratch).

Backs the notebooks in  "Nonparametric Bayesian Survival".

The survival arc fitted PARAMETRIC hazards (Weibull/exponential) and the Cox model as a
piecewise-exponential Poisson GLM. Here the baseline hazard itself is NONPARAMETRIC: chop time
into K intervals and let the hazard be a free constant lambda_k on each, with a prior that either
leaves the intervals independent or ties neighbours together.

    h(t) = lambda_k   for t in interval k,     x_i shifts it proportionally:  h_i(t) = lambda_k e^{x_i'beta}

Written on person-time, the piecewise-exponential likelihood is exactly Poisson: with exposure
e_ik (time subject i spent in interval k) and event indicator d_ik,
        d_ik ~ Poisson( lambda_k e^{x_i'beta} e_ik ).

Two priors on the baseline:
  * GAMMA PROCESS (Kalbfleisch 1978) -- independent increments of the cumulative hazard, i.e.
    lambda_k ~ Gamma(a,b) independently. This is CONJUGATE: lambda_k | . ~ Gamma(a + D_k, b + E_k),
    with D_k the events and E_k = sum_i e_ik e^{x_i'beta} the (covariate-weighted) exposure in
    interval k. The nonparametric baseline is then a direct Gibbs draw.
  * SMOOTHED -- a random walk on log lambda_k (log lambda_k - log lambda_{k-1} ~ N(0,sigma^2)),
    which borrows strength across intervals for a smooth hazard; updated by Metropolis.

The regression coefficient beta (proportional hazards) is updated by random-walk Metropolis on the
Poisson log-likelihood. Survival is S(t|x) = exp(-H_0(t) e^{x'beta}) from the sampled hazard, with
credible bands. Kaplan-Meier and Nelson-Aalen (the frequentist nonparametric estimators) are here too.
"""

import numpy as np


# --------------------------------------------------------------------------- #
#  person-time expansion                                                        #
# --------------------------------------------------------------------------- #

def expand_pwe(time, status, cuts):
    """exposure e_ik and event d_ik matrices for the piecewise-exponential model.
    cuts: increasing interval boundaries starting at 0. Returns E, D each (N, K)."""
    time = np.asarray(time, float); status = np.asarray(status, int)
    N = len(time); K = len(cuts) - 1
    E = np.zeros((N, K)); D = np.zeros((N, K))
    for i in range(N):
        for k in range(K):
            lo, hi = cuts[k], cuts[k + 1]
            if time[i] <= lo:
                break
            E[i, k] = min(time[i], hi) - lo
            if time[i] <= hi:
                if status[i] == 1:
                    D[i, k] = 1.0
                break
    return E, D


# --------------------------------------------------------------------------- #
#  Gibbs sampler                                                                #
# --------------------------------------------------------------------------- #

def npsurv_gibbs(time, status, cuts, rng, X=None, a=0.2, b=0.2, beta_sd=10.0,
                 smooth=False, draws=3000, burn=1500, step=0.3):
    """Piecewise-exponential Bayesian survival with a nonparametric baseline hazard.
    smooth=False -> independent Gamma-process prior (conjugate); True -> random-walk-smoothed
    log-hazard. X:(N,p) optional covariates (proportional hazards). Returns posterior draws of the
    interval hazards lambda:(draws,K), coefficients beta:(draws,p), and the interval widths."""
    time = np.asarray(time, float); status = np.asarray(status, int)
    E, D = expand_pwe(time, status, cuts); N, K = E.shape
    Dk = D.sum(0); di = D.sum(1); width = np.diff(cuts)
    p = 0 if X is None else X.shape[1]
    X = np.zeros((N, 0)) if X is None else np.asarray(X, float)
    beta = np.zeros(p); acc = 0
    lam = np.full(K, max(Dk.sum() / E.sum(), 1e-3))
    gamma = np.log(lam); sigma2 = 0.5
    LAM = np.empty((draws, K)); BETA = np.empty((draws, p))
    for it in range(draws + burn):
        w = np.exp(X @ beta) if p else np.ones(N)
        Ek = (E * w[:, None]).sum(0)
        # ---- baseline hazard ----
        if not smooth:
            lam = rng.gamma(a + Dk, 1.0 / (b + Ek))
        else:
            gamma = _update_gamma(gamma, Dk, Ek, sigma2, rng, step)
            # smoothing variance (InvGamma on the RW increments)
            dg = np.diff(gamma)
            sigma2 = 1.0 / rng.gamma(1e-2 + (K - 1) / 2, 1.0 / (1e-2 + 0.5 * (dg @ dg)))
            lam = np.exp(gamma)
        # ---- proportional-hazards coefficient (Metropolis) ----
        if p:
            EL = E @ lam
            def ll(bb):
                eta = X @ bb
                return np.sum(di * eta - np.exp(eta) * EL) - 0.5 * (bb @ bb) / beta_sd ** 2
            prop = beta + step * rng.standard_normal(p)
            if np.log(rng.random()) < ll(prop) - ll(beta):
                beta = prop; acc += 1
        if it >= burn:
            LAM[it - burn] = lam; BETA[it - burn] = beta
    return dict(lam=LAM, beta=BETA, width=width, cuts=np.asarray(cuts, float),
                accept=acc / (draws + burn))


def _update_gamma(gamma, Dk, Ek, sigma2, rng, step):
    """per-interval Metropolis for log-hazard under a random-walk prior."""
    K = len(gamma); g = gamma.copy()
    for k in range(K):
        prop = g[k] + step * rng.standard_normal()
        def lp(v, idx):
            like = Dk[idx] * v - np.exp(v) * Ek[idx]
            pri = 0.0
            if idx > 0: pri += -0.5 * (v - g[idx - 1]) ** 2 / sigma2
            if idx < K - 1: pri += -0.5 * (g[idx + 1] - v) ** 2 / sigma2
            return like + pri
        if np.log(rng.random()) < lp(prop, k) - lp(g[k], k):
            g[k] = prop
    return g


# --------------------------------------------------------------------------- #
#  posterior survival / hazard on a grid                                        #
# --------------------------------------------------------------------------- #

def survival_curves(res, tgrid, x=None):
    """posterior survival S(t|x) on tgrid from the sampled hazards. Returns (draws, len(tgrid))."""
    cuts = res["cuts"]; LAM = res["lam"]; draws, K = LAM.shape
    tgrid = np.asarray(tgrid, float)
    # cumulative baseline hazard at each grid time
    H = np.zeros((draws, len(tgrid)))
    for j, t in enumerate(tgrid):
        contrib = np.clip(np.minimum(t, cuts[1:]) - cuts[:-1], 0, None)   # time in each interval up to t
        H[:, j] = LAM @ contrib
    if x is not None and res["beta"].shape[1] > 0:
        H = H * np.exp(res["beta"] @ np.asarray(x, float))[:, None]
    return np.exp(-H)


def cum_hazard(res, tgrid):
    cuts = res["cuts"]; LAM = res["lam"]; tgrid = np.asarray(tgrid, float)
    H = np.zeros((LAM.shape[0], len(tgrid)))
    for j, t in enumerate(tgrid):
        H[:, j] = LAM @ np.clip(np.minimum(t, cuts[1:]) - cuts[:-1], 0, None)
    return H


# --------------------------------------------------------------------------- #
#  frequentist nonparametric estimators + simulation                            #
# --------------------------------------------------------------------------- #

def kaplan_meier(time, status):
    time = np.asarray(time, float); status = np.asarray(status, int)
    ts = np.sort(np.unique(time[status == 1])); S = np.empty(len(ts)); s = 1.0
    for i, t in enumerate(ts):
        d = np.sum((time == t) & (status == 1)); n = np.sum(time >= t)
        s *= (1 - d / n); S[i] = s
    return ts, S

def nelson_aalen(time, status):
    time = np.asarray(time, float); status = np.asarray(status, int)
    ts = np.sort(np.unique(time[status == 1])); H = np.empty(len(ts)); h = 0.0
    for i, t in enumerate(ts):
        d = np.sum((time == t) & (status == 1)); n = np.sum(time >= t)
        h += d / n; H[i] = h
    return ts, H

def simulate_survival(n, hazard, tmax, censor_time, rng):
    """simulate event times from a hazard function h(t) by inversion of the cumulative hazard on a
    fine grid, with administrative censoring at censor_time. hazard: callable t->h(t)."""
    tg = np.linspace(0, tmax, 4000); H = np.cumsum(hazard(tg)) * (tg[1] - tg[0])
    u = -np.log(rng.random(n))
    T = np.interp(u, H, tg, right=tmax)
    C = rng.uniform(0.4 * censor_time, censor_time, n)
    time = np.minimum(T, C); status = (T <= C).astype(int)
    return time, status