#!/usr/bin/env python3
"""Load results.json from multiple experiments and produce a comparison report.

Usage:
    python3 scripts/analysis/compare_experiments.py
    python3 scripts/analysis/compare_experiments.py --output comparison.md
"""

import argparse
import json
from collections import defaultdict
from pathlib import Path

REPO_ROOT = Path(__file__).resolve().parents[2]


def load_results():
    """Glob experiments/*/results/*/results.json and load them."""
    results = []
    for path in sorted(REPO_ROOT.glob("experiments/*/results/*/results.json")):
        try:
            with open(path, "r", encoding="utf-8") as f:
                data = json.load(f)
            data["_source"] = str(path.relative_to(REPO_ROOT))
            results.append(data)
        except (json.JSONDecodeError, OSError) as e:
            print(f"WARN: failed to load {path}: {e}")
    return results


def main():
    parser = argparse.ArgumentParser(description="Compare benchmark experiments.")
    parser.add_argument("--output", "-o", default="comparison_report.md", help="Output markdown file")
    args = parser.parse_args()

    results = load_results()
    if not results:
        print("No experiments found under experiments/*/results/*/results.json")
        return

    # Group by scenario name.
    by_scenario = defaultdict(list)
    for data in results:
        meta = data.get("metadata", {})
        for scenario in data.get("scenarios", []):
            key = scenario.get("name", "unknown")
            by_scenario[key].append((meta, scenario))

    output_path = Path(args.output)
    with open(output_path, "w", encoding="utf-8") as f:
        f.write("# Cross-Experiment Comparison\n\n")
        f.write("| Experiment | Chip | Engine | Scenario | Conc | In/Out | Req/s | OutTok/s | TTFT p99 | TPOT p99 | E2E p99 |\n")
        f.write("|---|---|---|---|---:|---:|---:|---:|---:|---:|---:|\n")

        for scenario_name in sorted(by_scenario.keys()):
            for meta, scenario in by_scenario[scenario_name]:
                lat = scenario.get("latencies", {})
                f.write(
                    f"| {meta.get('experiment', '')} "
                    f"| {meta.get('chip', '')} "
                    f"| {meta.get('engine', '')} "
                    f"| {scenario_name} "
                    f"| {scenario.get('concurrency', '')} "
                    f"| {scenario.get('input_len', '')}/{scenario.get('output_len', '')} "
                    f"| {scenario.get('request_throughput') or ''} "
                    f"| {scenario.get('output_token_throughput') or ''} "
                    f"| {lat.get('ttft_ms', {}).get('p99') or ''} "
                    f"| {lat.get('tpot_ms', {}).get('p99') or ''} "
                    f"| {lat.get('e2e_ms', {}).get('p99') or ''} |\n"
                )

    print(f"Loaded {len(results)} experiment runs.")
    print(f"Wrote {output_path}")


if __name__ == "__main__":
    main()