"""
Side-by-side comparison:
  [replay (collected actions)]  |  [original *_pred.mp4]  |  [DP-predicted actions]
All clipped to first 16 frames at 1280x720.
"""
import argparse
from pathlib import Path
import cv2
import numpy as np


def load_n(video, n):
    cap = cv2.VideoCapture(video)
    frames = []
    for _ in range(n):
        ret, f = cap.read()
        if not ret:
            break
        frames.append(f)
    cap.release()
    return frames


def label(img, text):
    out = img.copy()
    h = out.shape[0]
    cv2.rectangle(out, (0, 0), (out.shape[1], 60), (0, 0, 0), -1)
    cv2.putText(out, text, (20, 42), cv2.FONT_HERSHEY_SIMPLEX, 1.2, (255, 255, 255), 2)
    return out


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--replay_video", required=True)
    p.add_argument("--reference_video", required=True)
    p.add_argument("--dp_video", required=True)
    p.add_argument("--n", type=int, default=16)
    p.add_argument("--out", required=True)
    p.add_argument("--scale", type=float, default=0.5,
                   help="Scale factor per panel.")
    args = p.parse_args()

    rep = load_n(args.replay_video, args.n)
    ref = load_n(args.reference_video, args.n)
    dp = load_n(args.dp_video, args.n)

    n = min(len(rep), len(ref), len(dp))
    print(f"Loaded: replay={len(rep)} ref={len(ref)} dp={len(dp)} -> using n={n}")

    panels = []
    for i in range(n):
        r = label(rep[i], "replay (collected actions)")
        f = label(ref[i], "original _pred.mp4")
        d = label(dp[i], "DP-predicted actions")
        # Resize all to same height
        scale = args.scale
        new_w = int(r.shape[1] * scale)
        new_h = int(r.shape[0] * scale)
        r = cv2.resize(r, (new_w, new_h))
        f = cv2.resize(f, (new_w, new_h))
        d = cv2.resize(d, (new_w, new_h))
        panels.append(np.concatenate([r, f, d], axis=1))

    h, w = panels[0].shape[:2]
    Path(args.out).parent.mkdir(parents=True, exist_ok=True)
    out = cv2.VideoWriter(args.out, cv2.VideoWriter_fourcc(*"mp4v"), 5, (w, h))
    for p_ in panels:
        out.write(p_)
    out.release()
    print(f"Saved: {args.out}  ({w}x{h}, {n} frames @ 5fps)")


if __name__ == "__main__":
    main()