"""F2-conf-01: school-level programme effect with cluster-robust inference.
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
import pandas as pd
import statsmodels.formula.api as smf
from scipy import stats

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
df = pd.read_csv(os.path.join(ROOT, "data", "schools.csv"))
n_raw = len(df)

issues = []
excl = []
# Step 1: rows with missing values in required columns
n0 = len(df)
df = df.dropna(subset=["school_id", "student_id", "treated", "score"])
excl.append(("1", "drop rows with missing school_id/student_id/treated/score", n0, len(df), n0 - len(df),
             "required for estimation"))
# Step 2: duplicate student ids
n0 = len(df)
df = df.drop_duplicates(subset=["student_id"])
excl.append(("2", "drop duplicate student_id", n0, len(df), n0 - len(df), "student_id must be unique"))
for s, r, a, b, k, why in excl:
    if k > 0:
        issues.append({"issue": f"{k} rows removed: {r}", "handling": "excluded"})
# Design checks
within = df.groupby("school_id")["treated"].nunique()
if (within > 1).any():
    issues.append({"issue": f"treated varies within {int((within>1).sum())} schools", "handling": "reported; not altered"})
sizes = df.groupby("school_id").size()
if sizes.nunique() != 1:
    issues.append({"issue": "unbalanced school sizes", "handling": "kept; CR1 SE handles unequal clusters"})

pd.DataFrame(excl, columns=["step", "rule", "n_before", "n_after", "n_removed", "rationale"]).to_csv(
    os.path.join(ROOT, "out", "exclusions.csv"), index=False)

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

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

# Main: cluster-robust (CR1, Stata-style small-sample correction) by school, t with G-1 df
cl = smf.ols("score ~ treated", data=df).fit(cov_type="cluster", cov_kwds={"groups": df["school_id"]},
                                               use_t=True)
se_cl = float(cl.bse["treated"])
dof = G - 1
tcrit = float(stats.t.ppf(0.975, dof))
tstat = tau / se_cl
p = float(2 * stats.t.sf(abs(tstat), dof))
lo, hi = tau - tcrit * se_cl, tau + tcrit * se_cl

# Cross-check: two-sample Welch t-test on school means (aggregated analysis)
sm = df.groupby("school_id").agg(m=("score", "mean"), t=("treated", "first"))
a, b = sm.m[sm.t == 1], sm.m[sm.t == 0]
agg = stats.ttest_ind(a, b, equal_var=False)
se_agg = float(np.sqrt(a.var(ddof=1) / len(a) + b.var(ddof=1) / len(b)))

# ICC of residual (after removing arm means), one-way ANOVA estimator
df = df.assign(res=df["score"] - df.groupby("treated")["score"].transform("mean"))
grp = df.groupby("school_id")["res"]
ng = grp.size(); mg = grp.mean()
ssw = float(((df["res"] - grp.transform("mean")) ** 2).sum())
ssb = float((ng * (mg - df["res"].mean()) ** 2).sum())
msw = ssw / (N - G)
msb = ssb / (G - 2)  # two df used by arm means
n0_eff = (N - (ng ** 2).sum() / N) / (G - 1)
sb2 = max((msb - msw) / n0_eff, 0.0)
icc = sb2 / (sb2 + msw)
deff = 1 + (float(ng.mean()) - 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": 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
    "baseline_p_value_iid": float(ols.pvalues["treated"]),
    "baseline_ci_iid_low": float(ols.conf_int().loc["treated", 0]),
    "baseline_ci_iid_high": float(ols.conf_int().loc["treated", 1]),
    "t_stat_cluster": float(tstat),
    "df_cluster": int(dof),
    "se_ratio_cluster_to_iid": se_cl / se_iid,
    "design_effect_from_icc": float(deff),
    "school_means_welch_se": se_agg,
    "school_means_welch_p": float(agg.pvalue),
    "mean_score_treated": float(df.loc[df.treated == 1, "score"].mean()),
    "mean_score_control": float(df.loc[df.treated == 0, "score"].mean()),
}
with open(os.path.join(ROOT, "out", "results.json"), "w") as f:
    json.dump(res, f, indent=2)
print(json.dumps(res, indent=2))