"""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.
Randomness: one numpy default_rng(seed) per cell (seed from design.csv). Per replication, draws in
fixed order: u (G normals), e (N normals), W (B x G Rademacher via rng.integers(0,2)).
"""
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

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 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)
    sizes = np.array([int(x) for x in row.cluster_sizes.split("|")])
    assert len(sizes) == G and sizes.sum() == N
    gid = np.repeat(np.arange(G), sizes)
    var_exact = (icc * np.sum(sizes.astype(float) ** 2) + (1 - icc) * N) / N ** 2
    tcrit = stats.t.ppf(0.975, G - 1)
    zcrit = 1.96
    adj = G / (G - 1)
    rng = np.random.default_rng(int(row.seed))
    res = {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(icc), G)  # scale 0 when icc=0 -> u = 0
        e = rng.normal(0.0, np.sqrt(1 - icc), 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 = rng.integers(0, 2, size=(B, G)) * 2 - 1
        d = (W * S).sum(axis=1) / N                     # ybar* - ybar
        Sstar = W * S - np.outer(d, sizes)              # cluster sums of bootstrap residuals
        v1s = adj * np.sum(Sstar ** 2, axis=1) / N ** 2
        if np.any(v1s <= 0):
            n_degenerate += int(np.sum(v1s <= 0))
            v1s = np.where(v1s <= 0, np.nan, v1s)
        tabs = np.abs(d) / np.sqrt(v1s)
        q = np.nanquantile(tabs, 0.95)                  # numpy default (linear interpolation)
        res["ybar"][r], res["v0"][r], res["v1"][r], res["q"][r] = ybar, v0, v1, q
        res["v_iid"][r] = np.var(y, ddof=1) / N         # naive i.i.d. baseline
    hw = {"cr0": zcrit * np.sqrt(res["v0"]), "cr1": tcrit * np.sqrt(res["v1"]),
          "wcb": res["q"] * np.sqrt(res["v1"]), "iid": zcrit * np.sqrt(res["v_iid"])}
    out = {"var_exact": var_exact, "n_degenerate": n_degenerate, "R": R}
    for k, h in hw.items():
        cov = np.abs(res["ybar"] - mu) <= h
        out[k] = {"coverage": float(cov.mean()), "hw": float(h.mean()), "ind": cov.astype(float)}
    for k, v in [("cr0", res["v0"]), ("cr1", res["v1"])]:
        out[k]["rel_bias"] = float(v.mean() / var_exact)
        out[k]["rel_bias_mcse"] = float(v.std(ddof=1) / np.sqrt(R) / var_exact)
    out["mc_var_ybar_ratio"] = float(np.var(res["ybar"], ddof=1) / var_exact)
    out["mean_ybar"] = float(res["ybar"].mean())
    return out


def main():
    d = pd.read_csv(DATA)
    R_all, B_all = d.R.unique(), d.B.unique()
    assert len(R_all) == 1 and len(B_all) == 1
    R, B = int(R_all[0]), int(B_all[0])
    cov, se, rb, rb_se, base, chk = {}, {}, {}, {}, {}, {}
    rows, issues = [], []
    n_deg = 0
    ind_g5 = {k: [] for k in EST}
    for row in d.itertuples(index=False):
        o = simulate_cell(row)
        n_deg += o["n_degenerate"]
        tag = f"G{int(row.G)}_icc{float(row.icc):.1f}"
        for k in EST:
            key = f"{k}_{tag}"
            p = o[k]["coverage"]
            cov[key] = p
            se[key] = float(np.sqrt(p * (1 - p) / o["R"]))
            if k != "wcb":
                rb[key] = o[k]["rel_bias"]
                rb_se[key] = o[k]["rel_bias_mcse"]
            rows.append({"estimator": k, "cell_id": row.cell_id, "G": int(row.G), "icc": float(row.icc),
                         "coverage": p, "mc_se_coverage": se[key],
                         "rel_bias_var": o[k].get("rel_bias", None), "mean_ci_halfwidth": o[k]["hw"], "R": o["R"]})
        if int(row.G) == 5:
            for k in EST:
                ind_g5[k].append(o[k]["ind"])
        base[f"iid_{tag}"] = o["iid"]["coverage"]
        chk[tag] = {"mc_var_ybar_over_exact": o["mc_var_ybar_ratio"], "mean_ybar": o["mean_ybar"]}
        print(f"{row.cell_id} {tag}: " + ", ".join(f"{k}={cov[f'{k}_{tag}']:.4f}" for k in EST))
    table = pd.DataFrame(rows)
    table.to_csv(os.path.join(OUT, "table1.csv"), index=False)

    failure = [k for k, p in cov.items() if p < 0.90]
    nominal = [k for k, p in cov.items() if abs(p - 0.95) <= 2 * se[k]]
    g5 = {k: np.mean([cov[f"{k}_G5_icc{i:.1f}"] for i in (0.0, 0.2, 0.5)]) for k in EST}
    best = max(g5, key=g5.get)
    # paired MC SE of the difference in average G=5 coverage between best and each other procedure
    # (same replications -> paired; cells independent -> variance of mean over 3 cells = sum/9)
    g5_diff = {}
    for k in EST:
        if k == best:
            continue
        cell_vars = [np.var(a - b, ddof=1) / len(a) for a, b in zip(ind_g5[best], ind_g5[k])]
        g5_diff[f"{best}_minus_{k}"] = {"diff": float(g5[best] - g5[k]), "mc_se": float(np.sqrt(sum(cell_vars)) / len(cell_vars))}
    near = [k for k, p in cov.items() if abs(p - 0.90) <= 2 * se[k]]
    if 5 in d.G.values:
        issues.append({"issue": "With G=5 and Rademacher weights only 2^5=32 distinct weight vectors exist (16 distinct |t*| values, as |t*| is invariant to w -> -w); B=399 draws therefore repeat these few values and the 0.95 quantile of |t*| is coarse.",
                       "handling": "Kept as specified (B=399, Rademacher, random draws, numpy linear-interpolation quantile); discussed in report as a known feature of WCB with very few clusters."})
    if n_deg:
        issues.append({"issue": f"{n_deg} bootstrap draws had non-positive V_CR1*", "handling": "Excluded those draws from the |t*| quantile (nan-quantile)."})
    results = {
        "task_id": "T2-simulation", "R": R, "B": B, "n_cells": int(len(d)),
        "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": issues,
        "n_input_rows": {"design.csv": int(len(d))},
        "n_analysed": int(len(d) * R), "random_seed": int(d.loc[d.cell_id == "C1", "seed"].iloc[0]),
        # supplementary keys
        "avg_coverage_G5": {k: float(v) for k, v in g5.items()},
        "mc_se_rel_bias_var": rb_se,
        "avg_coverage_G5_paired_diff": g5_diff,
        "cells_within_2se_of_failure_line": near,
        "baseline_iid_coverage": base,
        "check_mc_var_ybar_over_exact": chk,
        "n_degenerate_bootstrap_draws": int(n_deg),
        "wcb_quantile_method": "numpy.quantile linear interpolation (type 7) of |t*_b|, b=1..B",
        "seeds_by_cell": {r.cell_id: int(r.seed) for r in d.itertuples(index=False)},
    }
    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(d.G.unique())
    style = {"cr0": ("o-", "CR0 (z)"), "cr1": ("s-", "CR1 (t, G-1)"), "wcb": ("^-", "WCB-t Rademacher")}
    for ax, icc in zip(axes, sorted(d.icc.unique())):
        for k in EST:
            p = np.array([cov[f"{k}_G{g}_icc{icc:.1f}"] for g in Gs])
            s = np.array([se[f"{k}_G{g}_icc{icc:.1f}"] for g in Gs])
            ax.errorbar(Gs, p, yerr=1.96 * s, fmt=style[k][0], capsize=3, label=style[k][1])
        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.minorticks_off(); ax.set_xticks(Gs); ax.set_xticklabels(Gs)
        ax.set_title(f"ICC = {icc:.1f}"); ax.set_xlabel("number of clusters G")
    axes[0].set_ylabel("coverage (R=2000; bars ±1.96 MC SE)")
    axes[0].legend(loc="lower right", fontsize=8)
    fig.tight_layout()
    fig.savefig(os.path.join(OUT, "figure1.png"), dpi=130)
    print("failure cells:", failure)
    print("best procedure at G=5:", best)


if __name__ == "__main__":
    main()