#!/usr/bin/env python3
"""benchmark_outputs/ 아래 실험들의 seedtts_en WER/SIM summary를 CSV로 모은다."""

from __future__ import annotations

import argparse
import csv
import re
from pathlib import Path

# CKPT_TAG 규칙: run_seed_for_loop.sh line 33
#   <experiment_name>_step_<step>_CFG<cfg>_STEPS<steps>_CFG_ZERO_STAR_<z>_CFG_SPK_<s>_ZERO_PATCH_CONTEXT_<p>_SIMULATE_NO_CONT_<c>
TAG_RE = re.compile(
    r"^(?P<experiment_name>.+?)"
    r"_step_(?P<step>\d+)"
    r"_CFG(?P<CFG>[\d.]+)"
    r"_STEPS(?P<STEPS>\d+)"
    r"_CFG_ZERO_STAR_(?P<CFG_ZERO_STAR>true|false)"
    r"_CFG_SPK_(?P<CFG_SPK>true|false)"
    r"_ZERO_PATCH_CONTEXT_(?P<ZERO_PATCH_CONTEXT>true|false)"
    r"_SIMULATE_NO_CONT_(?P<SIMULATE_NO_CONT>true|false)$"
)


def parse_wer_summary(path: Path) -> tuple[float | None, int]:
    """returns (wer_pct, n_utt). wer_pct None이면 summary 라인 미발견."""
    if not path.is_file():
        return None, 0
    lines = [ln for ln in path.read_text(encoding="utf-8", errors="replace").splitlines() if ln.strip()]
    if not lines:
        return None, 0
    last = lines[-1]
    wer_pct: float | None = None
    if last.startswith("WER:"):
        # "WER: 29.287%"
        raw = last[len("WER:"):].strip().rstrip("%").strip()
        try:
            wer_pct = float(raw)
        except ValueError:
            wer_pct = None
    # header 1 + summary 1 (있을 때만) 빼기
    n_utt = len(lines) - 1 - (1 if wer_pct is not None else 0)
    return wer_pct, max(n_utt, 0)


def parse_sim_summary(path: Path) -> tuple[float | None, float | None]:
    """returns (asv_mean, asv_var). 없는 값은 None."""
    if not path.is_file():
        return None, None
    asv_mean: float | None = None
    asv_var: float | None = None
    for ln in reversed(path.read_text(encoding="utf-8", errors="replace").splitlines()):
        ln = ln.strip()
        if not ln:
            continue
        if asv_var is None and ln.startswith("ASV-var:"):
            try:
                asv_var = float(ln[len("ASV-var:"):].strip())
            except ValueError:
                pass
        elif asv_mean is None and ln.startswith("ASV:"):
            try:
                asv_mean = float(ln[len("ASV:"):].strip())
            except ValueError:
                pass
        if asv_mean is not None and asv_var is not None:
            break
    return asv_mean, asv_var


def parse_tag(tag: str) -> dict[str, str]:
    """CKPT_TAG 디렉터리 이름을 분해. 매칭 실패하면 빈 dict."""
    m = TAG_RE.match(tag)
    if not m:
        return {}
    return m.groupdict()


def collect_rows(root: Path) -> list[dict]:
    rows: list[dict] = []
    for child in sorted(root.iterdir()):
        if not child.is_dir() or child.name == "viewer":
            continue
        seed_dir = child / "seedtts_en"
        wer_path = seed_dir / "seedtts_wer.txt"
        sim_path = seed_dir / "seedtts_sim.txt"
        wer_pct, n_utt = parse_wer_summary(wer_path)
        asv_mean, asv_var = parse_sim_summary(sim_path)

        if wer_pct is None and asv_mean is None:
            status = "missing" if not seed_dir.is_dir() else "incomplete"
        elif wer_pct is None or asv_mean is None:
            status = "incomplete"
        else:
            status = "ok"

        parsed = parse_tag(child.name)
        rows.append({
            "tag": child.name,
            "experiment_name": parsed.get("experiment_name", ""),
            "step": parsed.get("step", ""),
            "CFG": parsed.get("CFG", ""),
            "STEPS": parsed.get("STEPS", ""),
            "CFG_ZERO_STAR": parsed.get("CFG_ZERO_STAR", ""),
            "CFG_SPK": parsed.get("CFG_SPK", ""),
            "ZERO_PATCH_CONTEXT": parsed.get("ZERO_PATCH_CONTEXT", ""),
            "SIMULATE_NO_CONT": parsed.get("SIMULATE_NO_CONT", ""),
            "n_utt": n_utt if n_utt > 0 else "",
            "wer_pct": "" if wer_pct is None else f"{wer_pct:.3f}",
            "sim_asv": "" if asv_mean is None else f"{asv_mean:.3f}",
            "sim_asv_var": "" if asv_var is None else f"{asv_var:.3f}",
            "status": status,
        })
    return rows


def sort_key(row: dict):
    try:
        cfg = float(row["CFG"])
    except ValueError:
        cfg = float("inf")
    return (row["experiment_name"], row["step"], row["CFG_SPK"], cfg)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--root",
        type=Path,
        default=Path(__file__).resolve().parent.parent / "benchmark_outputs",
        help="benchmark_outputs 루트 경로",
    )
    parser.add_argument(
        "--out",
        type=Path,
        default=None,
        help="CSV 출력 경로 (default: <root>/benchmark_summary.csv)",
    )
    args = parser.parse_args()

    root: Path = args.root.resolve()
    if not root.is_dir():
        raise SystemExit(f"root not found: {root}")
    out_path: Path = (args.out or root / "benchmark_summary.csv").resolve()

    rows = collect_rows(root)
    rows.sort(key=sort_key)

    fieldnames = [
        "tag",
        "experiment_name",
        "step",
        "CFG",
        "STEPS",
        "CFG_ZERO_STAR",
        "CFG_SPK",
        "ZERO_PATCH_CONTEXT",
        "SIMULATE_NO_CONT",
        "n_utt",
        "wer_pct",
        "sim_asv",
        "sim_asv_var",
        "status",
    ]
    out_path.parent.mkdir(parents=True, exist_ok=True)
    with out_path.open("w", encoding="utf-8", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(rows)

    print(f"wrote {len(rows)} rows -> {out_path}")


if __name__ == "__main__":
    main()