#!/usr/bin/env python3
"""Lossless-image-sequence RIFE interpolation that preserves every source frame."""

import argparse
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("--multi", type=int, default=4)
    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.multi < 2:
        raise ValueError("--multi must be at least 2")
    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 for this workflow")

    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:
        # Practical-RIFE's warp grid follows the default tensor type. Setting it
        # before importing the model keeps both weights and grids consistently
        # in FP16 and avoids a mixed-type grid_sample operation.
        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 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_rgb_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])

    output_index = 0
    previous_bgr = first
    previous = to_tensor(previous_bgr)
    with torch.inference_mode():
        for path in tqdm(frames[1:], desc="RIFE pairs"):
            current_bgr = cv2.imread(str(path), cv2.IMREAD_COLOR)
            if current_bgr is None:
                raise RuntimeError(f"failed to read {path}")
            if current_bgr.shape[:2] != (height, width):
                raise RuntimeError(f"frame size changed at {path}")
            current = to_tensor(current_bgr)

            # Keep source pixels exact at every multiple of --multi.
            cv2.imwrite(str(args.output / f"{output_index:07d}.png"), previous_bgr)
            output_index += 1
            for step in range(1, args.multi):
                timestep = step / args.multi
                middle = model.inference(previous, current, timestep, args.scale)
                write_rgb_tensor(output_index, middle)
                output_index += 1

            previous_bgr = current_bgr
            previous = current

    cv2.imwrite(str(args.output / f"{output_index:07d}.png"), previous_bgr)
    print(f"wrote {output_index + 1} frames to {args.output}")


if __name__ == "__main__":
    main()