"""
Collect and visualize SV4D batch results.

For each sample, concatenate the 4 novel-view videos into a 2x2 grid,
and save all grids to a single folder for easy browsing.

Usage:
    python scripts/sampling/collect_results.py
"""

import os
import argparse
import numpy as np
import imageio
from PIL import Image


INPUT_DIR = "./outputs/outputs_4d"
OUTPUT_DIR = "./outputs/gallery"


def make_grid(frames_list, nrow=2):
    """
    frames_list: list of N lists, each containing T frames (H, W, 3).
    Returns: list of T grid frames.
    """
    n_views = len(frames_list)
    n_frames = min(len(f) for f in frames_list)
    h, w = frames_list[0][0].shape[:2]
    ncol = (n_views + nrow - 1) // nrow

    grids = []
    for t in range(n_frames):
        canvas = np.full((h * nrow, w * ncol, 3), 255, dtype=np.uint8)
        for idx in range(n_views):
            r, c = divmod(idx, ncol)
            canvas[r * h : (r + 1) * h, c * w : (c + 1) * w] = frames_list[idx][t]
        grids.append(canvas)
    return grids


def find_view_videos(directory):
    """Find novel-view mp4s (_v0XX.mp4) in a directory, handling both SV4D and SV4D2."""
    view_files = []
    for f in sorted(os.listdir(directory)):
        if not f.endswith(".mp4"):
            continue
        if "_v0" in f and "_v000" not in f:
            view_files.append(os.path.join(directory, f))
    return view_files


def process_one(sample_dir, output_dir):
    for sub in ["sv4d2", "sv4d", ""]:
        search_dir = os.path.join(sample_dir, sub) if sub else sample_dir
        if not os.path.isdir(search_dir):
            continue
        view_files = find_view_videos(search_dir)
        if view_files:
            break
    else:
        return False

    if not view_files:
        return False

    all_views = []
    for vf in view_files:
        frames = [np.array(f) for f in imageio.mimread(vf)]
        all_views.append(frames)

    grid_frames = make_grid(all_views, nrow=2)

    name = os.path.basename(sample_dir)
    out_path = os.path.join(output_dir, f"{name}.mp4")
    imageio.mimwrite(out_path, grid_frames, fps=4)
    return True


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--input_dir", default=INPUT_DIR)
    parser.add_argument("--output_dir", default=OUTPUT_DIR)
    args = parser.parse_args()

    os.makedirs(args.output_dir, exist_ok=True)

    samples = sorted([
        os.path.join(args.input_dir, d)
        for d in os.listdir(args.input_dir)
        if os.path.isdir(os.path.join(args.input_dir, d))
    ])

    done = 0
    for s in samples:
        name = os.path.basename(s)
        ok = process_one(s, args.output_dir)
        if ok:
            done += 1
            print(f"[{done}/{len(samples)}] {name}")
        else:
            print(f"[skip] {name} (no view videos found)")

    print(f"\nDone! {done} grid videos saved to {args.output_dir}")


if __name__ == "__main__":
    main()