# -*- coding: utf-8 -*-
"""
Image and Video Preprocessor for MiniCPM5-V with Adaptive Dynamic Slicing and OCR Grounding.
"""
import math
from typing import List, Tuple, Union, Optional
import numpy as np
from PIL import Image
import torch
from torchvision import transforms


class MiniCPM5VImageProcessor:
    r"""
    High-Resolution Dynamic Slicing Image and Video Processor for MiniCPM5-V.
    """
    def __init__(
        self,
        image_size: int = 448,
        max_slice_nums: int = 9,
        scale_resolution: int = 448,
        patch_size: int = 14,
        image_mean: Tuple[float, float, float] = (0.5, 0.5, 0.5),
        image_std: Tuple[float, float, float] = (0.5, 0.5, 0.5),
    ):
        self.image_size = image_size
        self.max_slice_nums = max_slice_nums
        self.scale_resolution = scale_resolution
        self.patch_size = patch_size
        self.image_mean = image_mean
        self.image_std = image_std

        self.transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize(mean=self.image_mean, std=self.image_std)
        ])

    def get_slice_grid(self, orig_w: int, orig_h: int) -> Tuple[int, int]:
        r"""
        Determine optimal grid layout (grid_w, grid_h) such that:
        - grid_w * grid_h <= max_slice_nums
        - Aspect ratio error is minimized
        """
        orig_aspect = orig_w / max(1, orig_h)
        best_grid = (1, 1)
        min_error = float("inf")

        for total_slices in range(1, self.max_slice_nums + 1):
            for gw in range(1, total_slices + 1):
                if total_slices % gw == 0:
                    gh = total_slices // gw
                    grid_aspect = gw / gh
                    # Log aspect ratio error
                    error = abs(math.log(orig_aspect / grid_aspect))
                    if error < min_error:
                        min_error = error
                        best_grid = (gw, gh)

        return best_grid

    def slice_image(self, image: Image.Image) -> List[Image.Image]:
        r"""
        Slice high-resolution image into adaptive grid patches + 1 overview thumbnail.
        Total slices returned: (gw * gh) + 1.
        """
        image = image.convert("RGB")
        w, h = image.size

        # 1. Global Overview Thumbnail (resized to scale_resolution x scale_resolution)
        overview = image.resize((self.scale_resolution, self.scale_resolution), Image.Resampling.BICUBIC)

        # 2. Determine grid
        gw, gh = self.get_slice_grid(w, h)

        if gw == 1 and gh == 1:
            # Single image does not need sub-slicing
            return [overview]

        # 3. High-res canvas resize to match exact grid size (gw * scale_res, gh * scale_res)
        target_w = gw * self.scale_resolution
        target_h = gh * self.scale_resolution
        resized_full = image.resize((target_w, target_h), Image.Resampling.BICUBIC)

        # 4. Crop each slice
        slices = []
        for j in range(gh):
            for i in range(gw):
                box = (
                    i * self.scale_resolution,
                    j * self.scale_resolution,
                    (i + 1) * self.scale_resolution,
                    (j + 1) * self.scale_resolution,
                )
                slice_patch = resized_full.crop(box)
                slices.append(slice_patch)

        # Return slices + global overview thumbnail at the end
        return slices + [overview]

    def preprocess_image(self, image: Image.Image) -> torch.Tensor:
        r"""
        Preprocesses a single PIL Image into a tensor of shape:
        [num_slices, 3, scale_resolution, scale_resolution]
        """
        slices = self.slice_image(image)
        tensors = [self.transform(s) for s in slices]
        return torch.stack(tensors, dim=0)

    def preprocess_video(
        self,
        frames: List[Image.Image],
        max_frames: int = 16
    ) -> torch.Tensor:
        r"""
        Uniformly sample up to max_frames from a video sequence.
        Preprocesses each frame as a 448x448 representation.
        Returns tensor: [num_frames, 3, scale_resolution, scale_resolution]
        """
        if len(frames) > max_frames:
            indices = np.linspace(0, len(frames) - 1, max_frames, dtype=int)
            sampled_frames = [frames[i] for i in indices]
        else:
            sampled_frames = frames

        tensors = []
        for f in sampled_frames:
            f_rgb = f.convert("RGB").resize(
                (self.scale_resolution, self.scale_resolution),
                Image.Resampling.BICUBIC
            )
            tensors.append(self.transform(f_rgb))

        return torch.stack(tensors, dim=0)

    def normalize_box(
        self,
        box: Tuple[int, int, int, int],
        orig_size: Tuple[int, int]
    ) -> Tuple[int, int, int, int]:
        r"""
        Normalize coordinate bounding box (ymin, xmin, ymax, xmax) to [0, 1000] integer range.
        Used for OCR visual grounding token prediction.
        """
        w, h = orig_size
        ymin, xmin, ymax, xmax = box
        norm_ymin = int(round(ymin / h * 1000))
        norm_xmin = int(round(xmin / w * 1000))
        norm_ymax = int(round(ymax / h * 1000))
        norm_xmax = int(round(xmax / w * 1000))

        # Clamp to [0, 1000]
        return (
            max(0, min(1000, norm_ymin)),
            max(0, min(1000, norm_xmin)),
            max(0, min(1000, norm_ymax)),
            max(0, min(1000, norm_xmax)),
        )

    def denormalize_box(
        self,
        norm_box: Tuple[int, int, int, int],
        orig_size: Tuple[int, int]
    ) -> Tuple[int, int, int, int]:
        r"""Convert normalized [0, 1000] coordinates back to original pixel coordinates."""
        w, h = orig_size
        ymin, xmin, ymax, xmax = norm_box
        return (
            int(round(ymin / 1000 * h)),
            int(round(xmin / 1000 * w)),
            int(round(ymax / 1000 * h)),
            int(round(xmax / 1000 * w)),
        )