"""
Multinomial probit (MNP) in PyMC via differenced-utility GHK.

Utilities:  U_ik = a_k + b * x_ik + e_ik ,  e_i ~ N(0, Sigma) ,  choose argmax_k U_ik
x_ik is a choice-specific covariate (here: log price).  a_k are alt-specific
intercepts (base alt's intercept = 0).

Identification (the hard part of MNP):
  - only utility *differences* are identified  -> work with w = utilities
    differenced vs a fixed base alternative (the last one).
  - scale not identified -> fix the (0,0) element of the differenced covariance
    Sigma* to 1 (Imai-van Dyk default).  We build Sigma* = diag(s) C diag(s)
    with s[0]=1 and C an LKJ correlation matrix.

Likelihood: choice c_i observed  <=>  all contrasts (U_c - U_k) > 0.  For each
possible chosen alternative we map w to those contrasts via a fixed matrix T_c,
giving g = T_c w ~ N(T_c mu*, T_c Sigma* T_c'); P(g>0) is evaluated by GHK with
auxiliary Uniforms (smooth -> NUTS-friendly).
"""
import numpy as np
import pymc as pm
import pytensor.tensor as pt

SQRT2 = np.sqrt(2.0)
_Phi = lambda x: 0.5 * pt.erfc(-x / SQRT2)
def _invPhi(p):
    p = pt.clip(p, 1e-12, 1 - 1e-12)
    return SQRT2 * pt.erfinv(2.0 * p - 1.0)


def _T(p, c):
    """Map w (base-differenced utilities, base=last) to contrasts U_c - U_k > 0."""
    pm1 = p - 1
    if c == p - 1:                       # base chosen: U_base - U_k = -w_k
        return -np.eye(pm1)
    rows = []
    e_c = np.zeros(pm1); e_c[c] = 1.0
    rows.append(e_c.copy())              # vs base: U_c - U_base = w_c
    for k in range(pm1):
        if k != c:
            r = np.zeros(pm1); r[c] = 1.0; r[k] = -1.0
            rows.append(r)               # vs k: U_c - U_k = w_c - w_k
    return np.array(rows)


def build_mnp(Xprice, choice, eta=3.0):
    """Xprice:(N,p) choice-specific covariate; choice:(N,) integer in 0..p-1
    with the base alternative coded as p-1."""
    Xprice = np.asarray(Xprice, float)
    choice = np.asarray(choice, int)
    N, p = Xprice.shape
    pm1 = p - 1
    dX = Xprice[:, :pm1] - Xprice[:, [p - 1]]     # (N, p-1) price diffs vs base
    Tmats = {c: _T(p, c) for c in range(p)}
    groups = {c: np.where(choice == c)[0] for c in range(p)}

    with pm.Model() as model:
        a = pm.Normal("a", 0.0, 5.0, shape=pm1)   # intercepts vs base
        b = pm.Normal("b", 0.0, 10.0)             # price coefficient

        _, C, _ = pm.LKJCholeskyCov("Cpack", n=pm1, eta=eta,
                                    sd_dist=pm.LogNormal.dist(0.0, 0.1),
                                    compute_corr=True)
        log_s = pm.Normal("log_s", 0.0, 0.5, shape=pm1 - 1)
        s = pt.concatenate([[1.0], pt.exp(log_s)])           # s[0]=1 fixed
        Sigma = (s[:, None] * C) * s[None, :]                # differenced cov
        pm.Deterministic("Sigma_star", Sigma)

        dX_t = pt.as_tensor_variable(dX)
        mu_star = a[None, :] + b * dX_t                      # (N, p-1)

        aux = pm.Uniform("aux", 0.0, 1.0, shape=(N, pm1))
        total = 0.0
        for c in range(p):
            idx = groups[c]
            if len(idx) == 0:
                continue
            Tc = pt.as_tensor_variable(Tmats[c])
            Oc = Tc @ Sigma @ Tc.T
            Lc = pt.linalg.cholesky(Oc)
            eta_g = mu_star[idx] @ Tc.T                      # (n_c, p-1)
            aux_g = aux[idx]
            u_cols, ll = [], pt.zeros((len(idx),))
            for k in range(pm1):
                partial = pt.zeros((len(idx),))
                for j in range(k):
                    partial = partial + Lc[k, j] * u_cols[j]
                lb = -(eta_g[:, k] + partial) / Lc[k, k]
                Phi_lb = pt.clip(_Phi(lb), 1e-12, 1 - 1e-12)
                u_cols.append(_invPhi(Phi_lb + aux_g[:, k] * (1.0 - Phi_lb)))
                ll = ll + pt.log1p(-Phi_lb)                  # P(g_k > 0)
            total = total + ll.sum()
        pm.Potential("ghk_loglik", total)
    return model


# --------------------------------------------------------------------------- #
def _simulate(N=1200, p=6, seed=0):
    rng = np.random.default_rng(seed)
    lp = rng.normal(-2.9, 0.25, size=(N, p))             # log prices
    a_true = np.array([0.8, 1.0, 0.3, 0.1, -0.4])        # intercepts vs base
    b_true = -7.0
    Ar = rng.normal(size=(p, p))
    Sig_full = Ar @ Ar.T / p + np.eye(p) * 0.4
    alpha = np.concatenate([a_true, [0.0]])
    U = alpha + b_true * lp + rng.multivariate_normal(np.zeros(p), Sig_full, N)
    choice = U.argmax(1)
    # identified truth: divide by sigma1 of the base-differenced covariance
    A = np.hstack([np.eye(p - 1), -np.ones((p - 1, 1))])
    Ss = A @ Sig_full @ A.T
    s1 = np.sqrt(Ss[0, 0])
    return lp, choice, a_true / s1, b_true / s1, Ss / s1**2


if __name__ == "__main__":
    import arviz as az
    lp, choice, a_id, b_id, Sig_id = _simulate(N=700)
    print("choice counts:", np.bincount(choice))
    model = build_mnp(lp, choice)
    import pymc.sampling.jax as pmjax
    with model:
        idata = pmjax.sample_numpyro_nuts(draws=400, tune=400, chains=2,
                                          target_accept=0.9, random_seed=3,
                                          progressbar=False, chain_method="sequential")
    print("\nintercepts a  true:", np.round(a_id, 2))
    print("intercepts a  est :",
          np.round(idata.posterior["a"].mean(("chain", "draw")).values, 2))
    print("price b       true:", round(float(b_id), 2),
          " est:", round(float(idata.posterior["b"].mean(("chain", "draw"))), 2))
    print("Sigma* diag true:", np.round(np.diag(Sig_id), 2))
    print("Sigma* diag est :",
          np.round(np.diag(idata.posterior["Sigma_star"].mean(("chain", "draw")).values), 2))
    print("divergences:", int(idata.sample_stats["diverging"].sum()),
          " max rhat:",
          round(float(az.rhat(idata, var_names=["a", "b"]).to_array().max()), 3))