"""
Bayesian hierarchical (multilevel) logistic regression — from-scratch Gibbs (Pólya–Gamma).

Model (varying intercept + varying slope, with a group-level predictor):
    y_i ~ Bernoulli(sigmoid(alpha_{j[i]} + beta_{j[i]} * x_i)),    i in group j[i]
    b_j = (alpha_j, beta_j)' ,   b_j ~ N(Delta' z_j, V_b)
where x_i is the individual covariate (e.g. income), z_j the group-level covariates
(e.g. [1, state_mean_income]), Delta (q x p) the group-level regression, V_b (p x p)
the random-effect covariance.

Pólya–Gamma augmentation (Polson–Scott–Windle 2013) makes every block conjugate:
  1. omega_i      ~ PG(1, eta_i)                              (data augmentation)
  2. b_j | .      ~ N( P_j^{-1}(X_j'kappa_j + Vb^{-1} Delta'z_j), P_j^{-1} ),
                    P_j = X_j' Omega_j X_j + Vb^{-1},  kappa = y - 1/2
  3. Delta | .    ~ MatrixNormal( (Z'Z+A0)^{-1} Z'B, (Z'Z+A0)^{-1}, V_b )
  4. V_b | .      ~ InverseWishart( nu0 + J, S0 + E'E ),  E = B - Z Delta

Reuses the PG sampler from binary_logit_mcmc. NumPy only. Companion to
binary_logit_mcmc.py (flat logit) and mts_gibbs.py.
"""
import numpy as np
from binary_logit_mcmc import rpg, sigmoid


def simulate_hier_logit(J=50, n_per=40, Delta=None, Vb=None, seed=0):
    """Simulate a 2-level varying-intercept/slope logit that exhibits the
    'rich state / poor state' pattern: positive within-group slope, negative
    group-level intercept–predictor relationship."""
    rng = np.random.default_rng(seed)
    u = rng.normal(size=J)                          # group-level predictor (e.g. state mean income)
    Z = np.column_stack([np.ones(J), u])            # (J, q=2)
    if Delta is None:
        # rows = group covariates (intercept, u); cols = REs (alpha, beta)
        Delta = np.array([[-0.10, 0.40],            # baseline alpha,   baseline beta (+ income slope)
                          [-0.70, 0.00]])           # effect of u on alpha (NEGATIVE), on beta (~0)
    if Vb is None:
        Vb = np.array([[0.30, 0.05], [0.05, 0.10]])
    B = Z @ Delta + rng.normal(size=(J, 2)) @ np.linalg.cholesky(Vb).T   # (J, 2) true REs
    state  = np.repeat(np.arange(J), n_per)
    income = rng.normal(size=J * n_per)             # individual predictor (centered)
    eta = B[state, 0] + B[state, 1] * income
    y = (rng.random(len(eta)) < sigmoid(eta)).astype(float)
    return dict(y=y, income=income, state=state, u=u, Z=Z,
                B_true=B, Delta_true=Delta, Vb_true=Vb)


def _riwishart(nu, S, rng):
    """Draw V ~ Inverse-Wishart(nu, scale=S) via a Bartlett-decomposition Wishart."""
    p = S.shape[0]
    L = np.linalg.cholesky(np.linalg.inv(S))        # chol of S^{-1}
    A = np.zeros((p, p))
    for i in range(p):
        A[i, i] = np.sqrt(rng.chisquare(nu - i))
        for k in range(i):
            A[i, k] = rng.standard_normal()
    W = L @ A @ A.T @ L.T                            # Wishart(nu, S^{-1})
    return np.linalg.inv(W)                           # Inverse-Wishart(nu, S)


def hier_logit_gibbs(y, income, state, Z, R=4000, burn=1000, seed=0,
                     A0=0.01, nu0=None, S0=None, n_terms=128):
    """Hierarchical PG-Gibbs. Returns posterior draws of B (J×2), Delta (q×2), V_b (2×2)."""
    rng = np.random.default_rng(seed)
    y = np.asarray(y, float); income = np.asarray(income, float); state = np.asarray(state, int)
    J, q = Z.shape; p = 2
    X = np.column_stack([np.ones(len(y)), income])   # individual design (1, income)
    kappa = y - 0.5
    idx = [np.where(state == j)[0] for j in range(J)]
    A0 = A0 * np.eye(q)
    nu0 = p + 2 if nu0 is None else nu0
    S0 = np.eye(p) if S0 is None else S0
    ZtZ = Z.T @ Z

    B = np.zeros((J, p)); Delta = np.zeros((q, p)); Vb = np.eye(p)
    keep = R - burn
    Bd = np.zeros((keep, J, p)); Dd = np.zeros((keep, q, p)); Vd = np.zeros((keep, p, p))

    for g in range(R):
        # 1. Polya-Gamma augmentation
        eta = (X * B[state]).sum(1)
        omega = rpg(1.0, eta, rng, n_terms)
        # 2. random effects b_j (Gaussian full conditional, per group)
        Vbinv = np.linalg.inv(Vb)
        for j in range(J):
            ij = idx[j]; Xj = X[ij]; oj = omega[ij]
            Pj = (Xj * oj[:, None]).T @ Xj + Vbinv
            rhs = Xj.T @ kappa[ij] + Vbinv @ (Delta.T @ Z[j])
            cov = np.linalg.inv(Pj)
            B[j] = cov @ rhs + np.linalg.cholesky(cov) @ rng.standard_normal(p)
        # 3. group-level regression Delta (matrix-normal full conditional)
        U = np.linalg.inv(ZtZ + A0)
        M = U @ (Z.T @ B)
        Delta = M + np.linalg.cholesky(U) @ rng.standard_normal((q, p)) @ np.linalg.cholesky(Vb).T
        # 4. RE covariance V_b (inverse-Wishart full conditional)
        E = B - Z @ Delta
        Vb = _riwishart(nu0 + J, S0 + E.T @ E, rng)
        if g >= burn:
            k = g - burn; Bd[k] = B; Dd[k] = Delta; Vd[k] = Vb
    return dict(B=Bd, Delta=Dd, Vb=Vd)


def summary(out):
    return dict(Delta=out['Delta'].mean(0), Vb=out['Vb'].mean(0), B=out['B'].mean(0))