#!/usr/bin/env python3
"""Verify GRL revision analysis outputs."""

from __future__ import annotations

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


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 require(path: Path, errors: list[str]) -> None:
    if not path.exists():
        errors.append(f"missing: {path}")


def numeric_ok(value: str) -> bool:
    if value == "":
        return True
    try:
        return math.isfinite(float(value))
    except ValueError:
        return True


def check_numeric(rows: list[dict[str, str]], path: Path, errors: list[str]) -> None:
    for i, row in enumerate(rows, start=2):
        for key, value in row.items():
            if not numeric_ok(value):
                errors.append(f"non-finite value in {path}:{i} column {key}: {value}")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--revision-dir", type=Path, default=DEFAULT_REVISION_DIR)
    parser.add_argument("--allow-missing-heavy", action="store_true", help="Allow missing per-sample/bootstrap outputs before heavy checkpoint evaluation has been run.")
    args = parser.parse_args()

    d = args.revision_dir
    errors: list[str] = []
    heavy_files = [
        d / "phase_picking" / "phase_per_window_outputs.csv.gz",
        d / "dispersion" / "dispersion_per_sample_metrics.csv.gz",
        d / "tables" / "phase_bootstrap_ci.csv",
        d / "tables" / "dispersion_bootstrap_ci.csv",
        d / "tables" / "phase_snr_stratified_metrics.csv",
        d / "tables" / "dispersion_snr_stratified_metrics.csv",
    ]
    light_files = [
        d / "tables" / "association_summary.csv",
        d / "tables" / "association_summary_for_manuscript.txt",
    ]
    for path in light_files:
        require(path, errors)
    if not args.allow_missing_heavy:
        for path in heavy_files:
            require(path, errors)

    phase_path = d / "phase_picking" / "phase_per_window_outputs.csv.gz"
    if phase_path.exists():
        rows = read_csv(phase_path)
        check_numeric(rows, phase_path, errors)
        sample_ids = {row["sample_id"] for row in rows}
        expected_rows = 10000 * 3 * 2
        if len(sample_ids) != 10000:
            errors.append(f"phase sample count expected 10000, found {len(sample_ids)}")
        if len(rows) != expected_rows:
            errors.append(f"phase row count expected {expected_rows}, found {len(rows)}")

    disp_path = d / "dispersion" / "dispersion_per_sample_metrics.csv.gz"
    if disp_path.exists():
        rows = read_csv(disp_path)
        check_numeric(rows, disp_path, errors)
        sample_ids = {row["sample_id"] for row in rows}
        expected_rows = 8292 * 3
        if len(sample_ids) != 8292:
            errors.append(f"dispersion sample count expected 8292, found {len(sample_ids)}")
        if len(rows) != expected_rows:
            errors.append(f"dispersion row count expected {expected_rows}, found {len(rows)}")

    assoc_path = d / "tables" / "association_summary.csv"
    if assoc_path.exists():
        rows = read_csv(assoc_path)
        check_numeric(rows, assoc_path, errors)
        by_condition = {row["condition"]: row for row in rows}
        expected = {"snr": ("1301", "2340"), "confidence": ("1561", "2340")}
        for condition, (tp, ref) in expected.items():
            row = by_condition.get(condition)
            if row is None:
                errors.append(f"association summary missing condition {condition}")
                continue
            if row["event_true_positive"] != tp or row["event_reference"] != ref:
                errors.append(f"association counts for {condition} expected {tp}/{ref}, found {row['event_true_positive']}/{row['event_reference']}")
            if row["retained_picks"] != "576875":
                errors.append(f"association retained picks for {condition} expected 576875, found {row['retained_picks']}")

    if errors:
        print("Verification failed:")
        for err in errors:
            print(f"  - {err}")
        raise SystemExit(1)
    print(f"Verification passed for {d}")


if __name__ == "__main__":
    main()