"""
patternmix.py -- PATTERN-MIXTURE MODELS and delta-sensitivity for MNAR missingness (from scratch).

Backs the notebooks in  "Pattern-Mixture Models -- MNAR by Sensitivity Analysis".

Selection models (Project 4) factor the joint density of outcome y and response indicator r as
p(y) p(r|y) -- a model for the outcome and a model for being observed. PATTERN-MIXTURE models take
the OTHER factorisation,

    p(y, r) = p(y | r) p(r),

stratifying by the missingness PATTERN and giving each stratum its own outcome distribution, then
mixing over the patterns. Its virtue is honesty. The outcome distribution among the MISSING
(r = 0) is never observed, so it cannot be estimated -- it must be ASSUMED. Pattern-mixture makes
that assumption an explicit, interpretable knob: a SENSITIVITY PARAMETER delta linking the missing
stratum to the observed one,

    E[y | missing]  =  E[y | observed, same covariates]  +  delta.

delta = 0 is the MAR ("missing at random", i.e. missing like the observed) assumption; delta != 0 is
MNAR -- the dropouts differ from completers by delta even after adjusting for what we saw. Because
delta is not identified by the data, the analysis is a SENSITIVITY ANALYSIS: sweep delta and watch
the conclusion. The TIPPING POINT is the delta at which the conclusion (e.g. a treatment benefit)
just reverses -- "how much worse would the dropouts have had to be for the result to go away?" This
is exactly the delta-adjustment / reference-based sensitivity analysis regulators ask for in trials
with dropout, and the transparent complement to the selection model's hard-to-interpret rho.

The engine is simple: fit a Bayesian regression to the COMPLETERS, impute each dropout from that
model shifted by delta (applied to a chosen arm), build multiple completed datasets, and combine
the estimand by Rubin's rules -- repeating across a grid of delta.
"""

import numpy as np


def _blr_predict(Xo, yo, Xm, rng):
    """Bayesian linear regression on completers (Xo,yo); return a posterior-predictive draw for Xm."""
    n, p = Xo.shape; XtXi = np.linalg.inv(Xo.T @ Xo + 1e-8 * np.eye(p))
    bhat = XtXi @ (Xo.T @ yo); resid = yo - Xo @ bhat
    s2 = max(resid @ resid, 1e-8) / rng.chisquare(max(n - p, 1))
    beta = bhat + np.linalg.cholesky(s2 * XtXi) @ rng.standard_normal(p)
    return Xm @ beta + np.sqrt(s2) * rng.standard_normal(Xm.shape[0])


def pm_impute(y, arm, Xbase, rng, delta=0.0, arms=(1,), m=50):
    """Pattern-mixture multiple imputation. y:(n,) endpoint with NaN for dropouts; arm:(n,) 0/1;
    Xbase:(n,k) fully observed baseline covariates (an intercept is added). Dropouts are imputed
    from the completer regression (fitted per arm); `delta` is ADDED to the imputed values of the
    dropouts in the arms listed in `arms` (the MNAR shift). Returns m completed y-vectors."""
    y = np.asarray(y, float); arm = np.asarray(arm); Xb = np.column_stack([np.ones(len(y)), Xbase])
    completed = []
    for _ in range(m):
        yc = y.copy()
        for a in (0, 1):
            inarm = arm == a; obs = inarm & ~np.isnan(y); mis = inarm & np.isnan(y)
            if mis.sum() == 0:
                continue
            pred = _blr_predict(Xb[obs], y[obs], Xb[mis], rng)
            if a in arms:
                pred = pred + delta
            yc[mis] = pred
        completed.append(yc)
    return completed


def effect(completed, arm):
    """treatment-arm minus control-arm mean of the completed endpoint, pooled over imputations by
    Rubin's rules. Returns (estimate, se)."""
    arm = np.asarray(arm); ests = np.array([yc[arm == 1].mean() - yc[arm == 0].mean() for yc in completed])
    # within-imputation variance of a difference of means
    within = np.mean([yc[arm == 1].var(ddof=1) / (arm == 1).sum() + yc[arm == 0].var(ddof=1) / (arm == 0).sum()
                      for yc in completed])
    m = len(completed); B = ests.var(ddof=1); T = within + (1 + 1 / m) * B
    return ests.mean(), np.sqrt(T)


def sensitivity(y, arm, Xbase, deltas, rng, arms=(1,), m=50):
    """sweep delta; return arrays of estimate and 95% CI bounds of the treatment effect."""
    est = np.empty(len(deltas)); lo = np.empty(len(deltas)); hi = np.empty(len(deltas))
    for i, dl in enumerate(deltas):
        e, se = effect(pm_impute(y, arm, Xbase, rng, delta=dl, arms=arms, m=m), arm)
        est[i] = e; lo[i] = e - 1.96 * se; hi[i] = e + 1.96 * se
    return est, lo, hi


def tipping_point(deltas, est, lo, hi):
    """smallest |delta| at which the 95% CI first includes 0 (the conclusion becomes non-significant)."""
    sig = (lo > 0) | (hi < 0)
    if sig.all():
        return None
    flip = np.where(~sig)[0]
    return deltas[flip[np.argmin(np.abs(deltas[flip]))]]


def complete_case_effect(y, arm):
    """difference in observed means (dropouts discarded)."""
    y = np.asarray(y, float); arm = np.asarray(arm); o = ~np.isnan(y)
    a1 = o & (arm == 1); a0 = o & (arm == 0)
    est = y[a1].mean() - y[a0].mean()
    se = np.sqrt(y[a1].var(ddof=1) / a1.sum() + y[a0].var(ddof=1) / a0.sum())
    return est, se


def simulate_trial(n, tau, rng, drop_mnar=1.2):
    """two-arm trial with MNAR dropout concentrated in the TREATMENT arm: the sicker (higher-y)
    treated patients drop out, so the observed treatment mean flatters the drug. Returns observed
    endpoint y (NaN=dropout), arm, baseline, and the full (pre-dropout) endpoint."""
    arm = (rng.random(n) < 0.5).astype(int); base = rng.normal(5, 1.5, n)
    yfull = 5 + tau * arm + 0.5 * (base - 5) + rng.normal(0, 1.5, n)          # lower = better; tau<0 = benefit
    pdrop = np.where(arm == 1, 1 / (1 + np.exp(-(yfull - 5) * drop_mnar)), 0.03)   # sicker treated -> drop
    return np.where(rng.random(n) < pdrop, np.nan, yfull), arm, base, yfull