"""Read the saved timing files and print the three formats side by side.

Each runner writes `bench/<tag>_timing.json`. This reads whatever is there, so
it works with two of three present -- which is the normal case while NCNN is
waiting on a Pi.

    python -m src.bench.eval_pt && python -m src.bench.eval_onnx
    python -m src.bench.compare

The row worth checking first is `depth`. It is the same numpy over the same
boxes whatever the backend, so it should be near-identical across columns. When
it is not, the run measured contention rather than the export format -- that is
exactly what happened here before onnxruntime's thread pool was pinned, where it
read 1.5 ms under PyTorch and 16.2 ms under ONNX.
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

STAGES = ("read", "preprocess", "inference", "postprocess", "depth", "total")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--bench-dir", type=Path, default=Path("bench"))
    parser.add_argument("--metric", default="p50", choices=["p50", "p95", "mean"])
    args = parser.parse_args()

    runs = {}
    for tag in ("pt", "onnx", "ncnn"):
        path = args.bench_dir / f"{tag}_timing.json"
        if path.exists():
            runs[tag] = json.loads(path.read_text())
    if not runs:
        raise SystemExit(f"no *_timing.json under {args.bench_dir}; run a benchmark first")

    tags = list(runs)
    print(f"\n{args.metric} milliseconds per frame\n")
    print(f"  {'stage':<14}" + "".join(f"{t:>12}" for t in tags))
    print("  " + "-" * (14 + 12 * len(tags)))
    for stage in STAGES:
        row = "".join(f"{runs[t]['stages'][stage][args.metric]:12.2f}" for t in tags)
        print(f"  {stage:<14}{row}")

    print("\n  " + "-" * (14 + 12 * len(tags)))
    print(f"  {'FPS':<14}"
          + "".join(f"{1000 / runs[t]['stages']['total'][args.metric]:12.2f}" for t in tags))
    print(f"  {'size (MB)':<14}" + "".join(f"{runs[t]['size_mb']:12.1f}" for t in tags))
    print(f"  {'load (ms)':<14}" + "".join(f"{runs[t]['load_ms']:12.0f}" for t in tags))
    print(f"  {'cores':<14}" + "".join(f"{str(runs[t].get('cpu_affinity','?')):>12}" for t in tags))

    # Relative to the first column, which is PyTorch when it is present.
    base = tags[0]
    print(f"\n  inference speedup vs {base}:")
    for tag in tags:
        ratio = (runs[base]["stages"]["inference"][args.metric]
                 / runs[tag]["stages"]["inference"][args.metric])
        print(f"    {tag:<8}{ratio:5.2f}x")

    spread = [runs[t]["stages"]["depth"][args.metric] for t in tags]
    if len(spread) > 1 and (max(spread) - min(spread)) > 0.25 * max(spread):
        print(f"\n  !! depth varies {min(spread):.2f}-{max(spread):.2f} ms between formats.")
        print(f"     It is identical code on the same boxes, so this run measured")
        print(f"     CPU contention, not the export format. Pin --threads.")


if __name__ == "__main__":
    main()