"""
Bayesian hierarchical multinomial logit (MNL) — from-scratch sampler (NumPy only).

Model:  respondent i makes choices over J alternatives in T_i tasks.
    P(choice = j | task) = exp(x_j' beta_i) / sum_l exp(x_l' beta_i)
    beta_i ~ N(Delta' z_i, V_beta)      (z_i = respondent-level covariates, e.g. [1, demos])

Sampler (RW-Metropolis within Gibbs) — MNL has no Pólya–Gamma/Gibbs step, so:
  1. beta_i   -- random-walk Metropolis, proposal scaled by each respondent's MNL Hessian
                 (the bayesm rhierMnlRwMixture move)
  2. Delta    -- matrix-normal (group-level regression of beta_i on z_i)
  3. V_beta   -- inverse-Wishart
Single-normal heterogeneity prior (bayesm's rhierMnlRwMixture uses a mixture-of-normals).

Data format: `data` = list, one per respondent, of (X, y) with
    X : (T_i, J, k) design tensor,  y : (T_i,) chosen-alternative index in 0..J-1.
Companion to binary_logit_mcmc.py (RW-Metropolis) and hier_logit_gibbs.py (the hierarchy).
"""
import numpy as np


def mnl_loglik(beta, X, y):
    """MNL log-likelihood for one respondent. X:(T,J,k), y:(T,)."""
    util = X @ beta                                   # (T, J)
    m = util.max(axis=1, keepdims=True)
    logden = m[:, 0] + np.log(np.exp(util - m).sum(axis=1))
    return float((util[np.arange(len(y)), y] - logden).sum())


def mnl_hess(beta, X):
    """Observed information sum_t X_t'(diag(p)-pp')X_t for one respondent."""
    util = X @ beta; util -= util.max(axis=1, keepdims=True)
    e = np.exp(util); p = e / e.sum(axis=1, keepdims=True)        # (T, J)
    k = X.shape[2]; H = np.zeros((k, k))
    for t in range(X.shape[0]):
        Xt = X[t]; pt = p[t]
        H += Xt.T @ (np.diag(pt) - np.outer(pt, pt)) @ Xt
    return H


def _pooled_beta(data, k, maxit=25):
    """Pooled MNL MLE (Newton) — center for the per-respondent RW proposals."""
    b = np.zeros(k)
    for _ in range(maxit):
        g = np.zeros(k); H = np.zeros((k, k))
        for X, y in data:
            util = X @ b; util -= util.max(axis=1, keepdims=True)
            e = np.exp(util); p = e / e.sum(axis=1, keepdims=True)
            for t in range(len(y)):
                g += X[t, y[t]] - p[t] @ X[t]
            H += mnl_hess(b, X)
        step = np.linalg.solve(H + 1e-6 * np.eye(k), g)
        b += step
        if np.max(np.abs(step)) < 1e-8:
            break
    return b


def _riwishart(nu, S, rng):
    p = S.shape[0]
    L = np.linalg.cholesky(np.linalg.inv(S)); A = np.zeros((p, p))
    for i in range(p):
        A[i, i] = np.sqrt(rng.chisquare(nu - i))
        for j in range(i):
            A[i, j] = rng.standard_normal()
    return np.linalg.inv(L @ A @ A.T @ L.T)


def simulate_hier_mnl(N=300, T=12, J=3, k=3, Delta=None, Vb=None, seed=0):
    """Simulate hierarchical MNL choice data. Returns dict incl. truth."""
    rng = np.random.default_rng(seed)
    u = rng.normal(size=(N, 1))
    Z = np.column_stack([np.ones(N), u])             # (N, q=2)
    if Delta is None:
        Delta = np.zeros((2, k))
        Delta[0] = np.array([1.0, -1.0, 0.5])[:k]    # baseline part-worths
        Delta[1] = np.array([0.4, 0.0, -0.3])[:k]    # demographic effect
    if Vb is None:
        Vb = 0.5 * np.eye(k)
    B = Z @ Delta + rng.normal(size=(N, k)) @ np.linalg.cholesky(Vb).T
    data = []
    for i in range(N):
        X = rng.normal(size=(T, J, k))
        util = X @ B[i]
        y = (util + rng.gumbel(size=(T, J))).argmax(axis=1)   # RUM choice
        data.append((X, y))
    return dict(data=data, Z=Z, B_true=B, Delta_true=Delta, Vb_true=Vb)


def hier_mnl_gibbs(data, Z, R=3000, burn=1000, seed=0, s=None,
                   A0=0.01, nu0=None, S0=None):
    """Hierarchical MNL via RW-Metropolis (beta_i) + Gibbs hierarchy (Delta, V_beta)."""
    rng = np.random.default_rng(seed)
    N = len(data); q = Z.shape[1]; k = data[0][0].shape[2]
    if s is None:
        s = 2.38 / np.sqrt(k)                        # RW scale
    nu0 = k + 3 if nu0 is None else nu0
    S0 = np.eye(k) if S0 is None else S0

    bpool = _pooled_beta(data, k)
    Hs = [mnl_hess(bpool, X) for X, _ in data]        # per-respondent MNL curvature at pooled MLE

    B = np.tile(bpool, (N, 1)).astype(float)
    Delta = np.zeros((q, k)); Vb = np.eye(k)
    llcur = np.array([mnl_loglik(B[i], data[i][0], data[i][1]) for i in range(N)])

    keep = R - burn
    Dd = np.zeros((keep, q, k)); Vd = np.zeros((keep, k, k)); Bm = np.zeros((N, k)); acc = 0
    for g in range(R):
        Vbinv = np.linalg.inv(Vb)
        for i in range(N):
            # proposal ~ N(0, (H_i + V_beta^{-1})^{-1}): the per-respondent posterior curvature
            # (data + prior precision), so the step adapts to the heterogeneity scale each sweep
            Lp = np.linalg.cholesky(Hs[i] + Vbinv)
            prop = B[i] + s * np.linalg.solve(Lp.T, rng.standard_normal(k))
            llp = mnl_loglik(prop, data[i][0], data[i][1])
            mu = Delta.T @ Z[i]
            d0 = B[i] - mu; d1 = prop - mu
            lpr = -0.5 * (d1 @ Vbinv @ d1 - d0 @ Vbinv @ d0)
            if np.log(rng.random()) < (llp - llcur[i] + lpr):
                B[i] = prop; llcur[i] = llp; acc += 1
        # Delta | B, Vb  (matrix-normal)
        U = np.linalg.inv(Z.T @ Z + A0 * np.eye(q)); M = U @ (Z.T @ B)
        Delta = M + np.linalg.cholesky(U) @ rng.standard_normal((q, k)) @ np.linalg.cholesky(Vb).T
        # V_beta | B, Delta  (inverse-Wishart)
        E = B - Z @ Delta; Vb = _riwishart(nu0 + N, S0 + E.T @ E, rng)
        if g >= burn:
            Dd[g - burn] = Delta; Vd[g - burn] = Vb; Bm += B
    return dict(Delta=Dd, Vb=Vd, B=Bm / keep, accept=acc / (R * N))


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