#!/usr/bin/env python
"""Frame statistics across the lambda experiments, read against the BF16 reference.

Mean RGB alone was what first flagged a regression, and it is not enough: a run can match the
average while having flattened the time axis or lost spatial contrast. Per-frame spread and
per-frame std are reported next to it, plus the distance to BF16 on each.

It is also not enough in the other direction. Every run here shares BF16's seed and therefore its
initial noise, so a quantization that stays faithful lands in the *same* sample; one that perturbs
the trajectory enough can land in a different, perfectly plausible one. That shows up in mean RGB
as a large delta which says nothing about image quality -- a darker scene is not a worse scene.
`PSNR` and `corr` are the columns that separate the two cases: they are computed frame-aligned
against BF16, so "same scene, slightly degraded" and "different scene" are distinguishable, and
only the first is a quantization-quality statement.
"""
import sys
from pathlib import Path
import numpy as np
import av


def stats(p):
    c = av.open(str(p))
    # Frame-threaded decode hands `to_ndarray` a frame whose buffer swscale can still be writing
    # to, which surfaces as EAGAIN from sws_scale rather than as anything decode-shaped.
    c.streams.video[0].thread_type = "NONE"
    c.streams.video[0].thread_count = 1
    m, s = [], []
    for f in c.decode(video=0):
        a = f.to_ndarray(format="rgb24").astype(np.float32)
        m.append(a.mean()); s.append(a.std())
    c.close()
    m, s = np.array(m), np.array(s)
    return dict(n=len(m), mean=m.mean(), lo=m.min(), hi=m.max(),
                span=m.max() - m.min(), std=s.mean())


def frames(p):
    c = av.open(str(p))
    c.streams.video[0].thread_type = "NONE"
    c.streams.video[0].thread_count = 1
    out = [f.to_ndarray(format="rgb24").astype(np.float32) for f in c.decode(video=0)]
    c.close()
    return out


def vs_ref(fs, rf):
    """Frame-aligned PSNR and grayscale correlation against the reference."""
    ps, cs = [], []
    for a, b in zip(fs, rf):
        if a.shape != b.shape:
            return float("nan"), float("nan")
        mse = float(((a - b) ** 2).mean())
        ps.append(10 * np.log10(255.0 ** 2 / max(mse, 1e-9)))
        x, y = a.mean(-1).ravel(), b.mean(-1).ravel()
        x, y = x - x.mean(), y - y.mean()
        cs.append(float((x @ y) / max(np.linalg.norm(x) * np.linalg.norm(y), 1e-9)))
    return float(np.mean(ps)), float(np.mean(cs))


def main():
    d = Path(sys.argv[1] if len(sys.argv) > 1 else "out/lambda_exps")
    files = sorted(d.glob("*.mp4"))
    ref = next((f for f in files if f.stem.startswith("0_")), None)
    R = stats(ref) if ref else None
    RF = frames(ref) if ref else None
    print(f"{'run':<22}{'n':>5}{'meanRGB':>9}{'per-frame lo-hi':>18}{'span':>7}{'std':>7}"
          f"{'Δmean':>8}{'PSNR':>8}{'corr':>7}")
    for f in files:
        st = stats(f)
        dm = f"{st['mean']-R['mean']:+.2f}" if R else "-"
        if RF and f is not ref:
            psnr, corr = vs_ref(frames(f), RF)
            pz, cz = f"{psnr:.2f}", f"{corr:.3f}"
        else:
            pz, cz = "-", "-"
        print(f"{f.stem:<22}{st['n']:>5}{st['mean']:>9.2f}"
              f"{st['lo']:>9.2f}-{st['hi']:<8.2f}{st['span']:>7.2f}{st['std']:>7.2f}"
              f"{dm:>8}{pz:>8}{cz:>7}")
    return 0


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