# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy", "opencv-python"]
# ///
"""
@effect fisheye
@version 1.0.0
@author artdag

@description
Barrel/fisheye lens distortion. Positive values bulge outward (fisheye),
negative values pinch inward (pincushion). Great for emphasis effects.

@param strength float
  @range -1 1
  @default 0.3
  Distortion strength. Positive = fisheye bulge, negative = pincushion pinch.

@param center_x float
  @range 0 1
  @default 0.5
  Horizontal center of distortion (0 = left, 1 = right).

@param center_y float
  @range 0 1
  @default 0.5
  Vertical center of distortion (0 = top, 1 = bottom).

@param zoom_correction bool
  @default true
  Automatically zoom to hide black edges on fisheye.

@example
  (effect fisheye :strength 0.5)

@example
  ;; Pulse fisheye on bass
  (effect fisheye :strength (bind bass :range [0 0.8] :transform sqrt))

@example
  ;; Pincushion effect
  (effect fisheye :strength -0.3)
"""

import numpy as np
import cv2


def process_frame(frame: np.ndarray, params: dict, state: dict) -> tuple:
    """
    Apply fisheye/barrel distortion to a video frame.

    Args:
        frame: Input frame as numpy array (H, W, 3) RGB uint8
        params: Effect parameters
            - strength: distortion strength -1 to 1 (default 0.3)
            - center_x: horizontal center 0-1 (default 0.5)
            - center_y: vertical center 0-1 (default 0.5)
            - zoom_correction: auto-zoom for fisheye (default True)
        state: Persistent state dict (unused)

    Returns:
        Tuple of (processed_frame, new_state)
    """
    strength = params.get("strength", 0.3)
    center_x = params.get("center_x", 0.5)
    center_y = params.get("center_y", 0.5)
    zoom_correction = params.get("zoom_correction", True)

    if strength == 0:
        return frame, state

    h, w = frame.shape[:2]

    # Calculate center in pixels
    cx = w * center_x
    cy = h * center_y

    # Create coordinate grids
    y_coords, x_coords = np.mgrid[0:h, 0:w].astype(np.float32)

    # Normalize coordinates relative to center
    x_norm = (x_coords - cx) / (w / 2)
    y_norm = (y_coords - cy) / (h / 2)

    # Calculate radius from center
    r = np.sqrt(x_norm**2 + y_norm**2)

    # Apply barrel/pincushion distortion
    if strength > 0:
        # Barrel distortion (fisheye)
        r_distorted = r * (1 + strength * r**2)
    else:
        # Pincushion distortion
        r_distorted = r / (1 - strength * r**2 + 0.001)

    # Calculate scale factor
    scale = np.where(r > 0, r_distorted / r, 1)

    # Apply zoom correction to hide black edges
    if zoom_correction and strength > 0:
        zoom = 1 + strength * 0.5
        scale = scale / zoom

    # Calculate new coordinates
    new_x = (x_norm * scale * (w / 2) + cx).astype(np.float32)
    new_y = (y_norm * scale * (h / 2) + cy).astype(np.float32)

    # Remap
    result = cv2.remap(
        frame, new_x, new_y,
        cv2.INTER_LINEAR,
        borderMode=cv2.BORDER_REFLECT
    )

    return result, state