import numpy as np
from PIL import Image
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.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 dither_pixel(img, xy, pal):
    global bayer
    pal2 = [np.array(rgb_to_linear(c)) for c in pal]
    counts = [0]*len(pal)
    x, y = xy
    def diff(c1, c2):
        return sum((c1-c2)**2)
    def nearest(rgb):
        cost, index = min((diff(rgb,c), i) for i,c in enumerate(pal2))
        return index
    color = np.array(rgb_to_linear(img.getpixel(xy)))
    err = np.array([0.0,0.0,0.0])
    for _ in range(bayer.size):
        err += color
        i = nearest(err)
        counts[i] += 1
        err -= pal2[i]
    thresh = bayer[y % bayer.shape[0]][x % bayer.shape[1]]
    for i, count in enumerate(counts):
        thresh -= count
        if thresh < 0:
            return i

dithered = Image.new('P', resized.size)
dithered.putpalette([x for rgb in PICO_RGB[:16] for x in rgb])
def do_it():
    global dithered, resized
    for x in tqdm.tqdm(range(128)):
        for y in range(128):
            i = dither_pixel(resized, (x,y), PICO_RGB[:16])
            dithered.putpixel((x,y), i)
# import cProfile
# cProfile.run('do_it()')
do_it()