"""
Seeds: binomial GLMM (logistic regression with group random effects + over-dispersion) -- from-scratch
Metropolis-within-Gibbs. BUGS Vol I "Seeds" (Crowder 1978; Breslow & Clayton 1993). The BINOMIAL twin of
the hierarchical-Poisson Epil model -- a plate random effect absorbs extra-binomial variation.

Model
-----
  r_i ~ Binomial(n_i, p_i),   logit(p_i) = x_i' beta + b_i,   b_i ~ N(0, sigma^2),  i = 1..I plates
  x_i = [1, seed-type, root-extract, seed-type*extract]  (a 2x2 factorial with interaction)
The plate random effect b_i captures OVER-DISPERSION: grouped binomial counts vary more than a single
logistic curve allows; b_i is the logistic-scale analogue of the NegBin's gamma or the Poisson-GLMM's
log-normal mixing. Self-contained (does not reuse/modify the older hierarchical-Poisson module).

Sampler (Metropolis-within-Gibbs)
---------------------------------
  1. beta   -- block RW-Metropolis (binomial-logistic loglik); proposal N(beta, s^2 (I_pool+prior_prec)^-1)
               from the pooled logistic Fisher info X' diag(n p(1-p)) X (computed once). s = 2.38/sqrt(k).
  2. b_i    -- per-plate RW-Metropolis, vectorised; step SD = s_b / sqrt(curvature_i),
               curvature_i = n_i p_i(1-p_i) + 1/sigma^2.  + centre b -> beta_0 (location move).
  3. sigma^2 -- Inverse-Gamma conjugate from {b_i}.
"""
import numpy as np
import pandas as pd
from scipy.special import expit


def seeds_data(csv='seeds.csv'):
    d = pd.read_csv(csv)
    X = np.column_stack([np.ones(len(d)), d.x1, d.x2, d.x1 * d.x2])   # 1, type, extract, interaction
    return d['r'].to_numpy(float), d['n'].to_numpy(float), X


def simulate_seeds(I=300, beta=(-0.5, 0.1, 1.3, -0.8), sigma=0.35, nbar=40, seed=0):
    rng = np.random.default_rng(seed); beta = np.asarray(beta, float)
    x1 = rng.integers(0, 2, I).astype(float); x2 = rng.integers(0, 2, I).astype(float)
    X = np.column_stack([np.ones(I), x1, x2, x1 * x2])
    b = rng.normal(0, sigma, I); p = expit(X @ beta + b)
    n = rng.integers(nbar // 2, nbar * 2, I).astype(float)
    r = rng.binomial(n.astype(int), p).astype(float)
    return r, n, X, b


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


def seeds_gibbs(r, n, X, R=20000, burn=5000, prior_sd=10.0, a0=0.001, b0=0.001,
                seed=0, s_beta=None, s_b=1.6):
    rng = np.random.default_rng(seed); I, k = X.shape
    s_beta = (2.38 / np.sqrt(k)) if s_beta is None else s_beta
    bp = _logit_mle(X, r, n); p0 = expit(X @ bp)
    Ipool = (X * (n * p0 * (1 - p0))[:, None]).T @ X
    Lb = np.linalg.cholesky(np.linalg.inv(Ipool + np.eye(k) / prior_sd ** 2))

    beta = bp.copy(); b = np.zeros(I); sigma = 0.3
    keep = R - burn
    B = np.zeros((keep, k)); SIG = np.zeros(keep); Bm = np.zeros(I); BRE = np.zeros((keep, I))
    accb = 0.0; acc_beta = 0
    for it in range(R):
        # 1. beta  (block RW-Metropolis)
        eta = X @ beta + b
        ll = np.sum(r * eta - n * np.logaddexp(0, eta)) - 0.5 * np.sum(beta ** 2) / prior_sd ** 2
        prop = beta + s_beta * (Lb @ rng.standard_normal(k)); ep = X @ prop + b
        llp = np.sum(r * ep - n * np.logaddexp(0, ep)) - 0.5 * np.sum(prop ** 2) / prior_sd ** 2
        if np.log(rng.random()) < llp - ll:
            beta = prop; acc_beta += 1
        eta0 = X @ beta
        # 2. b_i  (per-plate RW-Metropolis, vectorised, curvature-scaled)
        eta = eta0 + b; p = expit(eta)
        curv = n * p * (1 - p) + 1.0 / sigma ** 2
        prop_b = b + s_b * rng.standard_normal(I) / np.sqrt(curv)
        ep = eta0 + prop_b
        d = r * (prop_b - b) - n * (np.logaddexp(0, ep) - np.logaddexp(0, eta)) \
            - 0.5 * (prop_b ** 2 - b ** 2) / sigma ** 2
        ok = np.log(rng.random(I)) < d; b = np.where(ok, prop_b, b); accb += ok.mean()
        m = b.mean(); b -= m; beta[0] += m                       # centre random effects -> intercept
        # 3. sigma^2  (Inverse-Gamma)
        sigma = np.sqrt(1.0 / rng.gamma(a0 + I / 2.0, 1.0 / (b0 + 0.5 * np.sum(b ** 2))))
        if it >= burn:
            B[it - burn] = beta; SIG[it - burn] = sigma; Bm += b; BRE[it - burn] = b
    return dict(beta=B, sigma=SIG, b=Bm / keep, bdraws=BRE,
                accept_beta=acc_beta / R, accept_b=accb / R)


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