"""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

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
d = pd.read_csv(os.path.join(ROOT, "data", "schools.csv"))
n_raw = len(d)
issues = []
excl = []

def step(df, name, mask, rule, rationale):
    nb = len(df); df2 = df[mask].copy()
    excl.append(dict(step=name, rule=rule, n_before=nb, n_after=len(df2),
                     n_removed=nb - len(df2), rationale=rationale))
    return df2

d = step(d, "1", d[["school_id", "student_id", "treated", "score"]].notna().all(axis=1),
         "drop rows with any missing value", "outcome/treatment/cluster id required")
d = step(d, "2", ~d.student_id.duplicated(), "drop duplicate student_id", "one row per student")
d = step(d, "3", d.treated.isin([0, 1]), "treated must be 0/1", "binary assignment")
tconst = d.groupby("school_id").treated.transform("nunique") == 1
d = step(d, "4", tconst, "treated constant within school", "school-level assignment")
for e in excl:
    if e["n_removed"] > 0:
        issues.append({"issue": f"{e['rule']}: {e['n_removed']} rows", "handling": "excluded"})
sizes = d.groupby("school_id").size()
if sizes.nunique() != 1:
    issues.append({"issue": "unequal school sizes", "handling": "kept; CR1 handles unbalanced clusters"})

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

# Main: cluster-robust CR1 SE (small-sample factor G/(G-1)*(N-1)/(N-K)), clustered by school,
# inference with t distribution on G-1 degrees of freedom
G = int(d.school_id.nunique())
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"])
dfree = G - 1
tstat = tau / se_cl
p = float(2 * stats.t.sf(abs(tstat), dfree))
crit = float(stats.t.ppf(0.975, dfree))
lo, hi = tau - crit * se_cl, tau + crit * se_cl

# Cross-check: Welch t-test on school means (equal school sizes => identical point estimate)
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=False)
se_means = float((a.mean() - b.mean()) / tt.statistic)

# ICC of residual (score minus arm mean): one-way ANOVA estimator, schools as groups
d["resid"] = ols.resid
m = float(sizes.mean())
grp = d.groupby("school_id").resid
msb = float((sizes * grp.mean() ** 2).sum() / (G - 2))   # 2 arm means removed
msw = float(((d.resid - grp.transform("mean")) ** 2).sum() / (len(d) - G))
s2b = max((msb - msw) / m, 0.0)
icc = 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": float(icc),
    "n_schools": G, "n_treated_schools": int(sm.treated.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,
    # supplementary keys
    "t_stat_cluster": float(tstat), "df_cluster": dfree, "t_crit_975": crit,
    "se_ratio_cluster_over_iid": se_cl / se_iid,
    "p_value_iid": float(ols.pvalues["treated"]),
    "se_school_means_welch": se_means, "p_school_means_welch": float(tt.pvalue),
    "students_per_school": m, "design_effect_from_icc": float(deff),
    "sqrt_design_effect": float(np.sqrt(deff)),
    "sigma2_between_resid": float(s2b), "sigma2_within_resid": msw,
    "mean_score_treated": float(d.score[d.treated == 1].mean()),
    "mean_score_control": float(d.score[d.treated == 0].mean()),
}
with open(os.path.join(ROOT, "out", "results.json"), "w") as f:
    json.dump(res, f, indent=2)
pd.DataFrame(excl).to_csv(os.path.join(ROOT, "out", "exclusions.csv"), index=False)
print(json.dumps(res, indent=2))