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

@description
Posterize / Color Quantization effect. Reduces the number of colors
to create a poster/cartoon look. Great for stylized visuals.

@param levels int
  @range 2 32
  @default 8
  Number of color levels per channel. Lower = more stylized.

@param dither bool
  @default false
  Apply dithering to reduce color banding.

@example
  (effect posterize :levels 4)

@example
  ;; Beat-reactive posterization
  (effect posterize :levels (bind bass :range [2 16]))
"""

import numpy as np


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

    Args:
        frame: Input frame as numpy array (H, W, 3) RGB uint8
        params: Effect parameters
            - levels: color levels per channel (default 8)
            - dither: apply dithering (default False)
        state: Persistent state dict

    Returns:
        Tuple of (processed_frame, new_state)
    """
    levels = max(2, min(int(params.get("levels", 8)), 256))
    dither = params.get("dither", False)

    if state is None:
        state = {}

    step = 256 // levels

    if dither:
        # Floyd-Steinberg dithering
        result = frame.astype(np.float32).copy()
        h, w = result.shape[:2]

        for y in range(h - 1):
            for x in range(1, w - 1):
                for c in range(3):
                    old_val = result[y, x, c]
                    new_val = np.round(old_val / step) * step
                    result[y, x, c] = new_val
                    error = old_val - new_val

                    # Distribute error to neighbors
                    result[y, x + 1, c] += error * 7 / 16
                    result[y + 1, x - 1, c] += error * 3 / 16
                    result[y + 1, x, c] += error * 5 / 16
                    result[y + 1, x + 1, c] += error * 1 / 16

        return np.clip(result, 0, 255).astype(np.uint8), state
    else:
        # Simple quantization
        quantized = (frame // step) * step
        return quantized.astype(np.uint8), state