"""
Blur augmentation transforms.

Provides blur effects for document image augmentation:
- GaussianBlur: Simulates out-of-focus capture
- MotionBlur: Simulates camera/document movement during capture
"""

import cv2
import numpy as np

from shared.augmentation.base import AugmentationResult, BaseAugmentation


class GaussianBlur(BaseAugmentation):
    """
    Applies Gaussian blur to the image.

    Simulates out-of-focus capture or low-quality optics.
    Conservative defaults to preserve text readability.

    Parameters:
        kernel_size: Blur kernel size, int or (min, max) tuple (default: (3, 5)).
        sigma: Blur sigma, float or (min, max) tuple (default: (0.5, 1.5)).
    """

    name = "gaussian_blur"
    affects_geometry = False

    def _validate_params(self) -> None:
        kernel_size = self.params.get("kernel_size", (3, 5))
        if isinstance(kernel_size, int):
            if kernel_size < 1 or kernel_size % 2 == 0:
                raise ValueError("kernel_size must be a positive odd integer")
        elif isinstance(kernel_size, tuple):
            if kernel_size[0] < 1 or kernel_size[1] < kernel_size[0]:
                raise ValueError("kernel_size tuple must be (min, max) with min >= 1")

    def apply(
        self,
        image: np.ndarray,
        bboxes: np.ndarray | None = None,
        rng: np.random.Generator | None = None,
    ) -> AugmentationResult:
        rng = rng or np.random.default_rng()

        kernel_size = self.params.get("kernel_size", (3, 5))
        sigma = self.params.get("sigma", (0.5, 1.5))

        if isinstance(kernel_size, tuple):
            # Choose random odd kernel size
            min_k, max_k = kernel_size
            possible_sizes = [k for k in range(min_k, max_k + 1) if k % 2 == 1]
            if not possible_sizes:
                possible_sizes = [min_k if min_k % 2 == 1 else min_k + 1]
            kernel_size = rng.choice(possible_sizes)

        if isinstance(sigma, tuple):
            sigma = rng.uniform(sigma[0], sigma[1])

        # Ensure kernel size is odd
        if kernel_size % 2 == 0:
            kernel_size += 1

        # Apply Gaussian blur
        blurred = cv2.GaussianBlur(image, (kernel_size, kernel_size), sigma)

        return AugmentationResult(
            image=blurred,
            bboxes=bboxes.copy() if bboxes is not None else None,
            metadata={"kernel_size": kernel_size, "sigma": sigma},
        )

    def get_preview_params(self) -> dict:
        return {"kernel_size": 5, "sigma": 1.5}


class MotionBlur(BaseAugmentation):
    """
    Applies motion blur to the image.

    Simulates camera shake or document movement during capture.

    Parameters:
        kernel_size: Blur kernel size, int or (min, max) tuple (default: (5, 9)).
        angle_range: Motion angle range in degrees (default: (-45, 45)).
    """

    name = "motion_blur"
    affects_geometry = False

    def _validate_params(self) -> None:
        kernel_size = self.params.get("kernel_size", (5, 9))
        if isinstance(kernel_size, int):
            if kernel_size < 3:
                raise ValueError("kernel_size must be at least 3")
        elif isinstance(kernel_size, tuple):
            if kernel_size[0] < 3:
                raise ValueError("kernel_size min must be at least 3")

    def apply(
        self,
        image: np.ndarray,
        bboxes: np.ndarray | None = None,
        rng: np.random.Generator | None = None,
    ) -> AugmentationResult:
        rng = rng or np.random.default_rng()

        kernel_size = self.params.get("kernel_size", (5, 9))
        angle_range = self.params.get("angle_range", (-45, 45))

        if isinstance(kernel_size, tuple):
            kernel_size = rng.integers(kernel_size[0], kernel_size[1] + 1)

        angle = rng.uniform(angle_range[0], angle_range[1])

        # Create motion blur kernel
        kernel = np.zeros((kernel_size, kernel_size), dtype=np.float32)

        # Draw a line in the center of the kernel
        center = kernel_size // 2
        angle_rad = np.deg2rad(angle)

        for i in range(kernel_size):
            offset = i - center
            x = int(center + offset * np.cos(angle_rad))
            y = int(center + offset * np.sin(angle_rad))
            if 0 <= x < kernel_size and 0 <= y < kernel_size:
                kernel[y, x] = 1.0

        # Normalize kernel
        kernel = kernel / kernel.sum() if kernel.sum() > 0 else kernel

        # Apply motion blur
        blurred = cv2.filter2D(image, -1, kernel)

        return AugmentationResult(
            image=blurred,
            bboxes=bboxes.copy() if bboxes is not None else None,
            metadata={"kernel_size": kernel_size, "angle": angle},
        )

    def get_preview_params(self) -> dict:
        return {"kernel_size": 7, "angle_range": (-30, 30)}
