"""Gestalt: Large Multimodal Interplay Model — text-to-image demo.

Gestalt (GeWu-Lab/Gestalt) is a unified multimodal discrete-diffusion model.
This Space implements the reference text-to-image path from
https://github.com/GeWu-Lab/Gestalt `scripts/inference/run_t2i.py`:
an empty system prompt, the Stage-III T2I conversation template, and
MaskGit-style parallel sampling (64 steps, cfg 4.0, gumbel-max sampling with
confidence-based remasking). The resulting 32x32 VQVAE token grid is decoded
to a 512x512 image with the IBQ vision tokenizer
(TencentARC/IBQ-Tokenizer-16384) exactly as SEED-Voken decodes it.
"""

import os

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

import spaces  # noqa: E402  (must precede torch/CUDA imports)

import sys
from pathlib import Path

import numpy as np
import torch
import gradio as gr
from PIL import Image
from huggingface_hub import hf_hub_download
from transformers import PreTrainedTokenizerFast

sys.path.insert(0, str(Path(__file__).resolve().parent))

import gestalt.model.modeling_gestalt  # noqa: F401,E402  (registers GestaltConfig / GestaltModelLM)
from gestalt.model.config import SPECIAL_TOKENS, VOCAB_CONFIG  # noqa: E402
from gestalt.model.modeling_gestalt import GestaltModelLM  # noqa: E402
from ibq_decoder import IBQDecoder  # noqa: E402

MODEL_ID = "GeWu-Lab/Gestalt"
IBQ_ID = "TencentARC/IBQ-Tokenizer-16384"
LAT_H, LAT_W = 32, 32
MAX_STEPS = 64

DEFAULT_STEPS = 64
DEFAULT_CFG = 4.0
DEFAULT_TEMPERATURE = 1.0
DEFAULT_SEED = 0


def load_tokenizer() -> PreTrainedTokenizerFast:
    """Load and validate the Gestalt tokenizer (mirrors builder.load_tokenizer)."""
    tokenizer_file = hf_hub_download(MODEL_ID, "tokenizer.json")
    tokenizer = PreTrainedTokenizerFast(tokenizer_file=str(tokenizer_file))
    tokenizer.bos_token = "<|startoftext|>"
    # The EOS added-token (id 126081) has an unusual, visually-empty content
    # string; resolve it by ID so eos_token_id is correct across versions.
    tokenizer.eos_token = tokenizer.added_tokens_decoder[SPECIAL_TOKENS.EOS].content
    tokenizer.pad_token = "<padding>"

    expected_ids = {
        "<|startoftext|>": SPECIAL_TOKENS.BOS,
        "<system>": SPECIAL_TOKENS.SYSTEM_START,
        "</system>": SPECIAL_TOKENS.SYSTEM_END,
        "<user>": SPECIAL_TOKENS.USER_START,
        "</user>": SPECIAL_TOKENS.USER_END,
        "<|mdm_mask|>": SPECIAL_TOKENS.MASK,
        "<padding>": SPECIAL_TOKENS.PADDING,
        "<IMAGE>": SPECIAL_TOKENS.IMAGE_START,
        "</IMAGE>": SPECIAL_TOKENS.IMAGE_END,
        "<answer>": SPECIAL_TOKENS.ANSWER_START,
        "</answer>": SPECIAL_TOKENS.ANSWER_END,
    }
    for token, expected in expected_ids.items():
        if tokenizer.convert_tokens_to_ids(token) != expected:
            raise ValueError(f"Tokenizer special-token ID mismatch: {token}")
    if tokenizer.eos_token_id != SPECIAL_TOKENS.EOS:
        raise ValueError(
            f"Tokenizer EOS token id is {tokenizer.eos_token_id}; expected {SPECIAL_TOKENS.EOS}."
        )
    # The empty system prompt must tokenize to no tokens at all.
    if tokenizer("", add_special_tokens=False)["input_ids"] != []:
        raise ValueError("Empty string did not tokenize to an empty token list.")
    if len(tokenizer) != SPECIAL_TOKENS.VISUAL_TOKEN_END:
        raise ValueError(
            f"Tokenizer size is {len(tokenizer)}, expected {SPECIAL_TOKENS.VISUAL_TOKEN_END}."
        )
    return tokenizer


def validate_model_vocab(model: torch.nn.Module, expected_size: int) -> None:
    """Validate embedding/output rows against the expected extended vocab size."""
    input_rows = model.get_input_embeddings().weight.shape[0]
    if input_rows != expected_size:
        raise ValueError(f"Input embedding has {input_rows} rows; expected {expected_size}.")
    output_embeddings = model.get_output_embeddings()
    if output_embeddings is not None:
        output_rows = output_embeddings.weight.shape[0]
        if output_rows != expected_size:
            raise ValueError(f"Output projection has {output_rows} rows; expected {expected_size}.")


print("Loading Gestalt tokenizer...")
tokenizer = load_tokenizer()

print(f"Loading Gestalt model from {MODEL_ID}...")
model = GestaltModelLM.from_pretrained(
    MODEL_ID,
    dtype=torch.bfloat16,
    device_map=None,
    trust_remote_code=False,
)
model = model.to("cuda")
model.eval()
validate_model_vocab(model, VOCAB_CONFIG.EXTENDED_VOCAB_SIZE)
print(f"Model loaded (vocab_size={model.config.vocab_size})")

print("Loading IBQ vision tokenizer decoder...")
ibq_ckpt = hf_hub_download(IBQ_ID, "imagenet256_16384.ckpt")
ibq = IBQDecoder().load_from_checkpoint(ibq_ckpt).to("cuda").eval()
print("IBQ decoder loaded")


@spaces.GPU(duration=50)
def generate(
    prompt: str,
    seed: int = DEFAULT_SEED,
    steps: int = DEFAULT_STEPS,
    cfg_scale: float = DEFAULT_CFG,
    temperature: float = DEFAULT_TEMPERATURE,
    progress=gr.Progress(track_tqdm=True),
) -> Image.Image:
    """Generate a 512x512 image from a text prompt with Gestalt discrete diffusion.

    Args:
        prompt: Text prompt describing the image to generate.
        seed: Random seed (-1 for a random seed each run).
        steps: Number of MaskGit-style diffusion sampling steps (8-64).
        cfg_scale: Classifier-free guidance scale (0 disables CFG).
        temperature: Gumbel-max sampling temperature.

    Returns:
        The generated PIL image.
    """
    prompt = (prompt or "").strip()
    if not prompt:
        raise gr.Error("Please enter a text prompt.")
    seed = int(seed)
    if seed < 0:
        seed = int.from_bytes(os.urandom(4), "big") % (2**31)
    steps = int(min(max(int(steps), 8), MAX_STEPS))
    cfg_scale = float(max(cfg_scale, 0.0))
    temperature = float(max(temperature, 0.0))

    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

    # Reference: scripts/inference/run_t2i.py — empty system prompt, no special tokens.
    system_prompt_tokens = tokenizer("", add_special_tokens=False)["input_ids"]
    user_prompt_tokens = tokenizer(prompt, add_special_tokens=False)["input_ids"]

    with torch.no_grad():
        visual_tokens = model.generate_t2i(
            system_prompt_tokens=system_prompt_tokens,
            user_prompt_tokens=user_prompt_tokens,
            lat_h=LAT_H,
            lat_w=LAT_W,
            timesteps=steps,
            temperature=temperature,
            cfg_scale=cfg_scale,
        )

    raw_codes = (visual_tokens[0].cpu().numpy() - VOCAB_CONFIG.VISUAL_TOKEN_OFFSET).astype(np.int32)
    if np.any(raw_codes < 0) or np.any(raw_codes >= VOCAB_CONFIG.VQVAE_CODEBOOK_SIZE):
        raise RuntimeError("Model returned out-of-range VQVAE codebook indices.")

    with torch.no_grad():
        img = ibq.decode_indices(
            torch.from_numpy(raw_codes.astype(np.int64)).unsqueeze(0).to("cuda"),
            lat_h=LAT_H,
            lat_w=LAT_W,
        )[0].float()

    return IBQDecoder.to_pil(img)


EXAMPLES = [
    ["A cat sitting on a windowsill."],
    ["A cat wearing a pair of sunglasses, riding a skateboard, in space."],
]

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

with gr.Blocks(title="Gestalt · Text-to-Image") as demo:
    with gr.Column(elem_id="col-container"):
        gr.HTML(
            """
            <div style="text-align:center; max-width:700px; margin:0 auto;">
                <h1 style="font-size:2rem; margin-bottom:0.25rem;">Gestalt: Large Multimodal Interplay Model</h1>
                <p style="opacity:0.8; font-size:1rem;">
                    Unified multimodal discrete diffusion — text-to-image generation with the
                    <a href="https://huggingface.co/GeWu-Lab/Gestalt">GeWu-Lab/Gestalt</a> model,
                    decoded through the IBQ vision tokenizer.
                </p>
            </div>
            """
        )
        with gr.Row():
            with gr.Column():
                prompt = gr.Textbox(
                    label="Prompt",
                    placeholder="Describe the image you want...",
                    lines=2,
                    autofocus=True,
                )
                run_button = gr.Button("Generate", variant="primary")
            with gr.Column():
                image_out = gr.Image(label="Generated image", height=512)

        with gr.Accordion("Advanced options", open=False):
            seed = gr.Number(label="Seed (-1 = random)", value=DEFAULT_SEED, precision=0)
            steps = gr.Slider(label="Sampling steps", minimum=8, maximum=MAX_STEPS, step=1, value=DEFAULT_STEPS)
            cfg_scale = gr.Slider(label="CFG scale", minimum=0.0, maximum=10.0, step=0.1, value=DEFAULT_CFG)
            temperature = gr.Slider(label="Temperature", minimum=0.0, maximum=2.0, step=0.05, value=DEFAULT_TEMPERATURE)

        gr.Examples(
            examples=EXAMPLES,
            inputs=[prompt],
            outputs=[image_out],
            fn=generate,
            cache_examples=True,
            cache_mode="lazy",
        )

        gr.Markdown(
            """
            ---
            **Model**: [GeWu-Lab/Gestalt](https://huggingface.co/GeWu-Lab/Gestalt) (Apache-2.0) ·
            [Paper](https://arxiv.org/abs/2610.00576) · [Code](https://github.com/GeWu-Lab/Gestalt) ·
            Vision tokenizer: [TencentARC/IBQ-Tokenizer-16384](https://huggingface.co/TencentARC/IBQ-Tokenizer-16384) (Apache-2.0)
            """
        )

    run_button.click(
        fn=generate,
        inputs=[prompt, seed, steps, cfg_scale, temperature],
        outputs=[image_out],
        api_name="generate",
    )
    prompt.submit(
        fn=generate,
        inputs=[prompt, seed, steps, cfg_scale, temperature],
        outputs=[image_out],
        api_name=False,
    )

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