"""Bayesian variable selection for censored survival data -- from scratch.

Which covariates predict time-to-event, when many event times are only known to exceed
the follow-up (right-censoring)?  We fit a **Weibull proportional-hazards** model and put
a spike-and-slab prior on the covariate coefficients, exactly as in Part 1 -- the twist is
the likelihood, which must credit a censored patient with *surviving past* their last
observation rather than with an event.

Model (standardized covariates x, event indicator delta):
    hazard   h(t | x) = alpha t^{alpha-1} exp(b0 + x'beta)
    log-lik  sum_i [ delta_i (log alpha + (alpha-1) log t_i + b0 + x_i'beta) - t_i^alpha exp(b0 + x_i'beta) ]
The second term (the cumulative hazard) is subtracted for EVERY subject, censored or not;
only events add the log-hazard -- that is how censoring enters.

Because the Weibull likelihood is not conjugate, beta is updated by adaptive random-walk
Metropolis; but the inclusion step is unchanged from the Gaussian case -- gamma_j | beta_j
is a Bernoulli from the spike/slab density ratio, since that conditional never sees the
likelihood.  So this is Part-1 SSVS with a Metropolis coefficient update and a survival
log-likelihood:  beta (Metropolis) -> gamma (Bernoulli) -> baseline (b0, alpha; Metropolis).
"""
import numpy as np


def _loglik(t, logt, delta, eta, logalpha):
    a = np.exp(logalpha)
    return np.sum(delta * (logalpha + (a - 1.0) * logt + eta) - np.exp(a * logt + eta))


def weibull_ssvs(t, delta, X, tau=0.05, c=20.0, w=0.5, ndraw=15000, burn=5000, seed=0):
    """Weibull-PH spike-and-slab SSVS with right-censoring (Metropolis-within-Gibbs)."""
    rng = np.random.default_rng(seed)
    t = np.asarray(t, float); delta = np.asarray(delta, float); X = np.asarray(X, float)
    n, p = X.shape
    Xz = (X - X.mean(0)) / X.std(0)
    logt = np.log(t)

    beta = np.zeros(p); b0 = np.log(delta.sum() / (t ** 1).sum()); logalpha = 0.0
    gamma = np.ones(p, int)
    eta = b0 + Xz @ beta
    cur = _loglik(t, logt, delta, eta, logalpha)

    sb = np.full(p, 0.15); s0 = 0.15; sa = 0.10                    # adaptive RW scales
    ab = np.zeros(p); a0 = 0.0; aa = 0.0
    G = np.zeros((ndraw, p)); B = np.zeros((ndraw, p)); A = np.zeros(ndraw)

    for it in range(ndraw + burn):
        # ---- (1) beta_j | rest : adaptive RW-Metropolis, spike/slab prior variance ----
        for j in range(p):
            bp = beta[j] + sb[j] * rng.standard_normal()
            etp = eta + Xz[:, j] * (bp - beta[j])
            new = _loglik(t, logt, delta, etp, logalpha)
            sd2 = (c * tau) ** 2 if gamma[j] else tau ** 2
            la = new - cur - 0.5 * (bp ** 2 - beta[j] ** 2) / sd2
            if np.log(rng.random()) < la:
                beta[j] = bp; eta = etp; cur = new; ab[j] += 1
        # ---- (2) gamma_j | beta_j : Bernoulli spike/slab ratio ----
        l1 = np.log(w) - np.log(c * tau) - 0.5 * beta ** 2 / (c * tau) ** 2
        l0 = np.log(1 - w) - np.log(tau) - 0.5 * beta ** 2 / tau ** 2
        gamma = (rng.random(p) < 1.0 / (1.0 + np.exp(l0 - l1))).astype(int)
        # ---- (3) intercept b0 and log-shape logalpha : RW-Metropolis ----
        b0p = b0 + s0 * rng.standard_normal(); etp = eta + (b0p - b0)
        new = _loglik(t, logt, delta, etp, logalpha)
        if np.log(rng.random()) < new - cur - 0.5 * (b0p ** 2 - b0 ** 2) / 100.0:
            b0 = b0p; eta = etp; cur = new; a0 += 1
        lap = logalpha + sa * rng.standard_normal()
        new = _loglik(t, logt, delta, eta, lap)
        if np.log(rng.random()) < new - cur - 0.5 * (lap ** 2 - logalpha ** 2) / 1.0:
            logalpha = lap; cur = new; aa += 1
        # ---- adaptation during burn-in (target ~0.3 acceptance) ----
        if it < burn and it % 50 == 49:
            r = ab / 50.0; sb *= np.exp((r - 0.3) * 0.5); ab[:] = 0
            s0 *= np.exp((a0 / 50.0 - 0.3) * 0.5); a0 = 0
            sa *= np.exp((aa / 50.0 - 0.3) * 0.5); aa = 0
        if it >= burn:
            G[it - burn] = gamma; B[it - burn] = beta; A[it - burn] = np.exp(logalpha)

    return dict(incl=G.mean(0), beta=B.mean(0), alpha=A.mean(), n=n, p=p)