#!/usr/bin/env python3

from random import random
import sys

# Higher value = preserve color light level; lower = preserve color
FUDGE_FACTOR = 5.0

# E observer (equal energy):
REFERENCE_X = 100.0
REFERENCE_Y = 100.0
REFERENCE_Z = 100.0

def read_playpal(filename):
    """Read first palette from the given file containing a PLAYPAL lump."""
    with open(filename, "rb") as f:
        paldata = f.read(768)
        palette = []
        for i in range(256):
            off = i * 3
            palette.append((paldata[off], paldata[off+1], paldata[off+2]))
    return palette

def rgb_to_xyz(r, g, b):
    """Convert sRGB value to XYZ color space."""
    r = r / 255
    g = g / 255
    b = b / 255

    if r > 0.04045:
        r = ((r + 0.055) / 1.055) ** 2.4
    else:
        r = r / 12.92

    if g > 0.04045:
        g = ((g + 0.055) / 1.055) ** 2.4
    else:
        g = g / 12.92

    if b > 0.04045:
        b = ((b + 0.055) / 1.055) ** 2.4
    else:
        b = b / 12.92

    r = r * 100
    g = g * 100
    b = b * 100

    return (
        r * 0.4124 + g * 0.3576 + b * 0.1805,
        r * 0.2126 + g * 0.7152 + b * 0.0722,
        r * 0.0193 + g * 0.1192 + b * 0.9505,
    )

def xyz_to_cieluv(x, y, z):
    """Convert XYZ color value to CIE-LUV color space."""
    denom = max(x + 15*y + 3*z, 0.001)
    u = 4*x / denom
    v = 9*y / denom

    y = y / 100
    if y > 0.008856:
        y = y ** 0.3333
    else:
        y = 7.787*y + (16 / 116)

    ref_u = 4*REFERENCE_X / (REFERENCE_X + 15*REFERENCE_Y + 3*REFERENCE_Z)
    ref_v = 9*REFERENCE_Y / (REFERENCE_X + 15*REFERENCE_Y + 3*REFERENCE_Z)

    cie_l = 116*y - 16
    return (
        cie_l,
        13 * cie_l * (u - ref_u),  # CIE-u*
        13 * cie_l * (v - ref_v),  # CIE-v*
    )

def darken_palette(palette, frac):
    result = []
    for l, u, v in palette:
        frac2 = frac #+ random() * 0.1 - 0.05
        result.append((l * frac2, u * frac2, v * frac2))
    return result

def invert_palette(palette):
    white = xyz_to_cieluv(*rgb_to_xyz(255, 255, 255))
    return [
        (100 - l, 0, 0)
        for l, u, v in palette
    ]

def nearest_color(palette, color):
    best_index = None
    best_dist = 9999999
    for i, color2 in enumerate(palette):
        dist = (FUDGE_FACTOR * (color[0] - color2[0]) ** 2 +
                (color[1] - color2[1]) ** 2 +
                (color[2] - color2[2]) ** 2)
        if dist < best_dist:
            best_dist = dist
            best_index = i
    return best_index

def make_colormap(base_palette, palettes):
    result = []
    for palette in palettes:
        for color in palette:
            result.append(nearest_color(base_palette, color))
    return bytes(result)

def light_levels():
    for i in range(32):
        val = (31 - i) * 255 / 31
        l, _, _ = xyz_to_cieluv(*rgb_to_xyz(val, val, val))
        yield l / 100.0

palette = read_playpal(sys.argv[1])
luv_palette = [xyz_to_cieluv(*rgb_to_xyz(*rgb)) for rgb in palette]

colormap = make_colormap(luv_palette, [
    darken_palette(luv_palette, frac)
    for frac in light_levels()
] + [invert_palette(luv_palette), luv_palette])

with open("colormap.lmp", "wb") as f:
    f.write(colormap)