"""
Single-level count regression from scratch: Poisson and Negative Binomial.
MLE (Newton/IRLS, with asymptotic SEs) and Bayesian (random-walk Metropolis) for comparison.

Models
------
Poisson:           y_i ~ Poisson(mu_i),      mu_i = exp(x_i'beta)
                   E[y]=Var[y]=mu  (equidispersion)

Negative Binomial (NB2):  y_i ~ NB(mu_i, r),  mu_i = exp(x_i'beta)
                   Var[y] = mu + mu^2 / r     (r = size/dispersion; r -> inf gives Poisson)
                   Overdispersion when Var > mean, i.e. finite r.

Algorithms
----------
- MLE: Newton-Raphson / IRLS (Poisson uses Fisher scoring = IRLS; canonical log link so
  observed info = expected info = X' diag(mu) X). NB optimises (beta, log r) jointly.
  Asymptotic SEs from the inverse (negative) Hessian / Fisher information.
- Bayes: random-walk Metropolis with a proposal scaled by the posterior curvature at the
  MLE,  N(beta_cur, s^2 (H + prior_prec)^{-1}),  s = 2.38/sqrt(dim) (the Roberts optimal scale).
  Prior beta ~ N(0, prior_sd^2 I); for NB, log r ~ N(0, prior_sd_r^2).
"""
import numpy as np
from scipy import optimize
from scipy.special import gammaln


# ───────────────────────── data ─────────────────────────
def simulate_counts(n, beta, r=None, seed=0):
    """Simulate counts. r=None -> Poisson; r=finite -> NB2 (overdispersed)."""
    rng = np.random.default_rng(seed)
    beta = np.asarray(beta, float); k = len(beta)
    X = np.column_stack([np.ones(n)] + [rng.normal(size=n) for _ in range(k - 1)])
    mu = np.exp(X @ beta)
    if r is None:
        y = rng.poisson(mu)
    else:                                   # NB2 as a gamma-Poisson mixture
        lam = rng.gamma(shape=r, scale=mu / r)
        y = rng.poisson(lam)
    return X, y


# ───────────────────────── Poisson ─────────────────────────
def poisson_loglik(beta, X, y):
    eta = X @ beta
    return float(np.sum(y * eta - np.exp(eta) - gammaln(y + 1)))


def poisson_mle(X, y, tol=1e-10, maxit=100):
    """Newton-Raphson / IRLS. Returns (beta_hat, se, cov)."""
    k = X.shape[1]; beta = np.zeros(k)
    for _ in range(maxit):
        mu = np.exp(X @ beta)
        g = X.T @ (y - mu)                  # score
        I = (X * mu[:, None]).T @ X         # Fisher information = X' diag(mu) X
        step = np.linalg.solve(I, g); beta += step
        if np.max(np.abs(step)) < tol:
            break
    cov = np.linalg.inv((X * np.exp(X @ beta)[:, None]).T @ X)
    return beta, np.sqrt(np.diag(cov)), cov


def poisson_rwm(y, X, R=8000, burn=2000, prior_sd=10.0, seed=0, s=None):
    rng = np.random.default_rng(seed); n, k = X.shape
    b0, _, _ = poisson_mle(X, y)
    I = (X * np.exp(X @ b0)[:, None]).T @ X
    L = np.linalg.cholesky(np.linalg.inv(I + np.eye(k) / prior_sd ** 2))   # proposal scale
    s = (2.38 / np.sqrt(k)) if s is None else s
    beta = b0.copy()
    lp = poisson_loglik(beta, X, y) - 0.5 * np.sum(beta ** 2) / prior_sd ** 2
    keep = R - burn; B = np.zeros((keep, k)); acc = 0
    for r in range(R):
        prop = beta + s * (L @ rng.standard_normal(k))
        lpp = poisson_loglik(prop, X, y) - 0.5 * np.sum(prop ** 2) / prior_sd ** 2
        if np.log(rng.random()) < lpp - lp:
            beta, lp = prop, lpp; acc += 1
        if r >= burn:
            B[r - burn] = beta
    return dict(beta=B, accept=acc / R)


# ───────────────────────── Negative Binomial (NB2) ─────────────────────────
def negbin_loglik(beta, log_r, X, y):
    r = np.exp(log_r); mu = np.exp(X @ beta)
    return float(np.sum(gammaln(y + r) - gammaln(r) - gammaln(y + 1)
                        + r * np.log(r / (r + mu)) + y * np.log(mu / (r + mu))))


def _num_hess(f, x, eps=1e-4):
    """Numerical Hessian of scalar f at x via central differences (observed information)."""
    d = len(x); H = np.zeros((d, d)); E = np.eye(d) * eps; fx = f(x)
    for i in range(d):
        for j in range(i, d):
            if i == j:
                H[i, i] = (f(x + E[i]) - 2 * fx + f(x - E[i])) / eps ** 2
            else:
                H[i, j] = H[j, i] = (f(x + E[i] + E[j]) - f(x + E[i] - E[j])
                                     - f(x - E[i] + E[j]) + f(x - E[i] - E[j])) / (4 * eps ** 2)
    return H


def negbin_mle(X, y):
    """Joint MLE of (beta, log r) via BFGS. Returns (beta_hat, log_r_hat, se) with se over [beta, log r].
    SEs from a numerical Hessian of the negative log-likelihood at the optimum (BFGS hess_inv is an
    unreliable Hessian estimate and can badly understate SEs on real, correlated data)."""
    k = X.shape[1]
    nll = lambda th: -negbin_loglik(th[:k], th[k], X, y)
    th0 = np.r_[poisson_mle(X, y)[0], 0.0]
    res = optimize.minimize(nll, th0, method='BFGS')
    se = np.sqrt(np.diag(np.linalg.inv(_num_hess(nll, res.x))))
    return res.x[:k], float(res.x[k]), se


def negbin_rwm(y, X, R=8000, burn=2000, prior_sd=10.0, prior_sd_r=10.0, seed=0, s=None):
    rng = np.random.default_rng(seed); n, k = X.shape
    b0, lr0, _ = negbin_mle(X, y)
    th = np.r_[b0, lr0]; d = k + 1
    nll = lambda t: -negbin_loglik(t[:k], t[k], X, y)
    cov = np.linalg.inv(_num_hess(nll, th))          # full posterior-curvature cov at the MLE
    L = np.linalg.cholesky(cov + 1e-10 * np.eye(d))  # correlated proposal (handles param correlation)
    s = (2.38 / np.sqrt(d)) if s is None else s

    def logpost(t):
        return (negbin_loglik(t[:k], t[k], X, y)
                - 0.5 * np.sum(t[:k] ** 2) / prior_sd ** 2
                - 0.5 * t[k] ** 2 / prior_sd_r ** 2)
    lp = logpost(th); keep = R - burn
    B = np.zeros((keep, k)); LR = np.zeros(keep); acc = 0
    for r in range(R):
        prop = th + s * (L @ rng.standard_normal(d))
        lpp = logpost(prop)
        if np.log(rng.random()) < lpp - lp:
            th, lp = prop, lpp; acc += 1
        if r >= burn:
            B[r - burn] = th[:k]; LR[r - burn] = th[k]
    return dict(beta=B, log_r=LR, r=np.exp(LR), accept=acc / R)


def summary(draws, names=None):
    """Posterior mean / sd / 95% CrI for a (keep, k) draws array."""
    q = draws.shape[1]; names = names or [f'b{j}' for j in range(q)]
    out = []
    for j, nm in enumerate(names):
        d = draws[:, j]; lo, hi = np.percentile(d, [2.5, 97.5])
        out.append((nm, d.mean(), d.std(), lo, hi))
    return out