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

d = pd.read_csv("data/schools.csv")
n_rows = len(d)
issues = []
# ---- data checks ----
assert d[["school_id", "student_id", "treated", "score"]].notna().all().all(), "missing values"
assert not d.student_id.duplicated().any(), "duplicate student ids"
assert set(d.treated.unique()) <= {0, 1}
assert (d.groupby("school_id").treated.nunique() == 1).all(), "treatment varies within school"
sizes = d.groupby("school_id").size()
assert sizes.nunique() == 1, "unbalanced"

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

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

# ---- cross-check: two-sample t-test on school means (balanced design) ----
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=False)
se_means = float(np.sqrt(sm.y[sm.t == 1].var() / G1 + sm.y[sm.t == 0].var() / (G - G1)))

# ---- ICC of residual (one-way ANOVA estimator, balanced clusters) ----
d["resid"] = m_iid.resid
m = int(sizes.iloc[0])
grp = d.groupby("school_id").resid
msb = m * ((grp.mean() - d.resid.mean()) ** 2).sum() / (G - 1)
msw = ((d.resid - grp.transform("mean")) ** 2).sum() / (N - 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(N),
    "feasibility": "complete",
    "feasibility_reason": None,
    "data_issues": issues,
    "n_input_rows": {"schools.csv": int(n_rows)},
    "random_seed": None,
    "diagnostics": {
        "cluster_se_type": "CR1 (statsmodels cluster, G/(G-1)*(N-1)/(N-K) correction)",
        "inference_df": df,
        "t_stat_cluster": float(tstat),
        "t_stat_iid": float(tau / se_iid),
        "p_value_iid": float(m_iid.pvalues["treated"]),
        "school_means_welch_t": float(tt.statistic),
        "school_means_welch_p": float(tt.pvalue),
        "se_school_means": se_means,
        "students_per_school": m,
        "design_effect_1_plus_(m-1)icc": float(deff),
        "se_ratio_cluster_over_iid": se_cl / se_iid,
    },
}
os.makedirs("out", exist_ok=True)
with open("out/results.json", "w") as f:
    json.dump(res, f, indent=2)
print(json.dumps(res, indent=2))