"""T2-simulation: finite-sample coverage of CR0, CR1 and wild cluster bootstrap-t CIs for a mean.
Run from workspace root:  ~/wp29/venv/bin/python out/analysis.py
Reads only data/design.csv; writes out/results.json, out/table1.csv, out/figure1.png, out/exclusions.csv.
Random numbers: per cell numpy default_rng(seed from design.csv); draw order per cell:
  u = sqrt(icc)*N(0,1) of shape (R,G); e = sqrt(1-icc)*N(0,1) of shape (R,N);
  Rademacher weights W = 2*integers(0,2,(R,B,G))-1.
"""
import json, os
import numpy as np
import pandas as pd
from scipy import stats
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
OUT = os.path.join(ROOT, "out")
design = pd.read_csv(os.path.join(ROOT, "data", "design.csv"))
n_raw = len(design)

# ---------------- design audit ----------------
data_issues = []
cells = []
for _, row in design.iterrows():
    sizes = np.array([int(s) for s in str(row["cluster_sizes"]).split("|")])
    assert len(sizes) == row["G"], f"{row.cell_id}: len(cluster_sizes) != G"
    assert sizes.sum() == row["N"], f"{row.cell_id}: sum(cluster_sizes) != N"
    assert int(row["R"]) == 2000 and int(row["B"]) == 399
    cells.append(dict(cell_id=row["cell_id"], G=int(row["G"]), icc=float(row["icc"]), sizes=sizes,
                      N=int(row["N"]), R=int(row["R"]), B=int(row["B"]), seed=int(row["seed"]),
                      mu=float(row["mu"]), level=float(row["nominal_level"])))
assert design["cell_id"].is_unique and not design.isna().any().any()
if (design["G"] == 5).any():
    data_issues.append({
        "issue": "With G=5 Rademacher weights take only 2^5=32 distinct values (16 distinct |t*| up to sign symmetry, two of which (all weights equal) give t*=0); the B=399 bootstrap draws therefore repeat these few values and the WCB critical value is coarse.",
        "handling": "Kept as specified (B=399 i.i.d. Rademacher draws, not full enumeration); noted as a limitation in the report."})

# exclusion ledger: nothing is excluded
pd.DataFrame([{"step": 1, "rule": "keep all design rows passing consistency checks (G=len(sizes), N=sum(sizes), R=2000, B=399)",
               "n_before": n_raw, "n_after": n_raw, "n_removed": 0,
               "rationale": "all 9 cells consistent; no filtering"}]).to_csv(os.path.join(OUT, "exclusions.csv"), index=False)

def key(est, c):
    return f"{est}_G{c['G']}_icc{c['icc']:.1f}"

CHUNK = 250
rows, coverage, mcse, relbias, base_cov = [], {}, {}, {}, {}
relbias_se = {}
for c in cells:
    G, N, R, B, icc, mu, sizes = c["G"], c["N"], c["R"], c["B"], c["icc"], c["mu"], c["sizes"]
    rng = np.random.default_rng(c["seed"])
    gid = np.repeat(np.arange(G), sizes)
    u = np.sqrt(icc) * rng.standard_normal((R, G))
    e = np.sqrt(1 - icc) * rng.standard_normal((R, N))
    W = 2 * rng.integers(0, 2, size=(R, B, G)) - 1
    y = mu + u[:, gid] + e
    ybar = y.mean(axis=1)
    resid = y - ybar[:, None]
    S = np.zeros((R, G))
    for g in range(G):
        S[:, g] = resid[:, gid == g].sum(axis=1)       # cluster residual sums
    V0 = (S ** 2).sum(axis=1) / N ** 2
    V1 = G / (G - 1) * V0
    var_exact = (icc * (sizes ** 2).sum() + (1 - icc) * N) / N ** 2
    z = 1.96                                               # as specified in TASK.md
    tcrit = stats.t.ppf(0.5 + c["level"] / 2, G - 1)
    # wild cluster bootstrap-t: y* = ybar + w_g * resid; cluster sums of y*-ybar* are w_g S_g - n_g D
    q = np.empty(R)
    for s in range(0, R, CHUNK):
        Wc, Sc = W[s:s + CHUNK].astype(float), S[s:s + CHUNK]
        WS = Wc * Sc[:, None, :]                            # (r,B,G)
        D = WS.sum(axis=2) / N                              # ybar* - ybar
        Sstar = WS - sizes[None, None, :] * D[:, :, None]
        V1s = G / (G - 1) * (Sstar ** 2).sum(axis=2) / N ** 2
        with np.errstate(divide="ignore", invalid="ignore"):
            tstar = np.abs(D / np.sqrt(V1s))
        tstar = np.where(np.isfinite(tstar), tstar, 0.0)    # guard only (0/0 does not occur with continuous data)
        # 0.95 empirical quantile = inverted-CDF definition: the ceil(0.95*B)=380th order statistic
        q[s:s + CHUNK] = np.quantile(tstar, c["level"], axis=1, method="inverted_cdf")
    # naive i.i.d. baseline (ignores clustering)
    s2 = resid.var(axis=1, ddof=1)
    Viid = s2 / N
    hw = {"cr0": z * np.sqrt(V0), "cr1": tcrit * np.sqrt(V1), "wcb": q * np.sqrt(V1),
          "iid": z * np.sqrt(Viid)}
    err = np.abs(ybar - mu)
    for est in ["cr0", "cr1", "wcb", "iid"]:
        p = float(np.mean(err <= hw[est]))
        se = float(np.sqrt(p * (1 - p) / R))
        if est == "iid":
            base_cov[key("iid", c)] = p
            continue
        k = key(est, c)
        coverage[k], mcse[k] = p, se
        rb = None
        if est in ("cr0", "cr1"):
            rb = float(np.mean(V0 if est == "cr0" else V1) / var_exact)
            relbias[k] = rb
            Vk = V0 if est == "cr0" else V1
            relbias_se[k] = float(Vk.std(ddof=1) / np.sqrt(R) / var_exact)
        rows.append(dict(estimator=est, cell_id=c["cell_id"], G=G, icc=icc, coverage=p, mc_se_coverage=se,
                         rel_bias_var=rb, mean_ci_halfwidth=float(hw[est].mean()), R=R))
    print(f"{c['cell_id']} G={G} icc={icc}: " + ", ".join(f"{k2}={coverage[key(k2, c)]:.4f}" for k2 in ["cr0", "cr1", "wcb"])
          + f", iid={base_cov[key('iid', c)]:.4f}")

table = pd.DataFrame(rows)
table["_o"] = table["estimator"].map({"cr0": 0, "cr1": 1, "wcb": 2})
table = table.sort_values(["_o", "G", "icc"]).drop(columns="_o")
table.to_csv(os.path.join(OUT, "table1.csv"), index=False, float_format="%.6f")

# analytic E[V_CR0]/Var(ybar) (verification of the Monte Carlo bias estimates):
# S = A y with A = C - n 1'/N (C: G x N cluster indicator); E[S'S] = tr(A Sigma A'),
# Sigma = icc * C'C + (1-icc) I.
exact_rb = {}
for c in cells:
    G, N, icc, sizes = c["G"], c["N"], c["icc"], c["sizes"]
    C = (np.repeat(np.arange(G), sizes)[None, :] == np.arange(G)[:, None]).astype(float)
    A = C - np.outer(sizes, np.ones(N)) / N
    Sigma = icc * C.T @ C + (1 - icc) * np.eye(N)
    v_exact = (icc * (sizes ** 2).sum() + (1 - icc) * N) / N ** 2
    ev0 = np.trace(A @ Sigma @ A.T) / N ** 2 / v_exact
    exact_rb[key("cr0", c)] = float(ev0)
    exact_rb[key("cr1", c)] = float(G / (G - 1) * ev0)

failure = sorted(k for k, p in coverage.items() if p < 0.90)
nominal = sorted(k for k, p in coverage.items() if abs(p - 0.95) <= 2 * mcse[k])
small = {est: float(np.mean([coverage[key(est, c)] for c in cells if c["G"] == 5])) for est in ["cr0", "cr1", "wcb"]}
best_small = max(small, key=small.get)

results = {
    "task_id": "T2-simulation",
    "R": int(cells[0]["R"]), "B": int(cells[0]["B"]), "n_cells": len(cells),
    "coverage": coverage, "mc_se_coverage": mcse, "rel_bias_var": relbias,
    "failure_cells": failure, "n_failure_cells": len(failure), "nominal_cells": nominal,
    "best_procedure_small_G": best_small,
    "feasibility": "complete", "feasibility_reason": None,
    "data_issues": data_issues,
    "n_input_rows": {"design.csv": n_raw},
    "n_analysed": int(sum(c["R"] for c in cells)),
    "random_seed": int(cells[0]["seed"]),
    "baseline_coverage_iid": base_cov,
    "mean_coverage_small_G": small,
    "mc_se_rel_bias_var": relbias_se,
    "mean_ci_halfwidth": {f"{r['estimator']}_G{r['G']}_icc{r['icc']:.1f}": r["mean_ci_halfwidth"] for r in rows},
    "exact_rel_bias_var": exact_rb,
    "z_rel_bias_mc_vs_exact": {k: float((relbias[k] - exact_rb[k]) / relbias_se[k]) for k in relbias},
    "exact_var_ybar": {f"G{c['G']}_icc{c['icc']:.1f}": float((c['icc'] * (c['sizes'] ** 2).sum() + (1 - c['icc']) * c['N']) / c['N'] ** 2) for c in cells},
}
with open(os.path.join(OUT, "results.json"), "w") as f:
    json.dump(results, f, indent=2)

# figure
fig, axes = plt.subplots(1, 3, figsize=(13, 4.2), sharey=True)
Gs = sorted({c["G"] for c in cells})
for ax, icc in zip(axes, [0.0, 0.2, 0.5]):
    for est, mk in [("cr0", "o"), ("cr1", "s"), ("wcb", "^"), ("iid", "x")]:
        src = base_cov if est == "iid" else coverage
        p = np.array([src[f"{est}_G{G}_icc{icc:.1f}"] for G in Gs])
        se = np.sqrt(p * (1 - p) / 2000)
        ax.errorbar(Gs, p, yerr=2 * se, marker=mk, capsize=3,
                    label=est.upper() + (" (naive baseline)" if est == "iid" else ""),
                    linestyle=":" if est == "iid" else "-")
    ax.axhline(0.95, color="k", lw=1, label="nominal 0.95")
    ax.axhline(0.90, color="r", lw=1, ls="--", label="failure line 0.90")
    ax.set_xscale("log"); ax.set_xticks(Gs); ax.set_xticklabels(Gs); ax.minorticks_off()
    ax.set_title(f"ICC = {icc}"); ax.set_xlabel("number of clusters G")
axes[0].set_ylabel("coverage (±2 MC SE)")
axes[0].legend(fontsize=8, loc="lower right")
fig.suptitle("Coverage of nominal 95% intervals for the mean (R=2000 per cell)")
fig.tight_layout()
fig.savefig(os.path.join(OUT, "figure1.png"), dpi=130)
print("failure cells:", failure)
print("best procedure at G=5:", best_small)