#!/usr/bin/env python3
"""Compute paired bootstrap confidence intervals for GRL revision outputs."""

from __future__ import annotations

import argparse
import csv
import gzip
import math
from collections import defaultdict
from pathlib import Path

import numpy as np


DEFAULT_REVISION_DIR = Path("outputs/grl_revision_20260610")


def read_csv(path: Path) -> list[dict[str, str]]:
    opener = gzip.open if path.suffix == ".gz" else open
    with opener(path, "rt", newline="", encoding="utf-8") as f:
        return list(csv.DictReader(f))


def metric_from_counts(tp: np.ndarray, fp: np.ndarray, fn: np.ndarray) -> dict[str, float]:
    tp_sum = float(tp.sum())
    fp_sum = float(fp.sum())
    fn_sum = float(fn.sum())
    precision = tp_sum / (tp_sum + fp_sum) if tp_sum + fp_sum else 0.0
    recall = tp_sum / (tp_sum + fn_sum) if tp_sum + fn_sum else 0.0
    f1 = 2.0 * precision * recall / (precision + recall) if precision + recall else 0.0
    return {"precision": precision, "recall": recall, "f1": f1}


def percentile_ci(values: np.ndarray, alpha: float = 0.05) -> tuple[float, float]:
    return (float(np.percentile(values, 100 * alpha / 2)), float(np.percentile(values, 100 * (1 - alpha / 2))))


def phase_arrays(rows: list[dict[str, str]]):
    by_cond_phase: dict[tuple[str, str], dict[int, tuple[int, int, int]]] = defaultdict(dict)
    sample_ids = set()
    for row in rows:
        sid = int(row["sample_id"])
        sample_ids.add(sid)
        by_cond_phase[(row["condition"], row["phase"])][sid] = (
            int(row["tp"]),
            int(row["fp"]),
            int(row["fn"]),
        )
    ordered_ids = np.array(sorted(sample_ids), dtype=int)
    out = {}
    for key, vals in by_cond_phase.items():
        arr = np.array([vals[int(sid)] for sid in ordered_ids], dtype=np.int64)
        out[key] = {"tp": arr[:, 0], "fp": arr[:, 1], "fn": arr[:, 2]}
    return ordered_ids, out


def bootstrap_phase(input_csv: Path, output_csv: Path, summary_txt: Path, n_bootstrap: int, seed: int) -> None:
    rows = read_csv(input_csv)
    if not rows:
        raise RuntimeError(f"No rows in {input_csv}")
    sample_ids, arrays = phase_arrays(rows)
    n = len(sample_ids)
    rng = np.random.default_rng(seed)
    idx = rng.integers(0, n, size=(n_bootstrap, n), endpoint=False)
    conditions = ["full", "snr5", "snr10"]
    phases = ["P", "S"]
    estimates: dict[tuple[str, str], dict[str, float]] = {}
    boot: dict[tuple[str, str], dict[str, np.ndarray]] = {}
    for cond in conditions:
        for phase in phases:
            arr = arrays[(cond, phase)]
            estimates[(cond, phase)] = metric_from_counts(arr["tp"], arr["fp"], arr["fn"])
            boot[(cond, phase)] = {}
            for b in range(n_bootstrap):
                m = metric_from_counts(arr["tp"][idx[b]], arr["fp"][idx[b]], arr["fn"][idx[b]])
                for metric, value in m.items():
                    boot[(cond, phase)].setdefault(metric, np.empty(n_bootstrap, dtype=float))[b] = value
    output_csv.parent.mkdir(parents=True, exist_ok=True)
    with output_csv.open("w", newline="", encoding="utf-8") as f:
        fieldnames = ["task", "condition", "metric", "estimate", "ci_low", "ci_high", "diff_vs_full", "diff_ci_low", "diff_ci_high"]
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        for cond in conditions:
            for phase in phases:
                for metric in ["precision", "recall", "f1"]:
                    values = boot[(cond, phase)][metric]
                    lo, hi = percentile_ci(values)
                    full_est = estimates[("full", phase)][metric]
                    diff = estimates[(cond, phase)][metric] - full_est
                    if cond == "full":
                        dlo = dhi = 0.0
                    else:
                        diff_values = values - boot[("full", phase)][metric]
                        dlo, dhi = percentile_ci(diff_values)
                    writer.writerow(
                        {
                            "task": "phase_picking",
                            "condition": cond,
                            "metric": f"{phase}_{metric}",
                            "estimate": estimates[(cond, phase)][metric],
                            "ci_low": lo,
                            "ci_high": hi,
                            "diff_vs_full": diff,
                            "diff_ci_low": dlo,
                            "diff_ci_high": dhi,
                        }
                    )
            mean_est = 0.5 * (estimates[(cond, "P")]["f1"] + estimates[(cond, "S")]["f1"])
            mean_boot = 0.5 * (boot[(cond, "P")]["f1"] + boot[(cond, "S")]["f1"])
            lo, hi = percentile_ci(mean_boot)
            diff_values = mean_boot - 0.5 * (boot[("full", "P")]["f1"] + boot[("full", "S")]["f1"])
            dlo, dhi = (0.0, 0.0) if cond == "full" else percentile_ci(diff_values)
            writer.writerow(
                {
                    "task": "phase_picking",
                    "condition": cond,
                    "metric": "mean_f1",
                    "estimate": mean_est,
                    "ci_low": lo,
                    "ci_high": hi,
                    "diff_vs_full": mean_est - 0.5 * (estimates[("full", "P")]["f1"] + estimates[("full", "S")]["f1"]),
                    "diff_ci_low": dlo,
                    "diff_ci_high": dhi,
                }
            )
    summary_txt.parent.mkdir(parents=True, exist_ok=True)
    summary_txt.write_text(
        f"Phase-picking paired bootstrap: {n_bootstrap} resamples over {n} deterministic test windows. "
        f"Corpus-level TP/FP/FN were recomputed for each resample; per-window F1 was not averaged.\n",
        encoding="utf-8",
    )


def dispersion_arrays(rows: list[dict[str, str]]):
    by_cond: dict[str, dict[int, tuple[float, float, float]]] = defaultdict(dict)
    sample_ids = set()
    for row in rows:
        sid = int(row["sample_id"])
        sample_ids.add(sid)
        by_cond[row["condition"]][sid] = (
            float(row["abs_error_sum"]),
            float(row["squared_error_sum"]),
            float(row["valid_period_count"]),
        )
    ordered_ids = np.array(sorted(sample_ids), dtype=int)
    out = {}
    for cond, vals in by_cond.items():
        arr = np.array([vals[int(sid)] for sid in ordered_ids], dtype=float)
        out[cond] = {"abs": arr[:, 0], "sq": arr[:, 1], "n": arr[:, 2]}
    return ordered_ids, out


def disp_metrics(abs_sum: np.ndarray, sq_sum: np.ndarray, valid_n: np.ndarray) -> dict[str, float]:
    denom = float(valid_n.sum())
    mae = float(abs_sum.sum() / denom) if denom else 0.0
    rmse = float(math.sqrt(sq_sum.sum() / denom)) if denom else 0.0
    return {"mae": mae, "rmse": rmse}


def bootstrap_dispersion(input_csv: Path, output_csv: Path, summary_txt: Path, n_bootstrap: int, seed: int) -> None:
    rows = read_csv(input_csv)
    if not rows:
        raise RuntimeError(f"No rows in {input_csv}")
    sample_ids, arrays = dispersion_arrays(rows)
    n = len(sample_ids)
    rng = np.random.default_rng(seed)
    idx = rng.integers(0, n, size=(n_bootstrap, n), endpoint=False)
    conditions = ["full", "snr_q1", "snr_q2"]
    estimates = {cond: disp_metrics(arrays[cond]["abs"], arrays[cond]["sq"], arrays[cond]["n"]) for cond in conditions}
    boot = {cond: {"mae": np.empty(n_bootstrap), "rmse": np.empty(n_bootstrap)} for cond in conditions}
    for cond in conditions:
        arr = arrays[cond]
        for b in range(n_bootstrap):
            m = disp_metrics(arr["abs"][idx[b]], arr["sq"][idx[b]], arr["n"][idx[b]])
            boot[cond]["mae"][b] = m["mae"]
            boot[cond]["rmse"][b] = m["rmse"]
    output_csv.parent.mkdir(parents=True, exist_ok=True)
    with output_csv.open("w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=["task", "condition", "metric", "estimate", "ci_low", "ci_high", "diff_vs_full", "diff_ci_low", "diff_ci_high"])
        writer.writeheader()
        for cond in conditions:
            for metric in ["mae", "rmse"]:
                lo, hi = percentile_ci(boot[cond][metric])
                diff = estimates[cond][metric] - estimates["full"][metric]
                if cond == "full":
                    dlo = dhi = 0.0
                else:
                    dlo, dhi = percentile_ci(boot[cond][metric] - boot["full"][metric])
                writer.writerow(
                    {
                        "task": "dispersion",
                        "condition": cond,
                        "metric": metric,
                        "estimate": estimates[cond][metric],
                        "ci_low": lo,
                        "ci_high": hi,
                        "diff_vs_full": diff,
                        "diff_ci_low": dlo,
                        "diff_ci_high": dhi,
                    }
                )
    summary_txt.parent.mkdir(parents=True, exist_ok=True)
    summary_txt.write_text(f"Dispersion paired bootstrap: {n_bootstrap} resamples over {n} test samples.\n", encoding="utf-8")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--revision-dir", type=Path, default=DEFAULT_REVISION_DIR)
    parser.add_argument("--n-bootstrap", type=int, default=10000)
    parser.add_argument("--seed", type=int, default=20260609)
    args = parser.parse_args()
    tables = args.revision_dir / "tables"
    bootstrap_phase(
        args.revision_dir / "phase_picking" / "phase_per_window_outputs.csv.gz",
        tables / "phase_bootstrap_ci.csv",
        tables / "phase_bootstrap_summary.txt",
        args.n_bootstrap,
        args.seed,
    )
    bootstrap_dispersion(
        args.revision_dir / "dispersion" / "dispersion_per_sample_metrics.csv.gz",
        tables / "dispersion_bootstrap_ci.csv",
        tables / "dispersion_bootstrap_summary.txt",
        args.n_bootstrap,
        args.seed,
    )
    (tables / "bootstrap_summary_for_manuscript.txt").write_text(
        (tables / "phase_bootstrap_summary.txt").read_text(encoding="utf-8")
        + (tables / "dispersion_bootstrap_summary.txt").read_text(encoding="utf-8"),
        encoding="utf-8",
    )


if __name__ == "__main__":
    main()