import numpy as np
import scipy.spatial
import PIL
from PIL.Image import Image
from tqdm import tqdm

img = PIL.Image.open('peacock.jpg')
resized = img.resize((128,128))

RGB = tuple[int, int, int]
RGBLinear = tuple[float, float, float]

# sRGB values for the PICO-8 palette.
# The first 16 entries correspond with the default palette (0-15).
# The last 16 entries correspond with the "secret" palette (128-143).
PICO_RGB: list[RGB] = [
    (0, 0, 0),
    (29, 43, 83),
    (126, 37, 83),
    (0, 135, 81),
    (171, 82, 54),
    (95, 87, 79),
    (194, 195, 199),
    (255, 241, 232),
    (255, 0, 77),
    (255, 163, 0),
    (255, 236, 39),
    (0, 228, 54),
    (41, 173, 255),
    (131, 118, 156),
    (255, 119, 168),
    (255, 204, 170),
    (41, 24, 20),
    (17, 29, 53),
    (66, 33, 54),
    (18, 83, 89),
    (116, 47, 41),
    (73, 51, 59),
    (162, 136, 121),
    (243, 239, 125),
    (190, 18, 80),
    (255, 108, 36),
    (168, 231, 46),
    (0, 181, 67),
    (6, 90, 181),
    (117, 70, 101),
    (255, 110, 89),
    (255, 157, 129)
]

def component_to_linear(x: int) -> float:
    x /= 255
    if x <= 0.04045:
        return x/12.92
    else:
        return ((x + 0.055)/1.055)**2.4

def component_to_gamma(x: float) -> int:
    if x <= 0.04045/12.92:
        result = x * 12.92
    else:
        result = (x**(1/2.4) * 1.055) - 0.055
    result *= 255
    return round(result)

def rgb_to_linear(rgb: RGB) -> RGBLinear:
    return tuple(component_to_linear(c) for c in rgb)

def linear_to_rgb(linear: RGBLinear) -> RGB:
    return tuple(component_to_gamma(c) for c in linear)

PICO_RGB_LINEAR = [rgb_to_linear(x) for x in PICO_RGB]

def distance_sq(a: RGB, b: RGB) -> int:
    result = 0
    for i in range(3):
        result += (a[i] - b[i])**2
    return result

def make_bayer_matrix(size: int) -> np.ndarray:
    assert size > 0
    assert size & (size - 1) == 0 # power of two
    m = np.array([[0]])
    while m.shape[0] < size:
        m = np.block([[4*m, 4*m+3],[4*m+2, 4*m+1]])
    return m

bayer = make_bayer_matrix(4)

def pattern_dither(img, palette, *, pat_size=4, amount=0.75):
    assert img.mode == 'RGB'
    bayer = make_bayer_matrix(pat_size)

    # Reorder palette by luminance
    old_palette = palette
    reordered = list(zip(palette, range(len(palette))))
    reordered.sort(key=lambda c: 3*c[0][0] + 6*c[0][1] + c[0][2])
    palette = [c for c, _ in reordered]
    old_palette_indexes = [i for _, i in reordered]

    # Convert palette colors to linear RGB
    palette = [rgb_to_linear(c) for c in palette]
    palette_matrix = np.array(palette)

    def nearest(lc):
        """Takes a linear RGB color c and returns a palette index"""
        squared_errors = (palette_matrix - lc)**2
        return np.argmin(np.sum(squared_errors, axis=1))

    # Choose palette indexes for every pixel in the image
    output_indexes = []
    for y in tqdm(range(img.height)):
        for x in range(img.width):
            color = np.array(rgb_to_linear(img.getpixel((x, y))))
            err = np.zeros((3,)) + color
            counts = [0] * len(palette)
            for _ in range(bayer.size):
                i = nearest(err)
                counts[i] += 1
                err = (err - palette_matrix[i]) * amount + color
            thresh = bayer[y % pat_size][x % pat_size]
            for i, count in enumerate(counts):
                thresh -= count
                if thresh < 0:
                    output_indexes.append(i)
                    break

    # Convert palette indexes to an indexed image
    result = PIL.Image.new(mode='P', size=img.size)
    result.putpalette([c for rgb in old_palette for c in rgb])
    result.putdata([old_palette_indexes[i] for i in output_indexes])
    return result


def luminance(rgb: RGB) -> int:
    return 3*rgb[0] + 6*rgb[1] + rgb[2]

def pattern_dither_twotone(img: Image, palette: list[RGB]) -> Image:
    assert img.mode == 'RGB'
    bayer = make_bayer_matrix(4)

    # Sort palette indexes by luminance
    # Used to make sure blends consistently order their colors
    palette_indexes = list(range(len(palette)))
    palette_indexes.sort(key=lambda i: luminance(palette[i]))

    # Build lookup for all two-color blends
    linear_palette = [np.array(rgb_to_linear(c)) for c in palette]
    blends: list[RGB] = []
    recipes: list[tuple[int,int,int]] = [] # index 1, index 2, ratio/16
    def make_blend(a_idx: int, b_idx: int, ratio: int):
        if distance_sq(palette[a_idx], palette[b_idx]) > 3*128**2:
            return
        a = linear_palette[a_idx]
        b = linear_palette[b_idx]
        blend = (a*(16-ratio) + b*ratio)/16
        blend = linear_to_rgb(blend)
        blends.append(blend)
        recipes.append((a_idx, b_idx, ratio))

    for i in range(len(linear_palette)):
        make_blend(i, i, 0)
    for a_idx in range(len(linear_palette)):
        for b_idx in range(a_idx + 1, len(linear_palette)):
            for ratio in range(1, 16):
                make_blend(a_idx, b_idx, ratio)
    kd = scipy.spatial.KDTree(blends)
    print(len(blends), len(set(blends)))
    print(max(blends))
    print(min(blends))


    # Convert image pixels
    result_indexes = []
    for y in tqdm(range(img.height)):
        # row = [rgb_to_linear(img.getpixel((x,y))) for x in range(img.width)]
        # _, blends_indexes = kd.query(row)
        for x in range(img.width):
            color = img.getpixel((x, y))
            _dist, blends_idx = kd.query(color)
            recipe = recipes[blends_idx]
            thresh = bayer[y % 4][x % 4]
            if recipe[2] <= thresh:
                idx = recipe[0]
            else:
                idx = recipe[1]
            result_indexes.append(idx)

    # Output image
    result = PIL.Image.new(mode='P', size=img.size)
    result.putpalette([c for rgb in palette for c in rgb])
    result.putdata(result_indexes)
    return result