#!/usr/bin/env python3
"""Aggregate GRL SNR-transfer experiments across multiple random seeds."""

from __future__ import annotations

import argparse
import csv
import json
import math
from collections import defaultdict
from pathlib import Path
from statistics import mean, stdev


ROOT = Path(__file__).resolve().parents[1]


PHASE_METRICS = [
    "P_precision",
    "P_recall",
    "P_f1",
    "S_precision",
    "S_recall",
    "S_f1",
    "mean_f1",
]

DISP_METRICS = [
    "val_mae",
    "val_rmse",
    "val_certainty_f1",
]


def read_summary(path: Path) -> dict:
    if not path.exists():
        raise FileNotFoundError(path)
    with path.open(encoding="utf-8") as f:
        return json.load(f)


def stats(values: list[float]) -> dict[str, float]:
    if not values:
        return {"mean": math.nan, "std": math.nan, "n": 0}
    return {
        "mean": float(mean(values)),
        "std": float(stdev(values)) if len(values) > 1 else 0.0,
        "n": len(values),
    }


def aggregate_task(
    task: str,
    seeds: list[int],
    metrics: list[str],
    out_dir: Path,
    phase_balanced: bool = False,
    phase_prefix: str | None = None,
) -> tuple[list[dict], dict]:
    by_condition: dict[str, dict[str, list[float]]] = defaultdict(lambda: defaultdict(list))
    condition_labels: dict[str, str] = {}
    rows_long: list[dict] = []
    metadata: dict = {"seeds": seeds, "task": task, "summaries": []}

    if task == "phase":
        prefix = phase_prefix or ("snr_transfer_phase_balanced_seed" if phase_balanced else "snr_transfer_seed")
    else:
        prefix = "disp_snr_transfer_seed"
    for seed in seeds:
        path = ROOT / "outputs" / f"{prefix}{seed}" / "summary.json"
        summary = read_summary(path)
        metadata["summaries"].append(str(path.relative_to(ROOT)))
        for row in summary["rows"]:
            slug = row["slug"]
            condition_labels.setdefault(slug, row["label"])
            for metric in metrics:
                value = float(row[metric])
                by_condition[slug][metric].append(value)
                rows_long.append(
                    {
                        "task": task,
                        "seed": seed,
                        "condition_slug": slug,
                        "condition_label": row["label"],
                        "metric": metric,
                        "value": value,
                    }
                )

    rows_summary: list[dict] = []
    for condition, metric_values in by_condition.items():
        for metric, values in metric_values.items():
            s = stats(values)
            rows_summary.append(
                {
                    "task": task,
                    "condition_slug": condition,
                    "condition_label": condition_labels.get(condition, condition),
                    "metric": metric,
                    "mean": s["mean"],
                    "std": s["std"],
                    "n": s["n"],
                    "values": values,
                }
            )

    out_dir.mkdir(parents=True, exist_ok=True)
    long_path = out_dir / f"{task}_multiseed_long.csv"
    with long_path.open("w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=["task", "seed", "condition_slug", "condition_label", "metric", "value"])
        writer.writeheader()
        writer.writerows(rows_long)

    summary_path = out_dir / f"{task}_multiseed_summary.csv"
    with summary_path.open("w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(
            f,
            fieldnames=["task", "condition_slug", "condition_label", "metric", "mean", "std", "n", "values"],
        )
        writer.writeheader()
        for row in rows_summary:
            row = dict(row)
            row["values"] = json.dumps(row["values"])
            writer.writerow(row)

    return rows_summary, metadata


def lookup(rows: list[dict], task: str, condition: str, metric: str) -> dict:
    for row in rows:
        if row["task"] == task and row["condition_slug"] == condition and row["metric"] == metric:
            return row
    raise KeyError((task, condition, metric))


def has_task(all_rows: list[dict], task: str) -> bool:
    return any(row["task"] == task for row in all_rows)


def has_condition(all_rows: list[dict], task: str, condition: str) -> bool:
    return any(row["task"] == task and row["condition_slug"] == condition for row in all_rows)


def condition_label(all_rows: list[dict], task: str, condition: str, fallback: str) -> str:
    for row in all_rows:
        if row["task"] == task and row["condition_slug"] == condition:
            return row.get("condition_label") or fallback
    return fallback


def write_markdown(all_rows: list[dict], out_dir: Path, phase_balanced: bool = False) -> None:
    lines = [
        "# Multi-Seed SNR-Filtering Summary",
        "",
        "Values are mean +/- sample standard deviation across three seeds.",
        "",
    ]

    if has_task(all_rows, "phase"):
        lines.extend([
            "## Phase Picking",
            "",
            "| Training subset | Mean F1 | P F1 | S F1 |",
            "|---|---:|---:|---:|",
        ])
        if phase_balanced:
            if has_condition(all_rows, "phase", "finetune_full"):
                phase_labels = [
                    ("finetune_full", "Fine-tune full matched"),
                    ("finetune_p5_s_bal", "Fine-tune P>=5 dB, S-balanced matched"),
                    ("finetune_p10_s_bal", "Fine-tune P>=10 dB, S-balanced matched"),
                    ("scratch_full", "Scratch full matched"),
                    ("scratch_p5_s_bal", "Scratch P>=5 dB, S-balanced matched"),
                    ("scratch_p10_s_bal", "Scratch P>=10 dB, S-balanced matched"),
                ]
            else:
                phase_labels = [
                    ("full", "Full matched"),
                    ("p5_s_bal", "P>=5 dB, S-balanced matched"),
                    ("p10_s_bal", "P>=10 dB, S-balanced matched"),
                ]
        else:
            phase_labels = [
                ("full", "Full matched"),
                ("snr5", "SNR>5 dB matched"),
                ("snr10", "SNR>10 dB matched"),
            ]
        for slug, label in phase_labels:
            label = condition_label(all_rows, "phase", slug, label)
            mf = lookup(all_rows, "phase", slug, "mean_f1")
            pf = lookup(all_rows, "phase", slug, "P_f1")
            sf = lookup(all_rows, "phase", slug, "S_f1")
            lines.append(
                f"| {label} | {mf['mean']:.3f} +/- {mf['std']:.3f} | "
                f"{pf['mean']:.3f} +/- {pf['std']:.3f} | {sf['mean']:.3f} +/- {sf['std']:.3f} |"
            )
        lines.append("")

    if has_task(all_rows, "dispersion"):
        lines.extend([
            "",
            "## Dispersion Estimation",
            "",
            "| Training subset | MAE (km/s) | RMSE (km/s) |",
            "|---|---:|---:|",
        ])
        for slug, label in [("full", "Full matched"), ("snr_q1", "SNR>3.04 dB matched"), ("snr_q2", "SNR>6.77 dB matched")]:
            mae = lookup(all_rows, "dispersion", slug, "val_mae")
            rmse = lookup(all_rows, "dispersion", slug, "val_rmse")
            lines.append(
                f"| {label} | {mae['mean']:.4f} +/- {mae['std']:.4f} | "
                f"{rmse['mean']:.4f} +/- {rmse['std']:.4f} |"
            )
    (out_dir / "multiseed_summary.md").write_text("\n".join(lines) + "\n", encoding="utf-8")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--seeds", nargs="+", type=int, default=[20260609, 20260610, 20260611])
    parser.add_argument("--out-dir", default="outputs/grl_multiseed_seed20260609_20260611")
    parser.add_argument("--tasks", nargs="+", choices=["phase", "dispersion"], default=["phase", "dispersion"])
    parser.add_argument(
        "--phase-balanced",
        action="store_true",
        help="Aggregate phase-aware P/S-balanced SNR outputs for the phase task.",
    )
    parser.add_argument(
        "--phase-prefix",
        default=None,
        help="Override the phase output-directory prefix before the seed, e.g. snr_transfer_phase_complete_seed.",
    )
    args = parser.parse_args()

    out_dir = ROOT / args.out_dir
    all_rows: list[dict] = []
    metadata: dict[str, dict] = {}
    if "phase" in args.tasks:
        rows, meta = aggregate_task(
            "phase",
            args.seeds,
            PHASE_METRICS,
            out_dir,
            phase_balanced=args.phase_balanced,
            phase_prefix=args.phase_prefix,
        )
        all_rows.extend(rows)
        metadata["phase"] = meta
    if "dispersion" in args.tasks:
        rows, meta = aggregate_task("dispersion", args.seeds, DISP_METRICS, out_dir)
        all_rows.extend(rows)
        metadata["dispersion"] = meta

    with (out_dir / "multiseed_summary.json").open("w", encoding="utf-8") as f:
        json.dump({"metadata": metadata, "rows": all_rows}, f, indent=2)
    write_markdown(all_rows, out_dir, phase_balanced=args.phase_balanced)
    print((out_dir / "multiseed_summary.md").read_text(encoding="utf-8"))


if __name__ == "__main__":
    main()