"""
Zero-inflated counts (ZIP and ZINB) -- from-scratch Gibbs with data augmentation.
The PyMC / UCLA "fish" example and BUGS Vol II "Hearts" family. The new idea vs count_reg_mcmc.py
(plain Poisson/NegBin) is a MIXTURE likelihood with a latent regime indicator:

  with prob pi_i        : structural zero   (y_i = 0 for sure -- e.g. group never fished)
  with prob 1 - pi_i    : "at risk"         (y_i ~ Poisson(lambda_i) or NegBin -- may still be 0)

so   P(y=0) = pi + (1-pi) f(0),   P(y=k>0) = (1-pi) f(k).
Two linked regressions:
  log lambda_i = x_i' beta     (count model, on the at-risk regime)
  logit pi_i   = z_i' gamma    (zero-inflation model, who is a structural zero)

Sampler (Gibbs with augmentation)
---------------------------------
  1. s_i  -- latent structural-zero indicator. If y_i>0 then s_i=0 (must be at-risk). If y_i=0,
            P(s_i=1) = pi_i / (pi_i + (1-pi_i) f(0)); draw Bernoulli.
  2. beta -- count regression on the at-risk subset {s_i=0}: RW-Metropolis (Poisson / NB loglik).
  3. gamma-- logistic regression of s on Z: RW-Metropolis.
  4. (ZINB only) log r -- NB dispersion, updated jointly with beta as a block.

Priors: beta, gamma ~ N(0, prior_sd^2); log r ~ N(0, 4).
"""
import numpy as np
from scipy.special import gammaln, expit


def simulate_zip(n=3000, beta=(0.5, 0.8, -0.4), gamma=(-0.4, 1.1), r=None, seed=0):
    """Zero-inflated counts. X=[1,x1,x2] for the count part, Z=[1,z1] for the zero part."""
    rng = np.random.default_rng(seed)
    beta = np.asarray(beta, float); gamma = np.asarray(gamma, float)
    x1 = rng.normal(size=n); x2 = rng.normal(size=n); z1 = rng.normal(size=n)
    X = np.column_stack([np.ones(n), x1, x2]); Z = np.column_stack([np.ones(n), z1])
    pi = expit(Z @ gamma); lam = np.exp(X @ beta)
    struct = rng.random(n) < pi                       # structural zeros
    if r is None:
        cnt = rng.poisson(lam)                        # at-risk counts (Poisson)
    else:
        cnt = rng.poisson(lam * rng.gamma(r, 1.0 / r, size=n))  # NB via gamma mixing
    y = np.where(struct, 0, cnt)
    return y, X, Z, struct


def _pois_mle(X, y, maxit=100):
    b = np.zeros(X.shape[1])
    for _ in range(maxit):
        mu = np.exp(X @ b)
        b += np.linalg.solve((X * mu[:, None]).T @ X + 1e-8 * np.eye(X.shape[1]), X.T @ (y - mu))
    return b


def _nb_log0(lam, r):
    """log P(y=0) under NegBin(mean=lam, size=r)."""
    return r * (np.log(r) - np.log(r + lam))


def zip_gibbs(y, X, Z, R=8000, burn=2000, prior_sd=5.0, seed=0,
              s_beta=None, s_gamma=None, negbin=False, a_r=2.0):
    """Gibbs for ZIP (negbin=False) or ZINB (negbin=True)."""
    rng = np.random.default_rng(seed)
    y = np.asarray(y, float); n, k = X.shape; q = Z.shape[1]
    is0 = (y == 0)
    s_beta = (2.38 / np.sqrt(k)) if s_beta is None else s_beta
    s_gamma = (2.38 / np.sqrt(q)) if s_gamma is None else s_gamma

    # fixed proposal scales from crude full-data fits
    b0 = _pois_mle(X, y); Lb = np.linalg.cholesky(
        np.linalg.inv((X * np.exp(X @ b0)[:, None]).T @ X + np.eye(k) / prior_sd ** 2))
    w = 0.25
    Lg = np.linalg.cholesky(np.linalg.inv(w * Z.T @ Z + np.eye(q) / prior_sd ** 2))

    beta = b0.copy(); gamma = np.zeros(q); logr = 1.0
    keep = R - burn
    B = np.zeros((keep, k)); G = np.zeros((keep, q)); RR = np.zeros(keep)
    acc_b = 0; acc_g = 0
    for it in range(R):
        eta = X @ beta; lam = np.exp(eta); zg = Z @ gamma; pi = expit(zg)
        r = np.exp(logr) if negbin else None
        # 1. augment structural-zero indicator s (only ambiguous where y==0)
        logf0 = _nb_log0(lam, r) if negbin else -lam            # log P(at-risk count = 0)
        p_struct = pi / (pi + (1 - pi) * np.exp(logf0))         # P(s=1 | y=0)
        s = np.zeros(n)
        s[is0] = (rng.random(is0.sum()) < p_struct[is0]).astype(float)
        at = (s == 0)                                           # at-risk subset

        # 2. beta (+logr) on the at-risk subset
        Xa = X[at]; ya = y[at]
        def count_ll(bb, lr):
            e = Xa @ bb
            if not negbin:
                return np.sum(ya * e - np.exp(e)) - 0.5 * np.sum(bb ** 2) / prior_sd ** 2
            rr = np.exp(lr); mu = np.exp(e)
            return np.sum(gammaln(ya + rr) - gammaln(rr) + rr * np.log(rr) + ya * e
                          - (ya + rr) * np.log(rr + mu)) \
                - 0.5 * np.sum(bb ** 2) / prior_sd ** 2 - 0.5 * lr ** 2 / 4.0
        cur = count_ll(beta, logr)
        prop = beta + s_beta * (Lb @ rng.standard_normal(k))
        lrp = logr + (0.15 * rng.standard_normal() if negbin else 0.0)
        if np.log(rng.random()) < count_ll(prop, lrp) - cur:
            beta = prop; logr = lrp; acc_b += 1
        # 3. gamma: logistic regression of s on Z
        def logit_ll(gg):
            e = Z @ gg
            return np.sum(s * e - np.logaddexp(0, e)) - 0.5 * np.sum(gg ** 2) / prior_sd ** 2
        gp = gamma + s_gamma * (Lg @ rng.standard_normal(q))
        if np.log(rng.random()) < logit_ll(gp) - logit_ll(gamma):
            gamma = gp; acc_g += 1
        if it >= burn:
            i = it - burn; B[i] = beta; G[i] = gamma; RR[i] = np.exp(logr)
    out = dict(beta=B, gamma=G, accept_beta=acc_b / R, accept_gamma=acc_g / R)
    if negbin:
        out['r'] = RR
    return out


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