"""T2-simulation: coverage of CR0, CR1 and wild-cluster-bootstrap-t intervals for a mean.
Run from workspace root: ~/wp29/venv/bin/python out/analysis.py
"""
import json, os
import numpy as np, pandas as pd
from scipy import stats
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

OUT = "out"
design = pd.read_csv("data/design.csv")
issues = []

def simulate_cell(row):
    sizes = np.array([int(s) for s in str(row.cluster_sizes).split("|")])
    G, icc, R, B, mu = int(row.G), float(row.icc), int(row.R), int(row.B), float(row.mu)
    N = int(sizes.sum())
    if N != int(row.N) or len(sizes) != G:
        issues.append({"issue": f"{row.cell_id}: cluster_sizes inconsistent with G/N",
                       "handling": "used cluster_sizes as given"})
    rng = np.random.default_rng(int(row.seed))
    # data: all R replications at once (order: u then e, then bootstrap weights)
    u = rng.normal(0.0, np.sqrt(icc), size=(R, G)) if icc > 0 else np.zeros((R, G))
    e = rng.normal(0.0, np.sqrt(1.0 - icc), size=(R, N))
    gid = np.repeat(np.arange(G), sizes)
    y = mu + u[:, gid] + e                                   # R x N
    ybar = y.mean(axis=1)                                    # R
    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
    # wild cluster bootstrap-t (Rademacher, unrestricted); statistics depend only on cluster sums:
    # ybar* - ybar = sum_g w_g S_g / N ; bootstrap residual cluster sums = w_g S_g - n_g (ybar*-ybar)
    W = rng.choice(np.array([-1.0, 1.0]), size=(R, B, G))
    wS = W * S[:, None, :]                                   # R x B x G
    d = wS.sum(axis=2) / N                                   # R x B
    Sstar = wS - sizes[None, None, :] * d[:, :, None]
    V1star = G / (G - 1) * (Sstar ** 2).sum(axis=2) / N ** 2
    with np.errstate(divide="ignore", invalid="ignore"):
        tstar = np.abs(d) / np.sqrt(V1star)
    n_degen = int(np.sum(V1star <= 0))
    tstar = np.where(V1star > 0, tstar, 0.0)
    # 0.95 empirical quantile of |t*|: order statistic ceil(0.95*(B+1)) = 380th of 399
    k = int(np.ceil(0.95 * (B + 1))) - 1
    q = np.sort(tstar, axis=1)[:, k]
    z = stats.norm.ppf(0.975)  # 1.95996...; task says 1.96
    z = 1.96
    tc = stats.t.ppf(0.975, G - 1)
    hw = {"cr0": z * np.sqrt(V0), "cr1": tc * np.sqrt(V1), "wcb": q * np.sqrt(V1)}
    var_exact = (icc * (sizes ** 2).sum() + (1 - icc) * N) / N ** 2
    out = []
    for est, h in hw.items():
        cov = float(np.mean(np.abs(ybar - mu) <= h))
        se = float(np.sqrt(cov * (1 - cov) / R))
        rb = None
        if est == "cr0": rb = float(V0.mean() / var_exact)
        if est == "cr1": rb = float(V1.mean() / var_exact)
        out.append(dict(estimator=est, cell_id=row.cell_id, G=G, icc=icc, coverage=cov,
                        mc_se_coverage=se, rel_bias_var=rb, mean_ci_halfwidth=float(h.mean()), R=R))
    extra = dict(cell_id=row.cell_id, n_degenerate_boot=n_degen,
                 var_exact=var_exact, emp_var_ybar=float(ybar.var(ddof=1)),
                 sd_V0=float(V0.std(ddof=1)), sd_V1=float(V1.std(ddof=1)),
                 mean_q=float(q.mean()), t_crit=float(tc),
                 mcse_relbias_cr0=float(V0.std(ddof=1) / np.sqrt(R) / var_exact),
                 mcse_relbias_cr1=float(V1.std(ddof=1) / np.sqrt(R) / var_exact),
                 emp_var_ratio=float(ybar.var(ddof=1) / var_exact))
    return out, extra

rows, extras = [], []
for _, r in design.iterrows():
    o, x = simulate_cell(r)
    rows += o; extras.append(x)
    print(r.cell_id, [(d["estimator"], round(d["coverage"], 4)) for d in o])

tab = pd.DataFrame(rows)
tab.to_csv(os.path.join(OUT, "table1.csv"), index=False)
pd.DataFrame(extras).to_csv(os.path.join(OUT, "cell_diagnostics.csv"), index=False)

key = lambda d: f"{d['estimator']}_G{d['G']}_icc{d['icc']:.1f}"
coverage = {key(d): d["coverage"] for d in rows}
mcse = {key(d): d["mc_se_coverage"] for d in rows}
relb = {key(d): d["rel_bias_var"] for d in rows if d["estimator"] != "wcb"}
fail = [key(d) for d in rows if d["coverage"] < 0.90]
nominal = [key(d) for d in rows if abs(d["coverage"] - 0.95) <= 2 * d["mc_se_coverage"]]
small = tab[tab.G == 5].groupby("estimator").coverage.mean()
best = str(small.idxmax())

for x in extras:
    if x["n_degenerate_boot"] > 0:
        issues.append({"issue": f"{x['cell_id']}: {x['n_degenerate_boot']} bootstrap draws with zero V*_CR1",
                       "handling": "|t*| set to 0 for those draws"})
issues.append({"issue": "With G=5 Rademacher weights there are only 2^5=32 distinct weight vectors "
               "(16 distinct |t*| values up to sign), so the WCB critical value from B=399 draws is "
               "a quantile of a highly discrete distribution",
               "handling": "kept the specified procedure (B=399 random draws, no enumeration); noted in report"})

res = {
    "task_id": "T2-simulation", "R": int(design.R.iloc[0]), "B": int(design.B.iloc[0]),
    "n_cells": int(len(design)), "coverage": coverage, "mc_se_coverage": mcse,
    "rel_bias_var": relb, "failure_cells": fail, "n_failure_cells": len(fail),
    "nominal_cells": nominal, "best_procedure_small_G": best,
    "feasibility": "complete", "feasibility_reason": None, "data_issues": issues,
    "n_input_rows": {"design.csv": int(len(design))},
    "n_analysed": int(design.R.sum()), "random_seed": int(design.seed.iloc[0]),
}
with open(os.path.join(OUT, "results.json"), "w") as f:
    json.dump(res, f, indent=2)

fig, axes = plt.subplots(1, 3, figsize=(13, 4.2), sharey=True)
for ax, icc in zip(axes, sorted(tab.icc.unique())):
    for est, m in zip(["cr0", "cr1", "wcb"], ["o", "s", "^"]):
        s = tab[(tab.icc == icc) & (tab.estimator == est)].sort_values("G")
        ax.errorbar(s.G, s.coverage, yerr=1.96 * s.mc_se_coverage, marker=m, capsize=3, label=est.upper())
    ax.axhline(0.95, color="k", ls="--", lw=1, label="nominal 0.95")
    ax.axhline(0.90, color="r", ls=":", lw=1, label="failure 0.90")
    ax.set_xscale("log"); ax.set_xticks([5, 10, 30]); ax.set_xticklabels(["5", "10", "30"]); ax.minorticks_off()
    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(loc="lower right", fontsize=8)
fig.suptitle("Coverage of nominal 95% cluster-robust intervals for the mean (R=2000 per cell)")
fig.tight_layout(); fig.savefig(os.path.join(OUT, "figure1.png"), dpi=130)
print(json.dumps({k: res[k] for k in ["failure_cells", "nominal_cells", "best_procedure_small_G"]}, indent=1))