"""F2-conf-01: school-level programme effect with cluster-robust uncertainty.
Run from workspace root: ~/wp29/venv/bin/python out/analysis.py
Deterministic; no randomness used."""
import json, os, csv
import numpy as np, pandas as pd
import statsmodels.api as sm
from scipy import stats

D = "data/schools.csv"
d = pd.read_csv(D)
n_raw = len(d)
issues = []
ledger = [("0", "raw input", n_raw, n_raw, 0, "all rows read")]

# Audit-driven filters (none expected to remove rows; recorded for the ledger)
n0 = len(d); d = d.dropna(subset=["school_id", "treated", "score"])
ledger.append(("1", "drop rows with missing school_id/treated/score", n0, len(d), n0 - len(d), "cannot be used in OLS"))
n0 = len(d); d = d.drop_duplicates(subset=["student_id"])
ledger.append(("2", "drop duplicate student_id", n0, len(d), n0 - len(d), "student_id must be unique"))
if (d.groupby("school_id")["treated"].nunique() > 1).any():
    issues.append({"issue": "treated varies within school", "handling": "reported; not altered"})

y = d["score"].to_numpy(float); t = d["treated"].to_numpy(float)
X = sm.add_constant(t)
ols = sm.OLS(y, X).fit()
tau = float(ols.params[1]); se_iid = float(ols.bse[1])

# Cluster-robust (CR1, Stata-style small-sample correction), inference with t(G-1)
cl = sm.OLS(y, X).fit(cov_type="cluster", cov_kwds={"groups": d["school_id"].to_numpy(), "use_correction": True, "df_correction": True})
se_cl = float(cl.bse[1])
G = int(d["school_id"].nunique()); dfree = G - 1
tstat = tau / se_cl
p = float(2 * stats.t.sf(abs(tstat), dfree))
q = stats.t.ppf(0.975, dfree)
lo, hi = tau - q * se_cl, tau + q * se_cl

# Cross-check: two-sample t-test on school means (Welch & pooled, df=G-2)
sm_ = d.groupby("school_id").agg(m=("score", "mean"), tr=("treated", "first"))
a, b = sm_.loc[sm_.tr == 1, "m"], sm_.loc[sm_.tr == 0, "m"]
pooled = stats.ttest_ind(a, b, equal_var=True)
se_means = float(np.sqrt(a.var(ddof=1) / len(a) + b.var(ddof=1) / len(b)))

# ICC of residual: one-way ANOVA estimator on OLS residuals by school
r = pd.Series(ols.resid, index=d.index); gid = d["school_id"]
m = d.groupby("school_id").size(); n_bar = float(m.mean())
gm = r.groupby(gid).transform("mean")
ssw = float(((r - gm) ** 2).sum()); ssb = float((m * r.groupby(gid).mean() ** 2).sum())
msw = ssw / (len(d) - G); msb = ssb / (G - 1)
# note: residual df within schools unaffected since treated is constant within school
s2b = max((msb - msw) / n_bar, 0.0)
icc = s2b / (s2b + msw)
deff = 1 + (n_bar - 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": G, "n_treated_schools": int(sm_.tr.sum()), "n_analysed": int(len(d)),
    "feasibility": "complete", "feasibility_reason": None,
    "data_issues": issues, "n_input_rows": {"schools.csv": int(n_raw)},
    "random_seed": None,
    "cluster_df": dfree, "t_stat_cluster": float(tstat),
    "se_ratio_cluster_to_iid": se_cl / se_iid,
    "design_effect_from_icc": float(deff),
    "check_se_school_means": se_means,
    "check_p_school_means_pooled_t": float(pooled.pvalue),
    "p_value_iid": float(ols.pvalues[1]),
    "students_per_school": n_bar,
}
os.makedirs("out", exist_ok=True)
with open("out/results.json", "w") as f:
    json.dump(res, f, indent=2)
with open("out/exclusions.csv", "w", newline="") as f:
    w = csv.writer(f); w.writerow(["step", "rule", "n_before", "n_after", "n_removed", "rationale"]); w.writerows(ledger)
print(json.dumps(res, indent=2))