"""
Frailty survival: Weibull proportional hazards with a per-subject random effect (FRAILTY) and
right-censoring -- from-scratch Metropolis-within-Gibbs. BUGS Vol I "Kidney" (McGilchrist & Aisbett
1991): recurrence times of catheter-insertion infection, 2 per patient, with patient frailties.

This is the hierarchical (random-effects) extension of survival_mcmc.py -- and it is literally the
SURVIVAL LIKELIHOOD (from survival_mcmc.py) x the RANDOM-INTERCEPT SAMPLER (from hpois_gibbs.py).

Model
-----
  h(t_ij | x, w) = lambda_ij * k * t^{k-1},  lambda_ij = exp(x_ij' beta + w_{p(i)}),  w_j ~ N(0, sigma_w^2)
  events contribute h(t)S(t), censored only S(t); w_j is the frailty of patient j (shared by their obs).

Sampler (Metropolis-within-Gibbs)
---------------------------------
  1. (beta, log k)  -- block RW-Metropolis (Weibull censored loglik with the frailty as an offset);
                       proposal = chol(inv(Hessian)) from the no-frailty MLE.
  2. w_j            -- per-patient RW-Metropolis, VECTORISED. The frailty conditional is exactly the
                       hpois random-intercept form: "event count" A_j = sum_{i in j} delta_ij, and
                       "rate" S_j = sum_{i in j} exp(x_ij'beta) t_ij^k, so
                       loglik_j(w) = A_j w - S_j e^{w} - w^2/(2 sigma_w^2);
                       curvature-scaled step s_w / sqrt(S_j e^{w} + 1/sigma_w^2).
  +  centre w -> beta_0 (location move).
  3. sigma_w^2      -- Inverse-Gamma conjugate from {w_j}.
"""
import numpy as np
from survival_mcmc import weibull_mle, _negloglik, _num_hess


def frailty_weibull_gibbs(t, X, delta, patient, R=15000, burn=5000, seed=0,
                          prior_sd=10.0, logk_sd=2.0, a0=2.0, b0=1.0, s_w=1.3, s_beta=0.6):
    rng = np.random.default_rng(seed)
    t = np.asarray(t, float); delta = np.asarray(delta, float); patient = np.asarray(patient, int)
    n, p = X.shape; J = patient.max() + 1; logt = np.log(t)
    A = np.bincount(patient, weights=delta, minlength=J)        # events per patient (fixed)

    mle, _, _ = weibull_mle(X, t, delta)                        # no-frailty init + proposal
    H = _num_hess(lambda q: _negloglik(q, X, t, delta, False), mle)
    Lprop = np.linalg.cholesky(np.linalg.inv(H)) * s_beta
    beta = mle[:p].copy(); logk = mle[p]; w = np.zeros(J); sigma = 0.5

    def wll(b, lk, off):                                        # Weibull censored log-likelihood
        k = np.exp(lk); xb = X @ b + off
        return np.sum(delta * (xb + lk + (k - 1) * logt)) - np.sum(np.exp(xb) * t ** k)

    keep = R - burn
    B = np.zeros((keep, p)); K = np.zeros(keep); SIG = np.zeros(keep); W = np.zeros((keep, J))
    accb = 0; accw = 0.0
    for it in range(R):
        off = w[patient]
        # 1. (beta, log k) block RW-Metropolis
        cur = wll(beta, logk, off) - 0.5*np.sum(beta**2)/prior_sd**2 - 0.5*logk**2/logk_sd**2
        st = Lprop @ rng.standard_normal(p + 1); bp = beta + st[:p]; lkp = logk + st[p]
        pr = wll(bp, lkp, off) - 0.5*np.sum(bp**2)/prior_sd**2 - 0.5*lkp**2/logk_sd**2
        if np.log(rng.random()) < pr - cur:
            beta, logk = bp, lkp; accb += 1
        # 2. w_j vectorised RW-Metropolis (hpois random-intercept form)
        k = np.exp(logk); base = np.exp(X @ beta) * t ** k
        Sj = np.bincount(patient, weights=base, minlength=J)
        curv = Sj * np.exp(w) + 1.0 / sigma ** 2
        wp = w + s_w * rng.standard_normal(J) / np.sqrt(curv)
        d = A * (wp - w) - Sj * (np.exp(wp) - np.exp(w)) - 0.5 * (wp**2 - w**2) / sigma**2
        ok = np.log(rng.random(J)) < d; w = np.where(ok, wp, w); accw += ok.mean()
        m = w.mean(); w -= m; beta[0] += m                     # centre frailties -> intercept
        # 3. sigma_w^2 Inverse-Gamma
        sigma = np.sqrt(1.0 / rng.gamma(a0 + J/2.0, 1.0 / (b0 + 0.5*np.sum(w**2))))
        if it >= burn:
            i = it - burn; B[i] = beta; K[i] = np.exp(logk); SIG[i] = sigma; W[i] = w
    return dict(beta=B, k=K, sigma_w=SIG, w=W.mean(0), W=W,
                accept_beta=accb/R, accept_w=accw/R)


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)]