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

@description
Shatter effect. Explodes the image into flying pieces that move
outward from a center point. Great for beat drops and transitions.

@param intensity float
  @range 0 1
  @default 0.5
  Explosion force (0 = no effect, 1 = full explosion).

@param num_pieces int
  @range 10 200
  @default 50
  Number of shatter pieces.

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

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

@param rotation_speed float
  @range 0 5
  @default 1.0
  How fast pieces rotate as they fly.

@param gravity float
  @range 0 2
  @default 0.3
  Downward pull on pieces.

@param fade_out bool
  @default true
  Fade pieces as they fly away.

@param seed int
  @default 42
  Random seed for piece positions.

@state pieces list
  List of piece positions and velocities.

@example
  (effect shatter :intensity 0.7 :num_pieces 80)

@example
  ;; Beat-reactive explosion
  (effect shatter :intensity (bind onset :range [0 1]) :gravity 0.5)
"""

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 shatter effect to a video frame.

    Args:
        frame: Input frame as numpy array (H, W, 3) RGB uint8
        params: Effect parameters
        state: Persistent state dict

    Returns:
        Tuple of (processed_frame, new_state)
    """
    intensity = np.clip(params.get("intensity", 0.5), 0, 1)
    num_pieces = max(10, min(int(params.get("num_pieces", 50)), 200))
    center_x = params.get("center_x", 0.5)
    center_y = params.get("center_y", 0.5)
    rotation_speed = params.get("rotation_speed", 1.0)
    gravity = params.get("gravity", 0.3)
    fade_out = params.get("fade_out", True)
    seed = int(params.get("seed", 42))
    t = params.get("_time", 0)

    if state is None:
        state = {}

    if intensity < 0.01:
        return frame, state

    h, w = frame.shape[:2]
    cx, cy = int(center_x * w), int(center_y * h)

    # Initialize pieces
    if "pieces" not in state or len(state["pieces"]) != num_pieces:
        state["pieces"] = _init_pieces(w, h, num_pieces, cx, cy, seed)
        state["start_time"] = t

    pieces = state["pieces"]
    start_time = state.get("start_time", t)
    effect_t = t - start_time

    # Create output with black background
    result = np.zeros_like(frame)

    time_factor = effect_t * intensity * 3.0

    # Draw each piece
    for piece in pieces:
        px, py, pw, ph, vx, vy, rot = piece

        if pw <= 0 or ph <= 0:
            continue

        # Calculate current position with physics
        curr_x = px + vx * time_factor * w * 0.5
        curr_y = py + vy * time_factor * h * 0.5 + gravity * time_factor ** 2 * h * 0.2
        curr_rot = rot * rotation_speed * time_factor * 180

        # Calculate alpha (fade out over distance)
        distance = np.sqrt((curr_x - px)**2 + (curr_y - py)**2)
        alpha = 1.0 - (distance / max(w, h)) if fade_out else 1.0
        alpha = max(0, min(1, alpha))

        if alpha < 0.05:
            continue

        # Extract piece from original frame
        px1, py1 = max(0, int(px)), max(0, int(py))
        px2, py2 = min(w, int(px + pw)), min(h, int(py + ph))

        if px2 <= px1 or py2 <= py1:
            continue

        piece_img = frame[py1:py2, px1:px2].copy()

        # Rotate piece
        if abs(curr_rot) > 1:
            piece_img = ndimage.rotate(piece_img, curr_rot, reshape=False, mode='constant', cval=0)

        # Calculate destination
        dest_x = int(curr_x)
        dest_y = int(curr_y)
        piece_h, piece_w = piece_img.shape[:2]

        # Clip to frame bounds
        src_x1, src_y1 = 0, 0
        src_x2, src_y2 = piece_w, piece_h
        dst_x1, dst_y1 = dest_x, dest_y
        dst_x2, dst_y2 = dest_x + piece_w, dest_y + piece_h

        if dst_x1 < 0:
            src_x1 = -dst_x1
            dst_x1 = 0
        if dst_y1 < 0:
            src_y1 = -dst_y1
            dst_y1 = 0
        if dst_x2 > w:
            src_x2 -= (dst_x2 - w)
            dst_x2 = w
        if dst_y2 > h:
            src_y2 -= (dst_y2 - h)
            dst_y2 = h

        if dst_x2 <= dst_x1 or dst_y2 <= dst_y1:
            continue
        if src_x2 <= src_x1 or src_y2 <= src_y1:
            continue

        # Blend piece onto result
        piece_region = piece_img[src_y1:src_y2, src_x1:src_x2]
        if piece_region.size == 0:
            continue

        result_region = result[dst_y1:dst_y2, dst_x1:dst_x2]
        if result_region.shape != piece_region.shape:
            continue

        result[dst_y1:dst_y2, dst_x1:dst_x2] = (
            result_region * (1 - alpha) + piece_region * alpha
        ).astype(np.uint8)

    # Blend with original based on intensity
    final = ((1 - intensity) * frame + intensity * result).astype(np.uint8)

    return final, state


def _init_pieces(w: int, h: int, num_pieces: int, cx: int, cy: int, seed: int) -> list:
    """Initialize shatter pieces with random positions and velocities."""
    rng = DeterministicRNG(seed)

    pieces = []

    # Create grid of pieces
    cols = max(1, int(np.sqrt(num_pieces * w / h)))
    rows = max(1, int(num_pieces / cols))
    piece_w = w // cols
    piece_h = h // rows

    for row in range(rows):
        for col in range(cols):
            px = col * piece_w
            py = row * piece_h
            pw = piece_w + (w % cols if col == cols - 1 else 0)
            ph = piece_h + (h % rows if row == rows - 1 else 0)

            # Velocity away from center
            piece_cx = px + pw // 2
            piece_cy = py + ph // 2
            dx = piece_cx - cx
            dy = piece_cy - cy
            dist = max(1, np.sqrt(dx*dx + dy*dy))

            vx = dx / dist + rng.uniform(-0.3, 0.3)
            vy = dy / dist + rng.uniform(-0.3, 0.3)
            rot = rng.uniform(-2, 2)

            pieces.append((px, py, pw, ph, vx, vy, rot))

    return pieces