"""Grab the middle-frame of every *_pred_handobject.mp4 and tile into a single
image for quick visual QA of object-track placements across all episodes."""
import argparse, glob, os, math
import cv2
import numpy as np
import imageio.v2 as imageio


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--src_glob", required=True)
    ap.add_argument("--out_path", required=True)
    ap.add_argument("--cell_w", type=int, default=256)
    ap.add_argument("--cell_h", type=int, default=144)
    ap.add_argument("--cols", type=int, default=8)
    args = ap.parse_args()

    files = sorted(glob.glob(args.src_glob))
    n = len(files)
    rows = math.ceil(n / args.cols)
    print(f"{n} videos -> {rows}x{args.cols} grid")

    canvas = np.full((rows * args.cell_h, args.cols * args.cell_w, 3), 255, dtype=np.uint8)
    for i, p in enumerate(files):
        c = cv2.VideoCapture(p)
        nf = int(c.get(cv2.CAP_PROP_FRAME_COUNT))
        c.set(cv2.CAP_PROP_POS_FRAMES, nf // 2)
        r, frame = c.read()
        c.release()
        if not r:
            continue
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        cell = cv2.resize(frame, (args.cell_w, args.cell_h), interpolation=cv2.INTER_AREA)
        # tag with short name
        name = os.path.basename(p).replace("_pred_handobject.mp4", "").replace("table2_", "")
        cv2.putText(cell, name, (4, 14),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.4, (0, 0, 0), 1, cv2.LINE_AA)
        rr, cc = divmod(i, args.cols)
        canvas[rr * args.cell_h:(rr + 1) * args.cell_h,
               cc * args.cell_w:(cc + 1) * args.cell_w] = cell

    imageio.imwrite(args.out_path, canvas)
    print(f"saved {args.out_path}  ({canvas.shape[1]}x{canvas.shape[0]})")


if __name__ == "__main__":
    main()