"""
Longitudinal / mixed logistic regression -- binary GLMM with a SUBJECT random intercept shared across each
subject's repeated measurements -- from-scratch Metropolis-within-Gibbs. The toenail / onychomycosis trial
(De Backer 1998; Lesaffre & Spiessens 2001): a binary outcome measured repeatedly per patient over time.

This is the BINARY analogue of the Poisson GLMM hpois_gibbs.py (Epil): same group-random-intercept
structure (one b per subject, applied to all of that subject's rows), Bernoulli-logistic likelihood
instead of Poisson. It is NOT seeds_gibbs (whose random effect is per-row) -- here the effect is shared
across the subject's visits, which is what models the within-patient correlation of repeated measures.

Model
-----
  y_it ~ Bernoulli(p_it),  logit p_it = x_it' beta + b_{s(i)},  b_j ~ N(0, sigma^2)
  x_it = [1, treat, time, treat*time]  (treat constant within patient; time varies across visits)

Sampler: POLYA-GAMMA data augmentation (Polson, Scott & Windle 2013) -> EXACT Gibbs, no tuning, mixes
well even when sigma is large (the toenail sigma ~4 is notoriously hard for RW/quadrature).
  augment  omega_it ~ PG(1, eta_it),  eta_it = x_it'beta + b_{s(i)};   kappa_it = y_it - 1/2
  1. beta   | omega, b   ~ Normal (precision X'Omega X + prior_prec)
  2. b_j    | omega, beta ~ Normal, per subject (precision sum_{i in j} omega_i + 1/sigma^2)  + centre -> beta_0
  3. sigma^2 | b          ~ Inverse-Gamma
(A from-scratch RW-Metropolis version mixes too slowly here -- the estimates depend on the step size.)
"""
import numpy as np
import pandas as pd
from scipy.special import expit


def _rpg1(c, rng, n_terms=100):
    """Sample PG(1, c) (vectorised) via the sum-of-Gammas representation with an analytic TAIL-MEAN
    correction: omega = (1/2pi^2) [ full_mean + sum_{k<=K} (g_k - 1)/denom_k ], where full_mean =
    (pi/2a) tanh(pi a), a = |c|/2pi. Replacing the truncated tail by its exact expectation removes the
    truncation bias that otherwise inflates the variance estimate for near-separated units (large |c|),
    so a small K suffices. g_k ~ Exp(1)."""
    c = np.abs(np.asarray(c, float)); a = c / (2 * np.pi)
    k = np.arange(1, n_terms + 1)
    denom = (k - 0.5) ** 2 + a[:, None] ** 2                     # (n, K)
    g = rng.standard_gamma(1.0, size=(len(c), n_terms))
    full = np.where(a > 1e-8, (np.pi / (2 * np.where(a > 1e-8, a, 1.0))) * np.tanh(np.pi * a), np.pi ** 2 / 2)
    fluct = np.sum((g - 1.0) / denom, axis=1)                    # mean-zero fluctuation of the leading terms
    return (1.0 / (2 * np.pi ** 2)) * (full + fluct)


def toenail_data(csv='toenail.csv'):
    d = pd.read_csv(csv)
    uniq = np.unique(d['patient']); idx = {p: i for i, p in enumerate(uniq)}
    group = d['patient'].map(idx).to_numpy()
    X = np.column_stack([np.ones(len(d)), d['treat'], d['month'], d['treat'] * d['month']])
    return d['y'].to_numpy(float), X, group


def simulate_longlogit(J=300, ni=7, beta=(-1.5, -0.2, -0.4, -0.15), sigma=2.5, seed=0):
    rng = np.random.default_rng(seed); beta = np.asarray(beta, float)
    treat = np.repeat(rng.integers(0, 2, J), ni).astype(float)
    time = np.tile(np.arange(ni), J).astype(float)
    group = np.repeat(np.arange(J), ni)
    X = np.column_stack([np.ones(J*ni), treat, time, treat*time])
    b = rng.normal(0, sigma, J)
    y = (rng.random(J*ni) < expit(X @ beta + b[group])).astype(float)
    return y, X, group, b


def _logit_mle(X, y, maxit=100):
    beta = np.zeros(X.shape[1])
    for _ in range(maxit):
        p = expit(X @ beta); w = p * (1 - p)
        step = np.linalg.solve((X * w[:, None]).T @ X + 1e-8*np.eye(X.shape[1]), X.T @ (y - p)); beta += step
        if np.max(np.abs(step)) < 1e-10:
            break
    return beta


def longlogit_gibbs(y, X, group, R=8000, burn=2000, prior_sd=10.0, a0=0.001, b0=0.001,
                    seed=0, n_terms=200):
    rng = np.random.default_rng(seed)
    y = np.asarray(y, float); n, k = X.shape; group = np.asarray(group, int); J = group.max() + 1
    kappa = y - 0.5; Xt = X.T; P0 = np.eye(k) / prior_sd ** 2
    beta = _logit_mle(X, y); b = np.zeros(J); sigma = 1.0
    keep = R - burn
    B = np.zeros((keep, k)); SIG = np.zeros(keep); Bm = np.zeros(J); BRE = np.zeros((keep, J))
    for it in range(R):
        eta = X @ beta + b[group]
        om = _rpg1(eta, rng, n_terms)                            # Polya-Gamma augmentation
        # 1. beta | omega, b  (Gaussian)
        prec = (X * om[:, None]).T @ X + P0
        cov = np.linalg.inv(prec); mb = cov @ (Xt @ (kappa - om * b[group]))
        beta = mb + np.linalg.cholesky(cov) @ rng.standard_normal(k)
        eta0 = X @ beta
        # 2. b_j | omega, beta  (per-subject Gaussian, vectorised)
        prec_b = np.bincount(group, weights=om, minlength=J) + 1.0 / sigma ** 2
        rhs_b = np.bincount(group, weights=(kappa - om * eta0), minlength=J)
        b = rhs_b / prec_b + rng.standard_normal(J) / np.sqrt(prec_b)
        m = b.mean(); b -= m; beta[0] += m                        # centre subject effects -> intercept
        # 3. sigma^2 | b  (Inverse-Gamma)
        sigma = np.sqrt(1.0 / rng.gamma(a0 + J / 2.0, 1.0 / (b0 + 0.5 * np.sum(b ** 2))))
        if it >= burn:
            kk = it - burn; B[kk] = beta; SIG[kk] = sigma; Bm += b; BRE[kk] = b
    return dict(beta=B, sigma=SIG, b=Bm / keep, bdraws=BRE)


def summary(draws, names=None):
    q = draws.shape[1]; names = names or [f'b{j}' for j in range(q)]
    return [(nm, draws[:, j].mean(), draws[:, j].std(), *np.percentile(draws[:, j], [2.5, 97.5]))
            for j, nm in enumerate(names)]