"""PixelGen demo Space — a text-to-pixel-art model with 1,257 hand-written
parameters and zero training steps.

The `pixelgen/` package and the `weights/pixelgen.pt` checkpoint are vendored
from https://huggingface.co/DJLougen/pixelgen (MIT). The demo mirrors the
reference implementation (`PixelGen.sprite` / `generate.py`) 1:1: same
tokenizer, same forward pass, same nearest-neighbour upscale — it only adds a
UI. The model is tiny and CPU-bound, so inference runs directly in the web
process; `@spaces.GPU` stays as a portable no-op in case the hardware is ever
moved to ZeroGPU.
"""

import spaces  # imported first; a no-op off ZeroGPU

import gradio as gr
from PIL import Image

from pixelgen import weights as weights_io
from pixelgen.model import PixelGen
from pixelgen.palette import COLOR_WORDS, RAMP_NAMES
from pixelgen.templates import CLASSES
from pixelgen.tokenizer import CLASS_SYNONYMS

MODEL_ID = "DJLougen/pixelgen"
DEFAULT_SEED = 0
DEFAULT_SCALE = 8  # 16 x 8 -> 128x128 PNG, the model's default output


def _load_model() -> PixelGen:
    """Load the shipped checkpoint; fall back to rebuilding weights in code.

    `PixelGen()` sets every weight deterministically in `__init__`, so both
    paths produce bit-identical sprites (the repo's round-trip test covers
    this) — the fallback only guards against a missing/corrupt file.
    """
    try:
        model = weights_io.load()  # resolves to weights/pixelgen.pt
        print(f"loaded {MODEL_ID} from weights/pixelgen.pt")
    except Exception as exc:  # pragma: no cover - robustness fallback
        print(f"checkpoint load failed ({exc!r}); rebuilding weights from source")
        model = PixelGen()
    return model


MODEL = _load_model()
N_PARAMS = sum(p.numel() for p in MODEL.parameters())


@spaces.GPU(duration=30)
def generate(prompt: str, seed: int = DEFAULT_SEED, scale: int = DEFAULT_SCALE):
    """Generate a PICO-8 pixel-art sprite from a short text prompt.

    Args:
        prompt: A colour word (red, blue, green, yellow, purple, orange, pink,
            brown, white — plus synonyms such as gold, crimson, emerald) and a
            class word (potion, sword, shield, ship, invader, slime, skull,
            tree, key, heart, star — plus synonyms such as alien, blade).
            For example "red potion" or "golden key". Unknown words are
            ignored; a prompt with no known class renders the default sprite
            (a green potion).
        seed: Selects the hand-marked "?" pixels, giving small variations of
            the same sprite across seeds.
        scale: Nearest-neighbour upscale factor; the output PNG is
            16*scale x 16*scale pixels (default 8 -> 128x128).

    Returns:
        (sprite, seed_strip): the sprite as a PIL image, plus the same prompt
        rendered at seeds 0-7 to show the seed variation.
    """
    seed, scale = int(seed), int(scale)
    sprite = Image.fromarray(MODEL.generate([prompt], [seed], scale=scale)[0], "RGB")
    strip_arr = MODEL.generate([prompt] * 8, list(range(8)), scale=scale)
    strip = [Image.fromarray(arr, "RGB") for arr in strip_arr]
    return sprite, strip


def build_prompt(body: str, colour: str) -> str:
    """Compose a prompt like 'red potion' from the two dropdowns."""
    return f"{colour} {body}"


def vocab_markdown() -> str:
    """Render the model's fixed 59-word vocabulary as reference tables."""
    cls_rows = "\n".join(f"| {c} | {', '.join(ws)} |" for c, ws in CLASS_SYNONYMS.items())
    by_ramp: dict = {}
    for word, ramp in COLOR_WORDS.items():
        by_ramp.setdefault(ramp, []).append(word)
    col_rows = "\n".join(f"| {r} | {', '.join(ws)} |" for r, ws in by_ramp.items())
    return (
        "Prompts are a bag of words: every known word counts, everything else is "
        "ignored. A prompt with **no known class word** falls back to the default "
        "class (**potion**) and one with **no known colour word** to the default "
        "ramp (**green**) — so `\"dragon\"` gives you a green potion.\n\n"
        "### Sprite classes\n"
        "| class | words |\n|---|---|\n"
        f"{cls_rows}\n\n"
        "### Colour ramps\n"
        "| colour | words |\n|---|---|\n"
        f"{col_rows}"
    )


CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
.gradio-container img { image-rendering: pixelated; }
"""

EXAMPLES = [
    ["red potion", 0],
    ["golden key", 3],
    ["blue invader", 0],
    ["purple slime", 4],
]

with gr.Blocks(title="PixelGen") as demo:
    gr.Markdown(
        f"# 🕹️ PixelGen\n"
        f"A text-to-pixel-art model with **{N_PARAMS:,} parameters — zero of them "
        f"trained**. Every weight was typed in by hand: prompts like "
        f"`\"red potion\"` or `\"blue invader\"` become 16×16 PICO-8 sprites, "
        f"upscaled to crisp pixel art. "
        f"[Model card](https://huggingface.co/{MODEL_ID}) · "
        f"[weight-space programming](https://rikiyat.github.io/blog/weight-space-programming/)"
    )
    with gr.Column(elem_id="col-container"):
        with gr.Row():
            with gr.Column(scale=1):
                prompt = gr.Textbox(
                    label="Prompt", value="red potion", placeholder="e.g. blue invader"
                )
                with gr.Accordion("…or build one", open=False):
                    body_dd = gr.Dropdown(CLASSES, value="potion", label="sprite class")
                    colour_dd = gr.Dropdown(RAMP_NAMES, value="red", label="colour")
                    for dd in (body_dd, colour_dd):
                        dd.change(build_prompt, [body_dd, colour_dd], prompt)
                seed = gr.Slider(0, 99, value=DEFAULT_SEED, step=1, label="seed")
                scale = gr.Slider(
                    4, 16, value=DEFAULT_SCALE, step=1,
                    label="scale (output is 16×scale px)",
                )
                go = gr.Button("Generate", variant="primary")
            with gr.Column(scale=1):
                out = gr.Image(label="sprite", type="pil", height=300)
        gr.Markdown(
            "**Same prompt, seeds 0–7** — a seed only toggles the hand-marked "
            "`?` pixels, so seeds give small variations of the same sprite."
        )
        strip = gr.Gallery(columns=8, height=150, label="seeds")
        with gr.Accordion("Prompt vocabulary", open=False):
            gr.Markdown(vocab_markdown())
        gr.Examples(
            examples=EXAMPLES,
            inputs=[prompt, seed],
            outputs=[out, strip],
            fn=generate,
            cache_examples=True,
            cache_mode="lazy",
        )

    go.click(generate, [prompt, seed, scale], [out, strip], api_name="generate")
    prompt.submit(generate, [prompt, seed, scale], [out, strip])
    demo.load(generate, [prompt, seed, scale], [out, strip])

if __name__ == "__main__":
    # Gradio 6: theme/css are launch() parameters, not Blocks constructor ones.
    demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)