"""
Floyd-Steinberg dithering for E Ink Spectra 6 (6-color) display.

This is a 1:1 Python port of the C++ dithering in firmware/src/dither.h.
Same palette RGB values, same error diffusion coefficients, same color
distance metric. The output should be pixel-identical to the firmware.

Native panel color codes (what the controller expects over SPI):
    0x00 = Black
    0x01 = White
    0x02 = Yellow
    0x03 = Red
    0x04 = Orange  (not used in Spectra 6, but accepted by controller)
    0x05 = Blue
    0x06 = Green

Pixel packing: 4bpp, 2 pixels per byte.
    High nibble (bits 7-4) = LEFT pixel
    Low nibble  (bits 3-0) = RIGHT pixel
    Scan order: left-to-right, top-to-bottom

Data path on the ESP32:
    JPEG → decode to RGB888 → Floyd-Steinberg dither → 4bpp native buffer
    → writeNative(buf, invert=true) → SPI command 0x10 → panel refresh
"""

import numpy as np
from PIL import Image

# Palette: (R, G, B, native_panel_index)
# These RGB values must match firmware/src/dither.h PALETTE[] exactly.
PALETTE = [
    (  0,   0,   0, 0),  # Black
    (255, 255, 255, 1),  # White
    (200,  30,  30, 3),  # Red
    (  0, 145,   0, 6),  # Green
    (  0,   0, 160, 5),  # Blue
    (230, 210,   0, 2),  # Yellow
]

# Extract just the RGB values as a numpy array for vectorized distance calc
PALETTE_RGB = np.array([(r, g, b) for r, g, b, _ in PALETTE], dtype=np.int32)
PALETTE_INDICES = [idx for _, _, _, idx in PALETTE]

# Display RGB values — what the e-paper pigments actually look like.
# Used for rendering the simulator output. These may differ from the palette
# values above (which are what we compare against during dithering).
# For now they're the same, but can be tuned independently.
DISPLAY_RGB = {
    0: (  0,   0,   0),  # Black
    1: (255, 255, 255),  # White
    2: (230, 210,   0),  # Yellow
    3: (200,  30,  30),  # Red
    5: (  0,   0, 160),  # Blue
    6: (  0, 145,   0),  # Green
}


def find_closest_color(r: int, g: int, b: int) -> int:
    """Find the palette index with minimum squared Euclidean distance."""
    best_idx = 0
    best_dist = float('inf')
    for i, (pr, pg, pb, _) in enumerate(PALETTE):
        dr = r - pr
        dg = g - pg
        db = b - pb
        dist = dr * dr + dg * dg + db * db
        if dist < best_dist:
            best_dist = dist
            best_idx = i
    return best_idx


def dither_floyd_steinberg(img: Image.Image) -> tuple[np.ndarray, Image.Image]:
    """
    Apply Floyd-Steinberg dithering to an 800x480 RGB image.

    Returns:
        native_buffer: numpy array of shape (480, 400) dtype uint8
                       Each byte = 2 pixels packed as 4bpp native panel indices.
                       High nibble = left pixel, low nibble = right pixel.
        preview_img:   PIL Image showing what the e-paper display would look like.
    """
    assert img.size == (800, 480), f"Image must be 800x480, got {img.size}"

    width, height = 800, 480
    # Work with int16 to allow negative error values
    pixels = np.array(img, dtype=np.int16)

    # Output: what native panel index each pixel gets
    result = np.zeros((height, width), dtype=np.uint8)

    for y in range(height):
        for x in range(width):
            r = int(np.clip(pixels[y, x, 0], 0, 255))
            g = int(np.clip(pixels[y, x, 1], 0, 255))
            b = int(np.clip(pixels[y, x, 2], 0, 255))

            ci = find_closest_color(r, g, b)
            pr, pg, pb, native_idx = PALETTE[ci]
            result[y, x] = native_idx

            # Quantization error
            er = r - pr
            eg = g - pg
            eb = b - pb

            # Floyd-Steinberg error diffusion (same coefficients as C++)
            # Right: 7/16, Bottom-left: 3/16, Bottom: 5/16, Bottom-right: 1/16
            if x + 1 < width:
                pixels[y, x + 1, 0] += er * 7 // 16
                pixels[y, x + 1, 1] += eg * 7 // 16
                pixels[y, x + 1, 2] += eb * 7 // 16
            if y + 1 < height:
                if x - 1 >= 0:
                    pixels[y + 1, x - 1, 0] += er * 3 // 16
                    pixels[y + 1, x - 1, 1] += eg * 3 // 16
                    pixels[y + 1, x - 1, 2] += eb * 3 // 16
                pixels[y + 1, x, 0] += er * 5 // 16
                pixels[y + 1, x, 1] += eg * 5 // 16
                pixels[y + 1, x, 2] += eb * 5 // 16
                if x + 1 < width:
                    pixels[y + 1, x + 1, 0] += er * 1 // 16
                    pixels[y + 1, x + 1, 1] += eg * 1 // 16
                    pixels[y + 1, x + 1, 2] += eb * 1 // 16

    # Pack into 4bpp buffer (2 pixels per byte, high nibble = left)
    native_buffer = np.zeros((height, width // 2), dtype=np.uint8)
    for y in range(height):
        for x in range(0, width, 2):
            left = result[y, x]
            right = result[y, x + 1]
            native_buffer[y, x // 2] = (left << 4) | right

    # Render preview image using display RGB values
    preview = np.zeros((height, width, 3), dtype=np.uint8)
    for y in range(height):
        for x in range(width):
            preview[y, x] = DISPLAY_RGB[result[y, x]]
    preview_img = Image.fromarray(preview)

    return native_buffer, preview_img