#!/usr/bin/env python
"""Clustering/decile bake-off for NYC school peer groups & proclivity deciles.

Ten approaches, two families (cluster/peer-group + decile/need-index), scored on a
common predictive-validity battery against three outcomes (ELA/Math proficiency,
4-yr graduation rate, chronic absenteeism).

Unifying metric: leave-one-out "expected-outcome R^2". Each method yields an
EXPECTED outcome per school formed WITHOUT that school's own outcome, so cluster,
peer-set, and index methods are all scored on the same "how well does this
representation predict actual outcomes" scale.

Also: within-group homogeneity (clusters), temporal out-of-sample (fit year T ->
predict T+1), and residual ("punching above expectation") year-over-year stability.

Read-only on a cached CSV; writes results CSVs to scripts/analysis/data/.
"""
import warnings
import numpy as np
import pandas as pd
from sklearn.preprocessing import StandardScaler
from sklearn.neighbors import NearestNeighbors
from sklearn.cluster import KMeans, AgglomerativeClustering
from sklearn.mixture import GaussianMixture
from sklearn.decomposition import PCA
from sklearn.linear_model import LinearRegression
from sklearn.model_selection import KFold

warnings.filterwarnings("ignore")
RNG = 42
DATA = "scripts/analysis/data"

DEMOG = ["pct_black", "pct_hispanic", "pct_white", "pct_asian",
         "pct_ell", "pct_swd", "pct_econ_dis"]
GROUP_SIZE = 40  # target peer-group size, matches production K

# ----------------------------------------------------------------------------- data
def load_panel():
    df = pd.read_csv(f"{DATA}/school_year_panel.csv")
    df = df[df["include_in_default_comparisons"] == True].copy()
    # academic proficiency = mean of ELA & Math % proficient (0-100)
    df["prof"] = df[["ela_all_proficiency", "math_all_proficiency"]].mean(axis=1)
    df["grad"] = df["graduation_rate_4yr"]
    df["absent"] = df["chronic_absenteeism_rate"]
    # grade-band-median impute for econ_dis (only column with material missingness)
    for c in DEMOG:
        df[c] = df.groupby(["year", "grade_band"])[c].transform(
            lambda s: s.fillna(s.median())
        )
        df[c] = df[c].fillna(df.groupby("year")[c].transform("median"))
    df = df.dropna(subset=DEMOG)
    return df

# outcome -> (column, higher_is_better, eligible grade bands)
OUTCOMES = {
    "proficiency": ("prof", True, {"ES", "MS", "K8"}),
    "graduation": ("grad", True, {"HS"}),
    "absenteeism": ("absent", False, {"ES", "MS", "K8", "HS", "K12"}),
}

def r2(actual, expected):
    a = np.asarray(actual, float); e = np.asarray(expected, float)
    m = np.isfinite(a) & np.isfinite(e)
    if m.sum() < 10: return np.nan
    a, e = a[m], e[m]
    ss_res = np.sum((a - e) ** 2)
    ss_tot = np.sum((a - a.mean()) ** 2)
    return 1 - ss_res / ss_tot if ss_tot > 0 else np.nan

def zscaled(df, cols, scaler=None):
    if scaler is None:
        scaler = StandardScaler().fit(df[cols].values)
    return scaler.transform(df[cols].values), scaler

# ----------------------------------------------------------------------------- helpers
def loo_group_mean(labels, y):
    """Leave-one-out mean outcome within each label. y: array; labels: array."""
    y = np.asarray(y, float)
    out = np.full(len(y), np.nan)
    lab = pd.Series(labels)
    for g, idx in lab.groupby(lab).groups.items():
        idx = np.array(list(idx))
        yi = y[idx]
        valid = np.isfinite(yi)
        if valid.sum() <= 1:
            continue
        tot = np.nansum(yi); cnt = valid.sum()
        for j in idx:
            if np.isfinite(y[j]):
                out[j] = (tot - y[j]) / (cnt - 1)
            else:
                out[j] = tot / cnt
    return out

def ntile(series, n=10):
    """1..n decile labels by rank (ascending)."""
    r = series.rank(method="first")
    return np.minimum(n, (np.floor((r - 1) / len(series) * n) + 1)).astype(int)

# ----------------------------------------------------------------------------- methods
# Each method: fit on a per-band feature frame, return integer assignment labels
# (partition) OR a continuous index. Peer (k-NN) methods return None labels and
# are handled specially. The harness turns labels/index/peers into expected-outcome.

def k_for(n):
    return max(2, int(round(n / GROUP_SIZE)))

def assign_kmeans(X, n):
    return KMeans(n_clusters=k_for(n), random_state=RNG, n_init=10).fit_predict(X)

def assign_gmm(X, n):
    k = k_for(n)
    gm = GaussianMixture(n_components=k, random_state=RNG, covariance_type="diag",
                         reg_covar=1e-4, max_iter=200).fit(X)
    return gm.predict(X)

def assign_ward(X, n):
    return AgglomerativeClustering(n_clusters=k_for(n), linkage="ward").fit_predict(X)

# index methods (continuous "need" score; higher = more challenged)
PROCLIVITY3 = ["pct_econ_dis", "pct_ell", "pct_swd"]  # current production feature set

def index_meanpctrank(df):  # current proclivity: econ_dis, ell, swd (3 feat, no race)
    return df[PROCLIVITY3].rank(pct=True).mean(axis=1)

def index_meanpctrank7(df):  # same method, all 7 features (adds race/ethnicity)
    return df[DEMOG].rank(pct=True).mean(axis=1)

def index_pca3(df):  # PCA method on the 3 non-race proclivity features
    X, _ = zscaled(df, PROCLIVITY3)
    pc = PCA(n_components=1, random_state=RNG).fit(X)
    comp = pc.transform(X)[:, 0]
    if np.corrcoef(comp, df["pct_econ_dis"].values)[0, 1] < 0:
        comp = -comp
    return pd.Series(comp, index=df.index)

def index_pca(df):  # PCA on all 7 features (incl race/ethnicity)
    X, _ = zscaled(df, DEMOG)
    pc = PCA(n_components=1, random_state=RNG).fit(X)
    comp = pc.transform(X)[:, 0]
    # orient so higher = more challenged (positively correlated with econ_dis)
    if np.corrcoef(comp, df["pct_econ_dis"].values)[0, 1] < 0:
        comp = -comp
    return pd.Series(comp, index=df.index)

def index_eni(df):  # official-style weighted composite: econ need primary + ELL
    return 0.7 * df["pct_econ_dis"] + 0.3 * df["pct_ell"]

# ----------------------------------------------------------------------------- in-sample expected
def expected_partition(df, assign_fn, y):
    """Cluster within each grade band; LOO group-mean expected outcome."""
    exp = pd.Series(np.nan, index=df.index)
    lab_full = pd.Series("", index=df.index)
    for band, g in df.groupby("grade_band"):
        if len(g) < GROUP_SIZE:  # too few to cluster into ~40s -> single group
            labels = np.zeros(len(g), int)
        else:
            X, _ = zscaled(g, DEMOG)
            labels = assign_fn(X, len(g))
        lab = pd.Series([f"{band}:{l}" for l in labels], index=g.index)
        lab_full.loc[g.index] = lab
        exp.loc[g.index] = loo_group_mean(lab.values, g[y].values)
    return exp, lab_full

def expected_knn(df, y, feat_cols, feat_scaler_cols=None):
    """k-NN ego-group within band; mean outcome of K nearest peers (excl self)."""
    exp = pd.Series(np.nan, index=df.index)
    lab = pd.Series("", index=df.index)  # not a partition; left blank
    cols = feat_cols
    for band, g in df.groupby("grade_band"):
        gy = g[y].values
        have = np.isfinite(gy)
        if have.sum() <= GROUP_SIZE:
            continue
        Xall, scaler = zscaled(g, cols)
        cand = g[have]
        Xc = scaler.transform(cand[cols].values)
        k = min(GROUP_SIZE, len(cand) - 1)
        nn = NearestNeighbors(n_neighbors=k + 1).fit(Xc)
        _, ind = nn.kneighbors(Xall)
        cy = cand[y].values
        cand_pos = {pos: i for i, pos in enumerate(np.where(have)[0])}
        for row, gi in enumerate(g.index):
            self_cand = cand_pos.get(row, None)
            neigh = [p for p in ind[row] if p != self_cand][:k]
            exp.loc[gi] = np.nanmean(cy[neigh])
    return exp, lab

def expected_index(df, index_fn, y, n_dec=10):
    """Index -> decile bins; LOO decile-mean expected (the product 'decile' use)."""
    idx = index_fn(df)
    dec = pd.Series(ntile(idx, n_dec), index=df.index)
    exp = pd.Series(loo_group_mean(dec.values, df[y].values), index=df.index)
    return exp, dec, idx

def expected_regression(df, y, cols, kfolds=5):
    """Supervised: k-fold CV predictions of outcome from features. Expected = pred."""
    sub = df.dropna(subset=[y])
    if len(sub) < 50:
        return pd.Series(np.nan, index=df.index)
    X, _ = zscaled(sub, cols)
    yv = sub[y].values
    exp = pd.Series(np.nan, index=df.index)
    kf = KFold(n_splits=kfolds, shuffle=True, random_state=RNG)
    pred = np.full(len(sub), np.nan)
    for tr, te in kf.split(X):
        lr = LinearRegression().fit(X[tr], yv[tr])
        pred[te] = lr.predict(X[te])
    exp.loc[sub.index] = pred
    return exp

# ----------------------------------------------------------------------------- out-of-sample
def oos_partition(train, test, assign_fn, y):
    """Fit clusters per band on train; assign test to nearest centroid; train group-mean."""
    exp = pd.Series(np.nan, index=test.index)
    for band, gtr in train.groupby("grade_band"):
        gte = test[test["grade_band"] == band]
        if len(gte) == 0:
            continue
        if len(gtr) < GROUP_SIZE:
            exp.loc[gte.index] = gtr[y].mean()
            continue
        Xtr, scaler = zscaled(gtr, DEMOG)
        labels = assign_fn(Xtr, len(gtr))
        means = pd.Series(gtr[y].values).groupby(labels).mean()
        # centroids
        cents = np.vstack([Xtr[labels == l].mean(axis=0) for l in np.unique(labels)])
        clab = np.unique(labels)
        Xte = scaler.transform(gte[DEMOG].values)
        nn = NearestNeighbors(n_neighbors=1).fit(cents)
        _, ci = nn.kneighbors(Xte)
        exp.loc[gte.index] = [means.get(clab[c[0]], np.nan) for c in ci]
    return exp

def oos_knn(train, test, y, cols):
    exp = pd.Series(np.nan, index=test.index)
    for band, gtr in train.groupby("grade_band"):
        gte = test[test["grade_band"] == band]
        cand = gtr.dropna(subset=[y])
        if len(gte) == 0 or len(cand) <= GROUP_SIZE:
            continue
        Xc, scaler = zscaled(cand, cols)
        Xte = scaler.transform(gte[cols].values)
        k = min(GROUP_SIZE, len(cand))
        nn = NearestNeighbors(n_neighbors=k).fit(Xc)
        _, ind = nn.kneighbors(Xte)
        cy = cand[y].values
        exp.loc[gte.index] = [np.nanmean(cy[row]) for row in ind]
    return exp

def oos_index(train, test, index_fn, y, n_dec=10):
    itr = index_fn(train); ite = index_fn(test)
    # decile thresholds from train
    qs = np.quantile(itr, np.linspace(0, 1, n_dec + 1)[1:-1])
    dtr = np.digitize(itr, qs)
    dte = np.digitize(ite, qs)
    means = pd.Series(train[y].values).groupby(dtr).mean()
    return pd.Series([means.get(d, np.nan) for d in dte], index=test.index)

def oos_regression(train, test, y, cols):
    sub = train.dropna(subset=[y])
    if len(sub) < 50:
        return pd.Series(np.nan, index=test.index)
    Xtr, scaler = zscaled(sub, cols)
    lr = LinearRegression().fit(Xtr, sub[y].values)
    Xte = scaler.transform(test[cols].values)
    return pd.Series(lr.predict(Xte), index=test.index)

# ----------------------------------------------------------------------------- registry
# Methods 5 & 10b use prior-year outcome -> feature set includes 'prior'.
PARTITION_METHODS = {
    "02_kmeans": assign_kmeans,
    "03_gmm": assign_gmm,
    "04_ward": assign_ward,
}
INDEX_METHODS = {
    "07_meanpctrank": index_meanpctrank,
    "07b_meanpctrank7": index_meanpctrank7,
    "08_pca": index_pca,
    "08b_pca3": index_pca3,
    "09_eni": index_eni,
}

def build_prior(df, ycol):
    """Add prior-year (T-1) outcome per school as a feature column 'prior'."""
    years = sorted(df["year"].unique())
    prev = {years[i]: years[i - 1] for i in range(1, len(years))}
    key = df.set_index(["dbn", "year"])[ycol]
    def lookup(r):
        py = prev.get(r["year"])
        if py is None: return np.nan
        try: return key.loc[(r["dbn"], py)]
        except KeyError: return np.nan
    return df.apply(lookup, axis=1)