#!/usr/bin/env python3
"""Retime a numbered PNG sequence with RIFE while preserving both endpoints."""

import argparse
import math
import sys
from pathlib import Path

import cv2
import numpy as np
import torch
from torch.nn import functional as F
from tqdm import tqdm


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--repo", type=Path, required=True)
    parser.add_argument("--model", type=Path, required=True)
    parser.add_argument("--input", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--source-fps", type=float, required=True)
    parser.add_argument("--target-fps", type=float, required=True)
    parser.add_argument("--speed", type=float, required=True)
    parser.add_argument("--scale", type=float, default=1.0)
    parser.add_argument("--fp16", action="store_true")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    if args.source_fps <= 0 or args.target_fps <= 0 or args.speed <= 0:
        raise ValueError("frame rates and speed must be positive")
    if args.scale not in {0.25, 0.5, 1.0, 2.0, 4.0}:
        raise ValueError("unsupported --scale")
    if not torch.cuda.is_available():
        raise RuntimeError("CUDA GPU is required")

    args.repo = args.repo.resolve()
    args.model = args.model.resolve()
    args.input = args.input.resolve()
    args.output = args.output.resolve()
    sys.path.insert(0, str(args.repo))

    torch.set_grad_enabled(False)
    torch.backends.cudnn.enabled = True
    torch.backends.cudnn.benchmark = True
    if args.fp16:
        torch.set_default_tensor_type(torch.cuda.HalfTensor)

    from train_log.RIFE_HDv3 import Model

    model = Model()
    model.load_model(str(args.model), -1)
    model.eval()
    model.device()

    frames = sorted(args.input.glob("*.png"), key=lambda path: int(path.stem))
    if len(frames) < 2:
        raise RuntimeError("at least two numbered PNG frames are required")
    args.output.mkdir(parents=True, exist_ok=True)

    first = cv2.imread(str(frames[0]), cv2.IMREAD_COLOR)
    if first is None:
        raise RuntimeError(f"failed to read {frames[0]}")
    height, width = first.shape[:2]
    block = max(128, int(128 / args.scale))
    padded_height = ((height - 1) // block + 1) * block
    padded_width = ((width - 1) // block + 1) * block
    padding = (0, padded_width - width, 0, padded_height - height)

    def read_bgr(index: int) -> np.ndarray:
        frame = cv2.imread(str(frames[index]), cv2.IMREAD_COLOR)
        if frame is None:
            raise RuntimeError(f"failed to read {frames[index]}")
        if frame.shape[:2] != (height, width):
            raise RuntimeError(f"frame size changed at {frames[index]}")
        return frame

    def to_tensor(frame_bgr: np.ndarray) -> torch.Tensor:
        rgb = frame_bgr[:, :, ::-1].copy()
        tensor = torch.from_numpy(np.transpose(rgb, (2, 0, 1))).to(
            "cuda", non_blocking=True
        ).unsqueeze(0)
        tensor = tensor.half() if args.fp16 else tensor.float()
        return F.pad(tensor / 255.0, padding)

    def write_tensor(index: int, tensor: torch.Tensor) -> None:
        rgb = (
            tensor[0, :, :height, :width]
            .mul(255.0)
            .clamp(0.0, 255.0)
            .round()
            .to(torch.uint8)
            .cpu()
            .numpy()
            .transpose(1, 2, 0)
        )
        cv2.imwrite(str(args.output / f"{index:07d}.png"), rgb[:, :, ::-1])

    source_intervals = len(frames) - 1
    source_motion_duration = source_intervals / args.source_fps
    target_intervals = max(
        1,
        round(source_motion_duration / args.speed * args.target_fps),
    )
    positions = np.linspace(0.0, float(source_intervals), target_intervals + 1)

    cached_indices: tuple[int, int] | None = None
    cached_bgr: tuple[np.ndarray, np.ndarray] | None = None
    cached_tensors: tuple[torch.Tensor, torch.Tensor] | None = None

    with torch.inference_mode():
        for output_index, position in enumerate(tqdm(positions, desc="RIFE retime")):
            lower = min(int(math.floor(position)), source_intervals)
            upper = min(lower + 1, source_intervals)
            fraction = position - lower

            if fraction <= 1e-7 or lower == upper:
                cv2.imwrite(
                    str(args.output / f"{output_index:07d}.png"),
                    read_bgr(lower),
                )
                continue
            if 1.0 - fraction <= 1e-7:
                cv2.imwrite(
                    str(args.output / f"{output_index:07d}.png"),
                    read_bgr(upper),
                )
                continue

            pair = (lower, upper)
            if cached_indices != pair:
                lower_bgr = read_bgr(lower)
                upper_bgr = read_bgr(upper)
                cached_indices = pair
                cached_bgr = (lower_bgr, upper_bgr)
                cached_tensors = (to_tensor(lower_bgr), to_tensor(upper_bgr))
            assert cached_bgr is not None and cached_tensors is not None
            middle = model.inference(
                cached_tensors[0],
                cached_tensors[1],
                float(fraction),
                args.scale,
            )
            write_tensor(output_index, middle)

    print(
        f"wrote {target_intervals + 1} frames; "
        f"motion duration {target_intervals / args.target_fps:.6f}s; "
        f"effective speed {source_motion_duration / (target_intervals / args.target_fps):.6f}x"
    )


if __name__ == "__main__":
    main()