#!/usr/bin/env python
"""Step 8: SQNR summary table of every saved run (results/outputs_<name>.npy) against the FP32 reference
(and, as a second view, against the FP16 ONNX baseline). Writes results/summary.md and summary.json.

usage: 08_summary.py [--ref onnx_fp32] [--ref2 onnx_fp16] [--names a b c ...]
"""
import argparse, glob, json, os, sys
import numpy as np
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from metrics import compute_metrics, fmt_table

ROOT = os.environ.get("ROOT", "/data/users/logesh/Infernece_vision_Manual")
ap = argparse.ArgumentParser()
ap.add_argument("--ref", default="onnx_fp32")
ap.add_argument("--ref2", default="onnx_fp16")
ap.add_argument("--names", nargs="*", default=None)
ap.add_argument("--results_dir", default=f"{ROOT}/results")
a = ap.parse_args()

names = a.names or sorted(os.path.basename(p)[len("outputs_"):-4] for p in glob.glob(f"{a.results_dir}/outputs_*.npy"))
refs = {}
for r in (a.ref, a.ref2):
    p = f"{a.results_dir}/outputs_{r}.npy"
    if r and os.path.exists(p):
        refs[r] = np.load(p)
if not refs:
    sys.exit(f"no reference outputs found in {a.results_dir} (run 05_eval_onnx.py --name {a.ref} first)")

report, allj = [], {}
for rname, ref in refs.items():
    rows = {}
    for n in names:
        if n == rname:
            continue
        t = np.load(f"{a.results_dir}/outputs_{n}.npy")
        N = min(len(ref), len(t))
        rows[n] = compute_metrics(ref[:N], t[:N])
    report.append(f"### reference: {rname}   ({ref.shape[0]} images x {ref.shape[1]} tokens x {ref.shape[2]})\n\n" + fmt_table(rows))
    allj[rname] = rows
text = "\n\n".join(report)
print(text)
open(f"{a.results_dir}/summary.md", "w").write(text + "\n")
json.dump(allj, open(f"{a.results_dir}/summary.json", "w"), indent=2)
print(f"\nwritten {a.results_dir}/summary.md")