"""
catmiss.py -- MISSING CATEGORICAL DATA and nonresponse (from scratch).

Backs the notebooks in  "Categorical Missing Data & Nonresponse".

The whole arc has run on continuous outcomes; this final project moves to CATEGORICAL data --
contingency tables and surveys where some units are only PARTIALLY CLASSIFIED (one item answered,
another left blank). The Bayesian tool is the categorical twin of Project 1's Gaussian data
augmentation: a DIRICHLET-MULTINOMIAL Gibbs sampler.

For a two-way table with cell probabilities p (a 2x2 here), the conjugate loop is:

  I-step (augment): allocate each partially-classified unit to the full cells it is compatible with,
                    in proportion to the current p (a multinomial draw). A unit with Y1 = i but Y2
                    missing is split between (i,1) and (i,2); a unit missing both is spread over all
                    four cells.
  P-step (update):  with the table now complete, draw p ~ Dirichlet(counts + prior).

Under MAR this is exact and ignorable -- the partial units are allocated exactly as the observed
conditional distribution dictates, and using them SHARPENS the estimate over a complete-case
analysis that would throw them away.

NONIGNORABLE (MNAR) nonresponse -- the subject of Congdon's categorical missing-data chapter -- is
when the chance of answering DEPENDS on the true category (people with something to hide skip the
question). We encode it with a sensitivity parameter psi that TILTS the allocation of the missing
units away from the observed conditional: psi = 1 is MAR, psi != 1 makes the non-responders'
category distribution differ from the responders'. As in the pattern-mixture project, psi is not
identified by the data -- it is a knob for a sensitivity analysis.

CONNECTION TO LATENT CLASS ANALYSIS. An LCA (the Latent Class Analysis arc) handles a missing item on
exactly this principle: the item simply drops out of that subject's likelihood, i.e. the class
membership is inferred by marginalising over the unobserved response -- the same "sum over the cells
the unit is compatible with" that the I-step performs. Categorical missing data is where the
missing-data arc and the latent-class arc meet.
"""

import numpy as np


def augment_gibbs(nboth, ny2mis, ny1mis, nbothmis, rng, draws=4000, burn=2000,
                  psi1=1.0, psi2=1.0, alpha=0.5):
    """Dirichlet-multinomial augmentation for a 2x2 table with item nonresponse.
      nboth   : 2x2 complete-case counts (rows Y1, cols Y2)
      ny2mis  : length-2, units with Y1=i observed but Y2 missing
      ny1mis  : length-2, units with Y2=j observed but Y1 missing
      nbothmis: scalar, units missing both items
      psi1,psi2: MNAR sensitivity tilts (1 = MAR) for Y1, Y2 nonresponse
    Returns posterior draws of the 2x2 cell-probability matrix p."""
    nboth = np.asarray(nboth, float); ny2mis = np.asarray(ny2mis, float); ny1mis = np.asarray(ny1mis, float)
    t1 = np.array([1.0, psi1]); t2 = np.array([1.0, psi2])
    p = np.ones((2, 2)) / 4; P = np.empty((draws, 2, 2))
    for it in range(draws + burn):
        counts = nboth.copy()
        for i in (0, 1):                                    # Y1=i known, Y2 missing -> split over Y2
            w = p[i, :] * t2; counts[i, :] += rng.multinomial(int(ny2mis[i]), w / w.sum())
        for j in (0, 1):                                    # Y2=j known, Y1 missing -> split over Y1
            w = p[:, j] * t1; counts[:, j] += rng.multinomial(int(ny1mis[j]), w / w.sum())
        w = (p * np.outer(t1, t2)).ravel()                  # both missing -> spread over all 4
        counts += rng.multinomial(int(nbothmis), w / w.sum()).reshape(2, 2)
        g = rng.gamma(counts + alpha); p = g / g.sum()      # p ~ Dirichlet(counts + alpha)
        if it >= burn:
            P[it - burn] = p
    return P


def odds_ratio(P):
    """odds ratio of the 2x2 association from probability draws P:(draws,2,2)."""
    return (P[:, 0, 0] * P[:, 1, 1]) / (P[:, 0, 1] * P[:, 1, 0])


def pmarg(P, var, level):
    """marginal probability P(var = level); var in {0(Y1),1(Y2)}, level in {0,1}."""
    return P[:, level, :].sum(1) if var == 0 else P[:, :, level].sum(1)


def complete_case(nboth):
    """complete-case cell proportions (partial units discarded)."""
    nboth = np.asarray(nboth, float); return nboth / nboth.sum()


def simulate_table(N, p, resp2, rng):
    """simulate N units from a 2x2 p, then make Y2 missing with probability that depends on the true
    Y2 value (resp2 = [P(respond|Y2=1), P(respond|Y2=2)]; equal => MAR, unequal => MNAR).
    Returns the four partially-classified count groups."""
    p = np.asarray(p, float); flat = p.ravel()
    draw = rng.multinomial(N, flat).reshape(2, 2)
    nboth = np.zeros((2, 2)); ny2mis = np.zeros(2)
    for i in (0, 1):
        for j in (0, 1):
            nij = draw[i, j]; resp = rng.random(nij) < resp2[j]
            nboth[i, j] = resp.sum(); ny2mis[i] += (~resp).sum()
    return nboth, ny2mis, np.zeros(2), 0