"""Image processing transformations module for camera frames.

This module provides various image transformation functions that can be applied
to received camera frames, including lens distortion correction.
"""

from typing import Callable, Optional

import cv2
import numpy as np


class ImageTransformer:
    """Manages image transformations for camera frames.

    This class provides a pipeline for applying various transformations
    to images, such as lens distortion correction, scaling, and cropping.

    Attributes:
        transformation: The current transformation function to apply.
        fov_degrees: Field of view in degrees for the lens.
        k1: Primary radial distortion coefficient (tunable).
    """

    def __init__(self, fov_degrees: float = 150.0) -> None:
        """Initialize the image transformer with default (no-op) transformation.

        Args:
            fov_degrees: Field of view in degrees for the lens (default: 150.0).
        """
        self._transformation: Callable[[np.ndarray], np.ndarray] = self._passthrough
        self._transformation_name: str = "passthrough"
        self._enabled = False

        # Lens parameters
        self.fov_degrees = fov_degrees
        self.k1 = 0.72  # Primary radial distortion coefficient (tunable)
        self.k2 = 0.02  # Secondary radial distortion coefficient
        self.k3 = 0.0   # Tertiary radial distortion coefficient
        self.k4 = 0.0   # Quaternary radial distortion coefficient

        # Cache for distortion correction maps to avoid recomputation
        self._distortion_map_x: Optional[np.ndarray] = None
        self._distortion_map_y: Optional[np.ndarray] = None
        self._cached_size: Optional[tuple[int, int]] = None
        self._cached_k1: Optional[float] = None

    @property
    def transformation_name(self) -> str:
        """Get the name of the current transformation."""
        return self._transformation_name

    @property
    def is_enabled(self) -> bool:
        """Check if barrel correction is enabled."""
        return self._enabled

    @property
    def k1_value(self) -> float:
        """Get the current k1 coefficient value."""
        return self.k1

    def toggle(self) -> bool:
        """Toggle barrel correction on/off.

        Returns:
            True if correction is now enabled, False if disabled.
        """
        self._enabled = not self._enabled
        if self._enabled:
            self._transformation = self._barrel_correction
            self._transformation_name = "barrel_correction"
        else:
            self._transformation = self._passthrough
            self._transformation_name = "passthrough"
        return self._enabled

    def adjust_k1(self, delta: float) -> float:
        """Adjust the k1 coefficient by the given delta.

        Args:
            delta: Amount to adjust k1 by (can be positive or negative).

        Returns:
            The new k1 value after adjustment.
        """
        self.k1 += delta
        # Invalidate cache since k1 changed
        self._cached_k1 = None
        self._distortion_map_x = None
        self._distortion_map_y = None
        return self.k1

    def transform(self, frame: np.ndarray) -> np.ndarray:
        """Apply the current transformation to a frame.

        Args:
            frame: Input image frame (BGR format).

        Returns:
            Transformed frame.
        """
        return self._transformation(frame)

    def _passthrough(self, frame: np.ndarray) -> np.ndarray:
        """1:1 passthrough transformation - no changes to the frame.

        Args:
            frame: Input image frame.

        Returns:
            The same frame unchanged.
        """
        return frame

    def _barrel_correction(self, frame: np.ndarray) -> np.ndarray:
        """Apply barrel distortion correction for wide-angle lens.

        This correction is designed for wide-angle lenses that exhibit barrel
        distortion (fisheye effect). The correction uses OpenCV's fisheye
        undistortion model.

        Args:
            frame: Input image frame with barrel distortion.

        Returns:
            Corrected frame with reduced barrel distortion.
        """
        height, width = frame.shape[:2]
        size = (width, height)

        # Recompute maps only if frame size or k1 changes
        if (
            self._cached_size != size
            or self._cached_k1 != self.k1
            or self._distortion_map_x is None
        ):
            self._distortion_map_x, self._distortion_map_y = self._compute_barrel_maps(
                width, height
            )
            self._cached_size = size
            self._cached_k1 = self.k1

        # Apply remapping using precomputed maps
        return cv2.remap(
            frame,
            self._distortion_map_x,
            self._distortion_map_y,
            interpolation=cv2.INTER_LINEAR,
            borderMode=cv2.BORDER_CONSTANT,
            borderValue=(0, 0, 0),
        )

    def _compute_barrel_maps(
        self, width: int, height: int
    ) -> tuple[np.ndarray, np.ndarray]:
        """Compute distortion correction maps for barrel distortion.

        The K values control the radial distortion correction.

        Args:
            width: Image width in pixels.
            height: Image height in pixels.

        Returns:
            Tuple of (map_x, map_y) for cv2.remap.
        """
        # Camera matrix (assuming principal point at center)
        cx = width / 2.0
        cy = height / 2.0

        # Focal length estimate for wide-angle lens
        # f = (width / 2) / tan(FOV / 2)
        fov_radians = np.radians(self.fov_degrees)
        fx = fy = (width / 2.0) / np.tan(fov_radians / 2.0)

        camera_matrix = np.array(
            [[fx, 0.0, cx], [0.0, fy, cy], [0.0, 0.0, 1.0]], dtype=np.float64
        )

        # Distortion coefficients for barrel correction
        # K1, K2, K3, K4 for fisheye model
        # Positive values correct barrel distortion (pincushion effect)
        dist_coeffs = np.array(
            [self.k1, self.k2, self.k3, self.k4], dtype=np.float64
        )

        # Compute optimal new camera matrix to minimize black borders
        # Use balance=0.5 to keep some field of view while reducing black areas
        new_camera_matrix = cv2.fisheye.estimateNewCameraMatrixForUndistortRectify(
            camera_matrix,
            dist_coeffs,
            (width, height),
            np.eye(3),
            balance=0.5,
        )

        # Compute undistortion maps
        map_x, map_y = cv2.fisheye.initUndistortRectifyMap(
            camera_matrix,
            dist_coeffs,
            np.eye(3),
            new_camera_matrix,
            (width, height),
            cv2.CV_32FC1,
        )

        return map_x, map_y


def create_transformer(fov_degrees: float = 150.0) -> ImageTransformer:
    """Factory function to create a new ImageTransformer instance.

    Args:
        fov_degrees: Field of view in degrees for the lens (default: 150.0).

    Returns:
        New ImageTransformer with default passthrough transformation.
    """
    return ImageTransformer(fov_degrees=fov_degrees)
