"""
Degradation augmentation transforms.

Provides degradation effects for document image augmentation:
- Wrinkle: Paper wrinkle/crease simulation
- EdgeDamage: Damaged/torn edge effects
- Stain: Coffee stain/smudge effects
"""

import cv2
import numpy as np

from shared.augmentation.base import AugmentationResult, BaseAugmentation


class Wrinkle(BaseAugmentation):
    """
    Simulates paper wrinkles/creases using displacement mapping.

    Document-friendly: Uses subtle displacement to preserve text readability.

    Parameters:
        intensity: Wrinkle intensity (0-1) (default: 0.3).
        num_wrinkles: Number of wrinkles, int or (min, max) tuple (default: (2, 5)).
    """

    name = "wrinkle"
    affects_geometry = False

    def _validate_params(self) -> None:
        intensity = self.params.get("intensity", 0.3)
        if not (0 < intensity <= 1):
            raise ValueError("intensity must be between 0 and 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()

        h, w = image.shape[:2]
        intensity = self.params.get("intensity", 0.3)
        num_wrinkles = self.params.get("num_wrinkles", (2, 5))

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

        # Create displacement maps
        displacement_x = np.zeros((h, w), dtype=np.float32)
        displacement_y = np.zeros((h, w), dtype=np.float32)

        for _ in range(num_wrinkles):
            # Random wrinkle parameters
            angle = rng.uniform(0, np.pi)
            x0 = rng.uniform(0, w)
            y0 = rng.uniform(0, h)
            length = rng.uniform(0.3, 0.8) * min(h, w)
            width = rng.uniform(0.02, 0.05) * min(h, w)

            # Create coordinate grids
            xx, yy = np.meshgrid(np.arange(w), np.arange(h))

            # Distance from wrinkle line
            dx = (xx - x0) * np.cos(angle) + (yy - y0) * np.sin(angle)
            dy = -(xx - x0) * np.sin(angle) + (yy - y0) * np.cos(angle)

            # Gaussian falloff perpendicular to wrinkle
            mask = np.exp(-dy**2 / (2 * width**2))
            mask *= (np.abs(dx) < length / 2).astype(np.float32)

            # Displacement perpendicular to wrinkle
            disp_amount = intensity * rng.uniform(2, 8)
            displacement_x += mask * disp_amount * np.sin(angle)
            displacement_y += mask * disp_amount * np.cos(angle)

        # Create remap coordinates
        map_x = (np.arange(w)[np.newaxis, :] + displacement_x).astype(np.float32)
        map_y = (np.arange(h)[:, np.newaxis] + displacement_y).astype(np.float32)

        # Apply displacement
        augmented = cv2.remap(
            image, map_x, map_y, cv2.INTER_LINEAR, borderMode=cv2.BORDER_REFLECT
        )

        # Add subtle shading along wrinkles
        max_disp = np.max(np.abs(displacement_y)) + 1e-6
        shading = 1 - 0.1 * intensity * np.abs(displacement_y) / max_disp
        shading = shading[:, :, np.newaxis]
        augmented = (augmented.astype(np.float32) * shading).astype(np.uint8)

        return AugmentationResult(
            image=augmented,
            bboxes=bboxes.copy() if bboxes is not None else None,
            metadata={"num_wrinkles": num_wrinkles, "intensity": intensity},
        )

    def get_preview_params(self) -> dict:
        return {"intensity": 0.5, "num_wrinkles": 3}


class EdgeDamage(BaseAugmentation):
    """
    Adds damaged/torn edge effects to the image.

    Simulates worn or torn document edges.

    Parameters:
        max_damage_ratio: Maximum proportion of edge to damage (default: 0.05).
        edges: Which edges to potentially damage (default: all).
    """

    name = "edge_damage"
    affects_geometry = False

    def _validate_params(self) -> None:
        max_damage_ratio = self.params.get("max_damage_ratio", 0.05)
        if not (0 < max_damage_ratio <= 0.2):
            raise ValueError("max_damage_ratio must be between 0 and 0.2")

    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()

        h, w = image.shape[:2]
        max_damage_ratio = self.params.get("max_damage_ratio", 0.05)
        edges = self.params.get("edges", ["top", "bottom", "left", "right"])

        output = image.copy()

        # Select random edge to damage
        edge = rng.choice(edges)
        damage_size = int(max_damage_ratio * min(h, w))

        if edge == "top":
            # Create irregular top edge
            for x in range(w):
                depth = rng.integers(0, damage_size + 1)
                if depth > 0:
                    # Random color (white or darker)
                    color = rng.integers(200, 255) if rng.random() > 0.5 else rng.integers(100, 150)
                    output[:depth, x] = color

        elif edge == "bottom":
            for x in range(w):
                depth = rng.integers(0, damage_size + 1)
                if depth > 0:
                    color = rng.integers(200, 255) if rng.random() > 0.5 else rng.integers(100, 150)
                    output[h - depth:, x] = color

        elif edge == "left":
            for y in range(h):
                depth = rng.integers(0, damage_size + 1)
                if depth > 0:
                    color = rng.integers(200, 255) if rng.random() > 0.5 else rng.integers(100, 150)
                    output[y, :depth] = color

        else:  # right
            for y in range(h):
                depth = rng.integers(0, damage_size + 1)
                if depth > 0:
                    color = rng.integers(200, 255) if rng.random() > 0.5 else rng.integers(100, 150)
                    output[y, w - depth:] = color

        return AugmentationResult(
            image=output,
            bboxes=bboxes.copy() if bboxes is not None else None,
            metadata={"edge": edge, "damage_size": damage_size},
        )

    def get_preview_params(self) -> dict:
        return {"max_damage_ratio": 0.08}


class Stain(BaseAugmentation):
    """
    Adds coffee stain/smudge effects to the image.

    Simulates accidental stains on documents.

    Parameters:
        num_stains: Number of stains, int or (min, max) tuple (default: (1, 3)).
        max_radius_ratio: Maximum stain radius as ratio of image size (default: 0.1).
        opacity: Stain opacity, float or (min, max) tuple (default: (0.1, 0.3)).
    """

    name = "stain"
    affects_geometry = False

    def _validate_params(self) -> None:
        opacity = self.params.get("opacity", (0.1, 0.3))
        if isinstance(opacity, (int, float)):
            if not (0 < opacity <= 1):
                raise ValueError("opacity must be between 0 and 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()

        h, w = image.shape[:2]
        num_stains = self.params.get("num_stains", (1, 3))
        max_radius_ratio = self.params.get("max_radius_ratio", 0.1)
        opacity = self.params.get("opacity", (0.1, 0.3))

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

        output = image.astype(np.float32)
        max_radius = int(max_radius_ratio * min(h, w))

        for _ in range(num_stains):
            # Random stain position and size
            cx = rng.integers(max_radius, w - max_radius)
            cy = rng.integers(max_radius, h - max_radius)
            radius = rng.integers(max_radius // 3, max_radius)

            # Create stain mask with irregular edges
            yy, xx = np.ogrid[:h, :w]
            dist = np.sqrt((xx - cx) ** 2 + (yy - cy) ** 2)

            # Add noise to make edges irregular
            noise = rng.uniform(0.8, 1.2, (h, w))
            mask = (dist < radius * noise).astype(np.float32)

            # Blur for soft edges
            mask = cv2.GaussianBlur(mask, (21, 21), 0)

            # Random stain color (brownish/yellowish)
            stain_color = np.array([
                rng.integers(180, 220),  # R
                rng.integers(160, 200),  # G
                rng.integers(120, 160),  # B
            ], dtype=np.float32)

            # Apply stain
            mask_3d = mask[:, :, np.newaxis]
            output = output * (1 - mask_3d * opacity) + stain_color * mask_3d * opacity

        output = np.clip(output, 0, 255).astype(np.uint8)

        return AugmentationResult(
            image=output,
            bboxes=bboxes.copy() if bboxes is not None else None,
            metadata={"num_stains": num_stains, "opacity": opacity},
        )

    def get_preview_params(self) -> dict:
        return {"num_stains": 2, "max_radius_ratio": 0.1, "opacity": 0.25}
