#!/usr/bin/env python
"""Cross-outcome test: does adding PRIOR ACHIEVEMENT to the peer match improve
prediction of a DIFFERENT outcome than the one fed in? This isolates "better peer
group" from the trivial autocorrelation of predicting a metric from its own lag.

Feature added = prior-year PROFICIENCY (mean ELA/Math). Targets predicted =
absenteeism and graduation (and, as the autocorrelation reference, proficiency
itself). All methods scored on the SAME common subset (schools with prior
proficiency AND the target present)."""
import numpy as np
import pandas as pd
import bakeoff as B

panel = B.load_panel()

# prior-year proficiency, attached to every school-year (independent of target)
prof_panel = panel.copy()
prof_panel["prior_prof"] = B.build_prior(prof_panel, "prof")
key = prof_panel.set_index(["dbn", "year"])["prior_prof"]
panel = panel.merge(key.rename("prior_prof"), on=["dbn", "year"], how="left")

DEMO = B.DEMOG
DEMO_PLUS = B.DEMOG + ["prior_prof"]

def knn_expected(df, target, cols):
    e, _ = B.expected_knn(df, target, cols)
    return e

print("=" * 80)
print("CROSS-OUTCOME: peers built from demographics +/- prior PROFICIENCY,")
print("predicting a target that is NOT the fed-in metric. Common subset, LOO R^2.")
print("=" * 80)

# proficiency pool (ES/MS/K8) carries both proficiency and absenteeism + prior_prof
pool = panel[panel["grade_band"].isin({"ES", "MS", "K8"})].copy()

rows = []
for target, label in [("absent", "absenteeism (NOT the fed-in metric)"),
                      ("prof", "proficiency (autocorrelation reference)")]:
    for yr in sorted(pool["year"].unique()):
        d = pool[pool["year"] == yr].copy()
        mask = d["prior_prof"].notna() & d[target].notna()
        if mask.sum() < 50:
            continue
        # demographics-only peers
        e_demo = knn_expected(d, target, DEMO).reindex(d.index)
        # demographics + prior proficiency peers
        dd = d[d["prior_prof"].notna()].copy()
        e_plus = knn_expected(dd, target, DEMO_PLUS).reindex(d.index)
        r_demo = B.r2(d[target].values[mask.values], e_demo.values[mask.values])
        r_plus = B.r2(d[target].values[mask.values], e_plus.values[mask.values])
        rows.append(dict(target=label, year=yr, n=int(mask.sum()),
                         r2_demo_only=r_demo, r2_demo_plus_prior=r_plus,
                         delta=r_plus - r_demo))

res = pd.DataFrame(rows)
summary = (res.groupby("target")[["r2_demo_only", "r2_demo_plus_prior", "delta"]]
           .mean().reset_index())
print("\nPer-year:")
print(res.to_string(index=False, float_format=lambda v: f"{v:.3f}"))
print("\nMean across years:")
print(summary.to_string(index=False, float_format=lambda v: f"{v:.3f}"))
print("""
Read: if adding prior proficiency lifts R^2 for ABSENTEEISM (a different metric),
prior achievement genuinely improves the peer group. If it only lifts proficiency
(predicting proficiency from prior proficiency), that's mostly autocorrelation.""")