"""
Xector Primitives - Parallel array operations for GPU-style data parallelism.

Inspired by Connection Machine Lisp and hillisp. Xectors are parallel arrays
where operations automatically apply element-wise.

Usage in sexp:
    (require-primitives "xector")

    ;; Extract channels as xectors
    (let* ((r (red frame))
           (g (green frame))
           (b (blue frame))
           ;; Operations are element-wise on xectors
           (brightness (α+ (α* r 0.299) (α* g 0.587) (α* b 0.114))))
      ;; Reduce to scalar
      (βmax brightness))

    ;; Explicit α for element-wise, implicit also works
    (α+ r 10)     ;; explicit: add 10 to every element
    (+ r 10)      ;; implicit: same thing when r is a xector

    ;; β for reductions
    (β+ r)        ;; sum all elements
    (βmax r)      ;; maximum element
    (βmean r)     ;; average

Operators:
    α (alpha) - element-wise: (α+ x y) adds corresponding elements
    β (beta)  - reduce: (β+ x) sums all elements
"""

import numpy as np
from typing import Union, Callable, Any

# Try to use CuPy for GPU acceleration if available
try:
    import cupy as cp
    HAS_CUPY = True
except ImportError:
    cp = None
    HAS_CUPY = False


class Xector:
    """
    Parallel array type for element-wise operations.

    Wraps a numpy/cupy array and provides automatic broadcasting
    and element-wise operation semantics.
    """

    def __init__(self, data, shape=None):
        """
        Create a Xector from data.

        Args:
            data: numpy array, cupy array, scalar, or list
            shape: optional shape tuple (for coordinate xectors)
        """
        if isinstance(data, Xector):
            self._data = data._data
            self._shape = data._shape
        elif isinstance(data, np.ndarray):
            self._data = data.astype(np.float32)
            self._shape = shape or data.shape
        elif HAS_CUPY and isinstance(data, cp.ndarray):
            self._data = data.astype(cp.float32)
            self._shape = shape or data.shape
        elif isinstance(data, (list, tuple)):
            self._data = np.array(data, dtype=np.float32)
            self._shape = shape or self._data.shape
        else:
            # Scalar - will broadcast
            self._data = np.float32(data)
            self._shape = shape or ()

    @property
    def data(self):
        return self._data

    @property
    def shape(self):
        return self._shape

    def __len__(self):
        return self._data.size

    def __repr__(self):
        if self._data.size <= 10:
            return f"Xector({self._data})"
        return f"Xector(shape={self._shape}, size={self._data.size})"

    def to_numpy(self):
        """Convert to numpy array."""
        if HAS_CUPY and isinstance(self._data, cp.ndarray):
            return cp.asnumpy(self._data)
        return self._data

    def to_gpu(self):
        """Move to GPU if available."""
        if HAS_CUPY and not isinstance(self._data, cp.ndarray):
            self._data = cp.asarray(self._data)
        return self

    # Arithmetic operators - enable implicit element-wise ops
    def __add__(self, other):
        other_data = other._data if isinstance(other, Xector) else other
        return Xector(self._data + other_data, self._shape)

    def __radd__(self, other):
        return Xector(other + self._data, self._shape)

    def __sub__(self, other):
        other_data = other._data if isinstance(other, Xector) else other
        return Xector(self._data - other_data, self._shape)

    def __rsub__(self, other):
        return Xector(other - self._data, self._shape)

    def __mul__(self, other):
        other_data = other._data if isinstance(other, Xector) else other
        return Xector(self._data * other_data, self._shape)

    def __rmul__(self, other):
        return Xector(other * self._data, self._shape)

    def __truediv__(self, other):
        other_data = other._data if isinstance(other, Xector) else other
        return Xector(self._data / other_data, self._shape)

    def __rtruediv__(self, other):
        return Xector(other / self._data, self._shape)

    def __pow__(self, other):
        other_data = other._data if isinstance(other, Xector) else other
        return Xector(self._data ** other_data, self._shape)

    def __neg__(self):
        return Xector(-self._data, self._shape)

    def __abs__(self):
        return Xector(np.abs(self._data), self._shape)

    # Comparison operators - return boolean xectors
    def __lt__(self, other):
        other_data = other._data if isinstance(other, Xector) else other
        return Xector(self._data < other_data, self._shape)

    def __le__(self, other):
        other_data = other._data if isinstance(other, Xector) else other
        return Xector(self._data <= other_data, self._shape)

    def __gt__(self, other):
        other_data = other._data if isinstance(other, Xector) else other
        return Xector(self._data > other_data, self._shape)

    def __ge__(self, other):
        other_data = other._data if isinstance(other, Xector) else other
        return Xector(self._data >= other_data, self._shape)

    def __eq__(self, other):
        other_data = other._data if isinstance(other, Xector) else other
        return Xector(self._data == other_data, self._shape)

    def __ne__(self, other):
        other_data = other._data if isinstance(other, Xector) else other
        return Xector(self._data != other_data, self._shape)


def _unwrap(x):
    """Unwrap Xector to underlying data, or return as-is."""
    if isinstance(x, Xector):
        return x._data
    return x


def _wrap(data, shape=None):
    """Wrap result in Xector if it's an array."""
    if isinstance(data, (np.ndarray,)) or (HAS_CUPY and isinstance(data, cp.ndarray)):
        return Xector(data, shape)
    return data


# =============================================================================
# Frame/Xector Conversion
# =============================================================================
# NOTE: red, green, blue, gray are derived in derived.sexp using (channel frame n)

def xector_from_frame(frame):
    """Convert entire frame to xector (flattened RGB). (xector frame) -> Xector"""
    if isinstance(frame, np.ndarray):
        return Xector(frame.flatten().astype(np.float32), frame.shape)
    raise TypeError(f"Expected frame array, got {type(frame)}")


def xector_to_frame(x, shape=None):
    """Convert xector back to frame. (to-frame x) or (to-frame x shape) -> frame"""
    data = _unwrap(x)
    if shape is None and isinstance(x, Xector):
        shape = x._shape
    if shape is None:
        raise ValueError("Shape required to convert xector to frame")
    return np.clip(data, 0, 255).reshape(shape).astype(np.uint8)


# =============================================================================
# Coordinate Generators
# =============================================================================
# NOTE: x-coords, y-coords, x-norm, y-norm, dist-from-center are derived
# in derived.sexp using iota, tile, repeat primitives


# =============================================================================
# Alpha (α) - Element-wise Operations
# =============================================================================

def alpha_lift(fn):
    """Lift a scalar function to work element-wise on xectors."""
    def lifted(*args):
        # Check if any arg is a Xector
        has_xector = any(isinstance(a, Xector) for a in args)
        if not has_xector:
            return fn(*args)

        # Get shape from first xector
        shape = None
        for a in args:
            if isinstance(a, Xector):
                shape = a._shape
                break

        # Unwrap all args
        unwrapped = [_unwrap(a) for a in args]

        # Apply function
        result = fn(*unwrapped)

        return _wrap(result, shape)

    return lifted


# Element-wise math operations
def alpha_add(*args):
    """Element-wise addition. (α+ a b ...) -> Xector"""
    if len(args) == 0:
        return 0
    result = _unwrap(args[0])
    for a in args[1:]:
        result = result + _unwrap(a)
    return _wrap(result, args[0]._shape if isinstance(args[0], Xector) else None)


def alpha_sub(a, b=None):
    """Element-wise subtraction. (α- a b) -> Xector"""
    if b is None:
        return Xector(-_unwrap(a)) if isinstance(a, Xector) else -a
    shape = a._shape if isinstance(a, Xector) else (b._shape if isinstance(b, Xector) else None)
    return _wrap(_unwrap(a) - _unwrap(b), shape)


def alpha_mul(*args):
    """Element-wise multiplication. (α* a b ...) -> Xector"""
    if len(args) == 0:
        return 1
    result = _unwrap(args[0])
    for a in args[1:]:
        result = result * _unwrap(a)
    return _wrap(result, args[0]._shape if isinstance(args[0], Xector) else None)


def alpha_div(a, b):
    """Element-wise division. (α/ a b) -> Xector"""
    shape = a._shape if isinstance(a, Xector) else (b._shape if isinstance(b, Xector) else None)
    return _wrap(_unwrap(a) / _unwrap(b), shape)


def alpha_pow(a, b):
    """Element-wise power. (α** a b) -> Xector"""
    shape = a._shape if isinstance(a, Xector) else (b._shape if isinstance(b, Xector) else None)
    return _wrap(_unwrap(a) ** _unwrap(b), shape)


def alpha_sqrt(x):
    """Element-wise square root. (αsqrt x) -> Xector"""
    return _wrap(np.sqrt(_unwrap(x)), x._shape if isinstance(x, Xector) else None)


def alpha_abs(x):
    """Element-wise absolute value. (αabs x) -> Xector"""
    return _wrap(np.abs(_unwrap(x)), x._shape if isinstance(x, Xector) else None)


def alpha_sin(x):
    """Element-wise sine. (αsin x) -> Xector"""
    return _wrap(np.sin(_unwrap(x)), x._shape if isinstance(x, Xector) else None)


def alpha_cos(x):
    """Element-wise cosine. (αcos x) -> Xector"""
    return _wrap(np.cos(_unwrap(x)), x._shape if isinstance(x, Xector) else None)


def alpha_exp(x):
    """Element-wise exponential. (αexp x) -> Xector"""
    return _wrap(np.exp(_unwrap(x)), x._shape if isinstance(x, Xector) else None)


def alpha_log(x):
    """Element-wise natural log. (αlog x) -> Xector"""
    return _wrap(np.log(_unwrap(x)), x._shape if isinstance(x, Xector) else None)


# NOTE: alpha_clamp is derived in derived.sexp as (max2 lo (min2 hi x))

def alpha_min(a, b):
    """Element-wise minimum. (αmin a b) -> Xector"""
    shape = a._shape if isinstance(a, Xector) else (b._shape if isinstance(b, Xector) else None)
    return _wrap(np.minimum(_unwrap(a), _unwrap(b)), shape)


def alpha_max(a, b):
    """Element-wise maximum. (αmax a b) -> Xector"""
    shape = a._shape if isinstance(a, Xector) else (b._shape if isinstance(b, Xector) else None)
    return _wrap(np.maximum(_unwrap(a), _unwrap(b)), shape)


def alpha_mod(a, b):
    """Element-wise modulo. (αmod a b) -> Xector"""
    shape = a._shape if isinstance(a, Xector) else (b._shape if isinstance(b, Xector) else None)
    return _wrap(_unwrap(a) % _unwrap(b), shape)


def alpha_floor(x):
    """Element-wise floor. (αfloor x) -> Xector"""
    return _wrap(np.floor(_unwrap(x)), x._shape if isinstance(x, Xector) else None)


def alpha_ceil(x):
    """Element-wise ceiling. (αceil x) -> Xector"""
    return _wrap(np.ceil(_unwrap(x)), x._shape if isinstance(x, Xector) else None)


def alpha_round(x):
    """Element-wise round. (αround x) -> Xector"""
    return _wrap(np.round(_unwrap(x)), x._shape if isinstance(x, Xector) else None)


# NOTE: alpha_sq is derived in derived.sexp as (* x x)

# Comparison operators (return boolean xectors)
def alpha_lt(a, b):
    """Element-wise less than. (α< a b) -> Xector[bool]"""
    shape = a._shape if isinstance(a, Xector) else (b._shape if isinstance(b, Xector) else None)
    return _wrap(_unwrap(a) < _unwrap(b), shape)


def alpha_le(a, b):
    """Element-wise less-or-equal. (α<= a b) -> Xector[bool]"""
    shape = a._shape if isinstance(a, Xector) else (b._shape if isinstance(b, Xector) else None)
    return _wrap(_unwrap(a) <= _unwrap(b), shape)


def alpha_gt(a, b):
    """Element-wise greater than. (α> a b) -> Xector[bool]"""
    shape = a._shape if isinstance(a, Xector) else (b._shape if isinstance(b, Xector) else None)
    return _wrap(_unwrap(a) > _unwrap(b), shape)


def alpha_ge(a, b):
    """Element-wise greater-or-equal. (α>= a b) -> Xector[bool]"""
    shape = a._shape if isinstance(a, Xector) else (b._shape if isinstance(b, Xector) else None)
    return _wrap(_unwrap(a) >= _unwrap(b), shape)


def alpha_eq(a, b):
    """Element-wise equality. (α= a b) -> Xector[bool]"""
    shape = a._shape if isinstance(a, Xector) else (b._shape if isinstance(b, Xector) else None)
    return _wrap(_unwrap(a) == _unwrap(b), shape)


# Logical operators
def alpha_and(a, b):
    """Element-wise logical and. (αand a b) -> Xector[bool]"""
    shape = a._shape if isinstance(a, Xector) else (b._shape if isinstance(b, Xector) else None)
    return _wrap(np.logical_and(_unwrap(a), _unwrap(b)), shape)


def alpha_or(a, b):
    """Element-wise logical or. (αor a b) -> Xector[bool]"""
    shape = a._shape if isinstance(a, Xector) else (b._shape if isinstance(b, Xector) else None)
    return _wrap(np.logical_or(_unwrap(a), _unwrap(b)), shape)


def alpha_not(x):
    """Element-wise logical not. (αnot x) -> Xector[bool]"""
    return _wrap(np.logical_not(_unwrap(x)), x._shape if isinstance(x, Xector) else None)


# =============================================================================
# Beta (β) - Reduction Operations
# =============================================================================

def beta_add(x):
    """Sum all elements. (β+ x) -> scalar"""
    return float(np.sum(_unwrap(x)))


def beta_mul(x):
    """Product of all elements. (β* x) -> scalar"""
    return float(np.prod(_unwrap(x)))


def beta_min(x):
    """Minimum element. (βmin x) -> scalar"""
    return float(np.min(_unwrap(x)))


def beta_max(x):
    """Maximum element. (βmax x) -> scalar"""
    return float(np.max(_unwrap(x)))


def beta_mean(x):
    """Mean of all elements. (βmean x) -> scalar"""
    return float(np.mean(_unwrap(x)))


def beta_std(x):
    """Standard deviation. (βstd x) -> scalar"""
    return float(np.std(_unwrap(x)))


def beta_count(x):
    """Count of elements. (βcount x) -> scalar"""
    return int(np.size(_unwrap(x)))


def beta_any(x):
    """True if any element is truthy. (βany x) -> bool"""
    return bool(np.any(_unwrap(x)))


def beta_all(x):
    """True if all elements are truthy. (βall x) -> bool"""
    return bool(np.all(_unwrap(x)))


# =============================================================================
# Conditional / Selection
# =============================================================================

def xector_where(cond, true_val, false_val):
    """
    Conditional select. (where cond true-val false-val) -> Xector

    Like numpy.where - selects elements based on condition.
    """
    cond_data = _unwrap(cond)
    true_data = _unwrap(true_val)
    false_data = _unwrap(false_val)

    # Get shape from condition or values
    shape = None
    for x in [cond, true_val, false_val]:
        if isinstance(x, Xector):
            shape = x._shape
            break

    result = np.where(cond_data, true_data, false_data)
    return _wrap(result, shape)


# NOTE: fill, zeros, ones are derived in derived.sexp using iota

def xector_rand(size_or_frame):
    """Create xector of random values [0,1). (rand-x frame) -> Xector"""
    if isinstance(size_or_frame, np.ndarray):
        h, w = size_or_frame.shape[:2]
        size = h * w
        shape = (h, w)
    elif isinstance(size_or_frame, Xector):
        size = len(size_or_frame)
        shape = size_or_frame._shape
    else:
        size = int(size_or_frame)
        shape = (size,)

    return Xector(np.random.random(size).astype(np.float32), shape)


def xector_randn(size_or_frame, mean=0, std=1):
    """Create xector of normal random values. (randn-x frame) or (randn-x frame mean std) -> Xector"""
    if isinstance(size_or_frame, np.ndarray):
        h, w = size_or_frame.shape[:2]
        size = h * w
        shape = (h, w)
    elif isinstance(size_or_frame, Xector):
        size = len(size_or_frame)
        shape = size_or_frame._shape
    else:
        size = int(size_or_frame)
        shape = (size,)

    return Xector((np.random.randn(size) * std + mean).astype(np.float32), shape)


# =============================================================================
# Type checking
# =============================================================================

def is_xector(x):
    """Check if x is a Xector. (xector? x) -> bool"""
    return isinstance(x, Xector)


# =============================================================================
# CORE PRIMITIVES: gather, scatter, group-reduce, reshape
# These are the fundamental operations everything else builds on.
# =============================================================================

def xector_gather(data, indices):
    """
    Parallel index lookup. (gather data indices) -> Xector

    For each index in indices, look up the corresponding value in data.
    This is the fundamental operation for remapping/resampling.

    Example:
        (gather [10 20 30 40] [2 0 1 2])  ; -> [30 10 20 30]
    """
    data_arr = _unwrap(data)
    idx_arr = _unwrap(indices).astype(np.int32)

    # Flatten data for 1D indexing
    flat_data = data_arr.flatten()

    # Clip indices to valid range
    idx_clipped = np.clip(idx_arr, 0, len(flat_data) - 1)

    result = flat_data[idx_clipped]
    shape = indices._shape if isinstance(indices, Xector) else None
    return Xector(result, shape)


def xector_gather_2d(data, row_indices, col_indices):
    """
    2D parallel index lookup. (gather-2d data rows cols) -> Xector

    For each (row, col) pair, look up the value in 2D data.
    Essential for grid/cell operations.

    Example:
        (gather-2d image-lum cell-rows cell-cols)
    """
    data_arr = _unwrap(data)
    row_arr = _unwrap(row_indices).astype(np.int32)
    col_arr = _unwrap(col_indices).astype(np.int32)

    # Get data shape
    if isinstance(data, Xector) and data._shape and len(data._shape) >= 2:
        h, w = data._shape[:2]
        data_2d = data_arr.reshape(h, w)
    elif len(data_arr.shape) >= 2:
        h, w = data_arr.shape[:2]
        data_2d = data_arr.reshape(h, w) if data_arr.ndim == 1 else data_arr
    else:
        # Assume square
        size = int(np.sqrt(len(data_arr)))
        h, w = size, size
        data_2d = data_arr.reshape(h, w)

    # Clip indices
    row_clipped = np.clip(row_arr, 0, h - 1)
    col_clipped = np.clip(col_arr, 0, w - 1)

    result = data_2d[row_clipped.flatten(), col_clipped.flatten()]
    shape = row_indices._shape if isinstance(row_indices, Xector) else None
    return Xector(result, shape)


def xector_scatter(indices, values, size):
    """
    Parallel index write. (scatter indices values size) -> Xector

    Create a new xector of given size, writing values at indices.
    Later writes overwrite earlier ones at same index.

    Example:
        (scatter [0 2 4] [10 20 30] 5)  ; -> [10 0 20 0 30]
    """
    idx_arr = _unwrap(indices).astype(np.int32)
    val_arr = _unwrap(values)

    result = np.zeros(int(size), dtype=np.float32)
    idx_clipped = np.clip(idx_arr, 0, int(size) - 1)
    result[idx_clipped] = val_arr

    return Xector(result, (int(size),))


def xector_scatter_add(indices, values, size):
    """
    Parallel index accumulate. (scatter-add indices values size) -> Xector

    Like scatter, but adds to existing values instead of overwriting.
    Useful for histograms, pooling reductions.

    Example:
        (scatter-add [0 0 1] [1 2 3] 3)  ; -> [3 3 0] (1+2 at index 0)
    """
    idx_arr = _unwrap(indices).astype(np.int32)
    val_arr = _unwrap(values)

    result = np.zeros(int(size), dtype=np.float32)
    np.add.at(result, np.clip(idx_arr, 0, int(size) - 1), val_arr)

    return Xector(result, (int(size),))


def xector_group_reduce(values, group_indices, num_groups, op='mean'):
    """
    Reduce values by group. (group-reduce values groups num-groups op) -> Xector

    Groups values by group_indices and reduces each group.
    This is the primitive for pooling operations.

    Args:
        values: Xector of values to reduce
        group_indices: Xector of group assignments (integers)
        num_groups: Number of groups (output size)
        op: 'mean', 'sum', 'max', 'min'

    Example:
        ; Pool 4 values into 2 groups
        (group-reduce [1 2 3 4] [0 0 1 1] 2 "mean")  ; -> [1.5 3.5]
    """
    val_arr = _unwrap(values).flatten()
    grp_arr = _unwrap(group_indices).astype(np.int32).flatten()
    n = int(num_groups)

    if op == 'sum':
        result = np.zeros(n, dtype=np.float32)
        np.add.at(result, grp_arr, val_arr)
    elif op == 'mean':
        sums = np.zeros(n, dtype=np.float32)
        counts = np.zeros(n, dtype=np.float32)
        np.add.at(sums, grp_arr, val_arr)
        np.add.at(counts, grp_arr, 1)
        result = np.divide(sums, counts, out=np.zeros_like(sums), where=counts > 0)
    elif op == 'max':
        result = np.full(n, -np.inf, dtype=np.float32)
        np.maximum.at(result, grp_arr, val_arr)
        result[result == -np.inf] = 0
    elif op == 'min':
        result = np.full(n, np.inf, dtype=np.float32)
        np.minimum.at(result, grp_arr, val_arr)
        result[result == np.inf] = 0
    else:
        raise ValueError(f"Unknown reduce op: {op}")

    return Xector(result, (n,))


def xector_reshape(x, *dims):
    """
    Reshape xector. (reshape x h w) or (reshape x n) -> Xector

    Changes the logical shape of the xector without changing data.
    """
    data = _unwrap(x)
    if len(dims) == 1:
        new_shape = (int(dims[0]),)
    else:
        new_shape = tuple(int(d) for d in dims)

    return Xector(data.reshape(-1), new_shape)


def xector_shape(x):
    """Get shape of xector. (shape x) -> list"""
    if isinstance(x, Xector):
        return list(x._shape) if x._shape else [len(x)]
    if isinstance(x, np.ndarray):
        return list(x.shape)
    return []


def xector_len(x):
    """Get length of xector. (xlen x) -> int"""
    return len(_unwrap(x).flatten())


def xector_iota(n):
    """
    Generate indices 0 to n-1. (iota n) -> Xector

    Fundamental for generating coordinate xectors.

    Example:
        (iota 5)  ; -> [0 1 2 3 4]
    """
    return Xector(np.arange(int(n), dtype=np.float32), (int(n),))


def xector_repeat(x, n):
    """
    Repeat each element n times. (repeat x n) -> Xector

    Example:
        (repeat [1 2 3] 2)  ; -> [1 1 2 2 3 3]
    """
    data = _unwrap(x)
    result = np.repeat(data.flatten(), int(n))
    return Xector(result, (len(result),))


def xector_tile(x, n):
    """
    Tile entire xector n times. (tile x n) -> Xector

    Example:
        (tile [1 2 3] 2)  ; -> [1 2 3 1 2 3]
    """
    data = _unwrap(x)
    result = np.tile(data.flatten(), int(n))
    return Xector(result, (len(result),))


# =============================================================================
# 2D Grid Helpers (built on primitives above)
# =============================================================================

def xector_cell_indices(frame, cell_size):
    """
    Compute cell index for each pixel. (cell-indices frame cell-size) -> Xector

    Returns flat index of which cell each pixel belongs to.
    This is the bridge between pixel-space and cell-space.
    """
    h, w = frame.shape[:2]
    cell_size = int(cell_size)

    rows = h // cell_size
    cols = w // cell_size

    # For each pixel, compute its cell index
    y = np.repeat(np.arange(h), w)  # [0,0,0..., 1,1,1..., ...]
    x = np.tile(np.arange(w), h)    # [0,1,2..., 0,1,2..., ...]

    cell_row = y // cell_size
    cell_col = x // cell_size
    cell_idx = cell_row * cols + cell_col

    # Clip to valid range
    cell_idx = np.clip(cell_idx, 0, rows * cols - 1)

    return Xector(cell_idx.astype(np.float32), (h, w))


def xector_local_x(frame, cell_size):
    """
    X position within each cell [0, cell_size). (local-x frame cell-size) -> Xector
    """
    h, w = frame.shape[:2]
    x = np.tile(np.arange(w), h)
    local = (x % int(cell_size)).astype(np.float32)
    return Xector(local, (h, w))


def xector_local_y(frame, cell_size):
    """
    Y position within each cell [0, cell_size). (local-y frame cell-size) -> Xector
    """
    h, w = frame.shape[:2]
    y = np.repeat(np.arange(h), w)
    local = (y % int(cell_size)).astype(np.float32)
    return Xector(local, (h, w))


def xector_local_x_norm(frame, cell_size):
    """
    Normalized X within cell [0, 1]. (local-x-norm frame cell-size) -> Xector
    """
    h, w = frame.shape[:2]
    cs = int(cell_size)
    x = np.tile(np.arange(w), h)
    local = ((x % cs) / max(1, cs - 1)).astype(np.float32)
    return Xector(local, (h, w))


def xector_local_y_norm(frame, cell_size):
    """
    Normalized Y within cell [0, 1]. (local-y-norm frame cell-size) -> Xector
    """
    h, w = frame.shape[:2]
    cs = int(cell_size)
    y = np.repeat(np.arange(h), w)
    local = ((y % cs) / max(1, cs - 1)).astype(np.float32)
    return Xector(local, (h, w))


def xector_pool_frame(frame, cell_size, op='mean'):
    """
    Pool frame to cell values. (pool-frame frame cell-size) -> (r, g, b, lum) Xectors

    Returns tuple of xectors: (red, green, blue, luminance) for cells.
    """
    h, w = frame.shape[:2]
    cs = int(cell_size)
    rows = h // cs
    cols = w // cs
    num_cells = rows * cols

    # Compute cell indices for each pixel
    y = np.repeat(np.arange(h), w)
    x = np.tile(np.arange(w), h)
    cell_row = np.clip(y // cs, 0, rows - 1)
    cell_col = np.clip(x // cs, 0, cols - 1)
    cell_idx = cell_row * cols + cell_col

    # Extract channels
    r_flat = frame[:, :, 0].flatten().astype(np.float32)
    g_flat = frame[:, :, 1].flatten().astype(np.float32)
    b_flat = frame[:, :, 2].flatten().astype(np.float32)

    # Pool each channel
    def pool_channel(data):
        sums = np.zeros(num_cells, dtype=np.float32)
        counts = np.zeros(num_cells, dtype=np.float32)
        np.add.at(sums, cell_idx, data)
        np.add.at(counts, cell_idx, 1)
        return np.divide(sums, counts, out=np.zeros_like(sums), where=counts > 0)

    r_pooled = pool_channel(r_flat)
    g_pooled = pool_channel(g_flat)
    b_pooled = pool_channel(b_flat)
    lum = 0.299 * r_pooled + 0.587 * g_pooled + 0.114 * b_pooled

    shape = (rows, cols)
    return (Xector(r_pooled, shape),
            Xector(g_pooled, shape),
            Xector(b_pooled, shape),
            Xector(lum, shape))


def xector_cell_row(frame, cell_size):
    """
    Cell row index for each pixel. (cell-row frame cell-size) -> Xector
    """
    h, w = frame.shape[:2]
    cs = int(cell_size)
    rows = h // cs

    y = np.repeat(np.arange(h), w)
    cell_row = np.clip(y // cs, 0, rows - 1).astype(np.float32)
    return Xector(cell_row, (h, w))


def xector_cell_col(frame, cell_size):
    """
    Cell column index for each pixel. (cell-col frame cell-size) -> Xector
    """
    h, w = frame.shape[:2]
    cs = int(cell_size)
    cols = w // cs

    x = np.tile(np.arange(w), h)
    cell_col = np.clip(x // cs, 0, cols - 1).astype(np.float32)
    return Xector(cell_col, (h, w))


def xector_num_cells(frame, cell_size):
    """Number of cells. (num-cells frame cell-size) -> (rows, cols, total)"""
    h, w = frame.shape[:2]
    cs = int(cell_size)
    rows = h // cs
    cols = w // cs
    return (rows, cols, rows * cols)


# =============================================================================
# Scan (Prefix Operations) - cumulative reductions
# =============================================================================

def xector_scan_add(x, axis=None):
    """
    Cumulative sum (prefix sum). (scan+ x) or (scan+ x :axis 0)

    Returns array where each element is sum of all previous elements.
    Useful for integral images, cumulative effects.
    """
    data = _unwrap(x)
    shape = x._shape if isinstance(x, Xector) else None

    if axis is not None:
        # Reshape to 2D for axis operation
        if shape and len(shape) == 2:
            result = np.cumsum(data.reshape(shape), axis=int(axis)).flatten()
        else:
            result = np.cumsum(data, axis=int(axis))
    else:
        result = np.cumsum(data)

    return _wrap(result, shape)


def xector_scan_mul(x, axis=None):
    """Cumulative product. (scan* x) -> Xector"""
    data = _unwrap(x)
    shape = x._shape if isinstance(x, Xector) else None

    if axis is not None and shape and len(shape) == 2:
        result = np.cumprod(data.reshape(shape), axis=int(axis)).flatten()
    else:
        result = np.cumprod(data)

    return _wrap(result, shape)


def xector_scan_max(x, axis=None):
    """Cumulative maximum. (scan-max x) -> Xector"""
    data = _unwrap(x)
    shape = x._shape if isinstance(x, Xector) else None

    if axis is not None and shape and len(shape) == 2:
        result = np.maximum.accumulate(data.reshape(shape), axis=int(axis)).flatten()
    else:
        result = np.maximum.accumulate(data)

    return _wrap(result, shape)


def xector_scan_min(x, axis=None):
    """Cumulative minimum. (scan-min x) -> Xector"""
    data = _unwrap(x)
    shape = x._shape if isinstance(x, Xector) else None

    if axis is not None and shape and len(shape) == 2:
        result = np.minimum.accumulate(data.reshape(shape), axis=int(axis)).flatten()
    else:
        result = np.minimum.accumulate(data)

    return _wrap(result, shape)


# =============================================================================
# Outer Product - Cartesian operations
# =============================================================================

def xector_outer(x, y, op='*'):
    """
    Outer product. (outer x y) or (outer x y :op '+')

    Creates 2D result where result[i,j] = op(x[i], y[j]).
    Default is multiplication (*).

    Useful for generating 2D patterns from 1D vectors.
    """
    x_data = _unwrap(x)
    y_data = _unwrap(y)

    ops = {
        '*': np.multiply,
        '+': np.add,
        '-': np.subtract,
        '/': np.divide,
        'max': np.maximum,
        'min': np.minimum,
        'and': np.logical_and,
        'or': np.logical_or,
        'xor': np.logical_xor,
    }

    op_fn = ops.get(op, np.multiply)
    result = op_fn.outer(x_data.flatten(), y_data.flatten())

    # Return as xector with 2D shape
    h, w = len(x_data.flatten()), len(y_data.flatten())
    return _wrap(result.flatten(), (h, w))


def xector_outer_add(x, y):
    """Outer sum. (outer+ x y) -> result[i,j] = x[i] + y[j]"""
    return xector_outer(x, y, '+')


def xector_outer_mul(x, y):
    """Outer product. (outer* x y) -> result[i,j] = x[i] * y[j]"""
    return xector_outer(x, y, '*')


def xector_outer_max(x, y):
    """Outer max. (outer-max x y) -> result[i,j] = max(x[i], y[j])"""
    return xector_outer(x, y, 'max')


def xector_outer_min(x, y):
    """Outer min. (outer-min x y) -> result[i,j] = min(x[i], y[j])"""
    return xector_outer(x, y, 'min')


# =============================================================================
# Reduce with Axis - dimensional reductions
# =============================================================================

def xector_reduce_axis(x, op='sum', axis=0):
    """
    Reduce along an axis. (reduce-axis x :op 'sum' :axis 0)

    ops: 'sum', 'mean', 'max', 'min', 'prod', 'std'
    axis: 0 (rows), 1 (columns)

    For a frame-sized xector (H*W):
      axis=0: reduce across rows -> W values (one per column)
      axis=1: reduce across columns -> H values (one per row)
    """
    data = _unwrap(x)
    shape = x._shape if isinstance(x, Xector) else None

    if shape is None or len(shape) != 2:
        # Can't do axis reduction without 2D shape
        raise ValueError("reduce-axis requires 2D xector (with shape)")

    h, w = shape
    data_2d = data.reshape(h, w)
    axis = int(axis)

    ops = {
        'sum': lambda d, a: np.sum(d, axis=a),
        '+': lambda d, a: np.sum(d, axis=a),
        'mean': lambda d, a: np.mean(d, axis=a),
        'max': lambda d, a: np.max(d, axis=a),
        'min': lambda d, a: np.min(d, axis=a),
        'prod': lambda d, a: np.prod(d, axis=a),
        '*': lambda d, a: np.prod(d, axis=a),
        'std': lambda d, a: np.std(d, axis=a),
    }

    op_fn = ops.get(op, ops['sum'])
    result = op_fn(data_2d, axis)

    # Result shape: if axis=0, shape is (w,); if axis=1, shape is (h,)
    new_shape = (w,) if axis == 0 else (h,)
    return _wrap(result.flatten(), new_shape)


def xector_sum_axis(x, axis=0):
    """Sum along axis. (sum-axis x :axis 0)"""
    return xector_reduce_axis(x, 'sum', axis)


def xector_mean_axis(x, axis=0):
    """Mean along axis. (mean-axis x :axis 0)"""
    return xector_reduce_axis(x, 'mean', axis)


def xector_max_axis(x, axis=0):
    """Max along axis. (max-axis x :axis 0)"""
    return xector_reduce_axis(x, 'max', axis)


def xector_min_axis(x, axis=0):
    """Min along axis. (min-axis x :axis 0)"""
    return xector_reduce_axis(x, 'min', axis)


# =============================================================================
# Windowed Operations - sliding window computations
# =============================================================================

def xector_window(x, size, op='mean', stride=1):
    """
    Sliding window operation. (window x size :op 'mean' :stride 1)

    Applies reduction over sliding windows of given size.
    ops: 'sum', 'mean', 'max', 'min'

    For 1D: windows slide along the array
    For 2D (with shape): windows are size x size squares
    """
    data = _unwrap(x)
    shape = x._shape if isinstance(x, Xector) else None
    size = int(size)
    stride = int(stride)

    ops = {
        'sum': np.sum,
        'mean': np.mean,
        'max': np.max,
        'min': np.min,
        'std': np.std,
    }
    op_fn = ops.get(op, np.mean)

    if shape and len(shape) == 2:
        # 2D sliding window
        h, w = shape
        data_2d = data.reshape(h, w)

        # Use stride tricks for efficient windowing
        out_h = (h - size) // stride + 1
        out_w = (w - size) // stride + 1

        result = np.zeros((out_h, out_w))
        for i in range(out_h):
            for j in range(out_w):
                window = data_2d[i*stride:i*stride+size, j*stride:j*stride+size]
                result[i, j] = op_fn(window)

        return _wrap(result.flatten(), (out_h, out_w))
    else:
        # 1D sliding window
        n = len(data)
        out_n = (n - size) // stride + 1
        result = np.array([op_fn(data[i*stride:i*stride+size]) for i in range(out_n)])
        return _wrap(result, (out_n,))


def xector_window_sum(x, size, stride=1):
    """Sliding window sum. (window-sum x size)"""
    return xector_window(x, size, 'sum', stride)


def xector_window_mean(x, size, stride=1):
    """Sliding window mean. (window-mean x size)"""
    return xector_window(x, size, 'mean', stride)


def xector_window_max(x, size, stride=1):
    """Sliding window max. (window-max x size)"""
    return xector_window(x, size, 'max', stride)


def xector_window_min(x, size, stride=1):
    """Sliding window min. (window-min x size)"""
    return xector_window(x, size, 'min', stride)


def xector_integral_image(frame):
    """
    Compute integral image (summed area table). (integral-image frame)

    Each pixel contains sum of all pixels above and to the left.
    Enables O(1) box blur at any radius.

    Returns xector with same shape as frame's luminance.
    """
    if hasattr(frame, 'shape') and len(frame.shape) == 3:
        # Convert frame to grayscale
        gray = np.mean(frame, axis=2)
    else:
        data = _unwrap(frame)
        shape = frame._shape if isinstance(frame, Xector) else None
        if shape and len(shape) == 2:
            gray = data.reshape(shape)
        else:
            gray = data

    integral = np.cumsum(np.cumsum(gray, axis=0), axis=1)
    h, w = integral.shape
    return _wrap(integral.flatten(), (h, w))


def xector_box_blur_fast(integral, x, y, radius, width, height):
    """
    Fast box blur using integral image. (box-blur-fast integral x y radius w h)

    Given pre-computed integral image, compute average in box centered at (x,y).
    O(1) regardless of radius.
    """
    integral_data = _unwrap(integral)
    shape = integral._shape if isinstance(integral, Xector) else None

    if shape is None or len(shape) != 2:
        raise ValueError("box-blur-fast requires 2D integral image")

    h, w = shape
    integral_2d = integral_data.reshape(h, w)

    radius = int(radius)
    x, y = int(x), int(y)

    # Clamp coordinates
    x1 = max(0, x - radius)
    y1 = max(0, y - radius)
    x2 = min(w - 1, x + radius)
    y2 = min(h - 1, y + radius)

    # Sum in rectangle using integral image
    total = integral_2d[y2, x2]
    if x1 > 0:
        total -= integral_2d[y2, x1 - 1]
    if y1 > 0:
        total -= integral_2d[y1 - 1, x2]
    if x1 > 0 and y1 > 0:
        total += integral_2d[y1 - 1, x1 - 1]

    count = (x2 - x1 + 1) * (y2 - y1 + 1)
    return total / max(count, 1)


# =============================================================================
# PRIMITIVES Export
# =============================================================================

PRIMITIVES = {
    # Frame/Xector conversion
    # NOTE: red, green, blue, gray, rgb are derived in derived.sexp using (channel frame n)
    'xector': xector_from_frame,
    'to-frame': xector_to_frame,

    # Coordinate generators
    # NOTE: x-coords, y-coords, x-norm, y-norm, dist-from-center are derived
    # in derived.sexp using iota, tile, repeat primitives

    # Alpha (α) - element-wise operations
    'α+': alpha_add,
    'α-': alpha_sub,
    'α*': alpha_mul,
    'α/': alpha_div,
    'α**': alpha_pow,
    'αsqrt': alpha_sqrt,
    'αabs': alpha_abs,
    'αsin': alpha_sin,
    'αcos': alpha_cos,
    'αexp': alpha_exp,
    'αlog': alpha_log,
    # NOTE: αclamp is derived in derived.sexp as (max2 lo (min2 hi x))
    'αmin': alpha_min,
    'αmax': alpha_max,
    'αmod': alpha_mod,
    'αfloor': alpha_floor,
    'αceil': alpha_ceil,
    'αround': alpha_round,
    # NOTE: α² / αsq is derived in derived.sexp as (* x x)

    # Alpha comparison
    'α<': alpha_lt,
    'α<=': alpha_le,
    'α>': alpha_gt,
    'α>=': alpha_ge,
    'α=': alpha_eq,

    # Alpha logical
    'αand': alpha_and,
    'αor': alpha_or,
    'αnot': alpha_not,

    # ASCII fallbacks for α
    'alpha+': alpha_add,
    'alpha-': alpha_sub,
    'alpha*': alpha_mul,
    'alpha/': alpha_div,
    'alpha**': alpha_pow,
    'alpha-sqrt': alpha_sqrt,
    'alpha-abs': alpha_abs,
    'alpha-sin': alpha_sin,
    'alpha-cos': alpha_cos,
    'alpha-exp': alpha_exp,
    'alpha-log': alpha_log,
    'alpha-min': alpha_min,
    'alpha-max': alpha_max,
    'alpha-mod': alpha_mod,
    'alpha-floor': alpha_floor,
    'alpha-ceil': alpha_ceil,
    'alpha-round': alpha_round,
    'alpha<': alpha_lt,
    'alpha<=': alpha_le,
    'alpha>': alpha_gt,
    'alpha>=': alpha_ge,
    'alpha=': alpha_eq,
    'alpha-and': alpha_and,
    'alpha-or': alpha_or,
    'alpha-not': alpha_not,

    # Beta (β) - reduction operations
    'β+': beta_add,
    'β*': beta_mul,
    'βmin': beta_min,
    'βmax': beta_max,
    'βmean': beta_mean,
    'βstd': beta_std,
    'βcount': beta_count,
    'βany': beta_any,
    'βall': beta_all,

    # ASCII fallbacks for β
    'beta+': beta_add,
    'beta*': beta_mul,
    'beta-min': beta_min,
    'beta-max': beta_max,
    'beta-mean': beta_mean,
    'beta-std': beta_std,
    'beta-count': beta_count,
    'beta-any': beta_any,
    'beta-all': beta_all,

    # Convenience aliases
    'sum': beta_add,
    'product': beta_mul,
    'mean': beta_mean,

    # Conditional / Selection
    'where': xector_where,
    # NOTE: fill, zeros, ones are derived in derived.sexp using iota
    'rand-x': xector_rand,
    'randn-x': xector_randn,

    # Type checking
    'xector?': is_xector,

    # ===========================================
    # CORE PRIMITIVES - fundamental operations
    # ===========================================

    # Gather/Scatter - parallel indexing
    'gather': xector_gather,
    'gather-2d': xector_gather_2d,
    'scatter': xector_scatter,
    'scatter-add': xector_scatter_add,

    # Group reduce - pooling primitive
    'group-reduce': xector_group_reduce,

    # Shape operations
    'reshape': xector_reshape,
    'shape': xector_shape,
    'xlen': xector_len,

    # Index generation
    'iota': xector_iota,
    'repeat': xector_repeat,
    'tile': xector_tile,

    # Cell/Grid helpers (built on primitives)
    'cell-indices': xector_cell_indices,
    'cell-row': xector_cell_row,
    'cell-col': xector_cell_col,
    'local-x': xector_local_x,
    'local-y': xector_local_y,
    'local-x-norm': xector_local_x_norm,
    'local-y-norm': xector_local_y_norm,
    'pool-frame': xector_pool_frame,
    'num-cells': xector_num_cells,

    # Scan (prefix) operations - cumulative reductions
    'scan+': xector_scan_add,
    'scan*': xector_scan_mul,
    'scan-max': xector_scan_max,
    'scan-min': xector_scan_min,
    'scan-add': xector_scan_add,
    'scan-mul': xector_scan_mul,

    # Outer product - Cartesian operations
    'outer': xector_outer,
    'outer+': xector_outer_add,
    'outer*': xector_outer_mul,
    'outer-add': xector_outer_add,
    'outer-mul': xector_outer_mul,
    'outer-max': xector_outer_max,
    'outer-min': xector_outer_min,

    # Reduce with axis - dimensional reductions
    'reduce-axis': xector_reduce_axis,
    'sum-axis': xector_sum_axis,
    'mean-axis': xector_mean_axis,
    'max-axis': xector_max_axis,
    'min-axis': xector_min_axis,

    # Windowed operations - sliding window computations
    'window': xector_window,
    'window-sum': xector_window_sum,
    'window-mean': xector_window_mean,
    'window-max': xector_window_max,
    'window-min': xector_window_min,

    # Integral image - for fast box blur
    'integral-image': xector_integral_image,
    'box-blur-fast': xector_box_blur_fast,
}
