"""F2-conf-01: school-level programme effect with cluster-robust uncertainty.
Run from workspace root: ~/wp29/venv/bin/python out/analysis.py
Reads only data/schools.csv; writes out/results.json and out/exclusions.csv. No randomness."""
import json, os
import numpy as np, pandas as pd
import statsmodels.formula.api as smf
from scipy import stats

os.makedirs("out", exist_ok=True)
raw = pd.read_csv("data/schools.csv")
n_raw = len(raw)
issues, excl = [], []
d = raw.copy()

def step(name, rule, mask, why):
    global d
    nb = len(d); d = d[mask(d)].copy(); na = len(d)
    excl.append(dict(step=name, rule=rule, n_before=nb, n_after=na, n_removed=nb - na, rationale=why))

step("1", "drop rows with missing school_id/treated/score", lambda x: x[["school_id", "treated", "score"]].notna().all(axis=1), "cannot be used in OLS")
step("2", "drop exact duplicate rows", lambda x: ~x.duplicated(), "duplicate records")
step("3", "drop duplicated student_id (keep first)", lambda x: ~x.student_id.duplicated(), "student_id must be unique")
step("4", "treated must be 0/1", lambda x: x.treated.isin([0, 1]), "binary treatment")
# within-school treatment consistency
bad = d.groupby("school_id").treated.nunique()
bad = bad[bad > 1].index
step("5", "drop schools with non-constant treated", lambda x: ~x.school_id.isin(bad), "treatment must be constant within school")
for e in excl:
    if e["n_removed"] > 0:
        issues.append({"issue": f"{e['rule']}: {e['n_removed']} rows", "handling": "excluded"})
pd.DataFrame(excl).to_csv("out/exclusions.csv", index=False)

sizes = d.groupby("school_id").size()
if sizes.nunique() != 1:
    issues.append({"issue": "unequal school sizes", "handling": "kept; student-weighted OLS as defined"})

G = d.school_id.nunique()
G1 = int(d.groupby("school_id").treated.first().sum())
N = len(d)

# Baseline: conventional iid OLS
m_iid = smf.ols("score ~ treated", data=d).fit()
tau = float(m_iid.params["treated"]); se_iid = float(m_iid.bse["treated"])

# Main: cluster-robust (CR1, Liang-Zeger with small-sample factor G/(G-1)*(N-1)/(N-K)), t with G-1 df
m_cl = smf.ols("score ~ treated", data=d).fit(cov_type="cluster", cov_kwds={"groups": d.school_id}, use_t=True)
se_cl = float(m_cl.bse["treated"])
df_t = G - 1
tstat = tau / se_cl
p = float(2 * stats.t.sf(abs(tstat), df_t))
q = stats.t.ppf(0.975, df_t)
lo, hi = tau - q * se_cl, tau + q * se_cl

# Sensitivity: school-means Welch/pooled t-test (one obs per school)
sm = d.groupby("school_id").agg(y=("score", "mean"), t=("treated", "first"))
tt = stats.ttest_ind(sm.y[sm.t == 1], sm.y[sm.t == 0], equal_var=True)
se_means = float((sm.y[sm.t == 1].mean() - sm.y[sm.t == 0].mean()) / tt.statistic)

# ICC of residual (one-way ANOVA estimator, school as random factor)
d["res"] = m_iid.resid
k = N / G if sizes.nunique() == 1 else (N - (sizes**2).sum() / N) / (G - 1)
gm = d.groupby("school_id").res.agg(["mean", "size"])
ssb = float((gm["size"] * (gm["mean"] - d.res.mean())**2).sum())
ssw = float(((d.res - d.school_id.map(gm["mean"]))**2).sum())
msb = ssb / (G - 2)   # one df lost to treatment effect (between-school)
msw = ssw / (N - G)
s2b = max((msb - msw) / k, 0.0)
icc = s2b / (s2b + msw)
deff = 1 + (k - 1) * icc

res = {
    "task_id": "F2-conf-01",
    "tau_hat": tau,
    "se_iid": se_iid,
    "se_cluster": se_cl,
    "p_value": p,
    "tau_ci_low": float(lo),
    "tau_ci_high": float(hi),
    "effect_significant_5pct": bool(p < 0.05),
    "icc_hat": float(icc),
    "n_schools": int(G),
    "n_treated_schools": G1,
    "n_analysed": int(N),
    "feasibility": "complete",
    "feasibility_reason": None,
    "data_issues": issues,
    "n_input_rows": {"schools.csv": int(n_raw)},
    "random_seed": None,
    # supplementary
    "t_stat_cluster": float(tstat),
    "df_cluster": int(df_t),
    "p_value_iid": float(m_iid.pvalues["treated"]),
    "se_ratio_cluster_over_iid": se_cl / se_iid,
    "design_effect_from_icc": float(deff),
    "school_size": float(k),
    "sens_school_means_se": se_means,
    "sens_school_means_p": float(tt.pvalue),
    "mean_treated": float(d.score[d.treated == 1].mean()),
    "mean_control": float(d.score[d.treated == 0].mean()),
}
with open("out/results.json", "w") as f:
    json.dump(res, f, indent=2)
print(json.dumps(res, indent=2))