from __future__ import annotations

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


def _as_float(value: str | None) -> float | None:
    if value is None or value == "":
        return None
    try:
        return float(value)
    except ValueError:
        return None


def _load_rows(root: Path) -> list[dict[str, Any]]:
    rows: list[dict[str, Any]] = []
    for csv_path in sorted(root.glob("*/taonet_real_token_benchmark.csv")):
        variant = csv_path.parent.name
        with csv_path.open("r", newline="", encoding="utf-8") as handle:
            for row in csv.DictReader(handle):
                row = dict(row)
                row["variant"] = variant
                rows.append(row)
    return rows


def _best_forward_backward(rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
    candidates = [row for row in rows if row.get("mode") == "forward_backward"]
    grouped: dict[str, list[dict[str, Any]]] = {}
    for row in candidates:
        grouped.setdefault(row["variant"], []).append(row)

    best_rows = []
    for variant, items in grouped.items():
        items.sort(
            key=lambda row: (
                _as_float(row.get("eval_loss")) if _as_float(row.get("eval_loss")) is not None else float("inf"),
                -(_as_float(row.get("eval_accuracy")) or 0.0),
            )
        )
        best_rows.append(items[0])
    best_rows.sort(
        key=lambda row: (
            _as_float(row.get("eval_loss")) if _as_float(row.get("eval_loss")) is not None else float("inf"),
            -(_as_float(row.get("eval_accuracy")) or 0.0),
        )
    )
    return best_rows


def _project(row: dict[str, Any]) -> dict[str, Any]:
    keys = [
        "variant",
        "architecture",
        "hybrid_pattern",
        "batch_size",
        "seq_len",
        "total_params",
        "ssm_core",
        "ssm_hidden_dim",
        "ssm_mixer_dim",
        "ssm_num_lanes",
        "ssm_lane_mode",
        "ssm_split_mix",
        "tokens_per_s_mean",
        "eval_loss",
        "eval_perplexity",
        "eval_accuracy",
        "train_final_loss",
        "train_seconds",
        "peak_reserved_mb",
        "case_id",
        "checkpoint_path",
    ]
    return {key: row.get(key, "") for key in keys}


def _write_markdown(summary: list[dict[str, Any]], path: Path) -> None:
    headers = [
        "variant",
        "architecture",
        "batch",
        "params",
        "eval_loss",
        "eval_acc",
        "tok/s",
        "checkpoint",
    ]
    lines = [
        "# TaoNet Benchmark Suite Summary",
        "",
        "| " + " | ".join(headers) + " |",
        "| " + " | ".join(["---"] * len(headers)) + " |",
    ]
    for row in summary:
        lines.append(
            "| "
            + " | ".join(
                [
                    str(row["variant"]),
                    str(row["architecture"]),
                    str(row["batch_size"]),
                    str(row["total_params"]),
                    str(row["eval_loss"]),
                    str(row["eval_accuracy"]),
                    str(row["tokens_per_s_mean"]),
                    str(row["checkpoint_path"]),
                ]
            )
            + " |"
        )
    path.write_text("\n".join(lines) + "\n", encoding="utf-8")


def main() -> None:
    parser = argparse.ArgumentParser(description="Summarize a TaoNet benchmark suite output directory.")
    parser.add_argument("--suite-dir", required=True, help="Directory containing one subdirectory per benchmark variant.")
    parser.add_argument("--output-json", default="", help="Summary JSON path. Defaults to <suite-dir>/suite_summary.json.")
    parser.add_argument("--output-md", default="", help="Summary Markdown path. Defaults to <suite-dir>/suite_summary.md.")
    args = parser.parse_args()

    suite_dir = Path(args.suite_dir)
    rows = _load_rows(suite_dir)
    summary = [_project(row) for row in _best_forward_backward(rows)]
    json_path = Path(args.output_json) if args.output_json else suite_dir / "suite_summary.json"
    md_path = Path(args.output_md) if args.output_md else suite_dir / "suite_summary.md"
    json_path.write_text(json.dumps(summary, indent=2) + "\n", encoding="utf-8")
    _write_markdown(summary, md_path)
    print(f"Wrote {json_path}")
    print(f"Wrote {md_path}")


if __name__ == "__main__":
    main()