""" Color-by-Numbers generator Gradio app for HuggingFace Spaces

Upload an image (photo or pixel art)
Quantize colors to a limited palette (default 8 colors)
Create a printable "color-by-numbers" page where each pixel/block contains the color's number
Also export a legend showing each number and its representative color (hex)

Dependencies: pillow, numpy, gradio
Run: pip install pillow numpy gradio
then python huggingface_color_by_numbers_app.py
When deployed to Hugging Face Spaces, set runtime to use the same file.
"""

from PIL import Image, ImageDraw, ImageFont
import numpy as np
import gradio as gr
import io

def quantize_image(img: Image.Image, n_colors: int):
    """Quantize the image to n_colors using PIL adaptive palette.
    Returns: label_map (h,w) of 0..(k-1) and palette list of (r,g,b).
    """
    # convert to RGBA and remove alpha by compositing on white background
    img = img.convert("RGBA")
    background = Image.new("RGBA", img.size, (255, 255, 255, 255))
    img = Image.alpha_composite(background, img).convert("RGB")

    # Use PIL quantize to reduce to n_colors
    q = img.convert('P', palette=Image.ADAPTIVE, colors=n_colors)
    palette = q.getpalette()  # flat list r,g,b,r,g,b,... length >= n_colors*3

    # build palette list from used palette indices
    used_indices = list(sorted(set(q.getdata())))
    # extract first n_colors entries from palette
    pal = []
    for i in range(n_colors):
        r = palette[3 * i]
        g = palette[3 * i + 1]
        b = palette[3 * i + 2]
        pal.append((r, g, b))

    # Convert paletted image to numpy labels: each pixel index corresponds to 0..(n_colors-1)
    label_map = np.array(q, dtype=np.int32)

    return label_map, pal


def sort_palette_by_brightness(palette):
    # compute luminance and sort (brighter later) so numbers have consistent ordering
    lum = [0.299 * c[0] + 0.587 * c[1] + 0.114 * c[2] for c in palette]
    order = np.argsort(lum)
    sorted_palette = [palette[i] for i in order]
    return sorted_palette, order


def remap_labels(label_map, order):
    # order says old_index -> new position in sorted palette
    # produce mapping from old index to new index
    inv = np.empty_like(order)
    for new_idx, old_idx in enumerate(order):
        inv[old_idx] = new_idx
    remapped = inv[label_map]
    return remapped


def rgb_to_hex(rgb):
    return '#%02x%02x%02x' % rgb


def generate_color_by_numbers_page(label_map: np.ndarray, palette: list, block_size: int = 24, outline: bool = True, font_path: str = None):
    """Generate a color-by-numbers image where each block contains the number (1..k).
    - label_map: 2D array of labels 0..k-1
    - palette: list of RGB colors for legend ordering
    - block_size: pixel size of a single input pixel when scaled up
    """
    h, w = label_map.shape
    k = len(palette)

    # Create canvas
    canvas_w = w * block_size
    canvas_h = h * block_size
    canvas = Image.new('RGB', (canvas_w, canvas_h), (255, 255, 255))
    draw = ImageDraw.Draw(canvas)

    # Load font
    try:
        if font_path:
            font = ImageFont.truetype(font_path, int(block_size * 0.6))
        else:
            font = ImageFont.load_default()
    except Exception:
        font = ImageFont.load_default()

    # draw each cell
    for y in range(h):
        for x in range(w):
            label = int(label_map[y, x])  # 0..k-1
            left = x * block_size
            top = y * block_size
            right = left + block_size
            bottom = top + block_size

            # fill with white
            draw.rectangle([left, top, right, bottom], fill=(255, 255, 255))

            # draw outline grid
            if outline:
                draw.rectangle([left, top, right - 1, bottom - 1], outline=(200, 200, 200))

            # draw number centered
            number = str(label + 1)
            bbox = draw.textbbox((0, 0), number, font=font)
            tw = bbox[2] - bbox[0]
            th = bbox[3] - bbox[1]
            nx = left + (block_size - tw) / 2
            ny = top + (block_size - th) / 2 - 1
            draw.text((nx, ny), number, fill=(0, 0, 0), font=font)

    return canvas


def generate_palette_legend(palette: list, block_size: int = 48):
    k = len(palette)
    padding = 10
    sw = block_size * k + padding * (k + 1)
    sh = block_size + 2 * padding + 20  # extra space for hex text
    legend = Image.new('RGB', (sw, sh), (255, 255, 255))
    d = ImageDraw.Draw(legend)

    try:
        font = ImageFont.load_default()
    except Exception:
        font = ImageFont.load_default()

    for i, rgb in enumerate(palette):
        x = padding + i * (block_size + padding)
        y = padding
        d.rectangle([x, y, x + block_size, y + block_size], fill=rgb, outline=(0, 0, 0))
        label = str(i + 1)
        bbox = d.textbbox((0, 0), label, font=font)
        tw = bbox[2] - bbox[0]
        th = bbox[3] - bbox[1]
        d.text((x + (block_size - tw) / 2, y + (block_size - th) / 2), label, fill=(0, 0, 0), font=font)
        # hex text below
        hexcode = rgb_to_hex(rgb)
        d.text((x, y + block_size + 2), hexcode, fill=(0, 0, 0), font=font)

    return legend


def process_image_to_color_by_numbers(pil_img: Image.Image, n_colors: int = 8, block_size: int = 24, outline: bool = True):
    # 1) quantize
    label_map, pal = quantize_image(pil_img, n_colors)

    # 2) sort palette by brightness
    sorted_pal, order = sort_palette_by_brightness(pal)
    remapped = remap_labels(label_map, order)

    # 3) generate page
    page = generate_color_by_numbers_page(remapped, sorted_pal, block_size=block_size, outline=outline)

    # 4) generate legend
    legend = generate_palette_legend(sorted_pal, block_size=max(32, block_size * 2))

    return page, legend


# -------------------- Gradio interface --------------------

def process_and_return(inp_img, n_colors, block_size, show_legend):
    if inp_img is None:
        return None

    pil = Image.open(inp_img)  # inp_img is file path string
    page, legend = process_image_to_color_by_numbers(pil, n_colors=n_colors, block_size=block_size)

    # combine vertically
    if show_legend:
        combined_w = max(page.width, legend.width)
        combined_h = page.height + legend.height + 20
        combined = Image.new('RGB', (combined_w, combined_h), (255, 255, 255))
        combined.paste(page, (0, 0))
        combined.paste(legend, (0, page.height + 20))
    else:
        combined = page

    bio = io.BytesIO()
    combined.save(bio, format='PNG')
    bio.seek(0)
    return bio


demo = gr.Interface(
    fn=process_and_return,
    inputs=[
        gr.Image(type='filepath', label='Upload image'),  # Changed to 'filepath' for clarity
        gr.Slider(minimum=2, maximum=12, step=1, value=8, label='Number of colors (reduction)'),
        gr.Slider(minimum=8, maximum=64, step=1, value=24, label='Block size (pixels)'),
        gr.Checkbox(label='Show legend with color swatches and hex codes', value=True)
    ],
    outputs=[gr.File(label='Download Color-by-Numbers Page (PNG)')],  # Changed to File for download
    title='🎨 Color-by-Numbers Generator',
    description='Upload an image and get a printable color-by-numbers page. The app quantizes colors to a small palette and writes numbers in each pixel block so children (or adults) can color by numbers.',
    allow_flagging="never"
)

if __name__ == "__main__":
    demo.launch()