"""
selection.py -- HECKMAN SELECTION MODELS for nonignorable (MNAR) missingness (from scratch).

Backs the notebooks in  "Selection Models -- Missing Not at Random".

Projects 1-3 assumed IGNORABLE missingness (MCAR/MAR): the reason data are missing carries no
information beyond the observed data, so imputing from the data model is valid. This project crosses
into NONIGNORABLE (MNAR) territory, where whether a value is observed depends on the value ITSELF.
The canonical case is sample selection: a wage is seen only for people who work, and the same
unobserved traits (drive, ability) push both the decision to work and the wage -- so the observed
wages are a biased sample of all potential wages. No ignorable method can undo that; we must MODEL
the selection.

The SELECTION-MODEL factorisation writes the joint density of the outcome y and the response
indicator r as  p(y) p(r | y)  -- a model for the outcome times a model for BEING OBSERVED that is
allowed to depend on the (possibly missing) outcome. Heckman's Gaussian version is two equations
with correlated errors:

    outcome:    y_i   = x_i' beta  + eps_i                          (observed only when r_i = 1)
    selection:  S_i*  = w_i' gamma + u_i ,   r_i = 1{ S_i* > 0 }
    (eps_i, u_i) ~ N(0, [[sigma^2, lambda],[lambda, 1]])            var(u)=1 identifies the probit

The correlation rho = lambda/sigma is the crux: rho = 0 means selection is ignorable (MAR) and OLS
on the selected sample is fine; rho != 0 means the missingness is nonignorable and OLS is biased by
the classic inverse-Mills-ratio term. rho is what the model estimates -- but it is only weakly
identified by functional form, so an EXCLUSION RESTRICTION (a covariate in w but not in x) and a
frank sensitivity analysis are essential. This is the deep point of MNAR: the correction rests on an
assumption the data cannot fully check.

CONNECTION TO TOBIT. Tobit ("Bayesian Tobit -- Censored Gaussian Regression") is a
selection model too -- but one where selection is a
DETERMINISTIC function of the outcome (y observed iff its latent value clears a threshold), i.e. the
same equation drives value and observation (rho = 1, w = x). Heckman generalises it to a SEPARATE
selection equation with its own covariates and a free correlation. Tobit is Heckman with the
selection welded to the outcome.

The sampler is data augmentation, reusing the Albert-Chib truncated-normal machinery: augment the
latent index S_i* (truncated by r_i) and the missing outcomes y_i (for r=0), then draw beta, gamma
and the covariance (reparametrised as lambda = cov and tau^2 = var(eps|u), so sigma^2 = lambda^2 +
tau^2 and rho = lambda/sigma) from conjugate conditionals.
"""

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


def _rtn(mean, sd, positive, rng):
    """draw N(mean, sd^2) truncated to (0, inf) where `positive` is True, else to (-inf, 0).
    `positive` is a boolean array (per element)."""
    positive = np.asarray(positive); alpha = Phi(-mean / sd); uu = rng.random(mean.shape)
    p = np.where(positive, alpha + uu * (1 - alpha), uu * alpha)
    return mean + sd * Phinv(np.clip(p, 1e-12, 1 - 1e-12))


def simulate_heckman(n, beta, gamma, rho, sigma, rng):
    """Two-equation Heckman data with an exclusion restriction (z enters selection, not outcome).
    Returns observed y (NaN when not selected), r, X, W, and the full latent y for checking."""
    x1 = rng.standard_normal(n); x2 = rng.standard_normal(n); z = rng.standard_normal(n)
    X = np.column_stack([np.ones(n), x1, x2]); W = np.column_stack([np.ones(n), x1, z])
    u = rng.standard_normal(n)
    eps = rho * sigma * u + sigma * np.sqrt(1 - rho ** 2) * rng.standard_normal(n)   # corr(eps,u)=rho
    y = X @ np.asarray(beta, float) + eps; r = (W @ np.asarray(gamma, float) + u > 0).astype(float)
    return np.where(r > 0.5, y, np.nan), r, X, W, y


def selection_gibbs(yobs, r, X, W, rng, draws=3000, burn=1500, b_sd=10.0, g_sd=10.0):
    """Bayesian Heckman selection model by data augmentation. yobs:(n,) with NaN for r=0; r:(n,) 0/1;
    X,W design matrices (include intercept; W should have an exclusion restriction). Returns posterior
    draws of beta, gamma, sigma, and the selection correlation rho."""
    r = np.asarray(r, float); X = np.asarray(X, float); W = np.asarray(W, float)
    n, px = X.shape; pw = W.shape[1]; sel = r > 0.5
    yc = np.where(sel, np.nan_to_num(yobs), np.nan)
    # init
    b0, *_ = np.linalg.lstsq(X[sel], yobs[sel], rcond=None); beta = b0
    gamma = np.zeros(pw); gamma[0] = 0.0
    resid0 = yobs[sel] - X[sel] @ beta; tau2 = resid0 @ resid0 / sel.sum(); lam = 0.0
    yc[~sel] = (X[~sel] @ beta)
    XtXi = np.linalg.inv(X.T @ X + (1 / b_sd ** 2) * np.eye(px)); Lx = np.linalg.cholesky(XtXi)
    WtWi = np.linalg.inv(W.T @ W + (1 / g_sd ** 2) * np.eye(pw)); Lw = np.linalg.cholesky(WtWi)
    B = np.empty((draws, px)); G = np.empty((draws, pw)); SIG = np.empty(draws); RHO = np.empty(draws)
    for it in range(draws + burn):
        sig2 = lam ** 2 + tau2; eps = yc - X @ beta
        # ---- augment latent selection index S* | y, truncated by r ----
        mS = W @ gamma + (lam / sig2) * eps; vS = tau2 / sig2
        S = _rtn(mS, np.sqrt(vS), sel, rng); u = S - W @ gamma
        # ---- gamma | .  (probit-style, correlation-adjusted response, error var vS) ----
        resp_g = S - (lam / sig2) * eps
        gm = WtWi @ (W.T @ resp_g); gamma = gm + np.sqrt(vS) * (Lw @ rng.standard_normal(pw))
        u = S - W @ gamma
        # ---- augment missing outcomes y | u  (r=0) ----
        yc[~sel] = X[~sel] @ beta + lam * u[~sel] + np.sqrt(tau2) * rng.standard_normal((~sel).sum())
        # ---- beta | .  (regress y - lam*u on X, error var tau2) ----
        resp_b = yc - lam * u
        bm = XtXi @ (X.T @ resp_b); beta = bm + np.sqrt(tau2) * (Lx @ rng.standard_normal(px))
        eps = yc - X @ beta
        # ---- lambda, tau2 | .  (regress eps on u through the origin) ----
        uu = u @ u + 1e-6; ue = u @ eps; lam_hat = ue / uu
        rss = eps @ eps - lam_hat * ue
        tau2 = max(rss, 1e-6) / rng.chisquare(max(n - 1, 1))
        lam = lam_hat + np.sqrt(tau2 / uu) * rng.standard_normal()
        if it >= burn:
            i = it - burn; sig = np.sqrt(lam ** 2 + tau2)
            B[i] = beta; G[i] = gamma; SIG[i] = sig; RHO[i] = lam / sig
    return dict(beta=B, gamma=G, sigma=SIG, rho=RHO)


def naive_ols(yobs, X):
    """OLS on the selected sample only -- biased when the selection correlation rho != 0."""
    sel = ~np.isnan(yobs); Xs = X[sel]
    XtXi = np.linalg.inv(Xs.T @ Xs); b = XtXi @ (Xs.T @ yobs[sel])
    resid = yobs[sel] - Xs @ b; s2 = resid @ resid / (sel.sum() - Xs.shape[1])
    return b, np.sqrt(np.diag(s2 * XtXi))