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

@description
Deterministic random number generator for use in recipes and effects.
Given the same seed, produces the same sequence of values every time.

This effect doesn't modify the frame - it provides random values that
can be bound to other effect parameters. The random state persists
across frames for consistent sequences.

@param seed int
  @default 42
  Random seed for reproducibility. Same seed = same sequence.

@param min float
  @default 0
  Minimum output value.

@param max float
  @default 1
  Maximum output value.

@param mode string
  @enum uniform gaussian integer choice
  @default uniform
  Distribution type:
  - uniform: even distribution between min and max
  - gaussian: normal distribution (min=mean, max=stddev)
  - integer: random integers between min and max (inclusive)
  - choice: randomly pick from a list (use choices param)

@param choices list
  @default []
  List of values to choose from (for mode=choice).

@param step_every int
  @default 1
  Only generate new value every N frames (1 = every frame).

@state rng RandomState
  Numpy random state for deterministic sequence.

@state frame_count int
  Tracks frames for step_every.

@state current_value float
  Current random value (persists between steps).

@example
  ;; Random value 0-1 each frame
  (bind (random :seed 123))

@example
  ;; Random integer 1-10, changes every 5 frames
  (random :seed 42 :mode "integer" :min 1 :max 10 :step_every 5)

@example
  ;; Gaussian noise around 0.5
  (random :mode "gaussian" :min 0.5 :max 0.1)
"""

import numpy as np


def process_frame(frame: np.ndarray, params: dict, state: dict) -> tuple:
    """
    Generate deterministic random values.

    This effect passes through the frame unchanged but updates state
    with random values that can be used by the recipe/executor.

    Args:
        frame: Input frame (passed through unchanged)
        params: Effect parameters
            - seed: random seed (default 42)
            - min: minimum value (default 0)
            - max: maximum value (default 1)
            - mode: uniform/gaussian/integer/choice (default uniform)
            - choices: list for choice mode
            - step_every: frames between new values (default 1)
        state: Persistent state dict
            - rng: numpy RandomState
            - frame_count: frame counter
            - current_value: last generated value

    Returns:
        Tuple of (frame, state_with_random_value)
    """
    seed = int(params.get("seed", 42))
    min_val = params.get("min", 0)
    max_val = params.get("max", 1)
    mode = params.get("mode", "uniform")
    choices = params.get("choices", [])
    step_every = max(1, int(params.get("step_every", 1)))

    if state is None:
        state = {}

    # Initialize RNG on first call
    if "rng" not in state:
        state["rng"] = np.random.RandomState(seed)
        state["frame_count"] = 0
        state["current_value"] = None

    rng = state["rng"]
    frame_count = state["frame_count"]

    # Generate new value if needed
    if frame_count % step_every == 0 or state["current_value"] is None:
        if mode == "uniform":
            value = rng.uniform(min_val, max_val)
        elif mode == "gaussian":
            # min = mean, max = stddev
            value = rng.normal(min_val, max_val)
        elif mode == "integer":
            value = rng.randint(int(min_val), int(max_val) + 1)
        elif mode == "choice" and choices:
            value = choices[rng.randint(0, len(choices))]
        else:
            value = rng.uniform(min_val, max_val)

        state["current_value"] = value

    state["frame_count"] = frame_count + 1

    # Store value in state for recipe access
    state["value"] = state["current_value"]

    return frame, state


# Standalone RNG class for use in other effects
class DeterministicRNG:
    """
    Deterministic random number generator for use in effects.

    Usage in effects:
        from effects.random import DeterministicRNG

        def process_frame(frame, params, state):
            if "rng" not in state:
                state["rng"] = DeterministicRNG(params.get("seed", 42))
            rng = state["rng"]

            value = rng.uniform(0, 1)
            integer = rng.randint(0, 10)
            choice = rng.choice(["a", "b", "c"])
    """

    def __init__(self, seed: int = 42):
        """Initialize with seed for reproducibility."""
        self._rng = np.random.RandomState(seed)
        self._seed = seed

    def seed(self, seed: int):
        """Reset with new seed."""
        self._rng = np.random.RandomState(seed)
        self._seed = seed

    def uniform(self, low: float = 0, high: float = 1) -> float:
        """Random float in [low, high)."""
        return self._rng.uniform(low, high)

    def randint(self, low: int, high: int) -> int:
        """Random integer in [low, high]."""
        return self._rng.randint(low, high + 1)

    def gaussian(self, mean: float = 0, stddev: float = 1) -> float:
        """Random float from normal distribution."""
        return self._rng.normal(mean, stddev)

    def choice(self, items: list):
        """Random choice from list."""
        if not items:
            return None
        return items[self._rng.randint(0, len(items))]

    def shuffle(self, items: list) -> list:
        """Return shuffled copy of list."""
        result = list(items)
        self._rng.shuffle(result)
        return result

    def sample(self, items: list, n: int) -> list:
        """Random sample of n items without replacement."""
        if n >= len(items):
            return self.shuffle(items)
        indices = self._rng.choice(len(items), n, replace=False)
        return [items[i] for i in indices]

    def weighted_choice(self, items: list, weights: list):
        """Random choice with weights."""
        if not items or not weights:
            return None
        weights = np.array(weights, dtype=float)
        weights /= weights.sum()
        idx = self._rng.choice(len(items), p=weights)
        return items[idx]

    @property
    def state(self) -> dict:
        """Get RNG state for serialization."""
        return {"seed": self._seed, "state": self._rng.get_state()}

    @classmethod
    def from_state(cls, state: dict) -> 'DeterministicRNG':
        """Restore RNG from serialized state."""
        rng = cls(state["seed"])
        rng._rng.set_state(state["state"])
        return rng