import torch
import math

#===================================================================================
def srgb_to_linear(c: float) -> float:
    if c <= 0.04045: return c / 12.92
    return ((c + 0.055) / 1.055) ** 2.4

#===================================================================================
def srgb_to_linear_vec(x: torch.Tensor) -> torch.Tensor:
    """Vectorised sRGB to linear (matches standard IEC 61966‑2‑1)."""
    mask = x <= 0.04045
    out = torch.zeros_like(x)
    out[mask] = x[mask] / 12.92
    out[~mask] = ((x[~mask] + 0.055) / 1.055) ** 2.4
    return out

#===================================================================================
def linear_to_srgb_vec(x: torch.Tensor) -> torch.Tensor:
    """Vectorised linear to sRGB (inverse of above)."""
    mask = x <= 0.0031308
    out = torch.zeros_like(x)
    out[mask] = x[mask] * 12.92
    out[~mask] = 1.055 * (x[~mask] ** (1.0/2.4)) - 0.055
    return torch.clamp(out, 0.0, 1.0)

#===================================================================================
def rgb_to_lch_torch(image: torch.Tensor) -> torch.Tensor:
    """
    Convert RGB [0,1] to CIE LCH(ab) using D65 white point.
    Uses srgb_to_linear_vec (or your own) for gamma handling.
    """
    # 1. sRGB → linear RGB
    rgb_lin = srgb_to_linear_vec(image)

    # 2. linear RGB → XYZ (D65)
    M_rgb2xyz = torch.tensor([
        [0.4124, 0.3576, 0.1805],
        [0.2126, 0.7152, 0.0722],
        [0.0193, 0.1192, 0.9505]
    ], device=image.device, dtype=image.dtype)
    xyz = torch.einsum('ij,...j->...i', M_rgb2xyz, rgb_lin)

    # 3. XYZ → CIE L*a*b* (D65 white point)
    white = torch.tensor([0.95047, 1.0, 0.95883], device=image.device, dtype=image.dtype)
    xyz_n = xyz / white
    delta = 6.0 / 29.0
    delta2 = delta * delta
    delta3 = delta * delta * delta

    f = torch.zeros_like(xyz_n)
    mask_f = xyz_n > delta3
    f[mask_f] = xyz_n[mask_f] ** (1.0/3.0)
    f[~mask_f] = (xyz_n[~mask_f] / (3.0 * delta2)) + (4.0 / 29.0)

    L = 116.0 * f[..., 1] - 16.0
    a = 500.0 * (f[..., 0] - f[..., 1])
    b = 200.0 * (f[..., 1] - f[..., 2])

    # 4. Lab → LCH
    C = torch.sqrt(a*a + b*b)
    H = torch.atan2(b, a) * (180.0 / math.pi)
    H = H % 360.0
    H = torch.where(H >= 180.0, H - 360.0, H)   # map to [-180, 180)

    return torch.stack([L, C, H], dim=-1)

#===================================================================================
def lch_to_rgb_torch(lch: torch.Tensor) -> torch.Tensor:
    """
    Convert CIE LCH(ab) back to RGB [0,1].
    Uses linear_to_srgb_vec for final gamma compression.
    """
    L, C, H = lch[..., 0], lch[..., 1], lch[..., 2]
    H_rad = (H % 360.0) * (math.pi / 180.0)

    # LCH → Lab
    a = C * torch.cos(H_rad)
    b = C * torch.sin(H_rad)

    # Lab → XYZ
    delta = 6.0 / 29.0
    delta2 = delta * delta
    delta3 = delta * delta * delta

    fy = (L + 16.0) / 116.0
    fx = fy + a / 500.0
    fz = fy - b / 200.0

    xyz = torch.zeros_like(lch)
    # Inverse f(t)
    mask = fx > delta
    xyz[..., 0][mask] = fx[mask] ** 3
    xyz[..., 0][~mask] = 3.0 * delta2 * (fx[~mask] - 4.0/29.0)

    mask = fy > delta
    xyz[..., 1][mask] = fy[mask] ** 3
    xyz[..., 1][~mask] = 3.0 * delta2 * (fy[~mask] - 4.0/29.0)

    mask = fz > delta
    xyz[..., 2][mask] = fz[mask] ** 3
    xyz[..., 2][~mask] = 3.0 * delta2 * (fz[~mask] - 4.0/29.0)

    white = torch.tensor([0.95047, 1.0, 0.95883], device=lch.device, dtype=lch.dtype)
    xyz = xyz * white

    # XYZ → linear RGB
    M_xyz2rgb = torch.tensor([
        [ 3.2406, -1.5372, -0.4986],
        [-0.9689,  1.8758,  0.0415],
        [ 0.0557, -0.2040,  1.0570]
    ], device=lch.device, dtype=lch.dtype)
    rgb_lin = torch.einsum('ij,...j->...i', M_xyz2rgb, xyz)

    # linear → sRGB
    rgb_srgb = linear_to_srgb_vec(rgb_lin)

    return rgb_srgb

#===================================================================================
def hsl_to_rgb(h, s, l):
    hp = h / 60.0
    c = (1.0 - torch.abs(2.0 * l - 1.0)) * s
    x = c * (1.0 - torch.abs(hp % 2.0 - 1.0))
    m = l - 0.5 * c

    r = torch.zeros_like(l)
    g = torch.zeros_like(l)
    b = torch.zeros_like(l)

    m0 = (hp >= 0) & (hp < 1)
    m1 = (hp >= 1) & (hp < 2)
    m2 = (hp >= 2) & (hp < 3)
    m3 = (hp >= 3) & (hp < 4)
    m4 = (hp >= 4) & (hp < 5)
    m5 = (hp >= 5) & (hp < 6)

    r[m0], g[m0], b[m0] = c[m0], x[m0], 0
    r[m1], g[m1], b[m1] = x[m1], c[m1], 0
    r[m2], g[m2], b[m2] = 0, c[m2], x[m2]
    r[m3], g[m3], b[m3] = 0, x[m3], c[m3]
    r[m4], g[m4], b[m4] = x[m4], 0, c[m4]
    r[m5], g[m5], b[m5] = c[m5], 0, x[m5]

    return torch.clamp(torch.stack([r + m, g + m, b + m], dim=-1), 0.0, 1.0)

#===================================================================================
def rgb_to_hsl(rgb: torch.Tensor):
    maxc = rgb.max(dim=-1).values
    minc = rgb.min(dim=-1).values
    l = (maxc + minc) * 0.5

    delta = maxc - minc

    s = torch.zeros_like(l)
    nz = delta > 1e-8
    s[nz] = delta[nz] / (1.0 - torch.abs(2.0 * l[nz] - 1.0) + 1e-8)

    r, g, b = rgb[..., 0], rgb[..., 1], rgb[..., 2]

    h = torch.zeros_like(l)

    rmax = (maxc == r) & nz
    gmax = (maxc == g) & nz
    bmax = (maxc == b) & nz

    h_r = ((g - b) / (delta + 1e-8)) % 6.0
    h_g = ((b - r) / (delta + 1e-8)) + 2.0
    h_b = ((r - g) / (delta + 1e-8)) + 4.0

    h = torch.where(rmax, h_r, h)
    h = torch.where(gmax, h_g, h)
    h = torch.where(bmax, h_b, h)

    h = (h * 60.0) % 360.0

    return h, s, l

#===================================================================================
def rgb_to_hsv(rgb: torch.Tensor):
    """Convert RGB [0,1] to HSV. Returns (hue, saturation, value)."""
    r, g, b = rgb[..., 0], rgb[..., 1], rgb[..., 2]
    
    max_val, max_idx = torch.max(rgb, dim=-1)
    min_val, _ = torch.min(rgb, dim=-1)
    diff = max_val - min_val
    
    # Hue
    hue = torch.zeros_like(max_val)
    mask = diff != 0
    h = torch.zeros_like(max_val)
    h = torch.where((max_val == r) & mask, ((g - b) / diff) % 6, h)
    h = torch.where((max_val == g) & mask, ((b - r) / diff) + 2, h)
    h = torch.where((max_val == b) & mask, ((r - g) / diff) + 4, h)
    hue = (h / 6) % 1.0
    
    # Saturation
    sat = torch.zeros_like(max_val)
    sat = torch.where(max_val != 0, diff / max_val, sat)
    
    # Value
    val = max_val
    
    return hue, sat, val

#===================================================================================
def hsv_to_rgb(hue: torch.Tensor, sat: torch.Tensor, val: torch.Tensor):
    """Convert HSV to RGB [0,1]."""
    h = hue * 6.0
    i = torch.floor(h).long()
    f = h - i.float()
    p = val * (1 - sat)
    q = val * (1 - sat * f)
    t = val * (1 - sat * (1 - f))
    
    # Initialize output
    rgb = torch.zeros(*hue.shape, 3, device=hue.device)
    
    # Vectorized assignment using advanced indexing
    # This is simplified - in practice you'd use torch.where for each case
    mask0 = i == 0
    mask1 = i == 1
    mask2 = i == 2
    mask3 = i == 3
    mask4 = i == 4
    mask5 = i == 5
    
    rgb[..., 0] = torch.where(mask0, val, rgb[..., 0])
    rgb[..., 1] = torch.where(mask0, t, rgb[..., 1])
    rgb[..., 2] = torch.where(mask0, p, rgb[..., 2])
    
    rgb[..., 0] = torch.where(mask1, q, rgb[..., 0])
    rgb[..., 1] = torch.where(mask1, val, rgb[..., 1])
    rgb[..., 2] = torch.where(mask1, p, rgb[..., 2])
    
    rgb[..., 0] = torch.where(mask2, p, rgb[..., 0])
    rgb[..., 1] = torch.where(mask2, val, rgb[..., 1])
    rgb[..., 2] = torch.where(mask2, t, rgb[..., 2])
    
    rgb[..., 0] = torch.where(mask3, p, rgb[..., 0])
    rgb[..., 1] = torch.where(mask3, q, rgb[..., 1])
    rgb[..., 2] = torch.where(mask3, val, rgb[..., 2])
    
    rgb[..., 0] = torch.where(mask4, t, rgb[..., 0])
    rgb[..., 1] = torch.where(mask4, p, rgb[..., 1])
    rgb[..., 2] = torch.where(mask4, val, rgb[..., 2])
    
    rgb[..., 0] = torch.where(mask5, val, rgb[..., 0])
    rgb[..., 1] = torch.where(mask5, p, rgb[..., 1])
    rgb[..., 2] = torch.where(mask5, q, rgb[..., 2])
    
    return rgb

#===================================================================================
def rgb_to_lch(rgb: torch.Tensor):
    """Convert RGB to Lch (CIE L*a*b* with chroma)."""
    # Convert RGB to XYZ (D65)
    rgb_linear = rgb.clone()
    mask = rgb_linear <= 0.04045
    rgb_linear[mask] = rgb_linear[mask] / 12.92
    rgb_linear[~mask] = ((rgb_linear[~mask] + 0.055) / 1.055) ** 2.4
    
    # RGB to XYZ matrix
    x = 0.4124564 * rgb_linear[..., 0] + 0.3575761 * rgb_linear[..., 1] + 0.1804375 * rgb_linear[..., 2]
    y = 0.2126729 * rgb_linear[..., 0] + 0.7151522 * rgb_linear[..., 1] + 0.0721750 * rgb_linear[..., 2]
    z = 0.0193339 * rgb_linear[..., 0] + 0.1191920 * rgb_linear[..., 1] + 0.9503041 * rgb_linear[..., 2]
    
    # XYZ to Lab (D65 white point)
    xn, yn, zn = 0.95047, 1.0, 1.08883
    xr, yr, zr = x / xn, y / yn, z / zn
    
    f = lambda t: torch.where(t > 0.008856, t ** (1/3), 7.787 * t + 16/116)
    fx, fy, fz = f(xr), f(yr), f(zr)
    
    L = 116 * fy - 16
    a = 500 * (fx - fy)
    b = 200 * (fy - fz)
    
    # Chroma and hue
    C = torch.sqrt(a**2 + b**2)
    H = torch.atan2(b, a) % (2 * torch.pi)
    
    return L, C, H

#===================================================================================
def lch_to_rgb(L: torch.Tensor, C: torch.Tensor, H: torch.Tensor):
    """Convert Lch to RGB."""
    # Lch to Lab
    a = C * torch.cos(H)
    b = C * torch.sin(H)
    
    # Lab to XYZ (D65)
    fy = (L + 16) / 116
    fx = a / 500 + fy
    fz = fy - b / 200
    
    f_inv = lambda t: torch.where(t > 0.206893, t**3, (t - 16/116) / 7.787)
    xr, yr, zr = f_inv(fx), f_inv(fy), f_inv(fz)
    
    xn, yn, zn = 0.95047, 1.0, 1.08883
    x, y, z = xr * xn, yr * yn, zr * zn
    
    # XYZ to RGB
    r = 3.2404542 * x - 1.5371385 * y - 0.4985314 * z
    g = -0.9692660 * x + 1.8760108 * y + 0.0415560 * z
    b = 0.0556434 * x - 0.2040259 * y + 1.0572252 * z
    
    rgb = torch.stack([r, g, b], dim=-1)
    
    # Linear to sRGB
    mask = rgb <= 0.0031308
    rgb[mask] = 12.92 * rgb[mask]
    rgb[~mask] = 1.055 * rgb[~mask] ** (1/2.4) - 0.055
    
    return torch.clamp(rgb, 0, 1)

#===================================================================================
def rgb_to_yuv(rgb: torch.Tensor):
    """Convert RGB to Yuv (CIE Yuv)."""
    # BT.709 coefficients
    y = 0.2126 * rgb[..., 0] + 0.7152 * rgb[..., 1] + 0.0722 * rgb[..., 2]
    u = rgb[..., 0] - y
    v = rgb[..., 2] - y
    return y, u, v

#===================================================================================
def yuv_to_rgb(y: torch.Tensor, u: torch.Tensor, v: torch.Tensor):
    """Convert Yuv to RGB."""
    r = y + u
    g = y - 0.21482 * u - 0.38129 * v  # Simplified: approximate
    b = y + v
    rgb = torch.stack([r, g, b], dim=-1)
    return torch.clamp(rgb, 0, 1)

#===================================================================================
def rgb_to_lch_torch(image: torch.Tensor) -> torch.Tensor:
    """
    Convert RGB [0,1] to CIE LCH(ab) using D65 white point.
    Returns tensor with last dimension [L, C, H].
    """
    # 1. sRGB → linear RGB
    rgb_lin = srgb_to_linear_vec(image)

    # 2. linear RGB → XYZ (D65)
    M_rgb2xyz = torch.tensor([
        [0.4124, 0.3576, 0.1805],
        [0.2126, 0.7152, 0.0722],
        [0.0193, 0.1192, 0.9505]
    ], device=image.device, dtype=image.dtype)
    xyz = torch.einsum('ij,...j->...i', M_rgb2xyz, rgb_lin)

    # 3. XYZ → CIE L*a*b* (D65 white point)
    white = torch.tensor([0.95047, 1.0, 0.95883], device=image.device, dtype=image.dtype)
    xyz_n = xyz / white
    delta = 6.0 / 29.0
    delta2 = delta * delta
    delta3 = delta * delta * delta

    f = torch.zeros_like(xyz_n)
    mask_f = xyz_n > delta3
    f[mask_f] = xyz_n[mask_f] ** (1.0/3.0)
    f[~mask_f] = (xyz_n[~mask_f] / (3.0 * delta2)) + (4.0 / 29.0)

    L = 116.0 * f[..., 1] - 16.0
    a = 500.0 * (f[..., 0] - f[..., 1])
    b = 200.0 * (f[..., 1] - f[..., 2])

    # 4. Lab → LCH
    C = torch.sqrt(a*a + b*b)
    H = torch.atan2(b, a) * (180.0 / math.pi)
    H = H % 360.0
    H = torch.where(H >= 180.0, H - 360.0, H)   # map to [-180, 180)

    return torch.stack([L, C, H], dim=-1)

#===================================================================================
def lch_to_rgb_torch(lch: torch.Tensor) -> torch.Tensor:
    """
    Convert CIE LCH(ab) back to RGB [0,1].
    """
    L, C, H = lch[..., 0], lch[..., 1], lch[..., 2]
    H_rad = (H % 360.0) * (math.pi / 180.0)

    # LCH → Lab
    a = C * torch.cos(H_rad)
    b = C * torch.sin(H_rad)

    # Lab → XYZ
    delta = 6.0 / 29.0
    delta2 = delta * delta
    delta3 = delta * delta * delta

    fy = (L + 16.0) / 116.0
    fx = fy + a / 500.0
    fz = fy - b / 200.0

    xyz = torch.zeros_like(lch)
    # Inverse f(t)
    mask = fx > delta
    xyz[..., 0][mask] = fx[mask] ** 3
    xyz[..., 0][~mask] = 3.0 * delta2 * (fx[~mask] - 4.0/29.0)

    mask = fy > delta
    xyz[..., 1][mask] = fy[mask] ** 3
    xyz[..., 1][~mask] = 3.0 * delta2 * (fy[~mask] - 4.0/29.0)

    mask = fz > delta
    xyz[..., 2][mask] = fz[mask] ** 3
    xyz[..., 2][~mask] = 3.0 * delta2 * (fz[~mask] - 4.0/29.0)

    white = torch.tensor([0.95047, 1.0, 0.95883], device=lch.device, dtype=lch.dtype)
    xyz = xyz * white

    # XYZ → linear RGB
    M_xyz2rgb = torch.tensor([
        [ 3.2406, -1.5372, -0.4986],
        [-0.9689,  1.8758,  0.0415],
        [ 0.0557, -0.2040,  1.0570]
    ], device=lch.device, dtype=lch.dtype)
    rgb_lin = torch.einsum('ij,...j->...i', M_xyz2rgb, xyz)

    # linear → sRGB
    rgb_srgb = linear_to_srgb_vec(rgb_lin)
    return rgb_srgb

#===================================================================================
def pack_rgb(r, g, b):
    """r,g,b in 0..1, returns (packed_int, '#RRGGBB')"""
    r = max(0.0, min(1.0, r))
    g = max(0.0, min(1.0, g))
    b = max(0.0, min(1.0, b))
    ri = int(round(r * 255))
    gi = int(round(g * 255))
    bi = int(round(b * 255))
    packed = (ri << 16) | (gi << 8) | bi
    hex_str = f"#{ri:02X}{gi:02X}{bi:02X}"
    return packed, hex_str