import math
import PIL
from PIL.Image import Image
from functools import lru_cache

RGB = tuple[int, int, int]

# Load image, palette
img: Image = PIL.Image.open('peacock-bayer.png')

# Get palette info
assert img.mode == 'P'
raw_pal = img.getpalette()
palette: list[RGB] = []
for i in range(0, len(raw_pal), 3):
    palette.append(tuple(raw_pal[i:i+3]))
palette_indexes = {c: i for i, c in enumerate(palette)}

# Get darkest/lightest colors
@lru_cache
def brightness(c: RGB):
    return c[0]*0.3 + c[1]*0.6 + c[2]*0.1

darkest_color = min(palette, key=brightness)
lightest_color = max(palette, key=brightness)

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

def nearest_color(color: RGB, colors: list[RGB]) -> RGB:
    return min(colors, key=lambda c: distance_sq(c, color))

def choose_color_pair_kmeans(counts: dict[RGB,int]) -> tuple[RGB, RGB]:
    colors = [k for k, v in counts.items() if v > 0]
    total_count = sum(counts.values())
    if len(colors) == 0:
        raise ValueError('no colors present')
    elif len(colors) == 1:
        return colors[0], None

    def weighted_avg(colors: list[RGB]) -> RGB:
        result = [0, 0, 0]
        total_weight = 0
        for color in colors:
            weight = counts[color]
            total_weight += weight
            for i in range(3):
                result[i] += weight*color[i]
        if total_weight == 0:
            total_weight += 1 # Hack to handle division-by-zero
        return tuple(round(x/total_weight) for x in result)

    # k-means to choose the two colors
    old_result = (None, None)
    result = (min(colors, key=brightness), max(colors, key=brightness))
    print(f"k-means using {counts=}")
    print(f"starting value: {result}")
    while True:
        lists = ([], [])
        for color in colors:
            d_dark = distance_sq(color, result[0])
            d_light = distance_sq(color, result[1])
            if d_dark < d_light:
                lists[0].append(color)
            else:
                lists[1].append(color)
        old_result = result
        print(f"{lists=}")
        result = tuple(weighted_avg(cs) for cs in lists)
        print(f"updated: {result}")
        if old_result == result:
            break
    return tuple(nearest_color(x, palette) for x in result)


def choose_color_pair_luminance(counts: dict[RGB, int]) -> tuple[RGB,RGB]:
    # Calculate brightness threshold
    mean_brightness = 0
    total_count = 0
    for rgb, count in counts.items():
        mean_brightness += brightness(rgb)*count
        total_count += count
    mean_brightness = mean_brightness/total_count

    # Separate colors into two lists
    lists = ([], [])
    for color in counts.keys():
        if brightness(color) <= mean_brightness:
            lists[0].append(color)
        else:
            lists[1].append(color)
    print(lists)

    # Find least-bad approximations for each color
    def choose_color(cs: list[RGB]) -> RGB:
        best = None
        best_cost = math.inf
        for rgb in palette:
            cost = sum(distance_sq(rgb, c) * counts[c] for c in cs)
            if cost < best_cost:
                best = rgb
                best_cost = cost
        return best
    dark, light = tuple(choose_color(l) for l in lists)
    if dark == light:
        light = None
    print(dark, light)
    return dark, light


def two_tone(img: Image, xy1: tuple[int,int] = None, xy2: tuple[int,int] = None) -> None:
    # choose the two colors
    x1, y1 = xy1 or (0, 0)
    x2, y2 = xy2 or (img.width, img.height)
    x1 = max(x1, 0)
    y1 = max(y1, 0)
    x2 = min(x2, img.width)
    y2 = min(y2, img.height)
    counts = {}
    for x in range(x1, x2):
        for y in range(y1, y2):
            color = palette[img.getpixel((x,y))]
            if color not in counts:
                counts[color] = 0
            counts[color] += 1
    #dark, light = choose_color_pair_luminance(counts)
    print(f"choosing color pair for {xy1=}")
    dark, light = choose_color_pair_kmeans(counts)
    print(f"final: {dark=}, {light=}\n")
    if light == None:
        # Only one color, *ought* to be a no-op
        light = dark
    colors = [dark, light]

    # change each pixel's color to nearest
    for x in range(x1, x2):
        for y in range(y1, y2):
            color = palette[img.getpixel((x, y))]
            replacement = nearest_color(color, colors)
            img.putpixel((x, y), palette_indexes[replacement])


BLOCK_SIZE = 10
blocked = img.copy()
for x in range(0, img.width, BLOCK_SIZE):
    for y in range(0, img.height, BLOCK_SIZE):
        two_tone(blocked, (x, y), (x + BLOCK_SIZE, y + BLOCK_SIZE))
blocked.save("blocked.png")
for x in range(5):
    for y in range(5):
        blocked.putpixel((x, y), 0)