pipeline_miami/verify/residual_checks.py

"""Checks on the residual-z output (pipeline_miami/analysis/residual_z.py):

  R1  Independent refit: for every fitted outcome × year, re-derive the
      sampling N of every cell from the served October membership under the
      outcomes.yaml n_rule (own parser), re-derive eligibility (complete
      covariates, N ≥ min_n, year not excluded), require the eligible set to
      equal the engine's, then rebuild the design matrix and solve the
      weighted normal equations (X'WX)b = X'Wy with np.linalg.solve — a
      different solver from the engine's lstsq — and reproduce z to within
      1e-6 with the standardisation written out here independently:
        z = r / sqrt(tau² + se_i²), tau² = max(0, s²·n/(n−k) − Σw·se²/Σw).
      Before the 2026-09-16 review R1 copied N and the eligible set from the
      engine's CSV, so it could not catch a wrong n_rule; now it can.
  R2  Served = committed: every residual_z in miami_school_year_metrics equals
      the value in data/cities/miami/reference/residual_z.csv, and no served
      z exists for a cell absent from the CSV.
  R3  Sanity gates (from the reference scaffold's checks.py): every eligible z
      finite; flagged share below 50%; at least 30 eligible cells per fit.

Writes docs/qa_reports/miami/__residual_checks.{json,md}.
Run:  python -m pipeline_miami.verify.residual_checks
"""
from __future__ import annotations

import hashlib
import json
import re

import numpy as np
import pandas as pd

from pipeline_miami.analysis import residual_z as engine
from pipeline_miami.verify import common

GRADES = ["PK", "KG", "1", "2", "3", "4", "5", "6", "7", "8", "9", "10", "11", "12"]


def own_n(frame: pd.DataFrame, key: str, rule: str) -> pd.Series:
    """Sampling N under the outcomes.yaml n_rule, parsed here rather than by the engine."""
    if rule == "denominator":
        return frame[f"den__{key}"] if f"den__{key}" in frame.columns else pd.Series(np.nan, index=frame.index)
    if rule == "enrollment":
        return frame["enrollment_oct"]
    m = re.fullmatch(r"grades:([^;]+)(?:;share:([0-9.]+))?", rule)
    if not m:
        raise SystemExit(f"R1: cannot parse n_rule {rule}")
    grades: list[str] = []
    for part in m.group(1).split(","):
        if "-" in part:
            a, b = part.split("-")
            grades += [str(g) for g in range(int(a), int(b) + 1)]
        else:
            grades.append(part)
    share = float(m.group(2)) if m.group(2) else 1.0
    cols = [f"enrollment_oct_g{g}" for g in grades if f"enrollment_oct_g{g}" in frame.columns]
    return frame[cols].sum(axis=1, min_count=1) * share


def independent_z(e: pd.DataFrame, key: str, cov_cfg, scale: str, direction: str) -> np.ndarray:
    # Design: intercept, numeric (log where configured), drop-first dummies in sorted column order.
    parts = [np.ones(len(e))]
    for c in cov_cfg["numeric"]:
        x = e[c["key"]].to_numpy(dtype=float)
        parts.append(np.log(np.maximum(x, 1.0)) if c.get("transform") == "log" else x)
    for c in cov_cfg["categorical"]:
        d = pd.get_dummies(e[c["key"]].astype("category"), prefix=c["key"], drop_first=True)
        for col in sorted(d.columns):
            parts.append(d[col].to_numpy(dtype=float))
    X = np.column_stack(parts)
    y = e[key].to_numpy(dtype=float)
    w = e["n_own"].to_numpy(dtype=float)
    n_rows, k = X.shape
    # Solve the normal equations on centred, unit-scaled columns (the fitted
    # values are invariant to that reparameterisation); squaring an
    # unscaled design with log-enrollment next to percentages loses ~5
    # digits, which is a property of this solver, not of the engine.
    mu = X[:, 1:].mean(axis=0)
    sd = X[:, 1:].std(axis=0)
    sd = np.where(sd > 0, sd, 1.0)
    Xs = np.column_stack([np.ones(len(X)), (X[:, 1:] - mu) / sd])
    XtW = Xs.T * w
    beta = np.linalg.solve(XtW @ Xs, XtW @ y)
    resid = y - Xs @ beta
    s2 = np.sum(w * resid ** 2) / np.sum(w) * n_rows / (n_rows - k)
    n = w
    # scale here is the outcome's sampling model (outcomes.yaml sampling_error, else scale): percent → binomial on
    # N students, rate_per_100 → Poisson on the count, none → no per-cell noise term (attendance, a share of
    # student-days; review 2026-09-16b N11).
    if scale == "none":
        se = np.zeros_like(y, dtype=float)
    elif scale == "percent":
        p = np.clip(np.clip(y / 100.0, 0, 1), 0.5 / n, 1 - 0.5 / n)
        se = 100.0 * np.sqrt(p * (1 - p) / n)
    else:
        lam = np.maximum(y, 0.5 * 100.0 / n)
        se = np.sqrt(100.0 * lam / n)
    tau2 = max(0.0, s2 - np.sum(w * se ** 2) / np.sum(w))
    z = resid / np.sqrt(tau2 + se ** 2)
    return -z if direction == "LOWER_BETTER" else z


def main() -> None:
    import psycopg

    common.ensure_qa_dir()
    cov_cfg, out_cfg, comp = engine.load_config()
    packet = json.loads(engine.OUT_JSON.read_text())
    config_hash = hashlib.sha256("".join((engine.CONFIG_DIR / name).read_text() for name in
                                      ("covariates.yaml", "outcomes.yaml", "comparability.yaml")).encode()).hexdigest()[:12]
    packet_config_matches = packet.get("meta", {}).get("config_sha256_12") == config_hash
    csv = pd.read_csv(engine.OUT_CSV, dtype={"school_number": str, "year": str, "metric_key": str})
    with psycopg.connect(common.served_db_dsn()) as conn:
        conn.read_only = True
        frame, _, _ = engine.fetch_frame(conn, cov_cfg, list(out_cfg["outcomes"].keys()))
        served = pd.DataFrame(conn.execute("SELECT school_number, year, metric_key, residual_z::float8, outlier_flag FROM miami_school_year_metrics WHERE residual_z IS NOT NULL AND subgroup='ALL'").fetchall(),
                              columns=["school_number", "year", "metric_key", "z_served", "flag_served"])
    num_keys = [c["key"] for c in cov_cfg["numeric"]]
    cat_keys = [c["key"] for c in cov_cfg["categorical"]]
    min_n_default = out_cfg["defaults"]["min_n"]
    # R1 — independent N, eligibility and refit per outcome × year
    r1 = []
    for key, ocfg in out_cfg["outcomes"].items():
        if key not in frame.columns:
            continue
        excluded = set(str(y) for y in comp["regimes"][ocfg["comparability_regime"]].get("excluded_years", []))
        min_n = ocfg.get("min_n", min_n_default)
        f = frame[frame[key].notna()].copy()
        f["n_own"] = own_n(f, key, ocfg["n_rule"])
        f["eligible_own"] = f[num_keys + cat_keys].notna().all(axis=1) & (f.n_own >= min_n) & ~f.year.isin(excluded)
        sub = csv[(csv.metric_key == key) & (csv.eligible)]
        for year, g in sub.groupby("year"):
            own = f[(f.year == year) & f.eligible_own]
            own_set, csv_set = set(own.school), set(g.school_number)
            e = own.set_index("school").loc[sorted(own_set & csv_set)].reset_index()
            gi = g.set_index("school_number").loc[e.school]
            n_diff = float(np.nanmax(np.abs(e.n_own.to_numpy() - gi.n.to_numpy()))) if len(e) else float("nan")
            if len(e) >= 30 and own_set == csv_set:
                z_ind = independent_z(e, key, cov_cfg, ocfg.get("sampling_error", ocfg["scale"]), ocfg["direction"])
                independent_finite = bool(np.isfinite(z_ind).all())
                max_dz = float(np.max(np.abs(z_ind - gi.z.to_numpy())))
            else:
                independent_finite = False
                max_dz = float("nan")
            r1.append({"metric_key": key, "year": year, "n": int(len(g)), "eligible_set_equal": own_set == csv_set,
                       "own_only": int(len(own_set - csv_set)), "csv_only": int(len(csv_set - own_set)), "max_abs_dn": n_diff, "max_abs_dz": max_dz,
                       "independent_z_all_finite": independent_finite,
                       "pass": bool(independent_finite and own_set == csv_set and n_diff < 1e-9 and max_dz < 1e-6)})
    # R2 — served equals committed
    m = csv[csv.eligible].merge(served, on=["school_number", "year", "metric_key"], how="outer", indicator=True)
    both = m[m["_merge"] == "both"]
    r2 = {"csv_eligible": int(csv.eligible.sum()), "served_cells": int(len(served)), "matched": int(len(both)),
          "csv_only": int((m["_merge"] == "left_only").sum()), "served_only": int((m["_merge"] == "right_only").sum()),
          "all_z_finite": bool(len(both) and np.isfinite(both.z).all() and np.isfinite(both.z_served).all()),
          "max_abs_dz": float(np.max(np.abs(both.z - both.z_served))) if len(both) else None,
          "flag_mismatches": int((both.flag != both.flag_served).sum()) if len(both) else None}
    r2["pass"] = bool(r2["all_z_finite"] and r2["csv_only"] == 0 and r2["served_only"] == 0 and r2["max_abs_dz"] < 1e-6 and r2["flag_mismatches"] == 0)
    # R3 — sanity gates
    el = csv[csv.eligible]
    per_fit = el.groupby(["metric_key", "year"]).agg(n=("z", "size"), finite=("z", lambda s: bool(np.isfinite(s).all())), flagged=("flag", lambda s: float((s != "none").mean())))
    r3 = {"all_z_finite": bool(per_fit.finite.all()), "max_flagged_share": float(per_fit.flagged.max()), "min_eligible_per_fit": int(per_fit.n.min()),
          "pass": bool(per_fit.finite.all() and per_fit.flagged.max() < 0.5 and per_fit.n.min() >= 30)}
    verdict = "PASS" if packet_config_matches and r1 and all(r["pass"] for r in r1) and r2["pass"] and r3["pass"] else "REVIEW"
    out = {"verdict": verdict, "config_sha256_12": config_hash,
           "packet_generated_at": packet.get("meta", {}).get("generated_at"),
           "packet_config_matches": packet_config_matches,
           "R1_independent_refit": {"fits": len(r1), "all_pass": all(r["pass"] for r in r1), "max_abs_dz": max(r["max_abs_dz"] for r in r1) if r1 else None,
                                    "max_abs_dn": max(r["max_abs_dn"] for r in r1) if r1 else None, "eligible_sets_equal": all(r["eligible_set_equal"] for r in r1), "rows": r1},
           "R2_served_equals_committed": r2, "R3_sanity": r3}
    (common.QA_REPORTS_DIR / "__residual_checks.json").write_text(json.dumps(out, indent=2, default=str))
    md = [f"# Residual-z checks\n", f"Verdict: **{verdict}**\n",
          f"- Packet generation: {out['packet_generated_at']}; config {config_hash}; packet matches current config = {packet_config_matches}",
          f"- R1 independent N + eligibility + refit (normal equations vs lstsq): {len(r1)} outcome-years, eligible sets equal = {out['R1_independent_refit']['eligible_sets_equal']}, max |ΔN| = {out['R1_independent_refit']['max_abs_dn']:.2e}, max |Δz| = {out['R1_independent_refit']['max_abs_dz']:.2e}, all pass = {out['R1_independent_refit']['all_pass']}",
          f"- R2 served = committed: {r2['matched']} cells matched, {r2['csv_only']} CSV-only, {r2['served_only']} served-only, max |Δz| = {r2['max_abs_dz']}, flag mismatches = {r2['flag_mismatches']} → {'PASS' if r2['pass'] else 'REVIEW'}",
          f"- R3 sanity: all z finite = {r3['all_z_finite']}, max flagged share = {r3['max_flagged_share']:.3f}, min eligible per fit = {r3['min_eligible_per_fit']} → {'PASS' if r3['pass'] else 'REVIEW'}",
          "", "| outcome | year | n | eligible set equal | max abs dN | max abs dz |", "|---|---|---:|---|---:|---:|"]
    md += [f"| `{r['metric_key']}` | {r['year']} | {r['n']} | {r['eligible_set_equal']} | {r['max_abs_dn']:.1e} | {r['max_abs_dz']:.2e} |" for r in r1]
    (common.QA_REPORTS_DIR / "__residual_checks.md").write_text("\n".join(md))
    print(json.dumps({k: v for k, v in out.items() if k != "R1_independent_refit"} | {"R1": {k: v for k, v in out["R1_independent_refit"].items() if k != "rows"}}, indent=1, default=str))


if __name__ == "__main__":
    main()