#!/usr/bin/env python
"""Diagnostics for the bake-off report: production-vs-reimpl agreement, PCA
loadings / race contribution, index correlations, year-to-year robustness."""
import numpy as np
import pandas as pd
from scipy.stats import spearmanr
import bakeoff as B
from sklearn.preprocessing import StandardScaler
from sklearn.decomposition import PCA

panel = B.load_panel()

print("=" * 78)
print("1. Does our reimplementation match the PRODUCTION artifacts?")
print("=" * 78)
cur = pd.read_csv(f"{B.DATA}/school_current.csv")
cur = cur[cur["include_in_default_comparisons"] == True]
# production proclivity decile is computed on the LATEST demographic snapshot per
# school; reproduce method 07 on each school's latest year and correlate.
latest = panel.sort_values("year").groupby("dbn").tail(1).copy()
latest["mpr3"] = B.index_meanpctrank(latest)
latest["dec_reimpl"] = B.ntile(latest["mpr3"], 10)
m = latest.merge(cur[["dbn", "proclivity_decile"]], on="dbn", how="inner").dropna(
    subset=["proclivity_decile"])
rho, _ = spearmanr(m["dec_reimpl"], m["proclivity_decile"])
exact = (m["dec_reimpl"] == m["proclivity_decile"]).mean()
within1 = (np.abs(m["dec_reimpl"] - m["proclivity_decile"]) <= 1).mean()
print(f"  proclivity decile: n={len(m)}  Spearman rho={rho:.3f}  "
      f"exact match={exact:.1%}  within +/-1={within1:.1%}")
print("  (high agreement => method 07 faithfully reproduces production proclivity)")

print("\n" + "=" * 78)
print("2. PCA(7) loadings on proficiency pool (latest yr) — how much is RACE?")
print("=" * 78)
pool = panel[(panel["year"] == "2024-25") & panel["grade_band"].isin({"ES", "MS", "K8"})].copy()
X = StandardScaler().fit_transform(pool[B.DEMOG].values)
pc = PCA(n_components=2, random_state=42).fit(X)
comp = pc.transform(X)[:, 0]
if np.corrcoef(comp, pool["pct_econ_dis"])[0, 1] < 0:
    sign = -1
else:
    sign = 1
load = pd.Series(sign * pc.components_[0], index=B.DEMOG).sort_values()
print(f"  PC1 explains {pc.explained_variance_ratio_[0]:.1%} of demographic variance")
print("  PC1 loadings (oriented so higher = more challenged):")
for k, v in load.items():
    print(f"    {k:16s} {v:+.3f}")
race = ["pct_black", "pct_hispanic", "pct_white", "pct_asian"]
race_share = (load[race] ** 2).sum() / (load ** 2).sum()
print(f"  share of PC1 loading-magnitude from race/ethnicity: {race_share:.1%}")

print("\n" + "=" * 78)
print("3. Index correlations (proficiency pool, latest yr)")
print("=" * 78)
idx = pd.DataFrame({
    "mpr3 (current)": B.index_meanpctrank(pool),
    "mpr7": B.index_meanpctrank7(pool),
    "pca7": B.index_pca(pool),
    "pca3": B.index_pca3(pool),
    "eni": B.index_eni(pool),
})
print(idx.corr(method="spearman").round(3).to_string())

print("\n" + "=" * 78)
print("4. Year-to-year robustness of in-sample R^2 (std across years)")
print("=" * 78)
ins = pd.read_csv(f"{B.DATA}/results_insample.csv")
rob = (ins.groupby(["outcome", "method"])["r2_loo"]
       .agg(["mean", "std", "count"]).reset_index())
for oc in B.OUTCOMES:
    s = rob[rob["outcome"] == oc].sort_values("mean", ascending=False)
    print(f"\n  --- {oc} ---")
    for _, r in s.iterrows():
        sd = f"{r['std']:.3f}" if pd.notna(r['std']) else "  -  "
        print(f"    {r['method']:18s} mean={r['mean']:.3f}  sd={sd}  (n_yrs={int(r['count'])})")

print("\n" + "=" * 78)
print("5. Race contribution to predictiveness (proficiency, in-sample LOO R^2)")
print("=" * 78)
lb = pd.read_csv(f"{B.DATA}/leaderboard.csv")
p = lb[lb["outcome"] == "proficiency"].set_index("method")["r2_loo"]
print(f"  meanpctrank 3 non-race features : {p.get('07_meanpctrank', np.nan):.3f}")
print(f"  meanpctrank 7 incl race (eq wt) : {p.get('07b_meanpctrank7', np.nan):.3f}  (race HURTS equal-weighted)")
print(f"  PCA 3 non-race features         : {p.get('08b_pca3', np.nan):.3f}")
print(f"  PCA 7 incl race (cov-weighted)  : {p.get('08_pca', np.nan):.3f}  (race helps ONLY w/ optimal weighting)")