#!/usr/bin/env python3
"""Download a photo, letterbox to 800×480, Floyd-Steinberg dither to the
Waveshare 7.3" 6-colour palette, and write src/image.h as a PROGMEM array."""

import io
import math
import urllib.request
from pathlib import Path

import numpy as np
from PIL import Image

W, H = 480, 800   # portrait canvas; firmware rotates 90° CW onto landscape display
OUT = Path(__file__).parent / "src" / "image.h"

URL = (
    "https://d3eguztg5751m.cloudfront.net/as/assets-mem-com/cmi/3/0/0/7/11557003"
    "/20231128_115457370_0_orig.jpg/-/kenneth-edholm-fort-wayne-in-obituary.jpg"
    "?maxheight=650"
)

# Waveshare 6-colour palette: (display code, match RGB)
# RGB values tuned for photographic content — skin tones sit between
# white and yellow/red, so those two anchors matter most.
PALETTE = [
    (0x0, (  0,   0,   0)),  # Black
    (0x1, (255, 255, 255)),  # White
    (0x2, (255, 220,   0)),  # Yellow
    (0x3, (220,  40,  40)),  # Red
    (0x5, ( 20,  80, 200)),  # Blue
    (0x6, ( 50, 160,  50)),  # Green
]
CODES   = np.array([c for c, _ in PALETTE], dtype=np.uint8)
PAL_RGB = np.array([rgb for _, rgb in PALETTE], dtype=np.float32)  # (6,3)


def floyd_steinberg(img_rgb: np.ndarray) -> np.ndarray:
    """In-place F-S dither; returns (H,W) array of display colour codes."""
    arr = img_rgb.astype(np.float32)
    out = np.zeros((H, W), dtype=np.uint8)

    for y in range(H):
        for x in range(W):
            px = np.clip(arr[y, x], 0, 255)
            diffs = PAL_RGB - px
            idx = int(np.argmin((diffs ** 2).sum(axis=1)))
            out[y, x] = CODES[idx]

            err = px - PAL_RGB[idx]
            if x + 1 < W:
                arr[y, x + 1]     += err * (7 / 16)
            if y + 1 < H:
                if x > 0:
                    arr[y + 1, x - 1] += err * (3 / 16)
                arr[y + 1, x]     += err * (5 / 16)
                if x + 1 < W:
                    arr[y + 1, x + 1] += err * (1 / 16)

    return out


def main() -> None:
    print(f"Downloading {URL} ...")
    req = urllib.request.Request(URL, headers={"User-Agent": "pictureFrame/1.0"})
    with urllib.request.urlopen(req, timeout=30) as r:
        img_bytes = r.read()
    img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
    print(f"Downloaded: {img.size[0]}×{img.size[1]}")

    # Letterbox: fit inside 800×480, centre on white background
    scale = min(W / img.size[0], H / img.size[1])
    new_w = int(img.size[0] * scale)
    new_h = int(img.size[1] * scale)
    resized = img.resize((new_w, new_h), Image.LANCZOS)
    canvas = Image.new("RGB", (W, H), (255, 255, 255))
    canvas.paste(resized, ((W - new_w) // 2, (H - new_h) // 2))
    canvas.save("/tmp/picture_frame_preview.png")
    print("Preview → /tmp/picture_frame_preview.png")

    print(f"Dithering {W}×{H} to 6-colour palette (this takes ~60s) ...")
    arr = np.array(canvas, dtype=np.float32)
    codes = floyd_steinberg(arr)

    # Colour distribution
    names = ["Black", "White", "Yellow", "Red", "Blue", "Green"]
    for i, (code, name) in enumerate(zip(CODES, names)):
        cnt = int((codes == code).sum())
        print(f"  {name:7s} ({code:#04x}): {cnt:7d} px  ({100*cnt/(W*H):.1f}%)")

    # Pack 4bpp: high nibble = even column, low nibble = odd column
    high   = codes[:, 0::2].astype(np.uint8)
    low    = codes[:, 1::2].astype(np.uint8)
    packed = ((high << 4) | low)
    packed_bytes = packed.tobytes()
    assert len(packed_bytes) == H * (W // 2)

    OUT.parent.mkdir(parents=True, exist_ok=True)
    print(f"Writing {OUT} ...")
    ROW = W // 2
    with OUT.open("w") as f:
        f.write("#pragma once\n")
        f.write("#include <pgmspace.h>\n\n")
        f.write(f"// {W}×{H} photo, 4bpp Waveshare 7.3\" 6-colour\n")
        f.write("// Packing: byte[x/2] = (code[x]<<4)|code[x+1], even col = high nibble\n")
        f.write(f"#define IMAGE_ROW {ROW}  // bytes per display row\n\n")
        f.write("const uint8_t IMAGE_DATA[] PROGMEM = {\n")
        for y in range(H):
            row = packed_bytes[y * ROW:(y + 1) * ROW]
            for off in range(0, ROW, 16):
                chunk = row[off:off + 16]
                f.write("    " + ",".join(f"0x{b:02X}" for b in chunk) + ",\n")
        f.write("};\n")

    kb = len(packed_bytes) / 1024
    print(f"Done — {len(packed_bytes):,} bytes ({kb:.1f} KB) → {OUT}")


if __name__ == "__main__":
    main()