"""
Mixture-of-normals regression — from-scratch Gibbs (NumPy + scipy.norm).

The continuous sibling of the Geweke-Keane (1997) mixture-of-normals probit: instead of a
binary outcome with a latent-utility step, we observe y directly and let the regression error
follow a flexible mixture of normals (capturing skew AND fat tails in, e.g., earnings):

    y_t = x_t' beta + eps_t,   eps_t ~ sum_j p_j N(mu_j, sigma_j^2)

Identification: the error mixture is recentred to mean 0 each sweep (its location goes into the
intercept), so beta is the usual regression mean function and the mixture describes the
*shape* of the disturbance around it.

Gibbs blocks:
  1. s_t         -- component indicator, categorical  P(j) ∝ p_j N(r_t; mu_j, sigma_j^2)
  2. beta        -- weighted Gaussian (GLS, weights 1/sigma_{s_t}^2) on y - mu_{s_t}
  3. (mu_j,sig2_j)-- Normal-Inverse-Gamma per component
  4. p           -- Dirichlet
Closed-form density:  p(y|x) = sum_j p_j N(y; x'beta + mu_j, sigma_j^2)  -> used for WAIC.
"""
import numpy as np
from scipy.stats import norm


def mixreg_gibbs(y, X, m=2, R=6000, burn=2000, seed=0,
                 prior_sd_beta=10.0, a0=3.0, kappa0=0.01, alpha_dir=1.0):
    rng = np.random.default_rng(seed)
    n, k = X.shape
    Hb = np.eye(k) / prior_sd_beta ** 2
    beta = np.linalg.lstsq(X, y, rcond=None)[0]
    resid = y - X @ beta
    vy = float(np.var(resid)); b0 = 0.5 * vy * (a0 - 1)        # prior mean sigma2 ~ var(resid)
    sig2 = np.full(m, vy)
    mu = np.linspace(-1.0, 1.0, m) * np.std(resid) if m > 1 else np.zeros(1)
    p = np.ones(m) / m
    s = rng.integers(0, m, n)

    keep = R - burn
    B = np.zeros((keep, k)); P = np.zeros((keep, m)); MU = np.zeros((keep, m)); S2 = np.zeros((keep, m))
    LL = np.zeros((keep, n))
    for g in range(R):
        r = y - X @ beta
        # 1. component indicators
        logpj = (np.log(p)[None, :] - 0.5 * np.log(sig2)[None, :]
                 - 0.5 * (r[:, None] - mu[None, :]) ** 2 / sig2[None, :])
        logpj -= logpj.max(1, keepdims=True)
        pp = np.exp(logpj); pp /= pp.sum(1, keepdims=True)
        s = (rng.random(n)[:, None] > np.cumsum(pp, 1)).sum(1).clip(0, m - 1)
        # 2. beta | . (weighted Gaussian on y - mu_{s})
        w = 1.0 / sig2[s]
        yc = y - mu[s]
        V = np.linalg.inv((X * w[:, None]).T @ X + Hb)
        beta = V @ ((X * w[:, None]).T @ yc) + np.linalg.cholesky(V) @ rng.standard_normal(k)
        # 3. (mu_j, sigma2_j) | . (Normal-Inverse-Gamma)
        r = y - X @ beta
        for j in range(m):
            mk = s == j; nj = int(mk.sum()); rj = r[mk]
            kn = kappa0 + nj
            mn = rj.sum() / kn                                  # prior mean 0
            an = a0 + nj / 2.0
            bn = b0 + 0.5 * (np.sum((rj - rj.mean()) ** 2) if nj > 0 else 0.0) \
                 + 0.5 * (kappa0 * nj / kn) * (rj.mean() if nj > 0 else 0.0) ** 2
            sig2[j] = 1.0 / rng.gamma(an, 1.0 / bn)
            mu[j] = mn + np.sqrt(sig2[j] / kn) * rng.standard_normal()
        # 4. p | .
        p = rng.dirichlet(alpha_dir + np.bincount(s, minlength=m))
        # identification: recentre mixture to mean 0 (location -> intercept), then sort by mu
        mbar = np.sum(p * mu); beta[0] += mbar; mu -= mbar
        o = np.argsort(mu); mu = mu[o]; sig2 = sig2[o]; p = p[o]; s = np.argsort(o)[s]
        if g >= burn:
            gg = g - burn; B[gg] = beta; P[gg] = p; MU[gg] = mu; S2[gg] = sig2
            dens = (p[None, :] * norm.pdf(r[:, None], mu[None, :], np.sqrt(sig2)[None, :])).sum(1)
            LL[gg] = np.log(np.clip(dens, 1e-300, None))
    return dict(beta=B, p=P, mu=MU, sig2=S2, ll=LL, m=m)


def waic(out):
    ll = out['ll']; mx = ll.max(0)
    lppd = np.log(np.exp(ll - mx).mean(0)) + mx
    return float(-2 * (lppd.sum() - ll.var(0).sum()))


def error_density(out, grid):
    """Posterior-mean fitted error density on `grid`."""
    p = out['p'].mean(0); mu = out['mu'].mean(0); sd = np.sqrt(out['sig2'].mean(0))
    return sum(p[j] * norm.pdf(grid, mu[j], sd[j]) for j in range(out['m']))