#!/usr/bin/env python3
"""Summarize scored PlantGeneBench rows."""

from __future__ import annotations

import argparse
import csv
import json
from collections import Counter, defaultdict
from pathlib import Path
from typing import Any, Dict, Iterable, List


def load_jsonl(path: Path) -> List[Dict[str, Any]]:
    rows: List[Dict[str, Any]] = []
    with path.open("r", encoding="utf-8") as handle:
        for lineno, line in enumerate(handle, start=1):
            line = line.strip()
            if not line:
                continue
            try:
                rows.append(json.loads(line))
            except json.JSONDecodeError as exc:
                raise SystemExit(f"{path}:{lineno}: invalid JSON: {exc}") from exc
    return rows


def task_metadata(tasks_path: Path | None) -> Dict[str, Dict[str, Any]]:
    if not tasks_path:
        return {}
    if not tasks_path.exists():
        return {}
    return {row["task_id"]: row for row in load_jsonl(tasks_path)}


def mean(values: Iterable[float]) -> float:
    values = list(values)
    if not values:
        return 0.0
    return sum(values) / len(values)


def write_rows_csv(path: Path, rows: List[Dict[str, Any]], tasks: Dict[str, Dict[str, Any]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    fields = [
        "task_id",
        "category",
        "difficulty",
        "score",
        "stage_f1",
        "forbidden_ok",
        "params",
        "dependencies",
        "analysis_type",
        "errors",
        "expected_stages",
        "predicted_stages",
    ]
    with path.open("w", encoding="utf-8", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=fields)
        writer.writeheader()
        for row in rows:
            meta = tasks.get(row["task_id"], {})
            components = row.get("components") or {}
            writer.writerow({
                "task_id": row["task_id"],
                "category": meta.get("category", ""),
                "difficulty": meta.get("difficulty", ""),
                "score": row.get("score", 0.0),
                "stage_f1": components.get("stage_f1", 0.0),
                "forbidden_ok": components.get("forbidden_ok", 0.0),
                "params": components.get("params", 0.0),
                "dependencies": components.get("dependencies", 0.0),
                "analysis_type": components.get("analysis_type", 0.0),
                "errors": ";".join(row.get("errors") or []),
                "expected_stages": ",".join(row.get("expected_stages") or []),
                "predicted_stages": ",".join(row.get("predicted_stages") or []),
            })


def grouped_scores(rows: List[Dict[str, Any]], tasks: Dict[str, Dict[str, Any]], key: str) -> Dict[str, List[float]]:
    grouped: Dict[str, List[float]] = defaultdict(list)
    for row in rows:
        label = str(tasks.get(row["task_id"], {}).get(key) or "unknown")
        grouped[label].append(float(row.get("score") or 0.0))
    return grouped


def write_summary_md(path: Path, rows: List[Dict[str, Any]], tasks: Dict[str, Dict[str, Any]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    error_counter: Counter[str] = Counter()
    for row in rows:
        error_counter.update(row.get("errors") or [])

    lines: List[str] = []
    lines.append("# PlantGeneBench Summary")
    lines.append("")
    lines.append(f"- Tasks: {len(rows)}")
    lines.append(f"- Mean score: {mean(float(row.get('score') or 0.0) for row in rows):.3f}")
    lines.append(f"- Perfect tasks: {sum(1 for row in rows if float(row.get('score') or 0.0) >= 0.999)}")
    lines.append("")

    for group_key in ("category", "difficulty"):
        lines.append(f"## By {group_key}")
        lines.append("")
        lines.append("| Group | N | Mean score |")
        lines.append("|---|---:|---:|")
        for label, values in sorted(grouped_scores(rows, tasks, group_key).items()):
            lines.append(f"| {label} | {len(values)} | {mean(values):.3f} |")
        lines.append("")

    lines.append("## Error Types")
    lines.append("")
    lines.append("| Error | Count |")
    lines.append("|---|---:|")
    for error, count in error_counter.most_common():
        lines.append(f"| {error} | {count} |")
    if not error_counter:
        lines.append("| none | 0 |")
    lines.append("")

    path.write_text("\n".join(lines), encoding="utf-8")


def run(args: argparse.Namespace) -> int:
    rows = load_jsonl(args.scored)
    tasks = task_metadata(args.tasks)
    csv_path = Path(f"{args.out_prefix}.rows.csv")
    md_path = Path(f"{args.out_prefix}.summary.md")
    write_rows_csv(csv_path, rows, tasks)
    write_summary_md(md_path, rows, tasks)
    print(f"Wrote {csv_path}")
    print(f"Wrote {md_path}")
    return 0


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--scored", type=Path, required=True)
    parser.add_argument("--tasks", type=Path, default=Path("benchmarks/plantgenebench/tasks/plantgenebench_core_seed.jsonl"))
    parser.add_argument("--out-prefix", type=Path, required=True)
    return parser.parse_args()


if __name__ == "__main__":
    raise SystemExit(run(parse_args()))