from typing import Generator
from PIL import Image
from pathlib import Path
import numpy as np
import math

def load_palette(path: Path) -> np.ndarray:
    palette = []
    with open(path, 'r') as f:
        for line in f:
            hex_color = line.strip()
            # Convert hex to RGB tuple
            rgb_color = [int(hex_color[i:i+2], 16) for i in (0, 2, 4)]
            palette.append(rgb_color)
    return np.array(palette)

pico8_palette = load_palette('../pico-8-secret-palette.hex')

def load_image_128(path: Path) -> Image.Image:
    img = Image.open(path).convert('RGB')
    w, h = img.size
    if w > 128 or h > 128:
        if w % 128 == 0 and h % 128 == 0:
            resample = Image.Resampling.NEAREST
        else:
            resample = Image.Resampling.BICUBIC
        if w > h:
            w, h = 128, int(round(h * (128 / w)))
        else:
            w, h = int(round(w * (128 / h))), 128
        img = img.resize((w, h), resample=resample)
    return img


def quantize_image(img: Image.Image) -> Image.Image:
    pal_idxs = choose_subpalette_indexes(img, 16, pico8_palette)
    pal = np.array([pico8_palette[i] for i in pal_idxs]).astype(np.float32)
    img_np = np.array(img)
    pixels = img_np.reshape(-1, 3).astype(np.float32)
    distances = np.sum((pixels[:, None, :] - pal[None, :, :]) ** 2, axis=2)
    closest_idxs = np.argmin(distances, axis=1).astype(np.uint8).reshape(img_np.shape[1], -1)
    result = Image.fromarray(closest_idxs, mode="P")
    result.putpalette(pal.astype(np.uint8).reshape(-1).tolist())
    return result


def choose_subpalette_indexes(img: Image.Image, n: int, palette: np.ndarray):
    assert len(palette.shape) == 2;
    assert palette.shape[1] == 3
    img_np = np.array(img)
    result = list(range(n))
    num_steps = 0
    pal_idx = 0
    while num_steps < n:
        num_steps += 1
        old_val = result[pal_idx]
        result[pal_idx] = None
        best_err = math.inf
        best_val = None
        for val in [i for i in range(len(palette)) if i not in result]:
            result[pal_idx] = val
            err = _squared_quantization_error(img_np, [palette[i] for i in result])
            if err < best_err:
                best_err = err
                best_val = val
        result[pal_idx] = best_val
        if best_val != old_val:
            num_steps = 0
        pal_idx = (pal_idx + 1) % len(result)
    return result

def _squared_quantization_error(img_np: np.ndarray, palette: np.ndarray):
    pixels = img_np.reshape(-1, 3).astype(np.float32)
    pal = np.array(palette, dtype=np.float32)
    sq_distances = np.sum((pixels[:, None, :] - pal[None, :, :]) ** 2, axis=2)
    min_errs = np.min(sq_distances, axis=1)
    return np.sum(min_errs)

def walk_corpus(quantized=True) -> Generator[Path, None, None]:
    if quantized:
        return Path('../corpus/quantized/').glob('*.png')
    else:
        return Path('../corpus/').glob('*.*')