"""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, hashlib, os
import numpy as np, pandas as pd
import statsmodels.formula.api as smf
from scipy import stats

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

def step(name, rule, mask, why):
    global d
    nb = len(d); d = d[mask].copy(); excl.append([name, rule, nb, len(d), nb - len(d), why])

# Audit-driven filters (each logged even if it removes nothing)
step("1", "drop rows with missing school_id/treated/score", d[["school_id", "treated", "score"]].notna().all(axis=1), "required for estimation")
step("2", "drop duplicate student_id (keep first)", ~d.student_id.duplicated(), "one row per student")
nonconst = d.groupby("school_id").treated.transform("nunique") > 1
if nonconst.any():
    issues.append({"issue": "treated not constant within some schools", "handling": "schools dropped"})
step("3", "drop schools with non-constant treated", ~nonconst, "design: assignment at school level")
pd.DataFrame(excl, columns=["step", "rule", "n_before", "n_after", "n_removed", "rationale"]).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; CR1 cluster SE is valid regardless"})

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

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

# Main: CR1 cluster-robust SE by school, t reference with G-1 df
cl = smf.ols("score ~ treated", data=d).fit(cov_type="cluster", cov_kwds={"groups": d.school_id.values, "use_correction": True, "df_correction": True})
se_cl = float(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

# Diagnostic: two-sample t-test on school means (equivalent aggregated analysis)
sm = d.groupby("school_id").agg(score=("score", "mean"), treated=("treated", "first"))
a, b = sm.score[sm.treated == 1], sm.score[sm.treated == 0]
tt = 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 (after removing treatment-arm means): one-way ANOVA estimator
d["resid"] = ols.resid
m = sizes.mean()  # balanced
grp = d.groupby("school_id").resid
msb = m * ((grp.mean() - d.resid.mean()) ** 2).sum() / (G - 2)  # 2 df used by arm means
msw = ((d.resid - grp.transform("mean")) ** 2).sum() / (len(d) - G)
s2b = max((msb - msw) / m, 0.0)
icc = float(s2b / (s2b + msw))
deff = 1 + (m - 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": icc,
    "n_schools": int(G), "n_treated_schools": G1, "n_analysed": int(len(d)),
    "feasibility": "complete", "feasibility_reason": None,
    "data_issues": issues,
    "n_input_rows": {"schools.csv": int(n_raw)},
    "random_seed": None,
    "diagnostics": {
        "input_sha256": hashlib.sha256(open(IN, "rb").read()).hexdigest(),
        "cluster_se_type": "CR1 (Stata-style small-sample correction), t reference with G-1 df",
        "t_df": int(df_t), "t_stat_cluster": float(tstat),
        "p_value_iid": float(ols.pvalues["treated"]),
        "se_ratio_cluster_over_iid": se_cl / se_iid,
        "students_per_school": float(m),
        "design_effect_moulton": float(deff), "sqrt_design_effect": float(np.sqrt(deff)),
        "school_means_ttest_t": float(tt.statistic), "school_means_ttest_p_df38": float(tt.pvalue),
        "school_means_se": se_means,
        "msb": float(msb), "msw": float(msw), "sigma2_between": float(s2b),
    },
}
with open("out/results.json", "w") as f:
    json.dump(res, f, indent=2)
print(json.dumps({k: res[k] for k in ["tau_hat", "se_iid", "se_cluster", "p_value", "tau_ci_low", "tau_ci_high", "effect_significant_5pct", "icc_hat"]}, indent=2))