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

@description
Chromatic aberration effect. Creates color fringing by offsetting
RGB channels radially from the center (lens distortion simulation).

@param strength float
  @range 0 50
  @default 10
  Aberration strength. Bind to bass for reactive effect.

@param center_x float
  @range 0 1
  @default 0.5
  Aberration center X.

@param center_y float
  @range 0 1
  @default 0.5
  Aberration center Y.

@param radial bool
  @default true
  If true, aberration increases from center to edges.

@example
  (effect chromatic :strength 20)

@example
  ;; Beat-reactive chromatic aberration
  (effect chromatic :strength (bind bass :range [0 30]))
"""

import numpy as np
import cv2


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

    Args:
        frame: Input frame as numpy array (H, W, 3) RGB uint8
        params: Effect parameters
            - strength: aberration amount (default 10)
            - center_x: center X 0-1 (default 0.5)
            - center_y: center Y 0-1 (default 0.5)
            - radial: increase from center (default True)
        state: Persistent state dict

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

    if state is None:
        state = {}

    if strength == 0:
        return frame, state

    h, w = frame.shape[:2]
    r, g, b = frame[:, :, 0], frame[:, :, 1], frame[:, :, 2]

    if radial:
        # Create distance-from-center map
        y_coords, x_coords = np.ogrid[:h, :w]
        cx, cy = w * center_x, h * center_y
        dist = np.sqrt((x_coords - cx)**2 + (y_coords - cy)**2)
        max_dist = np.sqrt(cx**2 + cy**2)
        dist_normalized = (dist / max_dist).astype(np.float32)

        # Create coordinate maps for remapping
        map_x = np.tile(np.arange(w, dtype=np.float32), (h, 1))
        map_y = np.tile(np.arange(h, dtype=np.float32).reshape(-1, 1), (1, w))

        # Direction from center
        dx = (map_x - cx) / (dist + 1e-6)
        dy = (map_y - cy) / (dist + 1e-6)

        # Apply radial offset to red channel (outward)
        r_offset = strength * dist_normalized
        r_map_x = (map_x + dx * r_offset).astype(np.float32)
        r_map_y = (map_y + dy * r_offset).astype(np.float32)
        r_shifted = cv2.remap(r, r_map_x, r_map_y,
                              cv2.INTER_LINEAR, borderMode=cv2.BORDER_REPLICATE)

        # Apply radial offset to blue channel (inward)
        b_offset = -strength * dist_normalized
        b_map_x = (map_x + dx * b_offset).astype(np.float32)
        b_map_y = (map_y + dy * b_offset).astype(np.float32)
        b_shifted = cv2.remap(b, b_map_x, b_map_y,
                              cv2.INTER_LINEAR, borderMode=cv2.BORDER_REPLICATE)

        return np.stack([r_shifted, g, b_shifted], axis=-1).astype(np.uint8), state
    else:
        # Simple uniform offset
        offset = int(strength)
        M_r = np.float32([[1, 0, offset], [0, 1, 0]])
        M_b = np.float32([[1, 0, -offset], [0, 1, 0]])

        r_shifted = cv2.warpAffine(r, M_r, (w, h), borderMode=cv2.BORDER_REPLICATE)
        b_shifted = cv2.warpAffine(b, M_b, (w, h), borderMode=cv2.BORDER_REPLICATE)

        return np.stack([r_shifted, g, b_shifted], axis=-1).astype(np.uint8), state