docs/cities/miami/probe/race_imputation_prior.py

"""Does an adjacent-year prior beat the equal split for hidden race cells?

Follow-up to race_imputation.py after the 2026-09-16 review found the 2022-23
file marks zeros with '*' (17,464 tokens against ~5,900 in other years), so
the equal split over hidden cells inflates the small groups that year
(American Indian median share 0.000 → 0.958 → 0.000 across 2021-22 → 2023-24).

Diagnostic, derived from the file rather than chosen: tokens can only include
zeros when a school shows more hidden cells S than hidden students
U = total − visible (each genuinely hidden cell holds ≥ 1). Year-level S > U
(2022-23: 17,345 cells vs 17,077 students) says the file's '*' includes zeros;
school-level S > U picks the schools where that is certain.

Rules compared against the unsuppressed NCES CCD race subtotals:
  hybrid        : the loader's rule (fully-visible grade → zeros; equal split of U over the rest)
  prior_school  : hybrid, but where the school has S > U, split U over hidden groups in
                  proportion to the school's adjacent-year hybrid shares (years with S ≤ U only)
  prior_year    : the same prior for every school in a year whose file has S > U
  prior_all     : the prior wherever an adjacent trustworthy year exists (all years)
  prior_year_cap9: prior_year with no group receiving more than 9 students per hidden cell (the publisher's own
                  floor); the overflow is re-spread over the other hidden groups by prior, then equally
Fallback when no adjacent trustworthy year exists: hybrid. Error = rule − CCD, pp.

Held-out selection (added after the 2026-09-16b review, N9): the rule served
for 2022-23 (prior_year_cap9) was chosen on the same CCD file the harness now
reports it against. To measure the optimism of that in-sample choice, the
schools are split into two halves by the parity of their school number; the
best rule on each half (lowest mean |error| over the Black, Hispanic and
White shares) is applied to the other half and its held-out MAE reported
beside the in-sample one.

Run: pipeline_miami/.venv/bin/python docs/cities/miami/probe/race_imputation_prior.py [years]
"""
from __future__ import annotations

import sys
from pathlib import Path

import numpy as np
import pandas as pd
import psycopg

sys.path.insert(0, str(Path(__file__).parent))
import race_imputation as base  # noqa: E402

RACES = list(base.RACES)
ALL_YEARS = ["2018-19", "2021-22", "2022-23", "2023-24", "2024-25", "2025-26"]


def year_shares(year: str, con) -> tuple[pd.DataFrame, dict[str, float], bool]:
    """Per school×race: visible n, hidden cells, S; the equal-split (hybrid) share; and whether the year is trustworthy (S ≤ U)."""
    hy = base.fldoe_hybrid(year, con)
    enr = dict(con.execute("select school_number, value::float from miami_school_year_metrics where metric_key='enrollment_oct' and subgroup='ALL' and year=%s and value is not null", (year,)).fetchall())
    hy["total"] = hy.school.map(enr)
    hy = hy[hy.total.notna() & (hy.total > 0)].copy()
    vis = hy.groupby("school").n.sum()
    hy["U"] = (hy.total - hy.school.map(vis)).clip(lower=0)
    hy["hybrid"] = (hy.n + np.where(hy.S > 0, hy.U * hy.hid / hy.S.replace(0, np.nan), 0.0).astype(float)) / hy.total * 100
    hy["hybrid"] = hy.hybrid.fillna(hy.n / hy.total * 100)
    per = hy.groupby("school").agg(S=("S", "first"), U=("U", "first"))
    return hy, enr, int((per.S > per.U).sum()) == 0   # trustworthy = no school has more hidden cells than hidden students


def prior_alloc(hy: pd.DataFrame, prior: pd.DataFrame, apply_mask: pd.Series, cap: float | None = None) -> pd.Series:
    """Replace the equal split with a prior-proportional split for schools in apply_mask. With cap=9, no group
    receives more than 9 students per hidden cell (the publisher's own rule: a hidden cell holds 1–9); the excess
    is re-spread over the uncapped groups in proportion to their prior, and anything still left is spread
    equally over every hidden cell (the trustworthy-year rule)."""
    out = hy.hybrid.copy()
    p = prior.rename("prior")
    df = hy.join(p, on=["school", "race"])
    for school, g in df[apply_mask].groupby("school"):
        hidden = g[g.hid > 0]
        if not len(hidden) or g.U.iloc[0] <= 0:
            continue
        w = hidden.prior.fillna(np.nan)
        if w.isna().all():
            continue  # no adjacent trustworthy year → keep hybrid
        w = w.fillna(0.0)
        if w.sum() <= 0:
            continue  # every hidden group absent in adjacent years → keep hybrid (cannot all be zero when U > 0)
        U = float(g.U.iloc[0])
        if cap is None:
            alloc = U * w / w.sum()
        else:
            alloc = pd.Series(0.0, index=hidden.index)
            limit = cap * hidden.hid.astype(float)
            remaining, active = U, w.copy()
            for _ in range(len(hidden) + 1):
                if remaining <= 1e-9 or active.sum() <= 0:
                    break
                share = remaining * active / active.sum()
                room = (limit - alloc).clip(lower=0)
                take = np.minimum(share, room)
                alloc = alloc + take
                remaining -= float(take.sum())
                active = active.where(alloc < limit - 1e-9, 0.0)
            if remaining > 1e-9:   # every hidden cell already at 9: spread the rest equally over the hidden cells
                alloc = alloc + remaining * hidden.hid / hidden.hid.sum()
        out.loc[hidden.index] = (hidden.n + alloc) / hidden.total * 100
        zero_hidden = g[(g.hid > 0) & ~g.index.isin(hidden.index)]
        out.loc[zero_hidden.index] = zero_hidden.n / zero_hidden.total * 100
    return out


def main() -> None:
    con = psycopg.connect(base.DSN)
    years = sys.argv[1:] or ["2018-19", "2022-23", "2023-24"]
    shares: dict[str, tuple[pd.DataFrame, dict, bool]] = {}
    for y in ALL_YEARS:
        try:
            shares[y] = year_shares(y, con)
            print(f"[prior] {y}: tokens consistent with hidden-only (no school with S>U) = {shares[y][2]}; schools with S>U = {int((shares[y][0].groupby('school').S.first() > shares[y][0].groupby('school').U.first()).sum())}")
        except IndexError:
            pass
    for year in years:
        hy, enr, trusted = shares[year]
        i = ALL_YEARS.index(year)
        adj = [ALL_YEARS[j] for j in (i - 1, i + 1) if 0 <= j < len(ALL_YEARS) and ALL_YEARS[j] in shares and shares[ALL_YEARS[j]][2]]
        prior = pd.concat([shares[a][0].set_index(["school", "race"]).hybrid for a in adj], axis=1).mean(axis=1) if adj else pd.Series(dtype=float)
        S_school, U_school = hy.school.map(hy.groupby("school").S.first()), hy.school.map(hy.groupby("school").U.first())
        rules = {
            "hybrid": hy.hybrid,
            "prior_school": prior_alloc(hy, prior, S_school > U_school) if len(prior) else hy.hybrid,
            "prior_year": prior_alloc(hy, prior, pd.Series(True, index=hy.index)) if (len(prior) and not trusted) else hy.hybrid,
            "prior_all": prior_alloc(hy, prior, pd.Series(True, index=hy.index)) if len(prior) else hy.hybrid,
            "prior_year_cap9": prior_alloc(hy, prior, pd.Series(True, index=hy.index), cap=9.0) if (len(prior) and not trusted) else hy.hybrid,
        }
        ref = base.ccd(year)
        key = hy.set_index(["school", "race"]).index
        print(f"\n=== {year}: adjacent trustworthy years for the prior: {adj or 'none'}; year trustworthy = {trusted}")
        for race in RACES:
            m = (hy.race == race) & hy.school.isin(ref.index)
            truth = hy[m].school.map(ref[race])
            line = f"  {race:9s}"
            for name, col in rules.items():
                e = (col[m] - truth)
                line += f" | {name:12s} MAE {e.abs().mean():5.2f} bias {e.mean():+5.2f} p90 {e.abs().quantile(0.9):5.2f}"
            print(line)
        # process-bound test as the harness runs it: |rule − CCD| ≤ min(9·hid_r, U)/total·100 + 0.5
        for name, col in rules.items():
            m = hy.school.isin(ref.index)
            bound = (np.minimum(9.0 * hy.hid[m], hy.U[m]) / hy.total[m] * 100 + 0.5)
            truth = pd.Series([ref.loc[s_, r_] for s_, r_ in zip(hy.school[m], hy.race[m])], index=hy.index[m])
            ok = ((col[m] - truth).abs() <= bound)
            big3 = hy.race[m].isin(["black", "hispanic", "white"])
            print(f"  {name:14s} within process bound: all races {ok.mean() * 100:5.1f}%  black/hispanic/white {ok[big3].mean() * 100:5.1f}%")
        # held-out selection: choose on one half of the schools, score on the other
        big3 = ["black", "hispanic", "white"]
        m = hy.school.isin(ref.index) & hy.race.isin(big3)
        truth = pd.Series([ref.loc[s_, r_] for s_, r_ in zip(hy.school[m], hy.race[m])], index=hy.index[m])
        half = (hy.school[m].astype(int) % 2).rename("half")
        errs = pd.DataFrame({name: (col[m] - truth).abs() for name, col in rules.items()}).join(half)
        mae_by_half = errs.groupby("half").mean()
        for h in (0, 1):
            chosen = mae_by_half.loc[h].idxmin()
            other = 1 - h
            print(f"  held-out: chosen on half {h} (schools {'even' if h == 0 else 'odd'}) = {chosen:15s} in-sample MAE {mae_by_half.loc[h, chosen]:.2f}; on half {other} MAE {mae_by_half.loc[other, chosen]:.2f} (best rule there {mae_by_half.loc[other].idxmin()} {mae_by_half.loc[other].min():.2f})")
        # uniform-split school-years (every cell hidden) and the small-group medians
        allhid = hy.groupby("school").apply(lambda g: (g.n.sum() == 0) and (g.hid.sum() > 0))
        for name, col in rules.items():
            med = {r: float(np.median(col[hy.race == r])) for r in ("amerind", "hawaiian", "multi")}
            uni = int(sum(1 for s in allhid[allhid].index if np.allclose(col[hy.school == s].to_numpy(), 100 / 7, atol=1e-6)))
            print(f"  {name:12s} medians amerind {med['amerind']:.3f} hawaiian {med['hawaiian']:.3f} multi {med['multi']:.3f}; uniform-14.29% school-years {uni}")


if __name__ == "__main__":
    main()