"""
Geometric augmentation transforms.

Provides geometric transforms for document image augmentation:
- PerspectiveWarp: Subtle perspective distortion
"""

import cv2
import numpy as np

from shared.augmentation.base import AugmentationResult, BaseAugmentation


class PerspectiveWarp(BaseAugmentation):
    """
    Applies subtle perspective transformation to the image.

    Simulates viewing document at slight angle. Very conservative
    by default to preserve text readability.

    IMPORTANT: This transform affects bounding box coordinates.

    Parameters:
        max_warp: Maximum warp as proportion of image size (default: 0.02).
    """

    name = "perspective_warp"
    affects_geometry = True

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

        # Original corners
        src_pts = np.float32([
            [0, 0],
            [w, 0],
            [w, h],
            [0, h],
        ])

        # Add random perturbations to corners
        max_offset = max_warp * min(h, w)
        dst_pts = src_pts.copy()
        for i in range(4):
            dst_pts[i, 0] += rng.uniform(-max_offset, max_offset)
            dst_pts[i, 1] += rng.uniform(-max_offset, max_offset)

        # Compute perspective transform matrix
        transform_matrix = cv2.getPerspectiveTransform(src_pts, dst_pts)

        # Apply perspective transform
        warped = cv2.warpPerspective(
            image, transform_matrix, (w, h),
            borderMode=cv2.BORDER_REPLICATE
        )

        # Transform bounding boxes if present
        transformed_bboxes = None
        if bboxes is not None:
            transformed_bboxes = self._transform_bboxes(
                bboxes, transform_matrix, w, h
            )

        return AugmentationResult(
            image=warped,
            bboxes=transformed_bboxes,
            transform_matrix=transform_matrix,
            metadata={"max_warp": max_warp},
        )

    def _transform_bboxes(
        self,
        bboxes: np.ndarray,
        transform_matrix: np.ndarray,
        w: int,
        h: int,
    ) -> np.ndarray:
        """Transform bounding boxes using perspective matrix."""
        if len(bboxes) == 0:
            return bboxes.copy()

        transformed = []
        for bbox in bboxes:
            class_id, x_center, y_center, width, height = bbox

            # Convert normalized coords to pixel coords
            x_center_px = x_center * w
            y_center_px = y_center * h
            width_px = width * w
            height_px = height * h

            # Get corner points
            x1 = x_center_px - width_px / 2
            y1 = y_center_px - height_px / 2
            x2 = x_center_px + width_px / 2
            y2 = y_center_px + height_px / 2

            # Transform all 4 corners
            corners = np.float32([
                [x1, y1],
                [x2, y1],
                [x2, y2],
                [x1, y2],
            ]).reshape(-1, 1, 2)

            transformed_corners = cv2.perspectiveTransform(corners, transform_matrix)
            transformed_corners = transformed_corners.reshape(-1, 2)

            # Get bounding box of transformed corners
            new_x1 = np.min(transformed_corners[:, 0])
            new_y1 = np.min(transformed_corners[:, 1])
            new_x2 = np.max(transformed_corners[:, 0])
            new_y2 = np.max(transformed_corners[:, 1])

            # Convert back to normalized center format
            new_width = (new_x2 - new_x1) / w
            new_height = (new_y2 - new_y1) / h
            new_x_center = ((new_x1 + new_x2) / 2) / w
            new_y_center = ((new_y1 + new_y2) / 2) / h

            # Clamp to valid range
            new_x_center = np.clip(new_x_center, 0, 1)
            new_y_center = np.clip(new_y_center, 0, 1)
            new_width = np.clip(new_width, 0, 1)
            new_height = np.clip(new_height, 0, 1)

            transformed.append([class_id, new_x_center, new_y_center, new_width, new_height])

        return np.array(transformed, dtype=np.float32)

    def get_preview_params(self) -> dict:
        return {"max_warp": 0.03}
