"""
TimeLapseForge — Frame Interpolation Module
Creates smooth intermediate frames between panels using
alpha blending and optical flow warping.
"""

import cv2
import numpy as np
from PIL import Image
from typing import List, Optional, Callable


class FrameInterpolator:
    """Generate intermediate frames between timelapse panels for smooth video."""

    @staticmethod
    def pil_to_cv2(img: Image.Image) -> np.ndarray:
        """Convert PIL Image to OpenCV BGR format."""
        rgb = np.array(img)
        return cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)

    @staticmethod
    def cv2_to_pil(img: np.ndarray) -> Image.Image:
        """Convert OpenCV BGR image to PIL Image."""
        rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        return Image.fromarray(rgb)

    def blend_interpolate(
        self, img1: Image.Image, img2: Image.Image, num_frames: int = 4
    ) -> List[Image.Image]:
        """
        Simple alpha blending between two frames.
        Fast and reliable, good for most cases.
        """
        arr1 = np.array(img1).astype(np.float32)
        arr2 = np.array(img2).astype(np.float32)

        frames = []
        for i in range(1, num_frames + 1):
            alpha = i / (num_frames + 1)
            blended = ((1 - alpha) * arr1 + alpha * arr2).astype(np.uint8)
            frames.append(Image.fromarray(blended))

        return frames

    def flow_interpolate(
        self, img1: Image.Image, img2: Image.Image, num_frames: int = 4
    ) -> List[Image.Image]:
        """
        Optical flow based interpolation.
        Better quality than blending — objects move smoothly instead of fading.
        """
        cv_img1 = self.pil_to_cv2(img1)
        cv_img2 = self.pil_to_cv2(img2)

        gray1 = cv2.cvtColor(cv_img1, cv2.COLOR_BGR2GRAY)
        gray2 = cv2.cvtColor(cv_img2, cv2.COLOR_BGR2GRAY)

        # Calculate optical flow (Farneback method)
        try:
            flow = cv2.calcOpticalFlowFarneback(
                gray1, gray2, None,
                pyr_scale=0.5, levels=3, winsize=15,
                iterations=3, poly_n=5, poly_sigma=1.2, flags=0
            )
        except Exception:
            # Fallback to blend if flow calculation fails
            return self.blend_interpolate(img1, img2, num_frames)

        h, w = cv_img1.shape[:2]
        frames = []

        for i in range(1, num_frames + 1):
            alpha = i / (num_frames + 1)

            # Warp img1 towards img2 using scaled flow
            flow_scaled = flow * alpha
            map_x = np.float32(np.tile(np.arange(w), (h, 1)) + flow_scaled[..., 0])
            map_y = np.float32(np.tile(np.arange(h).reshape(-1, 1), (1, w)) + flow_scaled[..., 1])

            warped_1 = cv2.remap(cv_img1, map_x, map_y, cv2.INTER_LINEAR, borderMode=cv2.BORDER_REFLECT)

            # Also warp img2 backwards
            flow_back = flow * (alpha - 1)
            map_x_b = np.float32(np.tile(np.arange(w), (h, 1)) + flow_back[..., 0])
            map_y_b = np.float32(np.tile(np.arange(h).reshape(-1, 1), (1, w)) + flow_back[..., 1])

            warped_2 = cv2.remap(cv_img2, map_x_b, map_y_b, cv2.INTER_LINEAR, borderMode=cv2.BORDER_REFLECT)

            # Blend the two warped images
            blended = cv2.addWeighted(warped_1, 1 - alpha, warped_2, alpha, 0)
            frames.append(self.cv2_to_pil(blended))

        return frames

    def interpolate_sequence(
        self,
        images: List[Image.Image],
        multiplier: int = 4,
        method: str = "blend",
        progress_callback: Optional[Callable] = None,
    ) -> List[Image.Image]:
        """
        Interpolate an entire sequence of images.
        
        Args:
            images: List of panel images
            multiplier: How many intermediate frames between each panel pair
            method: "blend" for alpha blending, "flow" for optical flow
            progress_callback: Optional progress callback(current, total)
            
        Returns:
            Expanded list of frames with intermediate frames inserted
        """
        if len(images) < 2:
            return images

        interpolate_fn = self.flow_interpolate if method == "flow" else self.blend_interpolate
        result = [images[0]]
        total_pairs = len(images) - 1

        for i in range(total_pairs):
            intermediate = interpolate_fn(images[i], images[i + 1], num_frames=multiplier)
            result.extend(intermediate)
            result.append(images[i + 1])

            if progress_callback:
                progress_callback(i + 1, total_pairs)

        return result

    def create_morph_transition(
        self, img1: Image.Image, img2: Image.Image, num_frames: int = 8
    ) -> List[Image.Image]:
        """
        Create a morphing transition effect (blend + subtle zoom).
        Good for dramatic reveal moments.
        """
        frames = []
        arr1 = np.array(img1).astype(np.float32)
        arr2 = np.array(img2).astype(np.float32)
        h, w = arr1.shape[:2]

        for i in range(1, num_frames + 1):
            alpha = i / (num_frames + 1)

            # Subtle zoom effect (1.0 to 1.02 scale)
            scale = 1.0 + 0.02 * alpha
            center_x, center_y = w // 2, h // 2

            M = cv2.getRotationMatrix2D((center_x, center_y), 0, scale)
            zoomed_1 = cv2.warpAffine(arr1.astype(np.uint8), M, (w, h), borderMode=cv2.BORDER_REFLECT)

            blended = ((1 - alpha) * zoomed_1.astype(np.float32) + alpha * arr2).astype(np.uint8)
            frames.append(Image.fromarray(blended))

        return frames