import numpy as np

STACK_TWIST = 0.55
STACK_TWIST_CENTER = 0.50

A = np.array([0.2126, 0.7152, 0.0722], dtype=float)
A /= np.linalg.norm(A)
B = np.array([0.956622786415, -0.274519637823, -0.097528523834], dtype=float)
C = np.array([-0.066610303637,  0.119798469886, -0.990561151097], dtype=float)
P = np.column_stack([A, B, C])
W = float(np.dot([1.0, 1.0, 1.0], A))


def _rot(uv, angle):
    c = np.cos(angle)
    s = np.sin(angle)
    x = uv[..., 0]
    y = uv[..., 1]
    return np.stack([c*x - s*y, s*x + c*y], axis=-1)

def _gray_uv(z):
    return (z[..., None] * np.ones(3)) @ P[..., 1:3]

def rgb_to_zuv(rgb):
    rgb = np.asarray(rgb, dtype=float)
    t = rgb @ P
    z = t[..., 0] / W
    uv = t[..., 1:3] - _gray_uv(z)
    angle = STACK_TWIST * (z - STACK_TWIST_CENTER)
    uv = _rot(uv, angle)
    return np.concatenate([z[..., None], uv], axis=-1)

def zuv_to_rgb(zuv, clip=False):
    zuv = np.asarray(zuv, dtype=float)
    z = zuv[..., 0]
    uv = zuv[..., 1:3]
    angle = -STACK_TWIST * (z - STACK_TWIST_CENTER)
    uv = _rot(uv, angle) + _gray_uv(z)
    t = np.concatenate([(z * W)[..., None], uv], axis=-1)
    rgb = t @ P.T
    return np.clip(rgb, 0.0, 1.0) if clip else rgb