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

@description
Chromatic aberration / RGB channel separation. Offsets red and blue channels
in different directions, creating color fringing. Classic glitch aesthetic.

@param amount float
  @range 0 50
  @default 10
  Offset amount in pixels. Bind to bass for punchy glitch effect.

@param angle float
  @range 0 360
  @default 0
  Direction of split in degrees. 0 = horizontal, 90 = vertical.

@param red_offset float
  @range -50 50
  @default 0
  Override: specific red channel X offset (ignores amount/angle if set).

@param blue_offset float
  @range -50 50
  @default 0
  Override: specific blue channel X offset (ignores amount/angle if set).

@example
  (effect rgb_split :amount 15)

@example
  ;; Bass-reactive chromatic aberration
  (effect rgb_split :amount (bind bass :range [0 30] :transform sqrt))

@example
  ;; Vertical split
  (effect rgb_split :amount 20 :angle 90)
"""

import numpy as np
import cv2
import math


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

    Args:
        frame: Input frame as numpy array (H, W, 3) RGB uint8
        params: Effect parameters
            - amount: offset in pixels (default 10)
            - angle: split direction in degrees (default 0)
            - red_offset: override red X offset
            - blue_offset: override blue X offset
        state: Persistent state dict (unused)

    Returns:
        Tuple of (processed_frame, new_state)
    """
    amount = params.get("amount", 10)
    angle = params.get("angle", 0)
    red_override = params.get("red_offset")
    blue_override = params.get("blue_offset")

    # Calculate offsets
    if red_override is not None or blue_override is not None:
        # Use explicit offsets
        r_x = int(red_override or 0)
        r_y = 0
        b_x = int(blue_override or 0)
        b_y = 0
    else:
        # Calculate from amount and angle
        angle_rad = math.radians(angle)
        r_x = int(amount * math.cos(angle_rad))
        r_y = int(amount * math.sin(angle_rad))
        b_x = -r_x  # Blue goes opposite direction
        b_y = -r_y

    if r_x == 0 and r_y == 0 and b_x == 0 and b_y == 0:
        return frame, state

    h, w = frame.shape[:2]

    # Split channels
    r, g, b = frame[:, :, 0], frame[:, :, 1], frame[:, :, 2]

    # Create translation matrices
    M_r = np.float32([[1, 0, r_x], [0, 1, r_y]])
    M_b = np.float32([[1, 0, b_x], [0, 1, b_y]])

    # Translate red and blue channels
    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)

    # Merge channels
    result = np.stack([r_shifted, g, b_shifted], axis=-1)

    return result, state