"""
Scotland lip cancer -- Besag-York-Mollie (BYM) spatial Poisson, from-scratch Metropolis-within-Gibbs.
The classic disease-mapping model (Clayton & Kaldor 1987; Besag, York & Mollie 1991), WinBUGS/GeoBUGS
"Lip" example. 56 Scottish districts; the new ingredient vs the hierarchical-Poisson projects is a
STRUCTURED spatial random effect with an intrinsic CAR prior over the map's adjacency graph.

Model
-----
  O_i ~ Poisson(E_i * exp(eta_i)),   eta_i = alpha + beta * x_i + theta_i + phi_i
    O_i = observed lip-cancer cases,  E_i = age-standardised EXPECTED cases (offset),
    x_i = AFF/10  (% in agriculture/fishing/forestry -- outdoor sun exposure), divided by 10.
  theta_i ~ N(0, sigma_theta^2)                     -- UNSTRUCTURED heterogeneity (exchangeable)
  phi    ~ ICAR(sigma_phi^2)                          -- STRUCTURED spatial effect:
           phi_i | phi_{-i} ~ N( mean_{j~i} phi_j ,  sigma_phi^2 / n_i )
  alpha, beta ~ N(0, 100^2);  precisions 1/sigma^2 ~ Gamma(0.5, 0.0005)  (WinBUGS vague priors).

The ICAR prior penalises differences between neighbours: log p(phi) = -1/(2 sigma_phi^2) * sum_{i~j}
(phi_i - phi_j)^2. It is improper (invariant to adding a constant) -> phi is re-centred to sum-to-zero
each sweep, with the level absorbed into alpha (likewise theta). The fitted relative risk of area i is
exp(eta_i); the smoothed SMR.

Sampler (Metropolis-within-Gibbs)
---------------------------------
  1. (alpha, beta)        -- block RW-Metropolis.
  2. theta_i              -- VECTORISED RW-Metropolis (theta_i conditionally independent given the rest).
  3. phi_i                -- SEQUENTIAL RW-Metropolis (neighbours are coupled: must use current values).
  +  centre theta, phi -> alpha (location moves; phi's ICAR level is unidentified).
  4. sigma_theta^2        -- Inverse-Gamma (n/2 df).
     sigma_phi^2          -- Inverse-Gamma ((n-1)/2 df; ICAR has rank n-1).
"""
import numpy as np
import pandas as pd
import ast


def scotland_data(csv='scotland_lip.csv'):
    d = pd.read_csv(csv)
    O = d['CANCER'].to_numpy(float); E = d['CEXP'].to_numpy(float)
    x = d['AFF'].to_numpy(float) / 10.0                       # per 10% in AFF
    adj = [np.array(ast.literal_eval(a), int) - 1 for a in d['ADJ']]   # 0-indexed neighbour lists
    return O, E, x, adj, d['NAME'].tolist()


def bym_gibbs(O, E, x, adj, R=20000, burn=5000, seed=0, prior_fixed_sd=100.0,
              g_a=0.5, g_b=0.0005, s_ab=0.07, s_th=1.4, s_ph=1.4):
    rng = np.random.default_rng(seed); n = len(O)
    ni = np.array([len(a) for a in adj], float)
    alpha = np.log(O.sum() / E.sum()); beta = 0.0
    theta = np.zeros(n); phi = np.zeros(n); s2_th = 0.1; s2_ph = 0.1
    keep = R - burn
    A = np.zeros(keep); B = np.zeros(keep); STH = np.zeros(keep); SPH = np.zeros(keep)
    TH = np.zeros((keep, n)); PH = np.zeros((keep, n))
    acc_ab = 0; acc_th = 0.0; acc_ph = 0.0
    for it in range(R):
        # 1. (alpha, beta) block RW-Metropolis
        e0 = alpha + beta * x + theta + phi; mu = E * np.exp(e0)
        ll = np.sum(O * e0 - mu) - 0.5 * (alpha**2 + beta**2) / prior_fixed_sd**2
        ap = alpha + s_ab * rng.standard_normal(); bp = beta + s_ab * rng.standard_normal()
        ep = ap + bp * x + theta + phi; mup = E * np.exp(ep)
        llp = np.sum(O * ep - mup) - 0.5 * (ap**2 + bp**2) / prior_fixed_sd**2
        if np.log(rng.random()) < llp - ll:
            alpha, beta = ap, bp; acc_ab += 1
        # 2. theta_i  (vectorised; conditionally independent; curvature-scaled step)
        e0 = alpha + beta * x + theta + phi; mu = E * np.exp(e0)
        curv = mu + 1.0 / s2_th                              # Poisson curvature + prior precision
        prop = theta + s_th * rng.standard_normal(n) / np.sqrt(curv)
        mup = E * np.exp(e0 + (prop - theta))
        d = O * (prop - theta) - (mup - mu) - 0.5 * (prop**2 - theta**2) / s2_th
        ok = np.log(rng.random(n)) < d; theta = np.where(ok, prop, theta); acc_th += ok.mean()
        # 3. phi_i  (sequential; ICAR conditional N(neighbour mean, s2_ph/n_i); curvature-scaled)
        a_ph = 0
        for i in range(n):
            nbar = phi[adj[i]].mean(); prec = ni[i] / s2_ph
            cur = alpha + beta * x[i] + theta[i] + phi[i]; mui = E[i] * np.exp(cur)
            step = s_ph / np.sqrt(mui + prec)
            pp = phi[i] + step * rng.standard_normal(); muip = E[i] * np.exp(cur + (pp - phi[i]))
            d = O[i] * (pp - phi[i]) - (muip - mui) - 0.5 * prec * ((pp - nbar)**2 - (phi[i] - nbar)**2)
            if np.log(rng.random()) < d:
                phi[i] = pp; a_ph += 1
        acc_ph += a_ph / n
        # centre theta, phi -> alpha (location moves)
        m = theta.mean(); theta -= m; alpha += m
        m = phi.mean();   phi -= m;   alpha += m
        # 4. variances (Inverse-Gamma; precision ~ Gamma(g_a, g_b))
        s2_th = 1.0 / rng.gamma(g_a + n / 2.0, 1.0 / (g_b + 0.5 * np.sum(theta**2)))
        S = 0.5 * sum(np.sum((phi[i] - phi[adj[i]])**2) for i in range(n))   # sum over edges
        s2_ph = 1.0 / rng.gamma(g_a + (n - 1) / 2.0, 1.0 / (g_b + 0.5 * S))
        if it >= burn:
            k = it - burn
            A[k] = alpha; B[k] = beta; STH[k] = np.sqrt(s2_th); SPH[k] = np.sqrt(s2_ph)
            TH[k] = theta; PH[k] = phi
    return dict(alpha=A, beta=B, sd_theta=STH, sd_phi=SPH, theta=TH, phi=PH,
                accept_ab=acc_ab / R, accept_theta=acc_th / R, accept_phi=acc_ph / R)


def summary(draws, names=None):
    d = draws.reshape(draws.shape[0], -1); q = d.shape[1]
    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)]