#!/usr/bin/env python3
"""Aggregate raw benchmark samples across fresh device processes."""

from __future__ import annotations

import argparse
import hashlib
import json
import math
import statistics
from collections import defaultdict
from pathlib import Path
from typing import Iterable


def distribution(values: Iterable[float]) -> dict[str, float]:
    ordered = sorted(float(value) for value in values)
    if not ordered:
        raise SystemExit("cannot summarize an empty distribution")
    if any(not math.isfinite(value) or value <= 0 for value in ordered):
        raise SystemExit("timing samples must be finite and positive")

    def quantile(q: float) -> float:
        at = q * (len(ordered) - 1)
        lower = int(at)
        upper = min(lower + 1, len(ordered) - 1)
        return ordered[lower] + (ordered[upper] - ordered[lower]) * (at - lower)

    return {
        "mean": statistics.fmean(ordered),
        "median": quantile(0.5),
        "p25": quantile(0.25),
        "p75": quantile(0.75),
        "min": ordered[0],
        "max": ordered[-1],
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--records", required=True, type=Path)
    parser.add_argument("--metadata", required=True, type=Path)
    parser.add_argument("--output-json", required=True, type=Path)
    parser.add_argument("--output-markdown", required=True, type=Path)
    args = parser.parse_args()

    artifact = json.loads(args.records.read_text(encoding="utf-8-sig"))
    metadata = json.loads(args.metadata.read_text(encoding="utf-8-sig"))
    if int(metadata["processes"]) < 3:
        raise SystemExit("paper protocol requires at least three processes")
    expected_samples = int(metadata["retained_samples_per_workload"])
    if expected_samples < 20:
        raise SystemExit("paper protocol requires at least 20 samples per process")

    grouped: dict[str, list[dict[str, object]]] = defaultdict(list)
    for record in artifact.get("records", []):
        if record.get("kind") != "dinovision_benchmark":
            raise SystemExit("unexpected benchmark record kind")
        samples = record.get("samples_ms", [])
        if len(samples) != expected_samples:
            raise SystemExit(
                f"{record.get('label')} has {len(samples)} samples, expected {expected_samples}"
            )
        computed = distribution(float(value) for value in samples)
        for field, key in (
            ("median_ms", "median"),
            ("p25_ms", "p25"),
            ("p75_ms", "p75"),
            ("min_ms", "min"),
            ("max_ms", "max"),
            ("mean_ms", "mean"),
        ):
            if not math.isclose(
                float(record[field]), computed[key], rel_tol=0.0, abs_tol=1e-4
            ):
                raise SystemExit(
                    f"{record.get('label')} has inconsistent {field}: "
                    f"{record[field]} vs {computed[key]} from raw samples"
                )
        grouped[str(record["label"])].append(record)

    workloads: list[dict[str, object]] = []
    for label, records in sorted(grouped.items()):
        if len(records) != int(metadata["processes"]):
            raise SystemExit(
                f"{label!r} occurs in {len(records)} processes, expected {metadata['processes']}"
            )
        sources = {str(record.get("artifact_source")) for record in records}
        if len(sources) != len(records):
            raise SystemExit(f"{label!r} does not have one record per distinct process log")
        macs = {int(record["macs"]) for record in records}
        if len(macs) != 1:
            raise SystemExit(f"{label!r} has inconsistent MAC counts")
        mac_count = next(iter(macs))
        samples = [float(value) for record in records for value in record["samples_ms"]]
        timings = distribution(samples)
        process_medians = distribution(float(record["median_ms"]) for record in records)
        median_ms = timings["median"]
        workloads.append(
            {
                "label": label,
                "macs": mac_count,
                "processes": len(records),
                "samples": len(samples),
                "timing_ms": timings,
                "process_median_ms": process_medians,
                "median_gflops": (2.0 * mac_count) / (median_ms / 1000.0) / 1e9,
            }
        )

    private_metadata = {
        "device_serial",
        "binary",
        "recovery_manifest",
        "recovery_manifest_sha256",
    }
    public_metadata = {
        key: value for key, value in metadata.items() if key not in private_metadata
    }
    result = {
        "schema_version": 1,
        "benchmark_metadata": public_metadata,
        "records_sha256": hashlib.sha256(args.records.read_bytes()).hexdigest(),
        "workloads": workloads,
    }
    args.output_json.parent.mkdir(parents=True, exist_ok=True)
    args.output_markdown.parent.mkdir(parents=True, exist_ok=True)
    args.output_json.write_text(json.dumps(result, indent=2) + "\n", encoding="utf-8")

    lines = [
        "# Isolated Quest benchmark",
        "",
        f"Condition: `{metadata['condition']}`. Each workload has {metadata['processes']} processes x {expected_samples} retained samples after {metadata['warmups_per_process']} warmups/process.",
        "",
        "| Workload | Pooled median [IQR] (ms) | Process medians: median [min, max] (ms) | Pooled min-max (ms) | Median GFLOP/s |",
        "|---|---:|---:|---:|---:|",
    ]
    for workload in workloads:
        timing = workload["timing_ms"]
        lines.append(
            f"| {workload['label']} "
            f"| {timing['median']:.2f} [{timing['p25']:.2f}, {timing['p75']:.2f}] "
            f"| {workload['process_median_ms']['median']:.2f} "
            f"[{workload['process_median_ms']['min']:.2f}, {workload['process_median_ms']['max']:.2f}] "
            f"| {timing['min']:.2f}-{timing['max']:.2f} "
            f"| {workload['median_gflops']:.1f} |"
        )
    args.output_markdown.write_text("\n".join(lines) + "\n", encoding="utf-8")
    print(f"wrote {args.output_json} and {args.output_markdown}")


if __name__ == "__main__":
    main()