"""F2-conf-01: school-level programme effect with design-valid (cluster) uncertainty.
Run from workspace root: ~/wp29/venv/bin/python out/analysis.py"""
import json, hashlib
import numpy as np, pandas as pd
import statsmodels.formula.api as smf
from scipy import stats

DATA = "data/schools.csv"
raw = pd.read_csv(DATA)
n_raw = len(raw)
issues, excl = [], []

# Audit-based filters (each logged even if it removes nothing)
d = raw.copy()
def step(name, rule, mask, why):
    global d
    nb = len(d); d = d[mask(d)].copy()
    excl.append(dict(step=name, rule=rule, n_before=nb, n_after=len(d), n_removed=nb - len(d), rationale=why))
step("1", "drop rows with missing school_id/treated/score", lambda x: x[["school_id", "treated", "score"]].notna().all(axis=1), "estimand requires all three")
step("2", "drop duplicated student_id", lambda x: ~x.student_id.duplicated(), "one row per student")
step("3", "drop schools with non-constant treated", lambda x: x.groupby("school_id").treated.transform("nunique") == 1, "treatment must be school-level")
pd.DataFrame(excl).to_csv("out/exclusions.csv", index=False)

sizes = d.groupby("school_id").size()
if sizes.nunique() != 1:
    issues.append({"issue": "unbalanced school sizes", "handling": "kept; cluster-robust SE does not require balance"})

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

# Naive baseline: conventional OLS 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_t=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

# Robustness: OLS on school means (exact in balanced design), t with G-2 df
m = d.groupby("school_id").agg(score=("score", "mean"), treated=("treated", "first")).reset_index()
agg = smf.ols("score ~ treated", data=m).fit()
se_agg = float(agg.bse["treated"]); p_agg = float(agg.pvalues["treated"])

# ICC of residual (ANOVA estimator, treatment effect removed)
d["resid"] = ols.resid
msw = ((d.resid - d.groupby("school_id").resid.transform("mean")) ** 2).sum() / (N - G)
mbar = N / G
gm = d.groupby("school_id").resid.mean()
msb = (sizes * (gm - d.resid.mean()) ** 2).sum() / (G - 2)  # 2 df used by intercept+treatment
s2b = max((msb - msw) / mbar, 0.0)
icc = float(s2b / (s2b + msw))
deff = 1 + (mbar - 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(N),
    "feasibility": "complete", "feasibility_reason": None,
    "data_issues": issues, "n_input_rows": {"schools.csv": int(n_raw)}, "random_seed": None,
    # supplementary
    "cluster_se_type": "CR1 (statsmodels), t reference with G-1 df",
    "t_df": int(df_t), "t_stat_cluster": float(tstat),
    "se_ratio_cluster_to_iid": se_cl / se_iid, "design_effect_from_icc": float(deff),
    "school_means_se": se_agg, "school_means_p_value": p_agg,
    "mean_control": float(d.loc[d.treated == 0, "score"].mean()),
    "mean_treated": float(d.loc[d.treated == 1, "score"].mean()),
    "students_per_school": float(mbar),
    "input_sha256": hashlib.sha256(open(DATA, "rb").read()).hexdigest(),
}
json.dump(res, open("out/results.json", "w"), indent=2)
print(json.dumps(res, indent=2))