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

@description
Digital corruption / glitch block effect. Randomly corrupts rectangular
blocks by shifting, swapping, or duplicating from previous frames.
Simulates video compression artifacts.

@param block_size int
  @range 8 128
  @default 32
  Size of corruption blocks in pixels.

@param corruption float
  @range 0 1
  @default 0.3
  Probability of corrupting each block. Bind to energy for reactive glitch.

@param max_offset int
  @range 0 200
  @default 50
  Maximum pixel offset when shifting blocks.

@param color_corrupt bool
  @default true
  Also apply color channel shifts to blocks.

@param seed int
  @default 42
  Random seed for deterministic glitch patterns.

@state previous_frame ndarray
  Stores previous frame for frame-blending corruption.

@state rng DeterministicRNG
  Random number generator for reproducible results.

@example
  (effect datamosh :corruption 0.4)

@example
  ;; Heavy glitch on energy peaks
  (effect datamosh :corruption (bind energy :range [0 0.8]) :block_size 16)

@example
  ;; Reproducible glitch with seed
  (effect datamosh :corruption 0.5 :seed 12345)
"""

import numpy as np
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 datamosh/glitch block effect to a video frame.

    Args:
        frame: Input frame as numpy array (H, W, 3) RGB uint8
        params: Effect parameters
            - block_size: corruption block size (default 32)
            - corruption: probability 0-1 (default 0.3)
            - max_offset: max shift in pixels (default 50)
            - color_corrupt: apply color shifts (default True)
        state: Persistent state dict
            - previous_frame: last frame for duplication effect

    Returns:
        Tuple of (processed_frame, new_state)
    """
    block_size = max(8, min(int(params.get("block_size", 32)), 128))
    corruption = max(0, min(params.get("corruption", 0.3), 1))
    max_offset = int(params.get("max_offset", 50))
    color_corrupt = params.get("color_corrupt", True)
    seed = int(params.get("seed", 42))

    if state is None:
        state = {}

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

    if corruption == 0:
        state["previous_frame"] = frame.copy()
        return frame, state

    h, w = frame.shape[:2]
    result = frame.copy()
    prev_frame = state.get("previous_frame")

    # Process blocks
    for by in range(0, h, block_size):
        for bx in range(0, w, block_size):
            bh = min(block_size, h - by)
            bw = min(block_size, w - bx)

            if rng.uniform() < corruption:
                corruption_type = rng.choice(["shift", "duplicate", "color", "swap"])

                if corruption_type == "shift" and max_offset > 0:
                    ox = rng.randint(-max_offset, max_offset)
                    oy = rng.randint(-max_offset, max_offset)
                    src_x = max(0, min(bx + ox, w - bw))
                    src_y = max(0, min(by + oy, h - bh))
                    result[by:by+bh, bx:bx+bw] = frame[src_y:src_y+bh, src_x:src_x+bw]

                elif corruption_type == "duplicate" and prev_frame is not None:
                    if prev_frame.shape == frame.shape:
                        result[by:by+bh, bx:bx+bw] = prev_frame[by:by+bh, bx:bx+bw]

                elif corruption_type == "color" and color_corrupt:
                    block = result[by:by+bh, bx:bx+bw].copy()
                    shift = rng.randint(1, 3)
                    channel = rng.randint(0, 2)
                    block[:, :, channel] = np.roll(block[:, :, channel], shift, axis=0)
                    result[by:by+bh, bx:bx+bw] = block

                elif corruption_type == "swap":
                    other_bx = rng.randint(0, max(0, w - bw - 1))
                    other_by = rng.randint(0, max(0, h - bh - 1))
                    temp = result[by:by+bh, bx:bx+bw].copy()
                    result[by:by+bh, bx:bx+bw] = frame[other_by:other_by+bh, other_bx:other_bx+bw]
                    result[other_by:other_by+bh, other_bx:other_bx+bw] = temp

    state["previous_frame"] = frame.copy()
    return result, state