"""(iii-b) Cross-publisher check for the race-share covariates: the served
pct_black / pct_hispanic / pct_white (from FLDOE's suppressed race file, with
the allocation rule chosen in docs/cities/miami/probe/race_imputation.py)
against the unsuppressed NCES CCD race subtotals for the years both exist.

Three statements, all fixed before running:
  (a) Consistency with the process bound (NOT an accuracy test). Each hidden
      cell holds 1–9 students, so for race r the served share can miss the
      truth by at most min(9 × hidden_cells_r, U) / total, where U is the
      school's unallocated count. The re-read of the FLDOE file (independent
      Python) supplies hidden_cells_r and U per school, counting as hidden only
      the cells the loader also treats as hidden: a token in a grade whose
      by-grade total is already fully visible is a zero (loader rule 1), so it
      is not counted (review 2026-09-16, M7). Rule: |served − CCD| within the
      bound + 0.5 pp for ≥ 99% of school-years (the 0.5 pp allows for the two
      publishers' totals differing by a few students). Any allocation that
      keeps every hidden cell in [0, 9] and sums to the total satisfies this —
      the plain equal split passes it at 100% in 2022-23 with a 2.4 pp mean
      error (review 2026-09-16b, N9) — so a failure means a loader defect
      (a cell above the cap, a total that does not close), not a poor prior.
  (b) Uniform-split guard — no school-year may be served with all seven
      shares equal (the signature of an equal split over a school whose every
      cell is hidden; 10 such school-years were served in 2022-23 before the
      adjacent-year prior was adopted).
  (c) Accuracy against the no-information reference: per share, the served
      value's mean |Δ| from CCD is no larger than the equal split's (the
      unallocated students spread evenly over the hidden cells, recomputed
      here from the re-read file). The two coincide wherever no school
      exceeds the cap (every year but 2022-23), so the statement has power
      only where the capped adjacent-year prior is used. Caveat recorded with
      the verdict: that prior was chosen among five candidates on CCD
      2022-23, the same reference (in-sample selection; the split-half
      held-out check lives in docs/cities/miami/probe/race_imputation_prior.py).
  Informational — mean |Δ|, share within 2 pp and the signed median per year
      and per share, reported whatever they are. A first run of a
      95%-within-2-pp rule came in at 87–90% (2026-09-15) because small
      schools with many hidden cells cannot be resolved that finely; that
      rule is not the verdict.

Run:  python -m pipeline_miami.verify.cross_publisher_race
"""
from __future__ import annotations

import glob
import json
import re
import statistics

import pandas as pd

from pipeline_miami.verify import common
from pipeline_miami.verify.connectors import served_db, source_file

YEARS = ["2018-19", "2022-23", "2023-24"]
KEYS = {"pct_black": r"^black", "pct_hispanic": r"^hispanic", "pct_white": r"^white"}
ALL_RACES = [r"^white", r"^black", r"^hispanic", r"^asian", r"^native hawaiian", r"^american indian", r"^two or more"]
TOK = {"*", "**", "***"}


def _grade_key(v: str) -> str:
    t = str(v).strip().upper()
    return t if t in ("PK", "KG") else (str(int(t)) if t.isdigit() else t)


def served_grade_totals(year: str) -> dict[tuple[str, str], float]:
    """(school, grade) → served October membership for that grade (unsuppressed cells only)."""
    import psycopg

    out: dict[tuple[str, str], float] = {}
    with psycopg.connect(common.served_db_dsn()) as conn:
        conn.read_only = True
        for num, key, val in conn.execute("SELECT school_number, metric_key, value::float8 FROM miami_school_year_metrics WHERE year=%s AND subgroup='ALL' AND NOT suppressed AND value IS NOT NULL AND metric_key LIKE 'enrollment_oct_g%%'", (year,)):
            out[(str(num), _grade_key(key[len("enrollment_oct_g"):]))] = float(val)
    return out


def fldoe_hidden(year: str) -> dict[str, dict]:
    """Per school: visible count per race and hidden cells per race, where a token in a grade whose
    published by-grade total equals the row's visible sum is a zero, not a hidden cell (loader rule 1)."""
    path = glob.glob(str(common.SOURCES_DIR / f"fldoe_membership_by_school_grade_race_{year}" / "*" / "*.xlsx"))[0]
    raw = pd.read_excel(path, sheet_name="School", header=None, dtype=str)
    h = next(i for i in range(8) if str(raw.iloc[i, 0]).strip().lower().startswith("district"))
    hdr = [re.sub(r"[^a-z0-9]+", " ", str(c).lower()).strip() for c in raw.iloc[h].tolist()]
    cols = [next(i for i, c in enumerate(hdr) if re.match(p, c)) for p in ALL_RACES]
    ig = hdr.index("grade") if "grade" in hdr else None
    grade_totals = served_grade_totals(year) if ig is not None else {}
    d = raw.iloc[h + 1:]
    d = d[d.iloc[:, 0].astype(str).str.strip().str.replace(".0", "", regex=False) == "13"]
    out: dict[str, dict] = {}
    for _, row in d.iterrows():
        num = common.pad4(row.iloc[2])
        if not num:
            continue
        s = out.setdefault(num, {"visible": [0.0] * len(cols), "hidden": [0] * len(cols), "hidden_zero": 0})
        row_vis, row_hidden = 0.0, []
        for j, i in enumerate(cols):
            v = str(row.iloc[i]).strip()
            if v == "" or v.lower() == "nan":
                continue
            if v in TOK:
                row_hidden.append(j)
            else:
                x = float(v.replace(",", ""))
                s["visible"][j] += x
                row_vis += x
        g_total = grade_totals.get((num, _grade_key(row.iloc[ig]))) if ig is not None else None
        if g_total is not None and abs(g_total - row_vis) < 0.5:
            s["hidden_zero"] += len(row_hidden)      # fully visible grade: the tokens are zeros
            continue
        for j in row_hidden:
            s["hidden"][j] += 1
    return out


def main() -> None:
    common.ensure_qa_dir()
    enr = served_db.fetch("enrollment_oct")
    served = {k: served_db.fetch(k) for k in KEYS}
    per_key = {k: {"pairs": 0, "within_bound": 0, "within_2pp": 0, "deltas": [], "deltas_baseline": [], "largest": []} for k in KEYS}
    by_year: dict[str, dict] = {}
    all_shares = {k: served_db.fetch(k) for k in ["pct_white", "pct_black", "pct_hispanic", "pct_asian", "pct_hawaiian_pac_isl", "pct_amer_indian", "pct_two_or_more"]}
    for y in YEARS:
        try:
            ccd = {k: source_file.ccd_race(y, k) for k in KEYS}
        except SystemExit:
            continue
        hidden = fldoe_hidden(y)
        race_idx = {"pct_white": 0, "pct_black": 1, "pct_hispanic": 2}
        by_year[y] = {"hidden_cells": int(sum(sum(h["hidden"]) for h in hidden.values())), "tokens_resolved_as_zero": int(sum(h["hidden_zero"] for h in hidden.values())), "by_metric": {}}
        for k in KEYS:
            n_y, wb_y, w2_y, abs_y, absb_y = 0, 0, 0, [], []
            for (num, yy), ref in ccd[k].items():
                s = served[k].get((num, yy))
                e = enr.get((num, yy))
                hs = hidden.get(num)
                if not s or s["value"] is None or not e or e["value"] is None or not hs:
                    continue
                total = e["value"]
                U = max(0.0, total - sum(hs["visible"]))
                bound = min(9.0 * hs["hidden"][race_idx[k]], U) / total * 100.0 + 0.5
                d = s["value"] - ref["value"]
                # Equal-split baseline: the unallocated students spread evenly over the hidden cells. Every served
                # rule reduces to it where no school exceeds the cap, so it is the no-information reference the
                # served allocation must beat where it differs (2022-23, capped adjacent-year prior).
                n_hidden = sum(hs["hidden"])
                base_share = (hs["visible"][race_idx[k]] + (U * hs["hidden"][race_idx[k]] / n_hidden if n_hidden else 0.0)) / total * 100.0
                d_base = base_share - ref["value"]
                r = per_key[k]
                r["pairs"] += 1
                r["within_bound"] += int(abs(d) <= bound)
                r["within_2pp"] += int(abs(d) <= 2.0)
                r["deltas"].append(d)
                r["deltas_baseline"].append(d_base)
                r["largest"].append((num, yy, round(s["value"], 2), round(ref["value"], 2), round(bound, 2)))
                n_y += 1; wb_y += int(abs(d) <= bound); w2_y += int(abs(d) <= 2.0); abs_y.append(abs(d)); absb_y.append(abs(d_base))
            by_year[y]["by_metric"][k] = {"n_pairs": n_y, "within_bound_share": wb_y / n_y if n_y else None, "within_2pp_share": w2_y / n_y if n_y else None,
                                          "mean_abs_delta": statistics.fmean(abs_y) if abs_y else None, "mean_abs_delta_equal_split_baseline": statistics.fmean(absb_y) if absb_y else None,
                                          "lift_vs_equal_split": (1 - statistics.fmean(abs_y) / statistics.fmean(absb_y)) if absb_y and statistics.fmean(absb_y) > 0 else None}
    # (b) uniform-split guard over every served year
    uniform = []
    for (num, yy), cell in all_shares["pct_white"].items():
        vals = [all_shares[k].get((num, yy), {}).get("value") for k in all_shares]
        if all(v is not None for v in vals) and max(vals) - min(vals) < 1e-6 and vals[0] > 0:
            uniform.append(f"{num} {yy}")
    all_ok = not uniform
    summary_by = {}
    for k, r in per_key.items():
        n = r["pairs"]
        ok = n > 0 and r["within_bound"] / n >= 0.99
        all_ok = all_ok and ok
        mad_s = statistics.fmean(abs(d) for d in r["deltas"]) if r["deltas"] else None
        mad_b = statistics.fmean(abs(d) for d in r["deltas_baseline"]) if r["deltas_baseline"] else None
        beats = mad_s is not None and mad_b is not None and mad_s <= mad_b + 1e-9
        all_ok = all_ok and beats
        summary_by[k] = {"n_pairs": n, "within_bound": r["within_bound"], "within_bound_share": r["within_bound"] / n if n else None,
                         "within_2pp": r["within_2pp"], "within_2pp_share": r["within_2pp"] / n if n else None,
                         "median_delta": statistics.median(r["deltas"]) if r["deltas"] else None, "mean_abs_delta": mad_s,
                         "mean_abs_delta_equal_split_baseline": mad_b, "beats_equal_split": beats,
                         "verdict": "PASS" if ok and beats else "REVIEW", "largest": sorted(r["largest"], key=lambda t: -abs(t[2] - t[3]))[:8]}
    verdict = "PASS" if all_ok else "REVIEW"
    summary = {"summary": {"years": YEARS, "verdict": verdict, "uniform_split_school_years": uniform,
                           "verdict_note": "Served FLDOE-derived share vs NCES CCD race subtotal. (a) CONSISTENCY, not accuracy: |Δ| within the process bound min(9 × hidden cells, unallocated) / total + 0.5 pp for ≥ 99% of school-years, per share, with hidden cells counted as the loader counts them — any allocation that respects the cap and sums to the total satisfies it (the plain equal split included), so a failure means a loader defect, not a poor prior; (b) no school-year served with all seven shares equal; (c) ACCURACY against the no-information reference: the served share's mean |Δ| from CCD is no larger than the equal split's (the two coincide except where the capped adjacent-year prior applies, 2022-23). The prior was chosen among five candidates on CCD 2022-23, the same reference reported here (in-sample selection; the held-out split-half check is in docs/cities/miami/probe/race_imputation_prior.py). Mean |Δ| and the within-2-pp share are informational."},
               "by_metric": summary_by, "by_year": by_year}
    (common.QA_REPORTS_DIR / "__cross_publisher_race.json").write_text(json.dumps(summary, indent=2, default=str))
    out = ["# Cross-publisher (iii-b) — FLDOE race shares vs NCES CCD\n", f"Verdict: **{verdict}** ({', '.join(YEARS)}).\n",
           "| share | pairs | within process bound (consistency) | within 2 pp | median Δ (FLDOE − CCD) | mean |Δ| | mean |Δ| equal split | status |", "|---|---:|---:|---:|---:|---:|---:|---|"]
    for k, r in summary_by.items():
        out.append(f"| `{k}` | {r['n_pairs']} | {r['within_bound']} ({(r['within_bound_share'] or 0) * 100:.1f}%) | {r['within_2pp']} ({(r['within_2pp_share'] or 0) * 100:.1f}%) | {(r['median_delta'] or 0):+.2f} | {(r['mean_abs_delta'] or 0):.2f} | {(r['mean_abs_delta_equal_split_baseline'] or 0):.2f} | {r['verdict']} |")
    out.append("\n" + summary["summary"]["verdict_note"])
    out.append(f"\nUniform-split school-years served: {len(uniform)}" + (f" ({', '.join(uniform[:10])})" if uniform else ""))
    out += ["\n| year | hidden cells (loader rule) | tokens resolved as zero | share | pairs | within bound | within 2 pp | mean abs Δ | mean abs Δ equal split | lift |", "|---|---:|---:|---|---:|---:|---:|---:|---:|---:|"]
    for y, r in by_year.items():
        for k, v in r["by_metric"].items():
            out.append(f"| {y} | {r['hidden_cells']} | {r['tokens_resolved_as_zero']} | `{k}` | {v['n_pairs']} | {(v['within_bound_share'] or 0) * 100:.1f}% | {(v['within_2pp_share'] or 0) * 100:.1f}% | {(v['mean_abs_delta'] or 0):.2f} | {(v['mean_abs_delta_equal_split_baseline'] or 0):.2f} | {((v['lift_vs_equal_split'] or 0) * 100):+.1f}% |")
    for k, r in summary_by.items():
        out.append(f"\n## {k} — largest differences (school, year, FLDOE, CCD, bound)")
        out += [f"- {t[0]} {t[1]}: {t[2]} vs {t[3]} (bound ±{t[4]})" for t in r["largest"]]
    (common.QA_REPORTS_DIR / "__cross_publisher_race.md").write_text("\n".join(out))
    print(json.dumps({k: {kk: vv for kk, vv in v.items() if kk != "largest"} for k, v in summary_by.items()} | {"verdict": verdict}, indent=1, default=str))


if __name__ == "__main__":
    main()