"""T2-simulation: finite-sample coverage of CR0, CR1 and wild-cluster-bootstrap-t intervals 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.
RNG: per cell numpy default_rng(seed from design.csv); draw order per cell:
  u (R x G) ~ N(0, icc), e (R x N) ~ N(0, 1-icc), then Rademacher weights (B x G per replication)
  drawn chunk-wise (chunks of 200 replications, in replication order).
"""
import json, os
import numpy as np, pandas as pd
from scipy import stats
import matplotlib, matplotlib.ticker
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"))
CHUNK = 200
Z = 1.959963984540054  # 0.975 normal quantile ("1.96")


def simulate_cell(row):
    G, icc, N, R, B, mu = int(row.G), float(row.icc), int(row.N), int(row.R), int(row.B), float(row.mu)
    n = np.array([int(x) for x in row.cluster_sizes.split("|")])
    assert len(n) == G and n.sum() == N
    gid = np.repeat(np.arange(G), n)
    var_exact = (icc * np.sum(n.astype(float) ** 2) + (1 - icc) * N) / N ** 2
    rng = np.random.default_rng(int(row.seed))
    u = rng.normal(0.0, np.sqrt(icc), size=(R, G))
    e = rng.normal(0.0, np.sqrt(1 - icc), size=(R, N))
    y = mu + u[:, gid] + e                                   # R x N
    ybar = y.mean(axis=1)                                    # R
    # cluster sums of residuals S_g = sum_i (y_ig - ybar)
    csum = np.zeros((R, G))
    for g in range(G):
        csum[:, g] = y[:, gid == g].sum(axis=1)
    S = csum - n[None, :] * ybar[:, None]                    # R x G
    v0 = (S ** 2).sum(axis=1) / N ** 2
    c = G / (G - 1)
    v1 = c * v0
    # naive iid baseline (ignores clustering): s^2/N, normal critical value
    s2 = y.var(axis=1, ddof=1)
    vnaive = s2 / N
    # wild cluster bootstrap-t (unrestricted, Rademacher)
    k = int(np.ceil(0.95 * B))  # order statistic: smallest x with ECDF(x) >= 0.95 (=380th of 399)
    qwcb = np.empty(R)
    min_vstar = np.inf
    for start in range(0, R, CHUNK):
        stop = min(start + CHUNK, R)
        W = rng.integers(0, 2, size=(stop - start, B, G)).astype(np.float64) * 2 - 1
        A = W * S[start:stop, None, :]                       # w_g * S_g : bootstrap cluster sums of (y* - ybar)
        D = A.sum(axis=2) / N                                # ybar* - ybar
        Sstar = A - n[None, None, :] * D[:, :, None]         # cluster sums of bootstrap residuals
        vstar = c * (Sstar ** 2).sum(axis=2) / N ** 2
        min_vstar = min(min_vstar, float(vstar.min()))
        tabs = np.abs(D) / np.sqrt(vstar)
        qwcb[start:stop] = np.sort(tabs, axis=1)[:, k - 1]
    tcrit = stats.t.ppf(0.975, G - 1)
    hw = {"cr0": Z * np.sqrt(v0), "cr1": tcrit * np.sqrt(v1), "wcb": qwcb * np.sqrt(v1),
          "naive": Z * np.sqrt(vnaive)}
    err = np.abs(ybar - mu)
    res = {}
    for est, h in hw.items():
        p = float(np.mean(err <= h))
        res[est] = dict(coverage=p, mc_se=float(np.sqrt(p * (1 - p) / R)), mean_hw=float(h.mean()))
    res["cr0"]["rel_bias"] = float(v0.mean() / var_exact)
    res["cr1"]["rel_bias"] = float(v1.mean() / var_exact)
    res["cr0"]["rel_bias_mc_se"] = float(v0.std(ddof=1) / np.sqrt(R) / var_exact)
    res["cr1"]["rel_bias_mc_se"] = float(v1.std(ddof=1) / np.sqrt(R) / var_exact)
    # analytic E[V_CR0]: S_g = T_g - n_g*ybar, Var(T_g)=s2g, Cov(T_g,ybar)=s2g/N
    s2g = icc * n.astype(float) ** 2 + (1 - icc) * n
    exact_rb0 = float(np.sum(s2g - 2 * n * s2g / N + n.astype(float) ** 2 * var_exact) / N ** 2 / var_exact)
    res["_diag"] = dict(var_exact=var_exact, emp_var_ybar=float(ybar.var(ddof=1)),
                        min_vstar=min_vstar, exact_rel_bias_cr0=exact_rb0, exact_rel_bias_cr1=exact_rb0 * c, mean_q_wcb=float(qwcb.mean()), tcrit=float(tcrit),
                        n_distinct_q=int(len(np.unique(np.round(qwcb, 10)))))
    return res


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


rows, cov, se, rb, rbse = [], {}, {}, {}, {}
base_cov, base_se, diag = {}, {}, {}
for _, r in design.iterrows():
    res = simulate_cell(r)
    for est in ["cr0", "cr1", "wcb"]:
        kk = key(est, r.G, r.icc)
        cov[kk] = res[est]["coverage"]; se[kk] = res[est]["mc_se"]
        if est != "wcb":
            rb[kk] = res[est]["rel_bias"]; rbse[kk] = res[est]["rel_bias_mc_se"]
        rows.append(dict(estimator=est, cell_id=r.cell_id, G=int(r.G), icc=float(r.icc),
                         coverage=res[est]["coverage"], mc_se_coverage=res[est]["mc_se"],
                         rel_bias_var=res[est].get("rel_bias", np.nan),
                         mean_ci_halfwidth=res[est]["mean_hw"], R=int(r.R)))
    kb = key("naive", r.G, r.icc)
    base_cov[kb] = res["naive"]["coverage"]; base_se[kb] = res["naive"]["mc_se"]
    diag[r.cell_id] = res["_diag"]
    print(r.cell_id, {e: round(res[e]["coverage"], 4) for e in ["naive", "cr0", "cr1", "wcb"]})

table = pd.DataFrame(rows)
table.to_csv(os.path.join(OUT, "table1.csv"), index=False, float_format="%.6f")

failure = sorted([k for k, p in cov.items() if p < 0.90])
nominal = sorted([k for k, p in cov.items() if abs(p - 0.95) <= 2 * se[k]])
smallG = {est: float(np.mean([cov[key(est, 5, i)] for i in (0.0, 0.2, 0.5)])) for est in ["cr0", "cr1", "wcb"]}
best = max(smallG, key=smallG.get)
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": cov, "mc_se_coverage": se, "rel_bias_var": rb,
    "failure_cells": failure, "n_failure_cells": len(failure),
    "nominal_cells": nominal,
    "best_procedure_small_G": best,
    "feasibility": "complete", "feasibility_reason": None,
    "data_issues": [
        {"issue": "ICC=0 cells: cluster effect u_g has variance 0 (degenerate normal).",
         "handling": "Drawn with numpy normal(scale=0), which returns exact zeros; DGP is then i.i.d. as intended."},
        {"issue": "G=5 with Rademacher weights: only 2^5=32 distinct weight vectors (16 distinct |t*| values), so B=399 draws resample a tiny discrete set and the WCB critical value is coarse.",
         "handling": "Kept as specified (task forbids changing B or weights); documented as a limitation."},
    ],
    "n_input_rows": {"design.csv": int(len(design))},
    "n_analysed": int(len(design) * R),
    "random_seed": int(design.loc[design.cell_id == "C1", "seed"].iloc[0]),
    "rel_bias_var_mc_se": rbse,
    "baseline_naive_iid_coverage": base_cov,
    "baseline_naive_iid_mc_se": base_se,
    "avg_coverage_small_G": smallG,
    "wcb_quantile_rule": "order statistic ceil(0.95*B) of sorted |t*| (380th of 399)",
    "diagnostics": diag,
}
with open(os.path.join(OUT, "results.json"), "w") as f:
    json.dump(results, f, indent=1, sort_keys=False)

# figure
fig, axes = plt.subplots(1, 3, figsize=(13, 4.2), sharey=True)
style = {"cr0": ("o-", "CR0 + N(0,1)"), "cr1": ("s-", "CR1 + t(G-1)"), "wcb": ("^-", "WCB-t (Rademacher)")}
Gs = [5, 10, 30]
for ax, icc in zip(axes, [0.0, 0.2, 0.5]):
    for est, (m, lab) in style.items():
        p = np.array([cov[key(est, g, icc)] for g in Gs]); s = np.array([se[key(est, g, icc)] for g in Gs])
        ax.errorbar(Gs, p, yerr=1.96 * s, fmt=m, capsize=3, label=lab)
    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.set_xticks(Gs); ax.set_xticklabels(Gs); ax.xaxis.set_minor_formatter(matplotlib.ticker.NullFormatter())
    ax.set_title(f"ICC = {icc}"); ax.set_xlabel("number of clusters G")
axes[0].set_ylabel("coverage (bars: ±1.96 MC SE)")
axes[0].legend(fontsize=8, loc="lower right")
fig.suptitle(f"Coverage of nominal 95% intervals for the mean (R={R}, B={B})")
fig.tight_layout()
fig.savefig(os.path.join(OUT, "figure1.png"), dpi=130)
print("failure cells:", failure)
print("best procedure at G=5:", best, smallG)