"""Render a task video into one contact sheet for visual review.

The numeric check only reads final object poses, and it will pass an episode where the cup landed
upside down or a gripper hung empty in the air. Looking at the frames is the only way to catch
those, so every task gets a sheet before it is called done.
"""
import argparse, os
import imageio.v2 as iio
from PIL import Image, ImageDraw

ap = argparse.ArgumentParser()
ap.add_argument("videos", nargs="+")
ap.add_argument("--out", default="outputs/sheets")
ap.add_argument("--cols", type=int, default=6)
ap.add_argument("--w", type=int, default=340)
a = ap.parse_args()
os.makedirs(a.out, exist_ok=True)

for v in a.videos:
    if not os.path.exists(v):
        print(f"[sheet] MISSING {v}"); continue
    r = iio.get_reader(v)
    n = r.count_frames()
    idx = [min(n-1, int(n*f)) for f in [i/(a.cols-1) for i in range(a.cols)]]
    idx[0] = min(2, n-1)
    frames = [Image.fromarray(r.get_data(i)) for i in idx]
    fw, fh = frames[0].size
    h = int(a.w*fh/fw)
    sheet = Image.new("RGB", (a.w*len(frames), h), (20, 20, 22))
    for k, f in enumerate(frames):
        sheet.paste(f.resize((a.w, h), Image.LANCZOS), (a.w*k, 0))
        ImageDraw.Draw(sheet).text((a.w*k+6, h-16), f"{idx[k]}/{n}", fill=(255, 220, 60))
    p = os.path.join(a.out, os.path.basename(v).replace(".mp4", ".png"))
    sheet.save(p)
    print(f"[sheet] {p}  ({n} frames)")