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

@description
Scatter effect. Randomly redistributes pixels in local area creating
a dissolving, dispersed look.

@param amount float
  @range 0 100
  @default 10
  Scatter radius in pixels. Bind to energy for reactive dissolve.

@param randomize_per_frame bool
  @default false
  Different scatter pattern each frame (vs. static pattern).

@param grain float
  @range 0 1
  @default 0
  Add film grain to scattered result.

@param seed int
  @default 42
  Random seed for reproducible patterns.

@state rng DeterministicRNG
  Random number generator for displacement.

@state displacement_map ndarray
  Cached displacement map for static mode.

@example
  (effect scatter :amount 20)

@example
  ;; Reactive dissolve
  (effect scatter :amount (bind energy :range [5 50]) :randomize_per_frame true)
"""

import numpy as np
from scipy import ndimage
from pathlib import Path
import sys

# Import DeterministicRNG from same directory
_effects_dir = Path(__file__).parent
if str(_effects_dir) not in sys.path:
    sys.path.insert(0, str(_effects_dir))
from random import DeterministicRNG


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

    Args:
        frame: Input frame as numpy array (H, W, 3) RGB uint8
        params: Effect parameters
            - amount: scatter radius (default 10)
            - randomize_per_frame: animate pattern (default False)
            - grain: add film grain (default 0)
            - seed: random seed (default 42)
        state: Persistent state dict

    Returns:
        Tuple of (processed_frame, new_state)
    """
    amount = int(np.clip(params.get("amount", 10), 0, 100))
    randomize = params.get("randomize_per_frame", False)
    grain = params.get("grain", 0)
    seed = int(params.get("seed", 42))

    if state is None:
        state = {}

    if amount < 1:
        return frame, state

    # Initialize RNG
    if "rng" not in state:
        state["rng"] = DeterministicRNG(seed)
    rng = state["rng"]

    h, w = frame.shape[:2]

    # Generate or reuse displacement map
    last_size = state.get("last_size")
    if randomize or "displacement_map" not in state or last_size != (h, w):
        # Generate new displacement map
        displacement = np.zeros((h, w, 2), dtype=np.float32)
        for y in range(h):
            for x in range(w):
                displacement[y, x, 0] = rng.uniform(-amount, amount)
                displacement[y, x, 1] = rng.uniform(-amount, amount)
        state["displacement_map"] = displacement
        state["last_size"] = (h, w)

    displacement_map = state["displacement_map"]

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

    # Apply displacement
    new_y = y_coords + displacement_map[:, :, 0]
    new_x = x_coords + displacement_map[:, :, 1]

    # Clip to valid range
    new_y = np.clip(new_y, 0, h - 1)
    new_x = np.clip(new_x, 0, w - 1)

    # Sample from displaced positions
    result = np.zeros_like(frame)
    for c in range(frame.shape[2] if len(frame.shape) > 2 else 1):
        if len(frame.shape) > 2:
            result[:, :, c] = ndimage.map_coordinates(
                frame[:, :, c],
                [new_y, new_x],
                order=1,
                mode='reflect'
            )
        else:
            result = ndimage.map_coordinates(
                frame,
                [new_y, new_x],
                order=1,
                mode='reflect'
            )

    # Add grain
    if grain > 0:
        noise = np.array([[[rng.uniform(-grain * 20, grain * 20) for _ in range(3)]
                          for _ in range(w)]
                         for _ in range(h)])
        result = np.clip(result.astype(np.float32) + noise, 0, 255).astype(np.uint8)

    return result, state