"""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.
RNG: per cell numpy.random.default_rng(seed from design.csv). Within each replication, draws are
taken in the order: Z_u (G std normals), Z_e (N std normals), Rademacher weights (B x G, via rng.integers).
"""
import json, os
import numpy as np, 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"))


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


def run_cell(row):
    G, icc, N, R, B = int(row.G), float(row.icc), int(row.N), int(row.R), int(row.B)
    mu, level = float(row.mu), float(row.nominal_level)
    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
    z = 1.96                                         # as specified in TASK.md
    tcrit = stats.t.ppf(0.5 + level / 2, G - 1)
    adj = G / (G - 1)
    rng = np.random.default_rng(int(row.seed))
    res = {k: np.empty(R) for k in ["ybar", "v0", "v1", "q", "s2"]}
    for r in range(R):
        u = np.sqrt(icc) * rng.standard_normal(G)
        e = np.sqrt(1 - icc) * rng.standard_normal(N)
        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 = adj * v0
        # wild cluster bootstrap-t (unrestricted, Rademacher), vectorised over B
        W = 2.0 * rng.integers(0, 2, size=(B, G)) - 1.0      # B x G
        d = (W @ S) / N                                      # ybar* - ybar, length B
        # bootstrap residual cluster sums: sum_i (w_g ehat_ig - d) = w_g S_g - n_g d
        Sstar = W * S[None, :] - d[:, None] * n[None, :]
        vstar = adj * np.sum(Sstar ** 2, axis=1) / N ** 2
        tstar = np.abs(d) / np.sqrt(vstar)
        q = np.quantile(tstar, 0.95)                         # numpy default (linear interpolation)
        res["ybar"][r], res["v0"][r], res["v1"][r], res["q"][r] = ybar, v0, v1, q
        res["s2"][r] = ehat @ ehat / (N - 1)
    hw = {"cr0": z * np.sqrt(res["v0"]),
          "cr1": tcrit * np.sqrt(res["v1"]),
          "wcb": res["q"] * np.sqrt(res["v1"]),
          "iid": z * np.sqrt(res["s2"] / N)}                # baseline: naive i.i.d. interval
    out = {}
    for est, h in hw.items():
        cov = float(np.mean(np.abs(res["ybar"] - mu) <= h))
        out[est] = dict(coverage=cov, mc_se=float(np.sqrt(cov * (1 - cov) / R)),
                        mean_hw=float(np.mean(h)))
    out["cr0"]["rel_bias"] = float(np.mean(res["v0"]) / var_exact)
    out["cr1"]["rel_bias"] = float(np.mean(res["v1"]) / var_exact)
    out["iid"]["rel_bias"] = float(np.mean(res["s2"] / N) / var_exact)
    # MC SE of the relative-bias estimate (sd of V-hat / sqrt(R) / Var exact)
    out["cr0"]["rel_bias_se"] = float(np.std(res["v0"], ddof=1) / np.sqrt(R) / var_exact)
    out["cr1"]["rel_bias_se"] = float(np.std(res["v1"], ddof=1) / np.sqrt(R) / var_exact)
    # analytic E[V_CR0] (for validation): Var(T_g - n_g*ybar) summed over g, / N^2
    VT = icc * n.astype(float) ** 2 + (1 - icc) * n
    ev0 = np.sum(VT * (1 - 2 * n / N) + n.astype(float) ** 2 * var_exact) / N ** 2
    out["cr0"]["exact_rel_bias"] = float(ev0 / var_exact)
    out["cr1"]["exact_rel_bias"] = float(adj * ev0 / var_exact)
    out["meta"] = dict(var_exact=var_exact, emp_var_ybar=float(np.var(res["ybar"], ddof=1)),
                       mean_ybar=float(np.mean(res["ybar"])), tcrit=float(tcrit),
                       mean_q=float(np.mean(res["q"])), n_distinct_wcb_q=int(len(np.unique(np.round(res["q"], 10)))))
    return out


results = {"task_id": "T2-simulation"}
Rs, Bs = design.R.unique(), design.B.unique()
assert len(Rs) == 1 and len(Bs) == 1
R, B = int(Rs[0]), int(Bs[0])
cov, mcse, rb, rb_se, rb_exact = {}, {}, {}, {}, {}
base_cov, base_rb, cellmeta = {}, {}, {}
rows = []
for _, row in design.iterrows():
    o = run_cell(row)
    for est in ["cr0", "cr1", "wcb"]:
        k = key(est, row.G, row.icc)
        cov[k], mcse[k] = o[est]["coverage"], o[est]["mc_se"]
        if est != "wcb":
            rb[k], rb_se[k] = o[est]["rel_bias"], o[est]["rel_bias_se"]
            rb_exact[k] = o[est]["exact_rel_bias"]
        rows.append(dict(estimator=est, cell_id=row.cell_id, G=int(row.G), icc=float(row.icc),
                         coverage=o[est]["coverage"], mc_se_coverage=o[est]["mc_se"],
                         rel_bias_var=(o[est]["rel_bias"] if est != "wcb" else None),
                         mean_ci_halfwidth=o[est]["mean_hw"], R=int(row.R)))
    kb = key("iid", row.G, row.icc)
    base_cov[kb], base_rb[kb] = o["iid"]["coverage"], o["iid"]["rel_bias"]
    cellmeta[row.cell_id] = o["meta"]
    print(f"{row.cell_id} G={row.G} icc={row.icc}: cr0={o['cr0']['coverage']:.4f} "
          f"cr1={o['cr1']['coverage']:.4f} wcb={o['wcb']['coverage']:.4f} iid={o['iid']['coverage']:.4f}")

table = pd.DataFrame(rows)
table = table.sort_values(["estimator", "cell_id"]).reset_index(drop=True)
table.to_csv(os.path.join(OUT, "table1.csv"), index=False)

failure = sorted(k for k, v in cov.items() if v < 0.90)
nominal = sorted(k for k, v in cov.items() if abs(v - 0.95) <= 2 * mcse[k])
avg_small = {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(avg_small, key=avg_small.get)

data_issues = [
    {"issue": "ICC = 0.0 cells imply u_g ~ N(0, 0), a degenerate distribution.",
     "handling": "u_g drawn as sqrt(icc)*Z, which is identically 0 for icc = 0; DGP otherwise unchanged."},
    {"issue": "With G = 5 and Rademacher weights there are only 2^5 = 32 distinct weight vectors and |t*| is invariant to a global sign flip, so at most 16 distinct |t*| values exist; the B = 399 bootstrap quantile is highly discrete (draws repeat the same few values).",
     "handling": "Kept exactly as specified (no Webb weights or enumeration); flagged as a limitation in the report."},
    {"issue": "'0.95 empirical quantile' of |t*| is not uniquely defined for finite B.",
     "handling": "numpy.quantile default (linear interpolation, position 0.95*(B-1)); documented in report."},
]
results.update({
    "R": R, "B": B, "n_cells": int(len(design)),
    "coverage": cov, "mc_se_coverage": mcse, "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": data_issues,
    "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]),
    # supplementary keys
    "avg_coverage_small_G": avg_small,
    "mc_se_rel_bias_var": rb_se,
    "exact_rel_bias_var_analytic": rb_exact,
    "baseline_iid_coverage": base_cov,
    "baseline_iid_rel_bias_var": base_rb,
    "cell_diagnostics": cellmeta,
    "seeds": {r.cell_id: int(r.seed) for _, r in design.iterrows()},
})
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())
for ax, icc in zip(axes, [0.0, 0.2, 0.5]):
    for est, mk in [("cr0", "o"), ("cr1", "s"), ("wcb", "^")]:
        p = np.array([cov[key(est, G, icc)] for G in Gs])
        s = np.array([mcse[key(est, G, icc)] for G in Gs])
        ax.errorbar(Gs, p, yerr=1.96 * s, marker=mk, 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 line 0.90")
    ax.set_xscale("log"); ax.set_xticks(Gs); ax.set_xticklabels([str(g) for g in Gs]); ax.minorticks_off()
    ax.set_title(f"ICC = {icc:.1f}"); ax.set_xlabel("number of clusters G")
axes[0].set_ylabel("coverage of nominal 95% CI (±1.96 MC SE)")
axes[0].legend(fontsize=8, loc="lower right")
fig.suptitle(f"Coverage of cluster-robust CIs for a 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 (G=5):", best, avg_small)