"""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.
Deterministic: each cell uses numpy.random.default_rng(seed of that cell). Within a replication
the draw order is: u (G normals), e (N normals), then a B x G Rademacher weight matrix.
"""
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
import matplotlib.ticker

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")
EST = ["cr0", "cr1", "wcb"]


def key(est, G, icc):
    return f"{est}_G{int(G)}_icc{float(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, tv, level = float(row.mu), float(row.total_variance), float(row.nominal_level)
    sizes = np.array([int(s) for s in str(row.cluster_sizes).split("|")])
    assert len(sizes) == G and sizes.sum() == N
    cid = np.repeat(np.arange(G), sizes)
    var_u, var_e = icc * tv, (1.0 - icc) * tv
    var_ybar = (var_u * np.sum(sizes ** 2) + var_e * N) / N ** 2  # exact Var(ybar)
    z = stats.norm.ppf(1 - (1 - level) / 2)
    z = 1.96 if abs(z - 1.96) < 1e-3 else z  # TASK specifies 1.96
    tcrit = stats.t.ppf(1 - (1 - level) / 2, G - 1)
    cr1f = G / (G - 1)
    rng = np.random.default_rng(int(row.seed))

    out = {k: np.empty(R) for k in ["ybar", "v0", "v1", "q", "v_iid"]}
    n_degenerate = 0
    for r in range(R):
        u = rng.normal(0.0, np.sqrt(var_u), G)
        e = rng.normal(0.0, np.sqrt(var_e), N)
        y = mu + u[cid] + e
        ybar = y.mean()
        res = y - ybar
        S = np.bincount(cid, weights=res, minlength=G)  # cluster sums of residuals
        v0 = np.sum(S ** 2) / N ** 2
        v1 = cr1f * v0
        # wild cluster bootstrap-t (unrestricted, Rademacher)
        W = rng.choice(np.array([-1.0, 1.0]), size=(B, G))
        delta = (W @ S) / N                              # ybar* - ybar
        Sstar = W * S[None, :] - delta[:, None] * sizes[None, :]  # cluster sums of bootstrap residuals
        v1s = cr1f * np.sum(Sstar ** 2, axis=1) / N ** 2
        ok = v1s > 0
        n_degenerate += int((~ok).sum())
        tstar = np.abs(delta[ok]) / np.sqrt(v1s[ok])
        q = np.quantile(tstar, level, method="inverted_cdf")  # = ceil(0.95*B)-th order stat (380th of 399)
        out["ybar"][r], out["v0"][r], out["v1"][r], out["q"][r] = ybar, v0, v1, q
        out["v_iid"][r] = np.sum(res ** 2) / (N - 1) / N  # naive iid variance (baseline only)

    err = np.abs(out["ybar"] - mu)
    hw = {
        "cr0": z * np.sqrt(out["v0"]),
        "cr1": tcrit * np.sqrt(out["v1"]),
        "wcb": out["q"] * np.sqrt(out["v1"]),
        "iid": z * np.sqrt(out["v_iid"]),
    }
    res = {}
    for est, h in hw.items():
        cov = float(np.mean(err <= h))
        res[est] = dict(coverage=cov, mc_se=float(np.sqrt(cov * (1 - cov) / R)),
                        mean_hw=float(np.mean(h)))
    res["cr0"]["rel_bias"] = float(np.mean(out["v0"]) / var_ybar)
    res["cr1"]["rel_bias"] = float(np.mean(out["v1"]) / var_ybar)
    res["cr0"]["rel_bias_mcse"] = float(np.std(out["v0"], ddof=1) / np.sqrt(R) / var_ybar)
    res["cr1"]["rel_bias_mcse"] = float(np.std(out["v1"], ddof=1) / np.sqrt(R) / var_ybar)
    res["iid"]["rel_bias"] = float(np.mean(out["v_iid"]) / var_ybar)
    # exact E[V_CR0]/Var(ybar) (analytic check of the Monte Carlo bias estimate)
    Mm = np.eye(N) - 1.0 / N
    Sig = var_u * (cid[:, None] == cid[None, :]) + var_e * np.eye(N)
    A = np.zeros((G, N)); A[cid, np.arange(N)] = 1.0
    ev0 = float(np.trace(A @ Mm @ Sig @ Mm @ A.T) / N ** 2)
    res["cr0"]["rel_bias_exact"] = ev0 / var_ybar
    res["cr1"]["rel_bias_exact"] = cr1f * ev0 / var_ybar
    res["meta"] = dict(var_ybar=float(var_ybar), n_degenerate_boot=n_degenerate,
                       mean_ybar=float(out["ybar"].mean()),
                       var_ybar_mc=float(np.var(out["ybar"], ddof=1)),
                       n_distinct_q_min=int(len(np.unique(np.round(out["q"], 12)))))
    return res


def main():
    design = pd.read_csv(DATA)
    R_all, B_all = design.R.unique(), design.B.unique()
    assert len(R_all) == 1 and len(B_all) == 1
    R, B = int(R_all[0]), int(B_all[0])

    rows, cov, mcse, rb, extra = [], {}, {}, {}, {}
    base_cov, base_mcse = {}, {}
    n_degenerate = 0
    for _, row in design.iterrows():
        r = simulate_cell(row)
        n_degenerate += r["meta"]["n_degenerate_boot"]
        for est in EST:
            k = key(est, row.G, row.icc)
            cov[k], mcse[k] = r[est]["coverage"], r[est]["mc_se"]
            if est != "wcb":
                rb[k] = r[est]["rel_bias"]
                extra[k + "_mcse"] = r[est]["rel_bias_mcse"]
                extra[k + "_exact"] = r[est]["rel_bias_exact"]
            rows.append(dict(estimator=est, cell_id=row.cell_id, G=int(row.G), icc=float(row.icc),
                             coverage=r[est]["coverage"], mc_se_coverage=r[est]["mc_se"],
                             rel_bias_var=r[est].get("rel_bias", np.nan) if est != "wcb" else np.nan,
                             mean_ci_halfwidth=r[est]["mean_hw"], R=int(row.R)))
        kb = f"iid_G{int(row.G)}_icc{float(row.icc):.1f}"
        base_cov[kb], base_mcse[kb] = r["iid"]["coverage"], r["iid"]["mc_se"]
        extra[f"meta_{row.cell_id}"] = r["meta"]

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

    failure = [k for k, v in cov.items() if v < 0.90]
    nominal = [k for k, v in cov.items() if abs(v - 0.95) <= 2 * mcse[k]]
    g5 = {est: float(np.mean([cov[key(est, 5, i)] for i in sorted(design.icc.unique())])) for est in EST}
    best = max(g5, key=lambda e: g5[e])

    data_issues = [{
        "issue": "With G=5 and Rademacher weights only 2^5=32 distinct bootstrap weight vectors exist "
                 "(16 distinct |t*| values), so the B=399 draws resample a tiny support and the WCB "
                 "critical value is coarse and typically large.",
        "handling": "Procedure kept exactly as specified (no Webb weights, no enumeration); flagged as a "
                    "limitation in the report."}]
    if n_degenerate:
        data_issues.append({"issue": f"{n_degenerate} bootstrap draws had V*_CR1 = 0",
                            "handling": "dropped from that replication's |t*| set before taking the quantile"})

    results = {
        "task_id": "T2-simulation",
        "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]),
        "mean_coverage_G5": g5,
        "rel_bias_var_mcse": {k.replace("_mcse", ""): v for k, v in extra.items() if k.endswith("_mcse")},
        "rel_bias_var_exact": {k.replace("_exact", ""): v for k, v in extra.items() if k.endswith("_exact")},
        "baseline_iid_coverage": base_cov,
        "baseline_iid_mc_se_coverage": base_mcse,
        "cell_diagnostics": {k[5:]: v for k, v in extra.items() if k.startswith("meta_")},
        "seeds": {str(c): int(s) for c, s in zip(design.cell_id, design.seed)},
        "wcb_quantile_rule": "inverted_cdf 0.95 quantile of |t*| (380th order statistic of 399)",
    }
    with open(os.path.join(OUT, "results.json"), "w") as f:
        json.dump(results, f, indent=2)

    # figure
    iccs = sorted(design.icc.unique())
    Gs = sorted(design.G.unique())
    fig, axes = plt.subplots(1, len(iccs), figsize=(13, 4.2), sharey=True)
    styles = {"cr0": ("o-", "CR0 (z)"), "cr1": ("s-", "CR1 (t, G-1)"), "wcb": ("^-", "WCB-t (Rademacher)")}
    for ax, icc in zip(axes, iccs):
        for est in EST:
            p = np.array([cov[key(est, g, icc)] for g in Gs])
            se = np.array([mcse[key(est, g, icc)] for g in Gs])
            ax.errorbar(Gs, p, yerr=1.96 * se, fmt=styles[est][0], label=styles[est][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.set_xticks(Gs); ax.set_xticklabels([str(g) for g in Gs])
        ax.xaxis.set_minor_formatter(matplotlib.ticker.NullFormatter())
        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% CIs for the mean (R={R}, B={B})")
    fig.tight_layout()
    fig.savefig(os.path.join(OUT, "figure1.png"), dpi=120)
    plt.close(fig)

    print(table.to_string(index=False))
    print("failure cells:", failure)
    print("best procedure at G=5:", best, g5)


if __name__ == "__main__":
    main()