"""
Parametric survival regression (Weibull / exponential proportional hazards) with RIGHT-CENSORING --
from-scratch MLE + Bayesian RW-Metropolis. The BUGS Vol I "Leuk" example (Gehan leukemia, 6-MP vs
placebo) and the PyMC Bayesian-survival case study.

The new ingredient vs every earlier project is CENSORING: a subject still in remission at the end of
follow-up contributes only the survival probability S(t), not a density. With event indicator delta_i
(1 = relapse observed, 0 = censored):
    L = prod_i  h(t_i)^{delta_i} * S(t_i)
i.e. events contribute the hazard x survival, censored cases only the survival.

Weibull proportional-hazards parameterisation
---------------------------------------------
  h(t | x) = lambda_i * k * t^{k-1},   lambda_i = exp(x_i' beta),   S(t) = exp(-lambda_i t^k)
  log L = sum_i delta_i [ x_i'beta + log k + (k-1) log t_i ]  -  sum_i exp(x_i'beta) t_i^k
k = Weibull shape: k>1 increasing hazard, k<1 decreasing, k=1 = EXPONENTIAL (constant hazard).
exp(beta_j) is the HAZARD RATIO for covariate j. Note the log-lik is a Poisson GLM kernel in disguise
(delta as 0/1 "counts", offset log t^k) -- the piecewise-exponential/Poisson trick behind BUGS Leuk.

Priors (Bayes): beta ~ N(0, prior_sd^2),  log k ~ N(0, 1).
"""
import numpy as np
from scipy.optimize import minimize


def simulate_weibull(n=600, beta=(-1.0, -1.3), k=1.4, cens_scale=3.0, seed=0):
    """X = [1, treat]; right-censor by an independent exponential follow-up time."""
    rng = np.random.default_rng(seed)
    beta = np.asarray(beta, float)
    treat = rng.integers(0, 2, n).astype(float)
    X = np.column_stack([np.ones(n), treat])
    lam = np.exp(X @ beta)
    T = (-np.log(rng.random(n)) / lam) ** (1.0 / k)          # Weibull event times
    C = rng.exponential(cens_scale, n)                        # random right-censoring
    t = np.minimum(T, C); delta = (T <= C).astype(float)
    return t, X, delta


def _negloglik(params, X, t, delta, fix_k):
    k = 1.0 if fix_k else np.exp(params[-1])
    beta = params if fix_k else params[:-1]
    xb = X @ beta
    ll = np.sum(delta * (xb + np.log(k) + (k - 1) * np.log(t))) - np.sum(np.exp(xb) * t ** k)
    return -ll


def _num_hess(f, x, eps=1e-4):
    n = len(x); H = np.zeros((n, n)); fx = f(x)
    for i in range(n):
        for j in range(i, n):
            xpp = x.copy(); xpp[i] += eps; xpp[j] += eps
            xpm = x.copy(); xpm[i] += eps; xpm[j] -= eps
            xmp = x.copy(); xmp[i] -= eps; xmp[j] += eps
            xmm = x.copy(); xmm[i] -= eps; xmm[j] -= eps
            H[i, j] = H[j, i] = (f(xpp) - f(xpm) - f(xmp) + f(xmm)) / (4 * eps ** 2)
    return H


def weibull_mle(X, t, delta, fix_k=False):
    k0 = X.shape[1] + (0 if fix_k else 1)
    x0 = np.zeros(k0)
    res = minimize(_negloglik, x0, args=(X, t, delta, fix_k), method='Nelder-Mead',
                   options=dict(xatol=1e-8, fatol=1e-8, maxiter=20000))
    H = _num_hess(lambda p: _negloglik(p, X, t, delta, fix_k), res.x)
    se = np.sqrt(np.diag(np.linalg.inv(H)))
    return res.x, se, -res.fun


def weibull_rwm(X, t, delta, R=12000, burn=3000, prior_sd=10.0, fix_k=False, seed=0, s=1.0):
    rng = np.random.default_rng(seed)
    mle, _, _ = weibull_mle(X, t, delta, fix_k)
    H = _num_hess(lambda p: _negloglik(p, X, t, delta, fix_k), mle)
    L = np.linalg.cholesky(np.linalg.inv(H)) * s
    p = mle.copy(); d = len(p)

    def logpost(par):
        beta = par if fix_k else par[:-1]
        lp = -0.5 * np.sum(beta ** 2) / prior_sd ** 2
        if not fix_k:
            lp += -0.5 * par[-1] ** 2          # log k ~ N(0,1)
        return -_negloglik(par, X, t, delta, fix_k) + lp

    keep = R - burn; P = np.zeros((keep, d)); acc = 0; cur = logpost(p)
    for it in range(R):
        prop = p + L @ rng.standard_normal(d)
        lp = logpost(prop)
        if np.log(rng.random()) < lp - cur:
            p, cur = prop, lp; acc += 1
        if it >= burn:
            P[it - burn] = p
    return dict(draws=P, accept=acc / R, fix_k=fix_k)


def km_estimator(t, delta):
    """Kaplan-Meier product-limit estimate. Returns (times, survival)."""
    order = np.argsort(t); t, delta = t[order], delta[order]
    uniq = np.unique(t[delta == 1])
    n = len(t); S = 1.0; ts = [0.0]; Ss = [1.0]
    for ti in uniq:
        at_risk = np.sum(t >= ti)
        d = np.sum((t == ti) & (delta == 1))
        S *= (1 - d / at_risk); ts.append(ti); Ss.append(S)
    return np.array(ts), np.array(Ss)


def summary(draws, names):
    return [(nm, draws[:, j].mean(), draws[:, j].std(), *np.percentile(draws[:, j], [2.5, 97.5]))
            for j, nm in enumerate(names)]