"""Decode training videos to 15 fps frames (fit in 448 px) as uint8 arrays,
and mark scene cuts so no training triplet spans one."""
import json, subprocess, sys
from pathlib import Path
import numpy as np

src, out = Path(sys.argv[1]), Path(sys.argv[2])
out.mkdir(parents=True, exist_ok=True)
index = []
for video in sorted(src.glob("v*.webm")):
    probe = json.loads(subprocess.run(["ffprobe", "-v", "error", "-select_streams", "v:0", "-show_entries",
                                       "stream=width,height", "-of", "json", str(video)], capture_output=True, text=True).stdout)
    w, h = probe["streams"][0]["width"], probe["streams"][0]["height"]
    s = 448 / max(w, h)
    W, H = int(round(w * s / 8) * 8), int(round(h * s / 8) * 8)
    raw = subprocess.run(["ffmpeg", "-v", "error", "-i", str(video), "-vf", f"fps=15,scale={W}:{H}:flags=area",
                          "-f", "rawvideo", "-pix_fmt", "rgb24", "-"], capture_output=True).stdout
    frames = np.frombuffer(raw, np.uint8).reshape(-1, H, W, 3)
    small = frames[:, ::8, ::8].astype(np.float32).mean(-1)
    diff = np.abs(np.diff(small, axis=0)).mean((1, 2))
    cuts = np.where(diff > max(18.0, np.median(diff) * 6))[0] + 1   # first frame of each new shot
    np.save(out / f"{video.stem}.npy", frames)
    index.append({"name": video.stem, "frames": int(len(frames)), "width": W, "height": H, "cuts": cuts.tolist()})
    print(video.stem, frames.shape, "cuts:", len(cuts))
(out / "index.json").write_text(json.dumps(index, indent=1))