"""
seq_ord_probit_gibbs.py

Two Bayesian ordinal regression models via data-augmentation Gibbs sampling:

  1. Cumulative ordered probit  (Albert & Chib 1993)
       P(Y <= j | x) = Phi(gamma_j - x'beta)
       Single latent variable z* ~ N(x'beta, 1).
       Cut-points gamma_1 < ... < gamma_{J-1} drawn via uniform slice.

  2. Sequential ordinal probit  (Albert & Chib 2001)
       At stage k (k = 1,...,J-1), conditional on reaching stage k,
       a subject exits (Y = k-1) vs continues (Y >= k) via a binary probit:
           z_ik ~ N(x_i'alpha_k, 1),   s_ik = 1(z_ik > 0)
       J-1 independent binary probits — different alpha at every stage.

Joint probabilities (sequential, J = 3):
       P(Y=0|x) = Phi(-x'alpha_1)
       P(Y=1|x) = Phi(x'alpha_1) * Phi(-x'alpha_2)
       P(Y=2|x) = Phi(x'alpha_1) * Phi(x'alpha_2)

Design matrices
---------------
  Cumulative : X_c — NO intercept column (cut-points play the role of intercepts)
  Sequential : X_s — typically [1, covariates] (each stage has its own intercept)

References
----------
  Albert, J. H., & Chib, S. (1993).  JASA 88(422), 669-679.
  Albert, J. H., & Chib, S. (2001).  JASA 96(454), 415-425.
"""

import time
import numpy as np
from scipy.special import ndtr, ndtri


# ── truncated-Normal helper ────────────────────────────────────────────────────

def _rtn(lo, hi, mu, rng):
    """
    Vectorised draw from N(mu_i, 1) truncated to (lo_i, hi_i).
    Uses inverse-CDF; falls back to interval midpoint when mass < 1e-300.
    """
    p_lo = ndtr(lo - mu)
    p_hi = ndtr(hi - mu)
    gap  = p_hi - p_lo
    u    = rng.uniform(size=len(mu))
    p    = p_lo + u * gap
    tiny = gap < 1e-300
    p    = np.where(tiny, 0.5 * (p_lo + p_hi), p)
    return ndtri(np.clip(p, 1e-15, 1.0 - 1e-15)) + mu


# ── cumulative ordered probit ─────────────────────────────────────────────────

def cumulative_probit_gibbs(
    y, X_c,
    b0=None, B0=None,
    R=10000, burn=2000,
    seed=42, nprint=1000,
):
    """
    Albert & Chib (1993) Gibbs for the cumulative ordered probit.

    Model:  z_i ~ N(x_i'beta, 1)
            Y_i = j   iff   gamma[j-1] < z_i <= gamma[j]
            gamma[-1] = -inf,  gamma[J-1] = +inf  (sentinels)
            gamma[0] < ... < gamma[J-2]  — J-1 free cut-points

    X_c must NOT include an intercept column (cut-points absorb the intercept).

    Parameters
    ----------
    y   : (n,) int   labels in {0, ..., J-1}
    X_c : (n, k)     design matrix, no intercept
    b0  : (k,)       prior mean for beta       [default zeros]
    B0  : (k, k)     prior covariance          [default 100*I]
    R, burn : int    total / burn-in iterations
    seed, nprint : int

    Returns
    -------
    dict  'beta'  (R_kept, k) | 'gamma' (R_kept, J-1) | 'z' (R_kept, n)
    """
    rng = np.random.default_rng(seed)
    y   = np.asarray(y, int).ravel()
    X_c = np.asarray(X_c, float)
    n, k = X_c.shape
    J    = int(y.max()) + 1

    if b0 is None: b0 = np.zeros(k)
    if B0 is None: B0 = 100.0 * np.eye(k)
    b0        = np.asarray(b0, float)
    B0_inv    = np.linalg.inv(np.asarray(B0, float))
    B0_inv_b0 = B0_inv @ b0

    XtX  = X_c.T @ X_c
    Bbar = np.linalg.inv(XtX + B0_inv)
    L    = np.linalg.cholesky(Bbar)

    # initial cut-points: equally spaced between ±2
    gamma = np.linspace(-1.5, 1.5, J - 1)
    beta  = np.zeros(k)

    NINF = -1e10
    PINF =  1e10

    R_kept      = R - burn
    beta_draws  = np.empty((R_kept, k))
    gamma_draws = np.empty((R_kept, J - 1))
    z_draws     = np.empty((R_kept, n))
    kept = 0
    t0   = time.time()

    for r in range(1, R + 1):
        gext = np.r_[NINF, gamma, PINF]

        # Block 1: z_i | beta, gamma, y  — truncated Normal
        mu = X_c @ beta
        z  = _rtn(gext[y], gext[y + 1], mu, rng)

        # Block 2: beta | z  — Normal posterior (Bbar is constant)
        bbar = Bbar @ (X_c.T @ z + B0_inv_b0)
        beta = bbar + L @ rng.standard_normal(k)

        # Block 3: gamma_j | z, y  — uniform slice (independent per j)
        # Posterior: Uniform(max(z: y=j), min(z: y=j+1))
        # Ordering is automatically satisfied because z draws respect current gamma.
        for j in range(J - 1):
            lo_j = z[y == j].max()     if (y == j    ).any() else NINF
            hi_j = z[y == j + 1].min() if (y == j + 1).any() else PINF
            if lo_j < hi_j:
                gamma[j] = rng.uniform(lo_j, hi_j)

        if r > burn:
            beta_draws[kept]  = beta
            gamma_draws[kept] = gamma.copy()
            z_draws[kept]     = z
            kept += 1

        if nprint > 0 and r % nprint == 0:
            print(f'  [cumulative]  iter {r:6d}/{R}  ({time.time()-t0:.1f}s)')

    print(f'[cumulative] Done in {time.time()-t0:.1f}s  |  kept: {R_kept}')
    return {'beta': beta_draws, 'gamma': gamma_draws, 'z': z_draws}


# ── sequential ordinal probit ─────────────────────────────────────────────────

def sequential_probit_gibbs(
    y, X_s,
    b0=None, B0=None,
    R=10000, burn=2000,
    seed=42, nprint=1000,
):
    """
    Albert & Chib (2001) Gibbs for the sequential ordinal probit.

    Stage k (k = 1,...,J-1):  among subjects with Y_i >= k-1,
        define  s_ik = 1(Y_i >= k)   (binary: continue vs exit)
        z_ik ~ N(x_i'alpha_k, 1),   s_ik = 1(z_ik > 0)

    The J-1 stages share the same prior but have independent parameters.
    Each stage reduces to a binary probit (Albert & Chib 1993).

    Parameters
    ----------
    y   : (n,) int   labels in {0, ..., J-1}
    X_s : (n, k)     design matrix (typically includes intercept column)
    b0  : (k,)       prior mean for each alpha_k  [default zeros]
    B0  : (k, k)     prior covariance             [default 100*I]
    R, burn, seed, nprint : int

    Returns
    -------
    dict
      'alpha'    : list of J-1 arrays (R_kept, k)
      'z'        : list of J-1 arrays (R_kept, n_k)  — latent utilities per stage
      'at_risk'  : list of J-1 bool (n,) masks
    """
    rng = np.random.default_rng(seed)
    y   = np.asarray(y, int).ravel()
    X_s = np.asarray(X_s, float)
    n, k = X_s.shape
    J    = int(y.max()) + 1

    if b0 is None: b0 = np.zeros(k)
    if B0 is None: B0 = 100.0 * np.eye(k)
    b0        = np.asarray(b0, float)
    B0_inv    = np.linalg.inv(np.asarray(B0, float))
    B0_inv_b0 = B0_inv @ b0

    # build at-risk subsets for each stage
    at_risk = []   # bool masks length n
    s_list  = []   # binary outcome at each stage
    X_list  = []   # design submatrix
    for stage in range(1, J):              # stage = 1 .. J-1
        mask = (y >= stage - 1)
        at_risk.append(mask)
        s_list.append((y[mask] >= stage).astype(float))
        X_list.append(X_s[mask])

    # precompute posterior covariance for each stage (constant — sigma^2 = 1 fixed)
    Bbar_list = []
    L_list    = []
    for Xs in X_list:
        Bbar = np.linalg.inv(Xs.T @ Xs + B0_inv)
        Bbar_list.append(Bbar)
        L_list.append(np.linalg.cholesky(Bbar))

    alpha_list = [np.zeros(k) for _ in range(J - 1)]
    z_list     = [np.zeros(mask.sum()) for mask in at_risk]

    R_kept      = R - burn
    alpha_draws = [np.empty((R_kept, k)) for _ in range(J - 1)]
    z_draws_out = [np.empty((R_kept, mask.sum())) for mask in at_risk]
    kept = 0
    t0   = time.time()

    NINF = -1e10
    PINF =  1e10

    for r in range(1, R + 1):

        for q in range(J - 1):
            s    = s_list[q]
            Xs   = X_list[q]
            mu   = Xs @ alpha_list[q]
            lo   = np.where(s == 1, 0.0,  NINF)
            hi   = np.where(s == 1, PINF, 0.0)
            z    = _rtn(lo, hi, mu, rng)
            z_list[q] = z

            bbar = Bbar_list[q] @ (Xs.T @ z + B0_inv_b0)
            alpha_list[q] = bbar + L_list[q] @ rng.standard_normal(k)

        if r > burn:
            for q in range(J - 1):
                alpha_draws[q][kept]    = alpha_list[q]
                z_draws_out[q][kept]    = z_list[q]
            kept += 1

        if nprint > 0 and r % nprint == 0:
            print(f'  [sequential]  iter {r:6d}/{R}  ({time.time()-t0:.1f}s)')

    print(f'[sequential] Done in {time.time()-t0:.1f}s  |  kept: {R_kept}')
    return {'alpha': alpha_draws, 'z': z_draws_out, 'at_risk': at_risk}


# ── predicted probabilities ────────────────────────────────────────────────────

def cum_pred_probs(X_c, beta_draws, gamma_draws):
    """
    Posterior predictive probabilities for the cumulative model.

    Returns
    -------
    p_bar   : (n, J)    posterior mean P(Y=j | x_i)
    p_draws : (R, n, J) full draw array (for CPO / LPML)
    """
    n, R = X_c.shape[0], beta_draws.shape[0]
    J    = gamma_draws.shape[1] + 1
    p_draws = np.empty((R, n, J))

    for r in range(R):
        mu = X_c @ beta_draws[r]        # (n,)
        g  = gamma_draws[r]             # (J-1,)
        # cumulative CDF at each cut-point
        Fj = np.array([ndtr(g[j] - mu) for j in range(J - 1)]).T   # (n, J-1)
        p_draws[r, :, 0]        = Fj[:, 0]
        for j in range(1, J - 1):
            p_draws[r, :, j]    = Fj[:, j] - Fj[:, j - 1]
        p_draws[r, :, J - 1]   = 1.0 - Fj[:, J - 2]

    return p_draws.mean(axis=0), p_draws


def seq_pred_probs(X_s, alpha_draws):
    """
    Posterior predictive probabilities for the sequential model.

    Parameters
    ----------
    X_s         : (n, k)
    alpha_draws : list of J-1 arrays, each (R, k)

    Returns
    -------
    p_bar   : (n, J)
    p_draws : (R, n, J)
    """
    n  = X_s.shape[0]
    S  = len(alpha_draws)      # J - 1
    J  = S + 1
    R  = alpha_draws[0].shape[0]
    p_draws = np.empty((R, n, J))

    for r in range(R):
        # Phi(x'alpha_k) = P(continue at stage k)  ->  shape (n, S)
        phi_cont = np.column_stack(
            [ndtr( X_s @ alpha_draws[q][r]) for q in range(S)]
        )
        phi_exit = 1.0 - phi_cont

        # P(Y=0) = Phi(-x'alpha_1)
        p_draws[r, :, 0] = phi_exit[:, 0]
        # P(Y=j) = prod_{l=1}^{j} Phi(x'alpha_l) * Phi(-x'alpha_{j+1})
        for j in range(1, J - 1):
            p_draws[r, :, j] = phi_cont[:, :j].prod(axis=1) * phi_exit[:, j]
        # P(Y=J-1) = prod_{l=1}^{J-1} Phi(x'alpha_l)
        p_draws[r, :, J - 1] = phi_cont.prod(axis=1)

    return p_draws.mean(axis=0), p_draws


# ── LPML via harmonic-mean CPO ────────────────────────────────────────────────

def compute_lpml(y, p_draws):
    """
    LPML = sum_i log CPO_i,
    CPO_i = 1 / mean_r(1 / p(y_i | theta^r)).

    Parameters
    ----------
    y        : (n,) int
    p_draws  : (R, n, J)

    Returns
    -------
    lpml : float
    cpo  : (n,) array
    """
    y = np.asarray(y, int)
    R, n, J = p_draws.shape
    # lik[i, r] = p(y_i | theta^r)
    lik = p_draws[:, np.arange(n), y].T     # (n, R)
    lik = np.clip(lik, 1e-10, 1.0)
    cpo  = 1.0 / np.mean(1.0 / lik, axis=1)
    lpml = float(np.sum(np.log(np.clip(cpo, 1e-15, 1.0))))
    return lpml, cpo


# ── posterior summary ─────────────────────────────────────────────────────────

def posterior_summary(draws, param_names=None):
    """(R, p) draws -> DataFrame with mean, sd, q025, q975."""
    import pandas as pd
    p  = draws.shape[1]
    if param_names is None:
        param_names = [f'p{j}' for j in range(p)]
    lo = np.quantile(draws, 0.025, axis=0)
    hi = np.quantile(draws, 0.975, axis=0)
    return pd.DataFrame({
        'mean': draws.mean(0).round(4),
        'sd':   draws.std(0).round(4),
        'q025': lo.round(4),
        'q975': hi.round(4),
    }, index=param_names)