#!/usr/bin/env python
"""Summarise `score_modal_lambda.py` shards: error per modality under each lambda.

The row that matters is `text`. Every lambda was chosen to minimise an error measured on rows that
are 98.6% video, so the `video` row is close to the calibration's own objective and mostly repeats
what the artifacts already say. The `text` and `audio` rows are what that objective could not see.
"""

from __future__ import annotations

import argparse
import json
import statistics as st
from pathlib import Path

MODALITIES = ("video", "text", "audio")


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--dir", required=True)
    ap.add_argument("--top", type=int, default=8)
    args = ap.parse_args()

    layers: dict[str, dict] = {}
    tags: list[str] = []
    for p in sorted(Path(args.dir).glob("score_*.json")):
        blob = json.loads(p.read_text())
        tags = blob["tags"]
        layers.update(blob["layers"])

    print(f"relative L2 vs bf16, per modality, median over {len(layers)} layers "
          f"(rank-32 branch included)\n")
    head = "modality".ljust(10) + "".join(f"lam={t}".rjust(11) for t in tags)
    print(head)
    print("-" * len(head))
    med: dict[str, dict[str, float]] = {}
    for m in MODALITIES:
        vals = {t: sorted(r[m][t] for r in layers.values() if m in r and t in r[m])
                for t in tags}
        if not vals[tags[0]]:
            print(f"{m:<10}  (no rows sampled)")
            continue
        med[m] = {t: v[len(v) // 2] for t, v in vals.items()}
        print(f"{m:<10}" + "".join(f"{med[m][t]:.4f}".rjust(11) for t in tags))

    if "none" in tags:
        print("\nchange vs lambda=1 (negative = smoothing helps that modality)")
        print(head)
        print("-" * len(head))
        for m in MODALITIES:
            if m not in med:
                continue
            b = med[m]["none"]
            cells = ("--" if t == "none" else f"{(med[m][t]-b)/b*100:+.1f}%" for t in tags)
            print(f"{m:<10}" + "".join(c.rjust(11) for c in cells))

    if "video" in tags and "all" in tags:
        worst = sorted((r for r in layers.items() if "text" in r[1]),
                       key=lambda kv: -(kv[1]["text"]["video"] - kv[1]["text"]["all"]))
        print(f"\nlayers where lambda=video costs text the most, vs lambda=all:")
        for name, r in worst[: args.top]:
            print(f"  {name:<46} text {r['text']['all']:.4f} -> {r['text']['video']:.4f}   "
                  f"video {r['video']['all']:.4f} -> {r['video']['video']:.4f}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())