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

@description
Spiral/vortex distortion that twists the image around a center point.
Creates whirlpool-like effects. Great for psychedelic/hypnotic visuals.

@param strength float
  @range -10 10
  @default 1.0
  Swirl strength in radians. Positive = counter-clockwise, negative = clockwise.

@param radius float
  @range 0.1 2
  @default 0.5
  Effect radius as fraction of image size. Larger = wider swirl.

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

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

@param falloff string
  @enum linear quadratic gaussian
  @default quadratic
  How swirl strength decreases from center:
  - linear: constant decrease
  - quadratic: sharper center, softer edges
  - gaussian: smooth bell curve

@example
  (effect swirl :strength 2)

@example
  ;; Reactive swirl
  (effect swirl :strength (bind energy :range [0 5]))
"""

import numpy as np
import cv2


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

    Args:
        frame: Input frame as numpy array (H, W, 3) RGB uint8
        params: Effect parameters
            - strength: swirl amount in radians (default 1.0)
            - radius: effect radius as fraction (default 0.5)
            - center_x: horizontal center 0-1 (default 0.5)
            - center_y: vertical center 0-1 (default 0.5)
            - falloff: linear/quadratic/gaussian (default quadratic)
        state: Persistent state dict (unused)

    Returns:
        Tuple of (processed_frame, new_state)
    """
    strength = params.get("strength", 1.0)
    radius_frac = params.get("radius", 0.5)
    center_x = params.get("center_x", 0.5)
    center_y = params.get("center_y", 0.5)
    falloff = params.get("falloff", "quadratic")

    if strength == 0:
        return frame, state

    h, w = frame.shape[:2]

    # Calculate center and radius in pixels
    cx = w * center_x
    cy = h * center_y
    radius = max(w, h) * radius_frac

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

    # Calculate distance and angle from center
    dx = x_coords - cx
    dy = y_coords - cy
    dist = np.sqrt(dx**2 + dy**2)
    angle = np.arctan2(dy, dx)

    # Normalized distance for falloff
    norm_dist = dist / radius

    # Calculate falloff factor
    if falloff == "linear":
        factor = np.maximum(0, 1 - norm_dist)
    elif falloff == "gaussian":
        factor = np.exp(-norm_dist**2 * 2)
    else:  # quadratic
        factor = np.maximum(0, 1 - norm_dist**2)

    # Apply swirl rotation
    new_angle = angle + strength * factor

    # Calculate new coordinates
    new_x = (cx + dist * np.cos(new_angle)).astype(np.float32)
    new_y = (cy + dist * np.sin(new_angle)).astype(np.float32)

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

    return result, state