#!/usr/bin/env python
"""What each modality pays when lambda was calibrated on a different one.

The end-to-end runs put a number on the *outcome* -- lambda=video lands in a different sample than
BF16 does while lambda=all tracks it -- but not on the mechanism. This does, from the per-modality
statistics alone, with no GPU.

NVFP4 groups 16 consecutive input channels under one FP8 scale set by that group's absmax. A
channel whose own absmax is far below its group's absmax spends its mantissa on range it never
uses, and the loss is `log2(group_absmax / channel_absmax)` bits. That quantity is what smoothing
exists to reduce: `X/lambda` is the activation the kernel actually sees, so lambda reshapes exactly
this profile.

The point is that `W * lambda` is one weight and every modality's rows pass through it. Pick lambda
from video's absmax and video's profile flattens -- but text and audio are divided by a vector that
has nothing to do with their own profile, and an uncorrelated divisor makes a profile *sharper*,
not flatter. This prints the bits lost per modality under each lambda, so "video's gain" and
"text's loss" are on the same axis.

The `lambda = 1` column is the no-smoothing floor, i.e. what every modality pays if nobody is
favoured.
"""

from __future__ import annotations

import argparse
from pathlib import Path

import torch

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


def bits_lost(absmax: torch.Tensor, lam: torch.Tensor) -> float:
    """Mean bits a channel loses to its 16-wide group's scale, after dividing by lambda."""
    a = (absmax.double() / lam.double()).clamp_min(1e-12)
    n = (a.numel() // BLOCK) * BLOCK
    g = a[:n].reshape(-1, BLOCK)
    return float(torch.log2(g.amax(dim=1, keepdim=True) / g).mean())


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--stats", required=True, help="stats_modal_768p.pt")
    ap.add_argument("--calibs", nargs="+", required=True, help="tag=path.pt, e.g. all=...pt")
    ap.add_argument("--top", type=int, default=8)
    args = ap.parse_args()

    st = torch.load(args.stats, map_location="cpu", weights_only=False)["stats"]
    cal = {}
    for spec in args.calibs:
        tag, path = spec.split("=", 1)
        cal[tag] = torch.load(path, map_location="cpu", weights_only=False)["layers"]

    tags = ["1"] + list(cal)
    rows = {m: {t: [] for t in tags} for m in MODALITIES}
    per_layer = []

    for name, s in st.items():
        if name not in next(iter(cal.values())):
            continue
        lams = {"1": torch.ones_like(s["video"]["absmax"])}
        lams.update({t: cal[t][name]["lambda"].float() for t in cal})
        rec = {"layer": name}
        for m in MODALITIES:
            am = s[m]["absmax"].float()
            if float(am.max()) == 0.0:          # audio is absent from these clips for some layers
                continue
            for t in tags:
                v = bits_lost(am, lams[t])
                rows[m][t].append(v)
                rec[f"{m}/{t}"] = v
        per_layer.append(rec)

    print(f"bits lost to the block-16 scale, mean over channels, median over layers "
          f"({len(per_layer)} layers)\n")
    head = "modality".ljust(10) + "".join(f"lam={t}".rjust(12) for t in tags)
    print(head)
    print("-" * len(head))
    for m in MODALITIES:
        if not rows[m]["1"]:
            print(f"{m:<10}" + "  (no rows in calibration data)")
            continue
        cells = []
        for t in tags:
            v = sorted(rows[m][t])
            cells.append(f"{v[len(v)//2]:.3f}".rjust(12))
        print(f"{m:<10}" + "".join(cells))

    print("\nchange vs no smoothing (negative = smoothing helps that modality)")
    print(head)
    print("-" * len(head))
    for m in MODALITIES:
        if not rows[m]["1"]:
            continue
        base = sorted(rows[m]["1"]); base = base[len(base) // 2]
        cells = []
        for t in tags:
            v = sorted(rows[m][t]); v = v[len(v) // 2]
            cells.append(("--" if t == "1" else f"{v - base:+.3f}").rjust(12))
        print(f"{m:<10}" + "".join(cells))

    if "video" in cal:
        worst = sorted(per_layer, key=lambda r: -(r.get("text/video", 0) - r.get("text/all", 0)))
        print(f"\nlayers where lambda=video costs text the most, vs lambda=all:")
        for r in worst[: args.top]:
            print(f"  {r['layer']:<48} text {r.get('text/all', float('nan')):.3f} -> "
                  f"{r.get('text/video', float('nan')):.3f}   "
                  f"video {r.get('video/all', float('nan')):.3f} -> "
                  f"{r.get('video/video', float('nan')):.3f}")
    return 0


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