#!/usr/bin/env python3
"""Compare direct EasyCom floor-event probe result directories."""

from __future__ import annotations

import argparse
import csv
import json
from pathlib import Path
from typing import Any

import numpy as np


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--run",
        action="append",
        required=True,
        metavar="NAME=RESULT_DIR",
    )
    parser.add_argument("--output", type=Path, required=True)
    return parser.parse_args()


def parse_run(value: str) -> tuple[str, Path]:
    if "=" not in value:
        raise ValueError(f"Expected NAME=RESULT_DIR, got {value!r}")
    name, raw_path = value.split("=", 1)
    if not name.strip() or not raw_path.strip():
        raise ValueError(f"Expected NAME=RESULT_DIR, got {value!r}")
    return name.strip(), Path(raw_path.strip())


def mean_ap(rows: list[dict[str, Any]], event_type: str) -> float | None:
    values = [
        float(row["average_precision"])
        for row in rows
        if row["event_type"] == event_type
        and row.get("average_precision") is not None
    ]
    return float(np.mean(values)) if values else None


def horizon_row(
    rows: list[dict[str, Any]], event_type: str, horizon_index: int
) -> dict[str, Any]:
    matches = [
        row
        for row in rows
        if row["event_type"] == event_type
        and int(row["horizon_index"]) == horizon_index
    ]
    if len(matches) != 1:
        raise ValueError(
            f"Expected one {event_type} horizon {horizon_index}, got {len(matches)}"
        )
    return matches[0]


def main() -> None:
    args = parse_args()
    parsed = [parse_run(value) for value in args.run]
    if len({name for name, _ in parsed}) != len(parsed):
        raise ValueError("Duplicate run name")
    output_rows: list[dict[str, Any]] = []
    common_contract: tuple[Any, ...] | None = None
    for name, directory in parsed:
        path = directory / "summary.json"
        if not path.is_file():
            raise FileNotFoundError(path)
        payload = json.loads(path.read_text())
        if payload.get("protocol") != "easycom_direct_floor_event_probe_v1":
            raise ValueError(f"Unexpected protocol in {path}")
        contract = (
            tuple(payload["eval_sessions"]),
            payload["decision_timing"],
            float(payload["matched_recall"]),
        )
        if common_contract is None:
            common_contract = contract
        elif contract != common_contract:
            raise ValueError(
                f"Incomparable evaluation contract in {path}: {contract} vs "
                f"{common_contract}"
            )
        rows = payload["eval_metrics"]
        response_2s = horizon_row(rows, "response", 10)
        yield_2s = horizon_row(rows, "yield", 10)
        result = {
            "run": name,
            "oracle_condition": payload["oracle_condition"],
            "decision_timing": payload["decision_timing"],
            "eval_sessions": ",".join(payload["eval_sessions"]),
            "best_epoch": payload["best_epoch"],
            "mean_response_ap": mean_ap(rows, "response"),
            "mean_yield_ap": mean_ap(rows, "yield"),
            "response_2s_ap": response_2s["average_precision"],
            "response_2s_recall": response_2s["achieved_positive_recall"],
            "response_2s_false_entry_rate": response_2s[
                "negative_trigger_rate"
            ],
            "response_2s_short_resume_trigger_rate": response_2s[
                "short_resume_trigger_rate"
            ],
            "yield_2s_ap": yield_2s["average_precision"],
            "yield_2s_recall": yield_2s["achieved_positive_recall"],
            "yield_2s_false_stop_rate": yield_2s["negative_trigger_rate"],
        }
        output_rows.append(result)

    args.output.parent.mkdir(parents=True, exist_ok=True)
    with args.output.open("w", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=list(output_rows[0]))
        writer.writeheader()
        writer.writerows(output_rows)
    print(
        "run | mean response AP | response 2s AP/error | "
        "mean yield AP | yield 2s AP/error"
    )
    for row in output_rows:
        print(
            f"{row['run']} | {row['mean_response_ap']} | "
            f"{row['response_2s_ap']}/{row['response_2s_false_entry_rate']} | "
            f"{row['mean_yield_ap']} | "
            f"{row['yield_2s_ap']}/{row['yield_2s_false_stop_rate']}"
        )
    print(f"comparison: {args.output}")


if __name__ == "__main__":
    main()