"""
Array Primitives Library

Vectorized operations on numpy arrays for coordinate transformations.
"""
import numpy as np


# Arithmetic
def prim_arr_add(a, b):
    return np.add(a, b)


def prim_arr_sub(a, b):
    return np.subtract(a, b)


def prim_arr_mul(a, b):
    return np.multiply(a, b)


def prim_arr_div(a, b):
    return np.divide(a, b)


def prim_arr_mod(a, b):
    return np.mod(a, b)


def prim_arr_neg(a):
    return np.negative(a)


# Math functions
def prim_arr_sin(a):
    return np.sin(a)


def prim_arr_cos(a):
    return np.cos(a)


def prim_arr_tan(a):
    return np.tan(a)


def prim_arr_sqrt(a):
    return np.sqrt(np.maximum(a, 0))


def prim_arr_pow(a, b):
    return np.power(a, b)


def prim_arr_abs(a):
    return np.abs(a)


def prim_arr_exp(a):
    return np.exp(a)


def prim_arr_log(a):
    return np.log(np.maximum(a, 1e-10))


def prim_arr_atan2(y, x):
    return np.arctan2(y, x)


# Comparison / selection
def prim_arr_min(a, b):
    return np.minimum(a, b)


def prim_arr_max(a, b):
    return np.maximum(a, b)


def prim_arr_clip(a, lo, hi):
    return np.clip(a, lo, hi)


def prim_arr_where(cond, a, b):
    return np.where(cond, a, b)


def prim_arr_floor(a):
    return np.floor(a)


def prim_arr_ceil(a):
    return np.ceil(a)


def prim_arr_round(a):
    return np.round(a)


# Interpolation
def prim_arr_lerp(a, b, t):
    return a + (b - a) * t


def prim_arr_smoothstep(edge0, edge1, x):
    t = prim_arr_clip((x - edge0) / (edge1 - edge0), 0.0, 1.0)
    return t * t * (3 - 2 * t)


# Creation
def prim_arr_zeros(shape):
    return np.zeros(shape, dtype=np.float32)


def prim_arr_ones(shape):
    return np.ones(shape, dtype=np.float32)


def prim_arr_full(shape, value):
    return np.full(shape, value, dtype=np.float32)


def prim_arr_arange(start, stop, step=1):
    return np.arange(start, stop, step, dtype=np.float32)


def prim_arr_linspace(start, stop, num):
    return np.linspace(start, stop, num, dtype=np.float32)


def prim_arr_meshgrid(x, y):
    return np.meshgrid(x, y)


# Coordinate transforms
def prim_polar_from_center(map_x, map_y, cx, cy):
    """Convert Cartesian to polar coordinates centered at (cx, cy)."""
    dx = map_x - cx
    dy = map_y - cy
    r = np.sqrt(dx**2 + dy**2)
    theta = np.arctan2(dy, dx)
    return (r, theta)


def prim_cart_from_polar(r, theta, cx, cy):
    """Convert polar to Cartesian, adding center offset."""
    x = r * np.cos(theta) + cx
    y = r * np.sin(theta) + cy
    return (x, y)


PRIMITIVES = {
    # Arithmetic
    'arr+': prim_arr_add,
    'arr-': prim_arr_sub,
    'arr*': prim_arr_mul,
    'arr/': prim_arr_div,
    'arr-mod': prim_arr_mod,
    'arr-neg': prim_arr_neg,

    # Math
    'arr-sin': prim_arr_sin,
    'arr-cos': prim_arr_cos,
    'arr-tan': prim_arr_tan,
    'arr-sqrt': prim_arr_sqrt,
    'arr-pow': prim_arr_pow,
    'arr-abs': prim_arr_abs,
    'arr-exp': prim_arr_exp,
    'arr-log': prim_arr_log,
    'arr-atan2': prim_arr_atan2,

    # Selection
    'arr-min': prim_arr_min,
    'arr-max': prim_arr_max,
    'arr-clip': prim_arr_clip,
    'arr-where': prim_arr_where,
    'arr-floor': prim_arr_floor,
    'arr-ceil': prim_arr_ceil,
    'arr-round': prim_arr_round,

    # Interpolation
    'arr-lerp': prim_arr_lerp,
    'arr-smoothstep': prim_arr_smoothstep,

    # Creation
    'arr-zeros': prim_arr_zeros,
    'arr-ones': prim_arr_ones,
    'arr-full': prim_arr_full,
    'arr-arange': prim_arr_arange,
    'arr-linspace': prim_arr_linspace,
    'arr-meshgrid': prim_arr_meshgrid,

    # Coordinates
    'polar-from-center': prim_polar_from_center,
    'cart-from-polar': prim_cart_from_polar,
}