"""
AlignDRAW — 512x512 canvas adaptation.
Uses the original COCO checkpoint, adapting its initial canvas parameter.
IMPORTANT: This is not a newly trained 512x512 checkpoint.
"""

import os
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces
import pickle
import numpy as np
import torch
import torch.nn.functional as F
import gradio as gr
import imageio.v2 as imageio

from PIL import Image
from huggingface_hub import hf_hub_download
from model import AlignDRAW

REPO = "mediasynthesismuseum/aligndraw"
HERE = os.path.dirname(os.path.abspath(__file__))

IMAGE_SIZE = 512
DRAW_STEPS = 69

# Build the actual 512x512 model.
model = AlignDRAW(imgSize=IMAGE_SIZE, runSteps=DRAW_STEPS)

# Download the original checkpoint.
checkpoint_path = hf_hub_download(REPO, "aligndraw_coco.pt")
checkpoint = torch.load(checkpoint_path, map_location="cpu")

# Adapt the original learned 32x32 initial canvas to 512x512.
old_c0 = checkpoint["c0"].reshape(1, 3, 32, 32)

new_c0 = F.interpolate(
    old_c0,
    size=(IMAGE_SIZE, IMAGE_SIZE),
    mode="bilinear",
    align_corners=False
)

checkpoint["c0"] = new_c0.reshape(1, -1)

# Load the original weights, including the adapted canvas.
model.load_state_dict(checkpoint, strict=True)
model.eval().to("cuda")

# Original vocabulary.
DICT_PATH = hf_hub_download(REPO, "dictionary.pkl")

with open(DICT_PATH, "rb") as f:
    DICT = pickle.load(f, encoding="latin1")


def _tokens(caption):
    words = caption.strip().lower().replace(".", " .").split()

    unknown = sorted({
        w for w in words
        if w not in DICT and w != "."
    })

    toks = [
        DICT.get(w, DICT["UNK"])
        for w in words
    ] or [DICT["UNK"]]

    return toks, unknown


def _to_image(t):
    """Convert a native 512x512 tensor into a PIL image."""
    array = (
        t.detach()
        .cpu()
        .permute(1, 2, 0)
        .numpy() * 255
    ).clip(0, 255).astype(np.uint8)

    return Image.fromarray(array)


@spaces.GPU(duration=120)
def generate(caption, n_samples):
    if not caption.strip():
        raise gr.Error("Type a caption first.")

    toks, unknown = _tokens(caption)

    y = torch.tensor(
        [toks],
        dtype=torch.long,
        device="cuda"
    )

    count = int(n_samples)

    # Generate native 512x512 samples.
    imgs = model.generate(
        y.expand(count, -1),
        steps=DRAW_STEPS
    )

    gallery = [
        _to_image(imgs[i])
        for i in range(count)
    ]

    # Generate the 69-step drawing sequence.
    _, sequence = model.generate(
        y,
        steps=DRAW_STEPS,
        return_seq=True
    )

    frames = [
        np.array(_to_image(step[0]))
        for step in sequence
    ]

    frames += [frames[-1]] * 8

    gif_path = os.path.join(HERE, "draw.gif")
    imageio.mimsave(gif_path, frames, duration=0.09)

    note = (
        "💡 Not in the original vocabulary (treated as UNK): "
        + ", ".join(unknown)
    ) if unknown else ""

    note += (
        "\n\n⚠️ 512x512 canvas adaptation using original weights. "
        "This model has not been retrained at 512x512."
    )

    return gallery, gif_path, note


EXAMPLES = [
    "a red stop sign on a street corner .",
    "a herd of elephants walking in a field .",
    "a yellow school bus on the road .",
    "a bowl of oranges on a table .",
    "a plate with a sandwich and french fries .",
    "a large blue passenger jet in the sky .",
]

CSS = """
.gradio-container {
    max-width: 1100px !important;
    margin: 0 auto !important;
}
#ad-title {
    text-align: center;
}
"""

with gr.Blocks(title="AlignDRAW 512") as demo:

    gr.Markdown(
        "# 🖍️ AlignDRAW — 512×512",
        elem_id="ad-title"
    )

    gr.Markdown(
        "A 512×512 canvas adaptation of the 2015 AlignDRAW "
        "architecture, using the original COCO checkpoint. "
        "The drawing process runs for 69 steps."
    )

    with gr.Row():
        cap = gr.Textbox(
            label="Caption",
            value=EXAMPLES[0],
            scale=4
        )

        n = gr.Slider(
            1, 4, 1,
            step=1,
            label="Samples",
            scale=1
        )

    go = gr.Button("🖍️ Draw", variant="primary")
    note = gr.Markdown()

    with gr.Row():
        gallery = gr.Gallery(
            label="Native 512×512 samples",
            columns=2,
            height=520
        )

        anim = gr.Image(
            label="69-step drawing process",
            height=400
        )

    go.click(
        generate,
        inputs=[cap, n],
        outputs=[gallery, anim, note]
    )

    gr.Examples(EXAMPLES, inputs=cap)


if __name__ == "__main__":
    demo.queue().launch(
        theme=gr.themes.Citrus(),
        css=CSS
    )