"""
Description:
    given a video, save it as a filmstrip.
    downsample if necessary.
"""

import click
import numpy as np
import imageio
from PIL import Image
from typing import Tuple


"""
python -m sim.video2filmstrip \
    --video_path test.mp4 \
    --filmstrip_save_path test.pdf \
    --filmstrip_format pdf \
    --n_rows 1 \
    --n_frames 10 \
    --padding 5 \
    --background_color 255 255 255
"""
@click.command()
@click.option('--video_path', type=str, required=True)
@click.option('--filmstrip_save_path', type=str, required=True)
@click.option('--filmstrip_format', type=str, default='pdf')
@click.option('--n_rows', type=int, default=1)
@click.option('--n_frames', type=int, default=None)
@click.option('--padding', type=int, default=5)
@click.option('--background_color', type=click.Tuple([int, int, int]), default=(255, 255, 255))
def video2filmstrip(
    video_path: str, 
    filmstrip_save_path: str,
    filmstrip_format: str,
    n_rows: int,
    n_frames: int,
    padding: int,               # padding between frames
    background_color: Tuple[int, int, int]
):
    
    # load video
    video = imageio.get_reader(video_path)
    len_video = video.count_frames()
    if n_frames is None:
        n_frames = len_video
    else:
        n_frames = min(n_frames, len_video)
    assert n_frames % n_rows == 0
    print(f"video has {len_video} frames")

    # get video dimensions
    frame = video.get_data(0)
    h, w, _ = frame.shape

    # create filmstrip
    n_cols = n_frames // n_rows
    stride = len_video // n_frames
    print(f"creating a {n_rows}x{n_cols} filmstrip with {n_frames} frames")
    filmstrip = np.full((h * n_rows + padding * (n_rows - 1), w * n_cols + padding * (n_cols - 1), 3), background_color, dtype=np.uint8)
    for i in range(n_frames):
        row = i // n_cols
        col = i % n_cols
        frame = video.get_data(i * stride)
        filmstrip[
            row * h + row * padding : (row + 1) * h + row * padding,
            col * w + col * padding : (col + 1) * w + col * padding
        ] = frame

    # save filmstrip
    filmstrip = Image.fromarray(filmstrip)
    filmstrip.save(filmstrip_save_path, format=filmstrip_format)


if __name__ == '__main__':
    video2filmstrip()