"""
factorirt.py -- MULTIDIMENSIONAL IRT = the item factor-analysis model (from scratch).

Backs the notebooks in  "Multidimensional IRT & the Factor-Analysis Bridge".

A 2PL measures one ability. Real instruments often measure SEVERAL correlated traits at once
(extraversion AND neuroticism, verbal AND quantitative). Multidimensional IRT gives each person
a VECTOR of latent traits theta_i and each item a vector of DISCRIMINATIONS -- which are exactly
FACTOR LOADINGS. In the normal-ogive form,

    P(x_ij = 1 | theta_i) = Phi( a_j' theta_i + d_j ),     theta_i ~ N(0, Sigma),

a_j is item j's loading vector, d_j the easiness, and Sigma the correlation between the latent
factors. This IS confirmatory factor analysis for binary items (item factor analysis): the 2PL
is the one-factor special case, and the discrimination is a loading. A CONFIRMATORY model fixes
which items load on which factor (simple structure), which identifies the loadings without
rotation.

From scratch it is the ALBERT-CHIB augmentation again -- draw z_ij ~ N(a_j'theta_i + d_j, 1)
truncated by the sign of x_ij -- combined with the PARAMETER-EXPANDED Gibbs for factor analysis
(Ghosh & Dunson 2009): sample the factors and loadings with an UNCONSTRAINED factor covariance
(inverse-Wishart, conjugate), then STANDARDISE afterwards to unit-variance factors, which recovers
the identified loadings and the factor correlation. Per-factor sign is fixed by keeping each
factor's loadings summing positive.
"""

import numpy as np
from scipy.special import ndtr as Phi, ndtri as Phinv
from scipy.stats import invwishart


def _rtrunc_sign(mean, pos, rng):
    lo = np.where(pos, Phi(-mean), 0.0); hi = np.where(pos, 1.0, Phi(-mean))
    u = lo + rng.random(mean.shape) * (hi - lo)
    return mean + Phinv(np.clip(u, 1e-12, 1 - 1e-12))


# --------------------------------------------------------------------------- #
#  simulation                                                                   #
# --------------------------------------------------------------------------- #

def simulate_mirt(N, A, d, rng, Sigma=None):
    """A:(J,D) loadings, d:(J,) intercepts, Sigma:(D,D) factor correlation."""
    A = np.asarray(A, float); d = np.asarray(d, float); J, D = A.shape
    if Sigma is None:
        Sigma = np.eye(D)
    theta = rng.multivariate_normal(np.zeros(D), Sigma, N)
    z = theta @ A.T + d + rng.standard_normal((N, J))
    return (z > 0).astype(float), theta


# --------------------------------------------------------------------------- #
#  Gibbs sampler (Albert-Chib augmentation + parameter-expanded factor Gibbs)   #
# --------------------------------------------------------------------------- #

def mirt_gibbs(X, mask, rng, draws=2500, burn=1500, load_sd=2.0):
    """multidimensional 2PL / item factor analysis. mask:(J,D) boolean -- which factors each item
    loads on (a confirmatory simple structure). Returns standardised loadings A:(draws,J,D),
    intercepts d:(draws,J), factor correlations R:(draws,D,D) and factors theta:(draws,N,D)."""
    X = np.asarray(X, float); N, J = X.shape; D = mask.shape[1]
    A = mask.astype(float) * 0.7; d = np.zeros(J); theta = rng.standard_normal((N, D)); Sig = np.eye(D)
    AS = np.empty((draws, J, D)); DD = np.empty((draws, J)); RR = np.empty((draws, D, D)); TH = np.empty((draws, N, D))
    for it in range(draws + burn):
        z = _rtrunc_sign(theta @ A.T + d, X > 0.5, rng)
        # theta_i | z, A, d, Sig
        Vinv = A.T @ A + np.linalg.inv(Sig); V = np.linalg.inv(Vinv); Lv = np.linalg.cholesky(V)
        M = ((z - d) @ A) @ V.T
        theta = M + rng.standard_normal((N, D)) @ Lv.T
        # (a_j, d_j) | z, theta   -- masked regression of z_j on [theta_cols, 1]
        for j in range(J):
            cols = np.where(mask[j])[0]; p = len(cols)
            Xd = np.column_stack([theta[:, cols], np.ones(N)])
            prior = np.diag([1.0 / load_sd ** 2] * p + [1e-6])
            Vj = np.linalg.inv(Xd.T @ Xd + prior); mj = Vj @ (Xd.T @ z[:, j])
            dr = mj + np.linalg.cholesky(Vj) @ rng.standard_normal(p + 1)
            A[j, cols] = dr[:p]; d[j] = dr[-1]
        for k in range(D):                                # per-factor sign identification
            if A[:, k].sum() < 0:
                A[:, k] = -A[:, k]; theta[:, k] = -theta[:, k]
        # Sigma | theta  (unconstrained inverse-Wishart, then standardise to a correlation)
        Sig = np.atleast_2d(invwishart.rvs(df=N + D + 1, scale=np.eye(D) + theta.T @ theta, random_state=rng))
        s = np.sqrt(np.diag(Sig))
        if it >= burn:
            i = it - burn
            AS[i] = A * s[None, :]; DD[i] = d; RR[i] = Sig / np.outer(s, s); TH[i] = theta / s[None, :]
    return dict(A=AS, d=DD, R=RR, theta=TH)