import numpy as np
from PIL import Image
from tqdm import tqdm

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

# 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 = [
    (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):
    x /= 255
    if x <= 0.04045:
        return x/12.92
    else:
        return ((x + 0.055)/1.055)**2.4

def rgb_to_linear(rgb):
    return tuple(component_to_linear(c) for c in rgb)

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

def make_bayer_matrix(size):
    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 = 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