"""Quantitative arl2_patch vs frame-level ARL2 comparison on generated mp4s.

Same seed (0) and same prompts -> frame-aligned, directly comparable. Reuses
the same motion-magnitude / sharpness / rel-L2 diagnostics as
scripts/compare_gen_videos.py, adapted to the stage_b_{ckpt}_{split} directory
layout used by run_gen_arl2_6layer_3chunk_teacher_flow_all.sh and
run_gen_arl2_6layer_3chunk_patch_teacher_flow_all.sh.
"""
import argparse
import glob
import os

from torchvision.io import read_video


def vid_stats(path):
    v, _, _ = read_video(path, pts_unit="sec", output_format="TCHW")
    v = v.float() / 255.0
    dt = (v[1:] - v[:-1]).abs().mean().item()          # temporal motion magnitude
    dh = (v[..., 1:, :] - v[..., :-1, :]).pow(2).mean()  # spatial high-freq energy
    dw = (v[..., :, 1:] - v[..., :, :-1]).pow(2).mean()
    sharp = (dh + dw).item()
    return v, dt, sharp


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--frame_root", default="workspace/outputs/gen_arl2_6layer_3chunk_teacher_flow")
    ap.add_argument("--patch_root", default="workspace/outputs/gen_arl2_6layer_3chunk_patch_teacher_flow")
    ap.add_argument("--out", default="workspace/markdown/arl2_patch_vs_frame_comparison.md")
    ap.add_argument("--checkpoints", nargs="+", default=["best", "ema", "final"])
    args = ap.parse_args()

    L = ["# ARL2 patch (8-patch softmax+MLP inter) vs frame-level ARL2, teacher_flow objective\n",
         "Same seed=0, same prompts, same 832x480/21-latent-frame/81-decoded-frame protocol -> "
         "frame-aligned videos, directly comparable. rel-L2 = ||patch-frame||/||frame|| over the "
         "decoded video (higher = more different from the frame-level control).\n"]
    for ckpt in args.checkpoints:
        L.append(f"\n## checkpoint = {ckpt}\n")
        L.append("| set | prompt (truncated) | rel-L2 patch-vs-frame | motion-mag (patch / frame) | sharpness (patch / frame) |")
        L.append("|---|---|---|---|---|")
        for split in ["train", "ood"]:
            fdir = os.path.join(args.frame_root, f"stage_b_{ckpt}_{split}")
            pdir = os.path.join(args.patch_root, f"stage_b_{ckpt}_{split}")
            rels, dt_p, dt_f, sh_p, sh_f = [], [], [], [], []
            for fp in sorted(glob.glob(os.path.join(fdir, "*.mp4"))):
                name = os.path.basename(fp)
                pp = os.path.join(pdir, name)
                if not os.path.exists(pp):
                    continue
                vf, dtf, shf = vid_stats(fp)
                vp, dtp, shp = vid_stats(pp)
                n = min(vf.shape[0], vp.shape[0])
                rel = ((vp[:n] - vf[:n]).norm() / (vf[:n].norm() + 1e-8)).item()
                rels.append(rel); dt_p.append(dtp); dt_f.append(dtf); sh_p.append(shp); sh_f.append(shf)
                L.append(f"| {split} | {name[:42]} | {rel*100:.2f}% | {dtp:.4f} / {dtf:.4f} | {shp:.4f} / {shf:.4f} |")
            if rels:
                m = lambda x: sum(x) / len(x)
                L.append(f"| **{split} MEAN** | ({len(rels)} prompts) | **{m(rels)*100:.2f}%** | "
                          f"{m(dt_p):.4f} / {m(dt_f):.4f} | {m(sh_p):.4f} / {m(sh_f):.4f} |")

    L.append("\n**Read:** motion-mag (mean |frame_t+1 - frame_t|) higher = more/sharper motion, "
              "less = more static/blurred-together. sharpness (spatial gradient energy) higher = "
              "more spatial detail. rel-L2 just measures how much the two checkpoints diverge; it "
              "does not by itself say which is better -- read it together with motion-mag/sharpness "
              "and a visual check.\n")
    os.makedirs(os.path.dirname(args.out), exist_ok=True)
    with open(args.out, "w") as f:
        f.write("\n".join(L) + "\n")
    print("wrote", args.out)
    print("\n".join(L))


if __name__ == "__main__":
    main()