"""
Pumps: conjugate hierarchical Poisson (gamma-Poisson) -- from-scratch Gibbs.
BUGS Vol I "Pumps" (Gaver & O'Muircheartaigh 1987). Contrast with hpois_gibbs.py (Epil): there the
random effect is a NORMAL on the log-rate (non-conjugate -> Metropolis-within-Gibbs); here it is a
GAMMA on the rate (conjugate -> almost-closed-form Gibbs).

Model
-----
  x_i ~ Poisson(theta_i * t_i),      t_i = exposure (offset, in 1000s of hours)
  theta_i ~ Gamma(alpha, beta)       (failure rate of pump i; hierarchical)
  alpha ~ Exp(1),  beta ~ Gamma(a_beta=0.1, b_beta=1.0)

Conjugacy (the whole point)
---------------------------
  * theta_i | x_i, t_i, alpha, beta  ~  Gamma(alpha + x_i,  beta + t_i)      EXACT Gibbs
      (gamma is the conjugate prior for a Poisson rate; the offset t_i just adds to the rate param)
  * beta | theta, alpha              ~  Gamma(a_beta + n*alpha,  b_beta + sum theta_i)   EXACT Gibbs
      (beta is the rate of the Gamma(alpha,beta) on theta; gamma prior on beta is conjugate)
  * alpha | theta, beta              -- NOT conjugate (the Gamma shape sits inside Gamma(alpha)) ->
      a single random-walk Metropolis step on log(alpha).
So 2 of 3 blocks are exact draws; only the shape alpha needs Metropolis. (Epil needed Metropolis for
beta AND every random effect.)
"""
import numpy as np
import pandas as pd
from scipy.special import gammaln


def pumps_data(csv='pumps.csv'):
    """BUGS Vol I Pumps: x = failures, t = operating time (1000s of hours)."""
    d = pd.read_csv(csv)
    return d['x'].to_numpy(float), d['t'].to_numpy(float)


def pumps_gibbs(x, t, R=11000, burn=1000, seed=0, a_beta=0.1, b_beta=1.0, s_alpha=0.4):
    rng = np.random.default_rng(seed); n = len(x)
    alpha, beta = 1.0, 1.0
    keep = R - burn
    TH = np.zeros((keep, n)); AL = np.zeros(keep); BE = np.zeros(keep); acc = 0
    for r in range(R):
        # 1. theta_i | .  ~ Gamma(shape = alpha + x_i, rate = beta + t_i)   [EXACT, conjugate]
        theta = rng.gamma(alpha + x, 1.0 / (beta + t))
        # 2. beta | .      ~ Gamma(shape = a_beta + n*alpha, rate = b_beta + sum theta)  [EXACT]
        beta = rng.gamma(a_beta + n * alpha, 1.0 / (b_beta + theta.sum()))
        # 3. alpha | .     -- RW-Metropolis on log(alpha); prior alpha ~ Exp(1)
        st = np.sum(np.log(theta))
        def loglik_alpha(a):                       # up to a constant in alpha
            return n * (a * np.log(beta) - gammaln(a)) + a * st - a   # last term = Exp(1) log-prior
        la = np.log(alpha); lap = la + s_alpha * rng.standard_normal(); ap = np.exp(lap)
        if np.log(rng.random()) < loglik_alpha(ap) - loglik_alpha(alpha) + (lap - la):  # +Jacobian
            alpha = ap; acc += 1
        if r >= burn:
            TH[r - burn] = theta; AL[r - burn] = alpha; BE[r - burn] = beta
    return dict(theta=TH, alpha=AL, beta=BE, accept=acc / R)


def summary(draws, names=None):
    q = draws.shape[1] if draws.ndim == 2 else 1
    d = draws.reshape(-1, q)
    names = names or [f'p{j}' for j in range(q)]
    return [(nm, d[:, j].mean(), d[:, j].std(), *np.percentile(d[:, j], [2.5, 97.5]))
            for j, nm in enumerate(names)]