
import numpy as np

# Simple helper: build a 256x3 uint8 LUT from color stops in 0..1
# stops: list of (pos, (r,g,b)) with pos in [0,1], r,g,b in [0,255]
def _make_lut_from_stops(stops, size: int = 256) -> np.ndarray:
    stops = sorted(stops, key=lambda s: s[0])
    lut = np.zeros((size, 3), dtype=np.uint8)

    positions = [int(round(p * (size - 1))) for p, _ in stops]
    colors = [np.array(c, dtype=np.float32) for _, c in stops]

    for i in range(len(stops) - 1):
        x0 = positions[i]
        x1 = positions[i + 1]
        c0 = colors[i]
        c1 = colors[i + 1]

        if x1 <= x0:
            lut[x0] = c0.astype(np.uint8)
            continue

        length = x1 - x0
        for j in range(length):
            t = j / float(length)
            c = (1.0 - t) * c0 + t * c1
            lut[x0 + j] = c.astype(np.uint8)

    # last entry
    lut[positions[-1]] = colors[-1].astype(np.uint8)
    return lut


def _make_gray_lut() -> np.ndarray:
    lut = np.zeros((256, 3), dtype=np.uint8)
    for i in range(256):
        lut[i] = (i, i, i)
    return lut


def _make_fire_lut() -> np.ndarray:
    """Approximate 'fire' palette via HSL ramp: red → yellow, brightening."""
    import colorsys
    lut = np.zeros((256, 3), dtype=np.uint8)
    for x in range(256):
        # Hue 0..85 degrees
        h_deg = 85.0 * (x / 255.0)
        h = h_deg / 360.0
        s = 1.0
        # Lightness: 0..1 up to mid, then flat
        l = min(1.0, x / 128.0)
        r, g, b = colorsys.hls_to_rgb(h, l, s)
        lut[x] = (int(r * 255), int(g * 255), int(b * 255))
    return lut


def _make_doom_fire_lut() -> np.ndarray:
    """Classic Doom fire palette approximated as 256 RGB colors."""
    key_colors = np.array([
        [  0,   0,   0],
        [  7,   7,   7],
        [ 31,   7,   7],
        [ 47,  15,   7],
        [ 71,  15,   7],
        [ 87,  23,   7],
        [103,  31,   7],
        [119,  31,   7],
        [143,  39,   7],
        [159,  47,   7],
        [175,  63,   7],
        [191,  71,   7],
        [199,  71,   7],
        [223,  79,   7],
        [223,  87,   7],
        [223,  87,   7],
        [215,  95,   7],
        [215,  95,   7],
        [215, 103,  15],
        [207, 111,  15],
        [207, 119,  15],
        [207, 127,  15],
        [207, 135,  23],
        [199, 135,  23],
        [199, 143,  23],
        [199, 151,  31],
        [191, 159,  31],
        [191, 159,  31],
        [191, 167,  39],
        [191, 167,  39],
        [191, 175,  47],
        [183, 175,  47],
        [183, 183,  47],
        [183, 183,  55],
        [207, 207, 111],
        [223, 223, 159],
        [239, 239, 199],
        [255, 255, 255],
    ], dtype=np.uint8)

    stops = []
    n_keys = key_colors.shape[0]
    for i in range(n_keys):
        pos = i / (n_keys - 1)
        stops.append((pos, key_colors[i].tolist()))
    return _make_lut_from_stops(stops)


def _make_viridis_lut() -> np.ndarray:
    # Approximate viridis with a few key colors from the official palette
    stops = [
        (0.0,  (68,  1, 84)),
        (0.25, (59, 82,139)),
        (0.50, (33,145,140)),
        (0.75, (94,201, 98)),
        (1.0,  (253,231, 37)),
    ]
    return _make_lut_from_stops(stops)


def _make_inferno_lut() -> np.ndarray:
    stops = [
        (0.0,  (  0,   0,   4)),
        (0.25, ( 87,  15, 109)),
        (0.50, (187,  55,  84)),
        (0.75, (249, 142,   8)),
        (1.0,  (252, 255, 164)),
    ]
    return _make_lut_from_stops(stops)


def _make_plasma_lut() -> np.ndarray:
    stops = [
        (0.0,  ( 13,   8, 135)),
        (0.25, (126,   3, 167)),
        (0.50, (203,  71, 119)),
        (0.75, (248, 149,  64)),
        (1.0,  (240, 249,  33)),
    ]
    return _make_lut_from_stops(stops)


def _make_magma_lut() -> np.ndarray:
    stops = [
        (0.0,  (  0,   0,   4)),
        (0.25, ( 73,  18,  99)),
        (0.50, (150,  50,  98)),
        (0.75, (226, 102,  73)),
        (1.0,  (252, 253, 191)),
    ]
    return _make_lut_from_stops(stops)


def _make_turbo_lut() -> np.ndarray:
    # Approximate Google's Turbo colormap with a few key stops.
    stops = [
        (0.0,  ( 48,  18,  59)),
        (0.25, ( 31, 120, 180)),
        (0.50, ( 78, 181,  75)),
        (0.75, (241, 208,  29)),
        (1.0,  (133,  32,  26)),
    ]
    return _make_lut_from_stops(stops)


GRAY_LUT      = _make_gray_lut()
INFERNO_LUT   = _make_inferno_lut()
VIRIDIS_LUT   = _make_viridis_lut()
PLASMA_LUT    = _make_plasma_lut()
MAGMA_LUT     = _make_magma_lut()
TURBO_LUT     = _make_turbo_lut()
FIRE_LUT      = _make_fire_lut()
DOOM_FIRE_LUT = _make_doom_fire_lut()

COLOR_MAPS = {
    "Gray": GRAY_LUT,
    "Inferno": INFERNO_LUT,
    "Viridis": VIRIDIS_LUT,
    "Plasma": PLASMA_LUT,
    "Magma": MAGMA_LUT,
    "Turbo": TURBO_LUT,
    "Fire": FIRE_LUT,
    "Doom": DOOM_FIRE_LUT,
}

DEFAULT_CMAP_NAME = "Doom Fire"
