"""Does parent or student sentiment predict later enrollment loss, beyond what
utilization and the school's grade already say?

Pre-registered (2026-09-15, before running): for district-run schools open
in both years, regress the log change in October membership from the survey
year to 2025-26 on the survey year's "would recommend the school" share,
controlling for that year's utilization (headcount ÷ permanent capacity),
the school-grade letter index, school level and log enrollment. Report the
recommend coefficient with its 95% CI, the partial R² it adds over the
controls, and the same for the four survey composites. Two survey years are
used so the result is not one cohort's accident: 2021-22 (four years of
change) and 2022-23 (three). Expectation stated in advance: a positive
coefficient (higher recommend → less decline) that is small next to the
utilization and grade terms; an absent effect is reported as such.

Added after the 2026-09-16b review (N6), not pre-registered: the two cohorts
are also compared on the schools present in BOTH regressions (the 2022-23
workbook withholds fewer schools under 10 responses), the cohort difference
gets a paired school bootstrap, each predictor's year-to-year correlation is
reported, and the twelve predictor tests carry Holm-adjusted p-values. The
correlations do not distinguish sampling error from changes in sentiment or
respondents. Both windows end in 2025-26 and have different durations, so
their total-change coefficients are not independent replications of an
identical follow-up period.

Writes data/analysis/miami-survey-enrollment.json.
Run:  python -m pipeline_miami.analysis.survey_enrollment
"""
from __future__ import annotations

import datetime as dt
import json
import math

import numpy as np
import pandas as pd
from scipy import stats

from pipeline_miami.verify import common

OUT = common.REPO_ROOT / "data" / "analysis" / "miami-survey-enrollment.json"
END = "2025-26"
SURVEY_YEARS = ["2021-22", "2022-23"]
PREDICTORS = ["survey_parents_recommend", "survey_students_recommend", "survey_parents_teaching", "survey_students_safety", "survey_students_teaching", "survey_staff_leadership"]
CONTROLS = ["utilization_permanent_pct", "school_grade_letter_index", "log_enrollment"]
SCORING_VERSION = "favorable-response-v2"


def ols(X: np.ndarray, y: np.ndarray) -> dict:
    n, k = X.shape
    beta, *_ = np.linalg.lstsq(X, y, rcond=None)
    resid = y - X @ beta
    dof = n - k
    s2 = float(resid @ resid / dof)
    cov = s2 * np.linalg.pinv(X.T @ X)
    se = np.sqrt(np.diag(cov))
    r2 = 1 - float(resid @ resid) / float(((y - y.mean()) ** 2).sum())
    return {"beta": beta, "se": se, "r2": r2, "n": n, "dof": dof}


def design(df: pd.DataFrame, cols: list[str]) -> np.ndarray:
    parts = [np.ones(len(df))]
    for c in cols:
        parts.append(df[c].to_numpy(dtype=float))
    lv = pd.get_dummies(df["level"].astype("category"), drop_first=True)
    for c in sorted(lv.columns):
        parts.append(lv[c].to_numpy(dtype=float))
    return np.column_stack(parts)


def common_school_comparison(bases: dict[str, pd.DataFrame], years: list[str], n_boot: int = 2000, seed: int = 20260916) -> dict:
    """For each predictor: both cohorts refitted on the schools that have the predictor and the controls in BOTH
    survey years; the difference of the two per-pp coefficients with a paired school bootstrap (schools resampled
    with replacement, both cohorts refitted on the same resample); the predictor's year-to-year Pearson r on those
    schools (descriptive stability, not an identified measurement-error reliability);
    the median respondent count. The cohorts share an endpoint and differ in duration."""
    if len(years) != 2 or any(y not in bases for y in years):
        return {"skipped": "needs exactly two cohorts"}
    a, b = bases[years[0]], bases[years[1]]
    rng = np.random.default_rng(seed)
    out: dict = {"years": years, "n_boot": n_boot, "predictors": {}}
    for p in PREDICTORS:
        ca = a.dropna(subset=[p]).set_index("school")
        cb = b.dropna(subset=[p]).set_index("school")
        common = sorted(set(ca.index) & set(cb.index))
        if len(common) < 60:
            out["predictors"][p] = {"n_common": len(common), "skipped": "fewer than 60 common schools"}
            continue
        ca, cb = ca.loc[common].reset_index(), cb.loc[common].reset_index()
        j = len(CONTROLS) + 1

        def coef(df: pd.DataFrame) -> tuple[float, float]:
            f = ols(design(df, CONTROLS + [p]), df.dlog.to_numpy())
            return float(f["beta"][j]), float(f["se"][j])

        ba, sa = coef(ca)
        bb, sb = coef(cb)
        diffs = []
        for _ in range(n_boot):
            idx = rng.integers(0, len(common), len(common))
            ra, rb = ca.iloc[idx], cb.iloc[idx]
            try:
                diffs.append(coef(rb)[0] - coef(ra)[0])
            except (np.linalg.LinAlgError, ValueError):
                continue
        d = np.array(diffs)
        p_two = float(2 * min((d <= 0).mean(), (d >= 0).mean())) if len(d) else float("nan")
        rel = float(stats.pearsonr(ca[p], cb[p])[0])
        sd_pooled = float(pd.concat([ca[p], cb[p]]).std())
        out["predictors"][p] = {
            "n_common": len(common), "n_published": [int(len(a.dropna(subset=[p]))), int(len(b.dropna(subset=[p])))],
            years[0]: {"coef_per_pp": round(ba, 5), "se": round(sa, 5), "t": round(ba / sa, 3) if sa > 0 else None, "effect_of_1sd_in_pct_enrollment": round((math.exp(ba * sd_pooled) - 1) * 100, 3)},
            years[1]: {"coef_per_pp": round(bb, 5), "se": round(sb, 5), "t": round(bb / sb, 3) if sb > 0 else None, "effect_of_1sd_in_pct_enrollment": round((math.exp(bb * sd_pooled) - 1) * 100, 3)},
            "difference_per_pp": round(bb - ba, 5), "difference_ci95_bootstrap": [round(float(np.percentile(d, 2.5)), 5), round(float(np.percentile(d, 97.5)), 5)] if len(d) else None,
            "difference_p_bootstrap": p_two, "ratio_second_to_first": round(bb / ba, 3) if ba != 0 else None,
            "year_to_year_pearson_r": round(rel, 4), "predictor_sd_pooled_pp": round(sd_pooled, 2),
            "median_respondents_per_school": [float(ca[f"n__{p}"].median()) if f"n__{p}" in ca else None, float(cb[f"n__{p}"].median()) if f"n__{p}" in cb else None],
        }
    return out


def main() -> None:
    import psycopg

    keys = PREDICTORS + ["utilization_permanent_pct", "fish_capacity_permanent", "school_grade_letter_index", "enrollment_oct"]
    with psycopg.connect(common.served_db_dsn()) as conn:
        conn.read_only = True
        stale_sources = conn.execute("""
            SELECT DISTINCT COALESCE(l.source_id, '(missing source load)')
              FROM miami_school_year_metrics m LEFT JOIN miami_data_loads l ON l.id=m.source_load_id
             WHERE m.metric_key LIKE 'survey_%%' AND
                   (l.notes IS NULL OR l.notes NOT LIKE %s)
        """, (f"%scoring {SCORING_VERSION} %",)).fetchall()
        if stale_sources:
            raise SystemExit(f"survey scores need reloading with {SCORING_VERSION}: {', '.join(r[0] for r in stale_sources)}")
        rows = conn.execute("SELECT m.school_number, m.year, m.metric_key, m.value::float8, m.denominator::float8 FROM miami_school_year_metrics m WHERE m.subgroup='ALL' AND m.suppressed=false AND m.value IS NOT NULL AND m.metric_key = ANY(%s)", (keys,)).fetchall()
        sch = pd.DataFrame(conn.execute("SELECT school_number, school_level, is_charter, activity_code, closed_at FROM miami_schools").fetchall(), columns=["school", "level", "is_charter", "activity_code", "closed_at"])
    df = pd.DataFrame(rows, columns=["school", "year", "key", "value", "den"])
    wide = df.pivot_table(index=["school", "year"], columns="key", values="value", aggfunc="first").reset_index()
    dens = df[df.key.isin(PREDICTORS)].pivot_table(index=["school", "year"], columns="key", values="den", aggfunc="first")
    dens.columns = [f"n__{c}" for c in dens.columns]
    wide = wide.merge(dens.reset_index(), on=["school", "year"], how="left")
    end = wide[wide.year == END][["school", "enrollment_oct"]].rename(columns={"enrollment_oct": "enr_end"})
    results = {"meta": {"generated_at": dt.datetime.now(dt.timezone.utc).isoformat(), "scoring_version": SCORING_VERSION, "end_year": END, "survey_years": SURVEY_YEARS, "controls": CONTROLS + ["level dummies"],
                        "followup_years": {y: int(END[:4]) - int(y[:4]) for y in SURVEY_YEARS},
                        "cohort_comparison_note": "Total log changes over different follow-up durations, with the same endpoint and overlapping schools; not independent replications of an identical horizon.",
                        "year_to_year_correlation_note": "Descriptive stability; does not identify measurement reliability or explain the difference between coefficients.",
                        "universe": "district-run schools, active, with October membership in both years and all controls"}, "by_survey_year": {}}
    util_years = sorted(wide.loc[wide.utilization_permanent_pct.notna(), "year"].unique())
    missing_prereg = []
    bases: dict[str, pd.DataFrame] = {}
    for sy in SURVEY_YEARS:
        # Funnel counted BEFORE the join on the end year, so schools that closed or merged are visible as a drop.
        sy_rows = wide[wide.year == sy].merge(sch, on="school")
        funnel = {"schools_with_any_metric_row_in_survey_year": int(len(sy_rows)),
                  "with_survey_year_membership": int((sy_rows.enrollment_oct > 0).sum()),
                  "with_any_survey_predictor": int(sy_rows[PREDICTORS].notna().any(axis=1).sum())}
        base = sy_rows.merge(end, on="school")
        funnel["with_end_year_membership"] = int(((base.enrollment_oct > 0) & (base.enr_end > 0)).sum())
        funnel["dropped_no_end_year_membership_closed_or_merged"] = funnel["with_survey_year_membership"] - funnel["with_end_year_membership"]
        # Utilization needs a FISH capacity; the first report is 2022-23. For an earlier survey year the control is
        # that year's own October headcount over the first reported permanent capacity, so the survey-year
        # membership (the outcome's starting point) is not replaced by a later year's (review 2026-09-16b, N3:
        # the 2022-23 utilization value carried 2022-23 headcount and inflated the control's coefficient).
        uy = sy if sy in util_years else next((y for y in util_years if y > sy), None)
        if uy != sy:
            cap = wide[wide.year == uy][["school", "fish_capacity_permanent"]].rename(columns={"fish_capacity_permanent": "cap_first"})
            base = base.drop(columns=["utilization_permanent_pct"]).merge(cap, on="school", how="left")
            base["utilization_permanent_pct"] = (base.enrollment_oct / base.cap_first * 100.0).where(base.cap_first > 0)
            base = base.drop(columns=["cap_first"])
            results["meta"].setdefault("utilization_basis", {})[sy] = f"October {sy} headcount ÷ {uy} permanent capacity (first FISH report) × 100"
        else:
            results["meta"].setdefault("utilization_basis", {})[sy] = f"served utilization_permanent_pct for {sy}"
        base = base[(~base.is_charter) & (base.activity_code == "A") & base.closed_at.isna()].copy()
        funnel["district_run_active"] = int(len(base))
        base = base[(base.enrollment_oct > 0) & (base.enr_end > 0)].copy()
        funnel["with_membership_both_years"] = int(len(base))
        base["dlog"] = np.log(base.enr_end) - np.log(base.enrollment_oct)
        base["log_enrollment"] = np.log(base.enrollment_oct)
        for c in CONTROLS:
            base = base[base[c].notna()]
            funnel[f"with_{c}"] = int(len(base))
        base = base.dropna(subset=["dlog", "level"])
        out_sy = {"n_controls_only": int(len(base)), "funnel": funnel, "predictors": {}}
        bases[sy] = base
        Xc = design(base, CONTROLS)
        fc = ols(Xc, base.dlog.to_numpy())
        out_sy["controls_only"] = {"r2": round(fc["r2"], 4), "n": fc["n"], "coef": {c: round(float(b), 5) for c, b in zip(["intercept"] + CONTROLS, fc["beta"][: len(CONTROLS) + 1])},
                                  "se": {c: round(float(s), 5) for c, s in zip(["intercept"] + CONTROLS, fc["se"][: len(CONTROLS) + 1])}}
        for p in PREDICTORS:
            sub = base.dropna(subset=[p])
            if len(sub) < 60:
                out_sy["predictors"][p] = {"n": int(len(sub)), "skipped": "fewer than 60 schools with the predictor"}
                if p == PREDICTORS[0]:
                    missing_prereg.append(f"{p} {sy} (n={len(sub)})")
                continue
            y = sub.dlog.to_numpy()
            f0 = ols(design(sub, CONTROLS), y)
            f1 = ols(design(sub, CONTROLS + [p]), y)
            j = len(CONTROLS) + 1   # index of the predictor coefficient
            b, se = float(f1["beta"][j]), float(f1["se"][j])
            t = b / se if se > 0 else float("nan")
            pval = 2 * (1 - stats.t.cdf(abs(t), f1["dof"])) if se > 0 else float("nan")
            sd_p = float(sub[p].std())
            # raw (unadjusted) association for the decomposition
            rr = stats.pearsonr(sub[p], sub.dlog)
            n_resp = sub[f"n__{p}"] if f"n__{p}" in sub.columns else pd.Series(dtype=float)
            out_sy["predictors"][p] = {
                "n": int(len(sub)), "coef_per_pp": round(b, 5), "se": round(se, 5), "t": round(t, 3), "p": float(pval),
                "median_respondents_per_school": float(n_resp.median()) if n_resp.notna().any() else None,
                "ci95_per_pp": [round(b - 1.96 * se, 5), round(b + 1.96 * se, 5)],
                "effect_of_1sd_in_pct_enrollment": round((math.exp(b * sd_p) - 1) * 100, 3), "predictor_sd_pp": round(sd_p, 2),
                "r2_controls": round(f0["r2"], 4), "r2_with_predictor": round(f1["r2"], 4), "partial_r2": round((f1["r2"] - f0["r2"]) / max(1e-9, 1 - f0["r2"]), 4),
                "raw_pearson_r": round(float(rr[0]), 4), "raw_p": float(rr[1]),
            }
        results["by_survey_year"][sy] = out_sy
    # The pre-registration names two survey years for the headline predictor; a year where it is absent is a
    # loading failure, not a result (2022-23 parents' recommend was silently absent until 2026-09-16), so stop.
    if missing_prereg:
        raise SystemExit(f"pre-registered headline predictor missing: {'; '.join(missing_prereg)} — fix the survey load before publishing")
    # Holm over the family of predictor tests (6 predictors × 2 cohorts): the pre-registration named one headline
    # predictor but twelve coefficients are reported, so each carries its family-wise adjusted p as well.
    tests = [(sy, p, v["p"]) for sy, r in results["by_survey_year"].items() for p, v in r["predictors"].items() if "p" in v]
    order = sorted(range(len(tests)), key=lambda i: tests[i][2])
    m, running = len(tests), 0.0
    for rank, i in enumerate(order):
        running = max(running, min(1.0, (m - rank) * tests[i][2]))
        sy, p, _ = tests[i]
        results["by_survey_year"][sy]["predictors"][p]["p_holm"] = float(running)
        results["by_survey_year"][sy]["predictors"][p]["survives_holm_05"] = bool(running < 0.05)
    results["meta"]["multiple_testing"] = {"family_size": m, "method": "Holm step-down over all predictor × cohort tests", "survive_05": [f"{p} {sy}" for sy, r in results["by_survey_year"].items() for p, v in r["predictors"].items() if v.get("survives_holm_05")]}
    # Same schools in both cohorts (review 2026-09-16b, N6): the cohorts as published are different school sets
    # (the 2022-23 workbook withholds fewer schools under 10 responses), so the cohort comparison is repeated on
    # the intersection, with a paired school bootstrap of the difference and each predictor's year-to-year correlation.
    results["common_schools"] = common_school_comparison(bases, SURVEY_YEARS)
    # Coverage: which schools answer, by respondent group and year, and how many core items each composite uses.
    with psycopg.connect(common.served_db_dsn()) as conn:
        conn.read_only = True
        # with_value_active counts only schools active today, the same universe as the denominators below (review
        # 2026-09-16b, N13: closed schools with early-year values overstated coverage).
        cov = conn.execute("""
            SELECT m.metric_key, m.year, s.is_charter, count(*) FILTER (WHERE m.value IS NOT NULL) AS with_value,
                   count(*) FILTER (WHERE m.value IS NOT NULL AND s.activity_code = 'A' AND s.closed_at IS NULL) AS with_value_active, count(*) AS cells
              FROM miami_school_year_metrics m JOIN miami_schools s ON s.school_number = m.school_number
             WHERE m.metric_key IN ('survey_students_safety','survey_parents_safety','survey_staff_safety') AND m.subgroup = 'ALL'
             GROUP BY 1,2,3 ORDER BY 1,2,3""").fetchall()
        active = conn.execute("SELECT is_charter, count(*) FROM miami_schools WHERE activity_code='A' AND closed_at IS NULL GROUP BY 1").fetchall()
        # The like-for-like participation universe: active REGULAR schools (parents, staff) and, for students, those with
        # October membership in grades 4-12 in the latest survey year (review 2026-09-16, S5 / climate page denominators).
        latest_year = conn.execute("SELECT max(year) FROM miami_school_year_metrics WHERE metric_key='survey_students_safety'").fetchone()[0]
        regular = conn.execute("SELECT is_charter, count(*) FROM miami_schools WHERE activity_code='A' AND closed_at IS NULL AND school_type='REGULAR' GROUP BY 1").fetchall()
        regular_g4 = conn.execute("""
            SELECT s.is_charter, count(*) FROM miami_schools s WHERE s.activity_code='A' AND s.closed_at IS NULL AND s.school_type='REGULAR'
               AND EXISTS (SELECT 1 FROM miami_school_year_metrics g WHERE g.school_number=s.school_number AND g.year=%s AND g.subgroup='ALL' AND NOT g.suppressed AND g.value>0
                            AND g.metric_key IN ('enrollment_oct_g4','enrollment_oct_g5','enrollment_oct_g6','enrollment_oct_g7','enrollment_oct_g8','enrollment_oct_g9','enrollment_oct_g10','enrollment_oct_g11','enrollment_oct_g12'))
             GROUP BY 1""", (latest_year,)).fetchall()
        with_value_regular_g4 = conn.execute("""
            SELECT s.is_charter, count(DISTINCT m.school_number) FROM miami_school_year_metrics m JOIN miami_schools s ON s.school_number=m.school_number
             WHERE m.metric_key='survey_students_safety' AND m.year=%s AND m.value IS NOT NULL AND s.activity_code='A' AND s.closed_at IS NULL AND s.school_type='REGULAR'
               AND EXISTS (SELECT 1 FROM miami_school_year_metrics g WHERE g.school_number=s.school_number AND g.year=%s AND g.subgroup='ALL' AND NOT g.suppressed AND g.value>0
                            AND g.metric_key IN ('enrollment_oct_g4','enrollment_oct_g5','enrollment_oct_g6','enrollment_oct_g7','enrollment_oct_g8','enrollment_oct_g9','enrollment_oct_g10','enrollment_oct_g11','enrollment_oct_g12'))
             GROUP BY 1""", (latest_year, latest_year)).fetchall()
    xw = pd.read_csv(common.REPO_ROOT / "data" / "cities" / "miami" / "reference" / "survey_items.csv")
    core = xw[xw.core].drop_duplicates(["respondent", "canonical"]).groupby(["respondent", "domain"]).size()
    results["coverage"] = {
        "active_schools": {("charter" if c else "district"): int(n) for c, n in active},
        "active_regular": {("charter" if c else "district"): int(n) for c, n in regular},
        "students_universe": {"year": latest_year, "note": "active REGULAR schools with October membership in grades 4-12",
                              "eligible": {("charter" if c else "district"): int(n) for c, n in regular_g4},
                              "with_students_safety": {("charter" if c else "district"): int(n) for c, n in with_value_regular_g4}},
        "by_respondent_year": [{"respondent": k.split("_")[1], "year": y, "sector": "charter" if c else "district", "with_value": int(w), "with_value_active": int(wa), "cells": int(n)} for k, y, c, w, wa, n in cov],
        "core_items_per_domain": {f"{r}_{d}": int(n) for (r, d), n in core.items()},
        "items_per_respondent": {r: {"canonical_items": int(g.canonical.nunique()), "core_items": int(g[g.core].canonical.nunique()), "years": sorted(g.year.unique().tolist())} for r, g in xw.groupby("respondent")},
        "reverse_coded_items": int(xw[xw.reverse].drop_duplicates(["respondent", "canonical"]).shape[0]),
    }
    OUT.write_text(json.dumps(results, indent=1, default=str))
    for sy, r in results["by_survey_year"].items():
        print(f"== {sy}: controls-only R² {r['controls_only']['r2']} (n={r['controls_only']['n']}); util coef {r['controls_only']['coef']['utilization_permanent_pct']}, grade coef {r['controls_only']['coef']['school_grade_letter_index']}")
        for p, v in r["predictors"].items():
            if "skipped" in v:
                print(f"   {p:32s} {v['skipped']}")
            else:
                print(f"   {p:32s} n={v['n']:3d} coef/pp={v['coef_per_pp']:+.5f} (p={v['p']:.3g})  1 SD → {v['effect_of_1sd_in_pct_enrollment']:+.2f}% enrollment  partial R²={v['partial_r2']:.4f}  raw r={v['raw_pearson_r']:+.3f}")


if __name__ == "__main__":
    main()