"""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
Regenerates out/results.json, out/table1.csv, out/figure1.png from data/design.csv only.
"""
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

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

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

def simulate_cell(row):
    G, icc, R, B = int(row.G), float(row.icc), int(row.R), int(row.B)
    mu, level = float(row.mu), float(row.nominal_level)
    n = np.array([int(s) for s in str(row.cluster_sizes).split("|")])
    N = int(n.sum())
    assert len(n) == G, "cluster_sizes length != G"
    assert N == int(row.N), "sum(cluster_sizes) != N"
    var_exact = (icc * np.sum(n**2) + (1 - icc) * N) / N**2
    rng = np.random.default_rng(int(row.seed))
    # draw order (fixed, deterministic): U (R,G), E (R,N), bootstrap weights (R,B,G)
    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 - icc), size=(R, N))
    W = rng.choice(np.array([-1.0, 1.0]), size=(R, B, G))
    cl = np.repeat(np.arange(G), n)
    Y = mu + U[:, cl] + E                      # (R,N)
    ybar = Y.mean(axis=1)                      # (R,)
    ehat = Y - ybar[:, None]
    S = np.zeros((R, G))
    for g in range(G):
        S[:, g] = ehat[:, cl == g].sum(axis=1)  # cluster residual sums
    V0 = (S**2).sum(axis=1) / N**2
    V1 = G / (G - 1) * V0
    z = 1.96                                   # as specified in TASK.md
    tcrit = stats.t.ppf(1 - (1 - level) / 2, G - 1)
    hw0 = z * np.sqrt(V0)
    hw1 = tcrit * np.sqrt(V1)
    # Wild cluster bootstrap-t (unrestricted, Rademacher), using cluster sums:
    # ybar* - ybar = D = sum_g w_g S_g / N ; bootstrap residual cluster sums = w_g S_g - n_g D
    WS = W * S[:, None, :]                     # (R,B,G)
    D = WS.sum(axis=2) / N                     # (R,B)
    Sstar = WS - n[None, None, :] * D[:, :, None]
    V1s = G / (G - 1) * (Sstar**2).sum(axis=2) / N**2
    tstar = np.abs(D) / np.sqrt(V1s)
    q = np.quantile(tstar, level, axis=1)      # numpy default (linear interpolation) 0.95 quantile
    hwb = q * np.sqrt(V1)
    out = {}
    for est, hw, V in [("cr0", hw0, V0), ("cr1", hw1, V1), ("wcb", hwb, None)]:
        cov = np.mean(np.abs(ybar - mu) <= hw)
        out[est] = dict(coverage=float(cov), mc_se=float(np.sqrt(cov * (1 - cov) / R)),
                        rel_bias=None if V is None else float(V.mean() / var_exact),
                        rel_bias_mcse=None if V is None else float(V.std(ddof=1) / np.sqrt(R) / var_exact),
                        hw=float(hw.mean()))
    # analytic E[V_CR0]/Var(ybar): T_g = cluster sum of (y-mu), a_g = Var(T_g), S_g = T_g - (n_g/N) sum T
    a = icc * n**2 + (1 - icc) * n
    ES2 = a - 2 * (n / N) * a + (n / N) ** 2 * a.sum()
    rb0_exact = float(ES2.sum() / N**2 / var_exact)
    extra = dict(rel_bias_cr0_analytic=rb0_exact, rel_bias_cr1_analytic=rb0_exact * G / (G - 1), var_exact=float(var_exact), emp_var_ybar=float(ybar.var(ddof=1)),
                 n_distinct_abs_t_patterns=int(2 ** (G - 1)), mean_q=float(q.mean()), tcrit=float(tcrit))
    return out, extra

rows, cov, se, rb, rbse, extras = [], {}, {}, {}, {}, {}
for _, row in design.iterrows():
    res, extra = simulate_cell(row)
    extras[row.cell_id] = extra
    for est in ["cr0", "cr1", "wcb"]:
        r = res[est]; k = key(est, row.G, row.icc)
        cov[k] = r["coverage"]; se[k] = r["mc_se"]
        if r["rel_bias"] is not None:
            rb[k] = r["rel_bias"]; rbse[k] = r["rel_bias_mcse"]
        rows.append(dict(estimator=est, cell_id=row.cell_id, G=int(row.G), icc=float(row.icc),
                         coverage=r["coverage"], mc_se_coverage=r["mc_se"],
                         rel_bias_var=r["rel_bias"], mean_ci_halfwidth=r["hw"], R=int(row.R)))

if (design.G < 10).any():
    data_issues.append({"issue": "With G=5 and Rademacher weights there are only 2^5=32 distinct weight vectors "
                        "(16 distinct |t*| values), so the B=399 bootstrap draws repeat these few values and the "
                        "0.95 quantile of |t*| is coarse.",
                        "handling": "Kept the specified design (B=399, Rademacher) unchanged; flagged as a limitation in the report."})

table = pd.DataFrame(rows)
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 * se[k]])
g5 = {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(g5, key=g5.get)

results = {
    "task_id": "T2-simulation",
    "R": int(design.R.iloc[0]), "B": int(design.B.iloc[0]), "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": 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]),
    # supplementary diagnostics
    "supplementary": {"rel_bias_var_mc_se": rbse, "mean_coverage_G5": g5, "cell_diagnostics": extras},
}
with open(os.path.join(OUT, "results.json"), "w") as f:
    json.dump(results, f, indent=2)

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, m in [("cr0", "o"), ("cr1", "s"), ("wcb", "^")]:
        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, 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 line 0.90")
    ax.set_xscale("log"); ax.set_xticks(Gs); ax.set_xticklabels(Gs)
    ax.set_title(f"ICC = {icc}"); 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.tight_layout(); fig.savefig(os.path.join(OUT, "figure1.png"), dpi=130)
print(json.dumps({k: results[k] for k in ["coverage", "rel_bias_var", "failure_cells", "nominal_cells", "best_procedure_small_G"]}, indent=1))
print(json.dumps(extras, indent=1)); print(g5)