#!/usr/bin/env python3
"""Aggregate completed per-dataset metric CSVs without relabeling metric columns."""

from __future__ import annotations

import argparse
import csv
import json
from pathlib import Path

METRICS = [
    ("psnr", True), ("ssim", True), ("lpips", False), ("dists", False),
    ("niqe", False), ("musiq", True), ("maniqa_pipal", True),
    ("topiq_nr", True), ("clipiqa", True),
]
ARROW = {True: "↑", False: "↓"}


def read_mean(path: Path) -> dict[str, float | None] | None:
    if not path.is_file() or not path.stat().st_size:
        return None
    with path.open(newline="", encoding="utf-8") as handle:
        rows = csv.DictReader(handle)
        mean = next((row for row in rows if row.get("image") == "MEAN"), None)
    if mean is None:
        return None
    result: dict[str, float | None] = {}
    for metric, _ in METRICS:
        try:
            result[metric] = float(mean[metric])
        except (KeyError, TypeError, ValueError):
            result[metric] = None
    return result


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--metrics-root", type=Path, help="metrics directory paired with current comparison_run.json")
    parser.add_argument("--metric-csv", type=Path, action="append", default=[], help="current metric CSV; repeat per method")
    parser.add_argument("--dataset", required=True)
    parser.add_argument("--out-dir", type=Path, required=True)
    args = parser.parse_args()
    if bool(args.metrics_root) == bool(args.metric_csv):
        parser.error("provide exactly one of --metrics-root or --metric-csv")
    if args.metric_csv:
        metric_paths = [(path.parent.name, path) for path in args.metric_csv]
    else:
        if not args.metrics_root.is_dir():
            parser.error(f"metrics root does not exist: {args.metrics_root}")
        record_path = args.metrics_root.parent / "comparison_run.json"
        try:
            record = json.loads(record_path.read_text(encoding="utf-8"))
            completed = record["completed"]
        except (OSError, ValueError, KeyError, TypeError) as error:
            parser.error(f"current comparison record is required for --metrics-root: {record_path} ({error})")
        if not isinstance(completed, list):
            parser.error(f"invalid completed records in {record_path}")
        metrics_root = args.metrics_root.resolve()
        metric_paths = []
        for completion in completed:
            if not isinstance(completion, dict) or completion.get("dataset") != args.dataset:
                continue
            if not completion.get("parity_complete") or not isinstance(completion.get("method"), str):
                continue
            metric_csv = completion.get("metric_csv")
            if not isinstance(metric_csv, str):
                parser.error(f"invalid metric CSV in {record_path}")
            path = Path(metric_csv).resolve()
            try:
                path.relative_to(metrics_root)
            except ValueError:
                parser.error(f"recorded metric CSV is outside metrics root: {path}")
            metric_paths.append((completion["method"], path))
    rows = []
    seen_methods = set()
    for method, path in metric_paths:
        if method in seen_methods:
            parser.error(f"duplicate metric CSV for method: {method}")
        seen_methods.add(method)
        values = read_mean(path)
        if values is not None:
            rows.append({"method": method, **values})
    if not rows:
        raise SystemExit(f"no complete metric CSVs for {args.dataset}")
    args.out_dir.mkdir(parents=True, exist_ok=True)
    csv_path = args.out_dir / f"comparison_{args.dataset}.csv"
    with csv_path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.writer(handle)
        writer.writerow(["method", *[name for name, _ in METRICS]])
        for row in rows:
            writer.writerow([row["method"], *[row[name] for name, _ in METRICS]])
    table_rows = sorted(
        rows,
        key=lambda row: (row["lpips"] is None, row["lpips"] if row["lpips"] is not None else float("inf")),
    )
    method_width = max(len("method"), *(len(row["method"]) for row in rows))
    header = f"{'method':<{method_width}}  " + "  ".join(
        f"{name + ARROW[higher]:>14}" for name, higher in METRICS
    )
    lines = [header, "-" * len(header)]
    for row in table_rows:
        cells = [f"{row[name]:>14.4f}" if row[name] is not None else f"{'NA':>14}" for name, _ in METRICS]
        lines.append(f"{row['method']:<{method_width}}  " + "  ".join(cells))
    table_path = args.out_dir / f"comparison_{args.dataset}.txt"
    table_path.write_text("\n".join(lines) + "\n", encoding="utf-8")
    print(f"{len(rows)} methods -> {table_path}")


if __name__ == "__main__":
    main()