"""Commodore 64 (VIC-II) 16-colour palettes and colour-space helpers.

All colour distance work in the converter happens in CIELAB, which is far more
perceptually uniform than RGB, so the per-cell colour choices and dithering land
much closer to what a human eye judges as "the same colour".
"""

from __future__ import annotations

import numpy as np

# 16 fixed VIC-II colours, in canonical index order:
#  0 black      4 purple     8 orange     12 grey (medium)
#  1 white      5 green      9 brown      13 light green
#  2 red        6 blue      10 light red  14 light blue
#  3 cyan       7 yellow    11 dark grey  15 light grey

# "Colodore" (pepto's reworked, calibrated values) -- the modern default.
COLODORE = np.array([
    (0x00, 0x00, 0x00),
    (0xff, 0xff, 0xff),
    (0x81, 0x33, 0x38),
    (0x75, 0xce, 0xc8),
    (0x8e, 0x3c, 0x97),
    (0x56, 0xac, 0x4d),
    (0x2e, 0x2c, 0x9b),
    (0xed, 0xf1, 0x71),
    (0x8e, 0x50, 0x29),
    (0x55, 0x38, 0x00),
    (0xc4, 0x6c, 0x71),
    (0x4a, 0x4a, 0x4a),
    (0x7b, 0x7b, 0x7b),
    (0xa9, 0xff, 0x9f),
    (0x70, 0x6d, 0xeb),
    (0xb2, 0xb2, 0xb2),
], dtype=np.float64)

# "Pepto" (PAL) -- classic reference values, slightly more saturated.
PEPTO = np.array([
    (0x00, 0x00, 0x00),
    (0xff, 0xff, 0xff),
    (0x68, 0x37, 0x2b),
    (0x70, 0xa4, 0xb2),
    (0x6f, 0x3d, 0x86),
    (0x58, 0x8d, 0x43),
    (0x35, 0x28, 0x79),
    (0xb8, 0xc7, 0x6f),
    (0x6f, 0x4f, 0x25),
    (0x43, 0x39, 0x00),
    (0x9a, 0x67, 0x59),
    (0x44, 0x44, 0x44),
    (0x6c, 0x6c, 0x6c),
    (0x9a, 0xd2, 0x84),
    (0x6c, 0x5e, 0xb5),
    (0x95, 0x95, 0x95),
], dtype=np.float64)

PALETTES = {"colodore": COLODORE, "pepto": PEPTO}

COLOR_NAMES = [
    "black", "white", "red", "cyan", "purple", "green", "blue", "yellow",
    "orange", "brown", "light red", "dark grey", "grey", "light green",
    "light blue", "light grey",
]


def srgb_to_linear(rgb: np.ndarray) -> np.ndarray:
    """sRGB (0..255) -> linear-light (0..1)."""
    c = rgb.astype(np.float64) / 255.0
    return np.where(c <= 0.04045, c / 12.92, ((c + 0.055) / 1.055) ** 2.4)


def linear_to_srgb(lin: np.ndarray) -> np.ndarray:
    """linear-light (0..1) -> sRGB (0..255)."""
    c = np.clip(lin, 0.0, 1.0)
    s = np.where(c <= 0.0031308, c * 12.92, 1.055 * (c ** (1 / 2.4)) - 0.055)
    return np.clip(s * 255.0 + 0.5, 0, 255).astype(np.uint8)


# D65 reference white.
_XYZ_FROM_LIN = np.array([
    [0.4124564, 0.3575761, 0.1804375],
    [0.2126729, 0.7151522, 0.0721750],
    [0.0193339, 0.1191920, 0.9503041],
])
_WHITE = np.array([0.95047, 1.0, 1.08883])


def srgb_to_lab(rgb: np.ndarray) -> np.ndarray:
    """sRGB (0..255, last axis = RGB) -> CIELAB. Shape preserved except last axis."""
    lin = srgb_to_linear(rgb)
    xyz = lin @ _XYZ_FROM_LIN.T
    xyz = xyz / _WHITE
    eps = 216 / 24389
    kappa = 24389 / 27
    f = np.where(xyz > eps, np.cbrt(xyz), (kappa * xyz + 16) / 116)
    fx, fy, fz = f[..., 0], f[..., 1], f[..., 2]
    L = 116 * fy - 16
    a = 500 * (fx - fy)
    b = 200 * (fy - fz)
    return np.stack([L, a, b], axis=-1)


def get_palette(name: str = "colodore") -> np.ndarray:
    """Return the 16x3 sRGB palette (float64, 0..255)."""
    return PALETTES[name]


def palette_lab(name: str = "colodore") -> np.ndarray:
    """Return the 16 palette colours in CIELAB (16x3)."""
    return srgb_to_lab(get_palette(name))