"""Which rule for suppressed race cells best recovers the true shares?

FLDOE's Membership by School by Grade by Race hides cells under 10 students
('*' / '**'; blanks and '*' also mark grades not served in some years). The
October total is known exactly (enrollment_oct), so the unallocated count
U = total - sum(numeric cells) is known per school. Three rules for the
share of race r:
  zero          : numeric_r / total
  equal_split   : (numeric_r + U * suppressed_cells_r / suppressed_cells) / total
  grade_split   : per grade, U_g = N_g (by-grade file) - visible_g is exact; split over that grade's hidden cells
  adaptive      : hybrid, but where a school has more hidden cells than hidden students, weight cells by (visible_r + 1)
  hybrid        : the loader's rule — hidden cells in a fully-visible grade are zeros; school-level equal split over the rest
Blank cells are grades not served (zero), not hidden counts. In 2022-23 '*' marks both zeros and hidden
counts; the hybrid rule resolves the ones in fully-visible grades. Pass years on the command line
(default 2018-19 2023-24; 2022-23 once nces_ccd_membership_2022-23 is fetched).
Reference constructed independently: NCES CCD race counts (unsuppressed) for
2018-19 and 2023-24. Error = rule share - CCD share, in percentage points.

Run: pipeline_miami/.venv/bin/python docs/cities/miami/probe/race_imputation.py
"""
from __future__ import annotations

import glob
import re
import zipfile

import numpy as np
import pandas as pd
import psycopg

DSN = "postgresql://edu:edu@localhost:5432/edu_data_local"
RACES = {"white": r"^white", "black": r"^black", "hispanic": r"^hispanic", "asian": r"^asian", "hawaiian": r"^native hawaiian", "amerind": r"^american indian", "multi": r"^two or more"}
CCD_LABEL = {"white": "White", "black": "Black or African American", "hispanic": "Hispanic/Latino", "asian": "Asian", "hawaiian": "Native Hawaiian or Other Pacific Islander", "amerind": "American Indian or Alaska Native", "multi": "Two or more races"}
TOK = {"*", "**", "***", ""}


GRADE_KEY = {"PK": "PK", "KG": "KG", **{str(i): str(i) for i in range(1, 13)}}


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


def fldoe_grade_split(year: str, con) -> pd.DataFrame:
    """Per-grade equal split: U_g = N_g - visible_g from the by-grade file (exact), spread over that grade's hidden cells."""
    f = glob.glob(f"data/cities/miami/sources/fldoe_membership_by_school_grade_race_{year}/*/*.xlsx")[0]
    raw = pd.read_excel(f, 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 = {r: next(i for i, c in enumerate(hdr) if re.match(p, c)) for r, p in RACES.items()}
    ig = hdr.index("grade")
    d = raw.iloc[h + 1:]
    d = d[d.iloc[:, 0].astype(str).str.strip().str.replace(".0", "", regex=False) == "13"]
    ng = {}
    for k, v in con.execute("select school_number || '|' || substr(metric_key, 17), value::float from miami_school_year_metrics where metric_key like 'enrollment_oct_g%%' and subgroup='ALL' and year=%s and value is not null", (year,)).fetchall():
        sch, g = k.split("|"); ng[(sch, norm_grade(g))] = v
    recs = []
    for _, row in d.iterrows():
        num = str(row.iloc[2]).strip().zfill(4); g = norm_grade(row.iloc[ig])
        vis, hidden_cells, vals = 0.0, [], {}
        for r, i in cols.items():
            v = str(row.iloc[i]).strip(); v = "" if v.lower() == "nan" else v
            if v == "":
                vals[r] = 0.0
            elif v in TOK:
                hidden_cells.append(r); vals[r] = 0.0
            else:
                vals[r] = float(v.replace(",", "")); vis += vals[r]
        Ng = ng.get((num, g))
        Ug = max(0.0, Ng - vis) if Ng is not None else None
        for r in cols:
            add = (Ug / len(hidden_cells)) if (Ug is not None and r in hidden_cells and hidden_cells) else 0.0
            recs.append((num, r, vals[r] + add, Ug is None and r in hidden_cells))
    df = pd.DataFrame(recs, columns=["school", "race", "n", "unknown"])
    return df.groupby(["school", "race"]).agg(n=("n", "sum"), unknown=("unknown", "sum")).reset_index()


def fldoe_hybrid(year: str, con) -> pd.DataFrame:
    """The loader's rule (scripts/loaders/miami/race_membership.ts): hidden cells in a grade whose
    published total is already fully visible are zeros; the school's U is split equally over the rest."""
    f = glob.glob(f"data/cities/miami/sources/fldoe_membership_by_school_grade_race_{year}/*/*.xlsx")[0]
    raw = pd.read_excel(f, 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 = {r: next(i for i, c in enumerate(hdr) if re.match(p, c)) for r, p in RACES.items()}
    ig = hdr.index("grade") if "grade" in hdr else None   # 2025-26 has no grade column
    d = raw.iloc[h + 1:]
    d = d[d.iloc[:, 0].astype(str).str.strip().str.replace(".0", "", regex=False) == "13"]
    ng = {}
    for k, v in con.execute("select school_number || '|' || substr(metric_key, 17), value::float from miami_school_year_metrics where metric_key like 'enrollment_oct_g%%' and subgroup='ALL' and year=%s and value is not null", (year,)).fetchall():
        sch, g = k.split("|"); ng[(sch, norm_grade(g))] = v
    vis: dict[tuple[str, str], float] = {}
    hidden: dict[tuple[str, str], int] = {}
    for _, row in d.iterrows():
        num = str(row.iloc[2]).strip().zfill(4); g = norm_grade(row.iloc[ig]) if ig is not None else None
        rowvis, rowhidden = 0.0, []
        for r, i in cols.items():
            v = str(row.iloc[i]).strip(); v = "" if v.lower() == "nan" else v
            if v == "":
                continue
            if v in TOK:
                rowhidden.append(r)
            else:
                x = float(v.replace(",", "")); vis[(num, r)] = vis.get((num, r), 0.0) + x; rowvis += x
        Ng = ng.get((num, g)) if g is not None else None
        if Ng is not None and abs(Ng - rowvis) < 0.5:
            continue   # fully visible grade → hidden cells are zeros
        for r in rowhidden:
            hidden[(num, r)] = hidden.get((num, r), 0) + 1
    recs = []
    schools = {k[0] for k in vis} | {k[0] for k in hidden}
    for num in schools:
        S = sum(v for (n_, r), v in hidden.items() if n_ == num)
        for r in cols:
            recs.append((num, r, vis.get((num, r), 0.0), hidden.get((num, r), 0), S))
    return pd.DataFrame(recs, columns=["school", "race", "n", "hid", "S"])


def fldoe(year: str) -> pd.DataFrame:
    f = glob.glob(f"data/cities/miami/sources/fldoe_membership_by_school_grade_race_{year}/*/*.xlsx")[0]
    raw = pd.read_excel(f, 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 = {r: next(i for i, c in enumerate(hdr) if re.match(p, c)) for r, p in RACES.items()}
    d = raw.iloc[h + 1:]
    d = d[d.iloc[:, 0].astype(str).str.strip().str.replace(".0", "", regex=False) == "13"]
    recs = []
    for _, row in d.iterrows():
        num = str(row.iloc[2]).strip().zfill(4)
        for r, i in cols.items():
            v = str(row.iloc[i]).strip()
            if v.lower() == "nan":
                v = ""
            sup = v in TOK and v != ""   # blank = zero (grade not served); '*'/'**' = hidden 1..9
            recs.append((num, r, 0.0 if (sup or v == "") else float(v.replace(",", "")), sup))
    df = pd.DataFrame(recs, columns=["school", "race", "n", "sup"])
    return df.groupby(["school", "race"]).agg(n=("n", "sum"), sup=("sup", "sum")).reset_index()


def ccd(year: str) -> pd.DataFrame:
    f = glob.glob(f"data/cities/miami/sources/nces_ccd_membership_{year}/*/*.zip")[0]
    z = zipfile.ZipFile(f)
    member = [m for m in z.namelist() if m.lower().endswith(".csv")][0]
    parts = []
    with z.open(member) as fh:
        for chunk in pd.read_csv(fh, dtype=str, chunksize=500_000, encoding="latin1", low_memory=False):
            sel = chunk[(chunk.ST_LEAID == "FL-13") & chunk.TOTAL_INDICATOR.str.startswith("Derived - Subtotal by Race/Ethnicity and Sex")]
            parts.append(sel[["ST_SCHID", "RACE_ETHNICITY", "STUDENT_COUNT"]])
    d = pd.concat(parts)
    d["school"] = d.ST_SCHID.str.replace("FL-13-", "", regex=False).str.zfill(4)
    d["n"] = pd.to_numeric(d.STUDENT_COUNT, errors="coerce").fillna(0)
    g = d.groupby(["school", "RACE_ETHNICITY"])["n"].sum().unstack(fill_value=0)
    tot = g.sum(axis=1)
    out = pd.DataFrame({r: g.get(lbl, 0) / tot * 100 for r, lbl in CCD_LABEL.items()})
    return out[tot > 0]


def token_census(con) -> None:
    """How many suppressed cells vs how many hidden students, per year: if
    cells exceed hidden students, some tokens are zeros (unserved grades)."""
    print("year     | schools | '*' | '**' | blank | hidden students U | cells S | schools with S > U")
    for year in ["2018-19", "2021-22", "2022-23", "2023-24", "2024-25", "2025-26"]:
        f = glob.glob(f"data/cities/miami/sources/fldoe_membership_by_school_grade_race_{year}/*/*.xlsx")[0]
        raw = pd.read_excel(f, 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 RACES.values()]
        d = raw.iloc[h + 1:]
        d = d[d.iloc[:, 0].astype(str).str.strip().str.replace(".0", "", regex=False) == "13"]
        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())
        cnt = {"*": 0, "**": 0, "": 0}
        per_school: dict[str, list] = {}
        for _, row in d.iterrows():
            num = str(row.iloc[2]).strip().zfill(4)
            ps = per_school.setdefault(num, [0.0, 0])
            for i in cols:
                v = str(row.iloc[i]).strip()
                v = "" if v.lower() == "nan" else v
                if v in cnt:
                    cnt[v] += 1
                    ps[1] += 1
                elif v in TOK:
                    ps[1] += 1
                else:
                    ps[0] += float(v.replace(",", ""))
        U = sum(max(0.0, enr[k] - v[0]) for k, v in per_school.items() if k in enr)
        S = sum(v[1] for k, v in per_school.items() if k in enr)
        s_gt_u = sum(1 for k, v in per_school.items() if k in enr and v[1] > max(0.0, enr[k] - v[0]))
        print(f"{year} | {len(per_school):7d} | {cnt['*']:5d} | {cnt['**']:4d} | {cnt['']:5d} | {U:17.0f} | {S:7d} | {s_gt_u}")


def main() -> None:
    con = psycopg.connect(DSN)
    token_census(con)
    import sys
    years = sys.argv[1:] or ["2018-19", "2023-24"]
    for year in years:
        fl = fldoe(year)
        gs = fldoe_grade_split(year, con).set_index(["school", "race"])
        hy = fldoe_hybrid(year, con).set_index(["school", "race"])
        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())
        ref = ccd(year)
        rows = []
        for school, g in fl.groupby("school"):
            total = enr.get(school)
            if not total or school not in ref.index:
                continue
            numeric = g.n.sum()
            U = max(0.0, total - numeric)
            S = g.sup.sum()
            for _, r in g.iterrows():
                zero = r.n / total * 100
                eq = (r.n + (U * r.sup / S if S else 0)) / total * 100
                # capped: each hidden cell gets min(9, U/S); any remainder follows the visible mix
                per_cell = min(9.0, U / S) if S else 0.0
                rem = max(0.0, U - per_cell * S)
                capped = (r.n + per_cell * r.sup + (rem * r.n / numeric if numeric > 0 else 0)) / total * 100
                gsplit = gs.loc[(school, r.race), "n"] / total * 100 if (school, r.race) in gs.index else np.nan
                if (school, r.race) in hy.index:
                    hrow = hy.loc[(school, r.race)]
                    hyb = (hrow.n + (U * hrow.hid / hrow.S if hrow.S else 0)) / total * 100
                    # adaptive: when a school has more hidden cells than hidden students (so some cells are zeros),
                    # weight each hidden cell by the race's visible presence (+1 smoothing) instead of equally
                    if hrow.S > U:
                        hs = hy.loc[school]
                        wts = {rr: (hs.loc[rr, "n"] + 1.0) * hs.loc[rr, "hid"] for rr in hs.index}
                        W = sum(wts.values())
                        adapt = (hrow.n + (U * wts[r.race] / W if W else 0)) / total * 100
                    else:
                        adapt = hyb
                else:
                    hyb = np.nan; adapt = np.nan
                rows.append((school, r.race, total, U / total * 100, zero - ref.loc[school, r.race], eq - ref.loc[school, r.race], gsplit - ref.loc[school, r.race], hyb - ref.loc[school, r.race], adapt - ref.loc[school, r.race]))
        e = pd.DataFrame(rows, columns=["school", "race", "total", "unalloc_pp", "err_zero", "err_equal", "err_grade", "err_hybrid", "err_adaptive"])
        print(f"\n=== {year}: {e.school.nunique()} schools matched to CCD; unallocated pp median {e.groupby('school').unalloc_pp.first().median():.2f}")
        for race in ["black", "hispanic", "white", "asian", "multi"]:
            s = e[e.race == race]
            line = f"  {race:9s}"
            for c in ["err_zero", "err_equal", "err_grade", "err_hybrid", "err_adaptive"]:
                x = s[c].dropna()
                line += f" | {c[4:]:6s} MAE {x.abs().mean():5.2f} p90 {x.abs().quantile(0.9):5.2f} bias {x.mean():+5.2f} within2pp {(x.abs() <= 2).mean() * 100:5.1f}%"
            print(line)
        big = e[e.unalloc_pp >= 20]
        print(f"  schools with ≥20 pp unallocated: {big.school.nunique()} (median total {big.groupby('school').total.first().median():.0f}); their black/hispanic MAE zero {big[big.race.isin(['black','hispanic'])].err_zero.abs().mean():.2f} equal {big[big.race.isin(['black','hispanic'])].err_equal.abs().mean():.2f} grade {big[big.race.isin(['black','hispanic'])].err_grade.abs().mean():.2f} hybrid {big[big.race.isin(['black','hispanic'])].err_hybrid.abs().mean():.2f} adaptive {big[big.race.isin(['black','hispanic'])].err_adaptive.abs().mean():.2f}")


if __name__ == "__main__":
    main()