"""
Interval-censored survival -- Weibull proportional hazards from scratch (MLE + Bayes RW-Metropolis).
The natural sequel to survival_mcmc.py (right-censoring only): here the event time is known only to lie
in an interval (l, u]. One likelihood covers every censoring type:

    contribution_i = S(l_i) - S(u_i),     S(t) = exp(-lambda_i t^k),  lambda_i = exp(x_i' beta)

    interval (l<u<inf): S(l)-S(u)                         (event between two visits)
    left   (l=0)       : S(0)-S(u) = 1 - S(u) = F(u)       (event before first visit)
    right  (u=inf)     : S(l)-0    = S(l)                  (no event by last visit)
    exact  (l=u)       : handled with the density f(l) = lambda k l^{k-1} S(l)

(Kept as a NEW module rather than editing survival_mcmc.py, which does right-censoring only.)
"""
import numpy as np
from scipy.optimize import minimize


def simulate_icweib(n=1500, beta=(-7.0, 0.8), k=2.0, n_visits=6, t_max=None, seed=0):
    """Weibull event times, then INTERVAL-censor by a grid of inspection visits (current-status-like)."""
    rng = np.random.default_rng(seed); beta = np.asarray(beta, float)
    grp = rng.integers(0, 2, n).astype(float); X = np.column_stack([np.ones(n), grp])
    lam = np.exp(X @ beta)
    T = (-np.log(rng.random(n)) / lam) ** (1.0 / k)
    t_max = (np.quantile(T, 0.9)) if t_max is None else t_max
    visits = np.sort(rng.uniform(0, t_max, size=(n, n_visits)), axis=1)
    l = np.zeros(n); u = np.full(n, np.inf)
    for i in range(n):
        below = visits[i][visits[i] < T[i]]; above = visits[i][visits[i] >= T[i]]
        l[i] = below[-1] if below.size else 0.0
        u[i] = above[0] if above.size else np.inf
    exact = np.zeros(n, bool)
    return l, u, X, exact, T


def _ic_negloglik(par, X, l, u, exact):
    k = np.exp(par[-1]); beta = par[:-1]; lam = np.exp(X @ beta)
    ll = 0.0
    if exact.any():
        e = exact
        ll += np.sum(np.log(lam[e] * k * l[e] ** (k - 1)) - lam[e] * l[e] ** k)
    c = ~exact
    if c.any():
        a = -lam[c] * l[c] ** k                                   # log S(l)
        bu = np.where(np.isinf(u[c]), -np.inf, -lam[c] * np.where(np.isinf(u[c]), 0.0, u[c]) ** k)  # log S(u)
        ll += np.sum(a + np.log(-np.expm1(bu - a)))               # log(S(l) - S(u)), stable
    return -ll


def _num_hess(f, x, eps=1e-4):
    n = len(x); H = np.zeros((n, n));
    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 icweib_mle(X, l, u, exact):
    k0 = X.shape[1]
    x0 = np.concatenate([np.zeros(k0), [0.0]]); x0[0] = -np.log(np.mean(l[l > 0]) + 1.0)
    res = minimize(_ic_negloglik, x0, args=(X, l, u, exact), method='Nelder-Mead',
                   options=dict(maxiter=40000, xatol=1e-8, fatol=1e-8))
    H = _num_hess(lambda p: _ic_negloglik(p, X, l, u, exact), res.x)
    se = np.sqrt(np.diag(np.linalg.inv(H)))
    return res.x, se


def icweib_rwm(X, l, u, exact, R=15000, burn=4000, prior_sd=100.0, logk_sd=5.0, seed=0, s=0.6):
    rng = np.random.default_rng(seed)
    mle, _ = icweib_mle(X, l, u, exact)
    H = _num_hess(lambda p: _ic_negloglik(p, X, l, u, exact), mle)
    L = np.linalg.cholesky(np.linalg.inv(H)) * s; d = len(mle)

    def lp(par):
        beta = par[:-1]
        return -_ic_negloglik(par, X, l, u, exact) - 0.5 * np.sum(beta ** 2) / prior_sd ** 2 \
            - 0.5 * par[-1] ** 2 / logk_sd ** 2

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


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