"""T2-simulation: finite-sample coverage of CR0, CR1 and wild-cluster-bootstrap-t CIs for a mean.

Run from the workspace root:  ~/wp29/venv/bin/python out/analysis.py
Reads only data/design.csv; writes out/results.json, out/table1.csv, out/figure1.png.

RNG: for each cell, rng = numpy.random.default_rng(cell seed). Per replication the draw order is
  u = rng.standard_normal(G); e = rng.standard_normal(N); W = rng.integers(0, 2, (B, G))*2-1.
Baseline (P2): naive i.i.d. interval ybar +- 1.96*s/sqrt(N), ignoring clustering (baseline_* keys).
"""
import json
import 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__)))
DATA = os.path.join(ROOT, "data", "design.csv")
OUT = os.path.join(ROOT, "out")

design = pd.read_csv(DATA)
data_issues = []


def key(est, G, icc):
    return f"{est}_G{int(G)}_icc{icc:.1f}"


def simulate_cell(row):
    G, icc, N, R, B = int(row.G), float(row.icc), int(row.N), int(row.R), int(row.B)
    mu = float(row.mu)
    n = np.array([int(x) for x in str(row.cluster_sizes).split("|")])
    assert len(n) == G and n.sum() == N, f"inconsistent design row {row.cell_id}"
    gid = np.repeat(np.arange(G), n)
    var_exact = (icc * np.sum(n ** 2) + (1 - icc) * N) / N ** 2
    z = 1.96
    tcrit = stats.t.ppf(0.975, G - 1)
    cfac = G / (G - 1)
    qidx = int(np.ceil((B + 1) * 0.95)) - 1  # 0-based index of the 380th order statistic for B=399
    rng = np.random.default_rng(int(row.seed))
    rec = {k: np.empty(R) for k in ["ybar", "v0", "v1", "q", "s2", "qalt"]}
    n_zero_vstar = 0
    for r in range(R):
        u = rng.standard_normal(G) * np.sqrt(icc)
        e = rng.standard_normal(N) * np.sqrt(1 - icc)
        W = rng.integers(0, 2, size=(B, G)) * 2 - 1
        y = mu + u[gid] + e
        ybar = y.mean()
        ehat = y - ybar
        S = np.bincount(gid, weights=ehat, minlength=G)  # cluster residual sums
        v0 = np.sum(S ** 2) / N ** 2
        v1 = cfac * v0
        # wild cluster bootstrap: y* = ybar + w_g ehat_ig
        d = (W @ S) / N                               # ybar* - ybar, shape (B,)
        Sstar = W * S[None, :] - d[:, None] * n[None, :]  # cluster sums of bootstrap residuals
        vstar = cfac * np.sum(Sstar ** 2, axis=1) / N ** 2
        n_zero_vstar += int(np.sum(vstar <= 0))
        tabs = np.abs(d) / np.sqrt(vstar)
        q = np.sort(tabs)[qidx]
        rec["qalt"][r] = np.quantile(tabs, 0.95)  # sensitivity: numpy linear-interpolation quantile
        rec["ybar"][r], rec["v0"][r], rec["v1"][r], rec["q"][r] = ybar, v0, v1, q
        rec["s2"][r] = ehat.var(ddof=1)
    ybar, v0, v1, q = rec["ybar"], rec["v0"], rec["v1"], rec["q"]
    hw = {"cr0": z * np.sqrt(v0), "cr1": tcrit * np.sqrt(v1), "wcb": q * np.sqrt(v1),
          "iid": z * np.sqrt(rec["s2"] / N), "wcb_alt": rec["qalt"] * np.sqrt(v1)}
    out = {}
    for est, h in hw.items():
        cov = float(np.mean(np.abs(ybar - mu) <= h))
        out[est] = {"coverage": cov, "mc_se": float(np.sqrt(cov * (1 - cov) / R)),
                    "hw": float(np.mean(h))}
    out["cr0"]["rel_bias"] = float(np.mean(v0) / var_exact)
    out["cr1"]["rel_bias"] = float(np.mean(v1) / var_exact)
    out["cr0"]["rel_bias_se"] = float(np.std(v0, ddof=1) / np.sqrt(R) / var_exact)
    out["cr1"]["rel_bias_se"] = float(np.std(v1, ddof=1) / np.sqrt(R) / var_exact)
    out["meta"] = {"var_exact": float(var_exact), "mc_var_ybar": float(np.var(ybar, ddof=1)),
                   "n_zero_vstar": n_zero_vstar, "n_distinct_q": int(len(np.unique(np.round(q, 12))))}
    return out


ests = ["cr0", "cr1", "wcb"]
rows, coverage, mcse, relbias = [], {}, {}, {}
base_cov, base_se, meta, rb_se, wcb_alt = {}, {}, {}, {}, {}
for _, row in design.iterrows():
    res = simulate_cell(row)
    G, icc = int(row.G), float(row.icc)
    for est in ests:
        k = key(est, G, icc)
        coverage[k] = res[est]["coverage"]
        mcse[k] = res[est]["mc_se"]
        rb = res[est].get("rel_bias")
        if rb is not None:
            relbias[k] = rb
            rb_se[k] = res[est]["rel_bias_se"]
        rows.append({"estimator": est, "cell_id": row.cell_id, "G": G, "icc": icc,
                     "coverage": res[est]["coverage"], "mc_se_coverage": res[est]["mc_se"],
                     "rel_bias_var": rb, "mean_ci_halfwidth": res[est]["hw"], "R": int(row.R)})
    bk = f"G{G}_icc{icc:.1f}"
    base_cov[bk] = res["iid"]["coverage"]
    base_se[bk] = res["iid"]["mc_se"]
    wcb_alt[key("wcb", G, icc)] = res["wcb_alt"]["coverage"]
    meta[row.cell_id] = res["meta"]
    if res["meta"]["n_zero_vstar"] > 0:
        data_issues.append({"issue": f"{row.cell_id}: {res['meta']['n_zero_vstar']} bootstrap draws with zero V*",
                            "handling": "t* set to inf/nan by numpy; kept as is"})

table = pd.DataFrame(rows)
table["est_order"] = table.estimator.map({e: i for i, e in enumerate(ests)})
table = table.sort_values(["est_order", "cell_id"]).drop(columns="est_order")
table.to_csv(os.path.join(OUT, "table1.csv"), index=False, float_format="%.6f")

failure = sorted([k for k, v in coverage.items() if v < 0.90])
nominal = sorted([k for k, v in coverage.items() if abs(v - 0.95) <= 2 * mcse[k]])
smallG = {e: float(np.mean([coverage[key(e, 5, i)] for i in [0.0, 0.2, 0.5]])) for e in ests}
best = max(smallG, key=smallG.get)

if (design.G == 5).any():
    data_issues.append({"issue": "G=5 cells: Rademacher weights give only 2^5=32 distinct weight vectors (16 distinct |t*| values), "
                                 "so the WCB critical value takes few distinct values and B=399 draws mostly repeat them",
                        "handling": "procedure kept exactly as specified (no Webb weights / enumeration); reported as a limitation"})
data_issues.append({"issue": "ICC=0.0 cells: cluster effect variance is zero (degenerate N(0,0))",
                    "handling": "u_g drawn as sqrt(0)*N(0,1)=0 so the RNG stream structure is the same in every cell"})

R = int(design.R.iloc[0]); B = int(design.B.iloc[0])
results = {
    "task_id": "T2-simulation",
    "R": R, "B": B, "n_cells": int(len(design)),
    "coverage": coverage, "mc_se_coverage": mcse, "rel_bias_var": relbias,
    "mc_se_rel_bias_var": rb_se,
    "failure_cells": failure, "n_failure_cells": len(failure), "nominal_cells": nominal,
    "best_procedure_small_G": best,
    "mean_coverage_G5": smallG,
    "feasibility": "complete", "feasibility_reason": None,
    "data_issues": data_issues,
    "n_input_rows": {"design.csv": int(len(design))},
    "n_analysed": int(design.R.sum()),
    "random_seed": int(design.loc[design.cell_id == "C1", "seed"].iloc[0]),
    "baseline_iid_coverage": base_cov, "baseline_iid_mc_se": base_se,
    "sensitivity_wcb_coverage_numpy_quantile": wcb_alt,
    "wcb_quantile_rule": "380th order statistic of B=399 |t*| values (ceil((B+1)*0.95))",
    "cell_diagnostics": meta,
}
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(design.G.unique())
style = {"cr0": ("o-", "CR0 (z)"), "cr1": ("s-", "CR1 (t, G-1)"), "wcb": ("^-", "WCB-t (Rademacher)")}
for ax, icc in zip(axes, sorted(design.icc.unique())):
    for e in ests:
        c = np.array([coverage[key(e, g, icc)] for g in Gs])
        s = np.array([mcse[key(e, g, icc)] for g in Gs])
        ax.errorbar(Gs, c, yerr=1.96 * s, fmt=style[e][0], label=style[e][1], capsize=3)
    ax.axhline(0.95, color="k", ls="--", lw=1, label="nominal 0.95")
    ax.axhline(0.90, color="r", ls=":", lw=1, label="failure line 0.90")
    ax.set_xscale("log"); ax.minorticks_off(); ax.set_xticks(Gs); ax.set_xticklabels([str(g) for g in Gs])
    ax.set_title(f"ICC = {icc:.1f}"); ax.set_xlabel("number of clusters G")
axes[0].set_ylabel("coverage (±1.96 MC SE)")
axes[0].legend(fontsize=8, loc="lower right")
fig.suptitle(f"Coverage of nominal 95% cluster-robust CIs for the mean (R={R}, B={B})")
fig.tight_layout()
fig.savefig(os.path.join(OUT, "figure1.png"), dpi=130)

print(table.to_string(index=False))
print("failure cells:", failure)
print("baseline iid coverage:", base_cov)
print("WCB coverage, numpy-quantile sensitivity:", wcb_alt)
print("best procedure at G=5:", best, smallG)