import os

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

import spaces  # noqa: E402  (must be imported before any torch/CUDA usage)
import time  # noqa: E402
from pathlib import Path  # noqa: E402

import gradio as gr  # noqa: E402
import torch  # noqa: E402
from jinja2 import Template  # noqa: E402
from PIL import Image  # noqa: E402
from transformers import (  # noqa: E402
    AutoProcessor,
    Qwen2_5_VLForConditionalGeneration,
)

# ---------------------------------------------------------------------------
# Model — mirrors the authors' reference implementation in
# SpatialBlock/src/demo.py (Apache-2.0). The reason model wraps the question
# in a chain-of-thought prompt and answers between <answer> </answer> tags.
# ---------------------------------------------------------------------------

MODEL_ID = "rsoohyun/SpatialBlock-7B-reason"
MAX_PIXELS = 128 * 28 * 28  # authors' default (demo.py)
MIN_PIXELS = 16 * 28 * 28
MAX_NEW_TOKENS = 1024  # authors' demo.sh default for the reason variant

IMAGE_TOKEN = "<image>"
HERE = Path(__file__).parent
COT_TEMPLATE_FILE = HERE / "rl_cot.jinja"

model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    MODEL_ID,
    dtype=torch.bfloat16,
    attn_implementation="sdpa",
)
model.eval()
model = model.to("cuda")

processor = AutoProcessor.from_pretrained(
    MODEL_ID,
    use_fast=True,
    max_pixels=MAX_PIXELS,
    min_pixels=MIN_PIXELS,
)


def apply_cot_template(text: str) -> str:
    """Exactly the authors' apply_cot_template() from src/demo.py."""
    format_prompt = COT_TEMPLATE_FILE.read_text(encoding="utf-8")
    format_prompt = Template(format_prompt.strip())
    return format_prompt.render(content=text + "\n")


def build_messages(text: str, images: list) -> list:
    """Exactly the authors' build_messages() from src/demo.py."""
    chunks = text.split(IMAGE_TOKEN)
    content = []
    if len(chunks) == 1:
        content += [{"type": "image", "image": image} for image in images]
        content.append({"type": "text", "text": text})
    else:
        if len(chunks) - 1 != len(images):
            raise ValueError(
                f"{len(chunks) - 1} '{IMAGE_TOKEN}' markers but {len(images)} images"
            )
        for i, chunk in enumerate(chunks):
            if chunk:
                content.append({"type": "text", "text": chunk})
            if i < len(images):
                content.append({"type": "image", "image": images[i]})
    return [{"role": "user", "content": content}]


def _parse_answer_letter(text: str) -> str:
    """Extract the option letter from '<answer>X</answer>' (or a bare letter)."""
    if "<answer>" in text:
        inner = text.split("<answer>")[-1].split("</answer>")[0]
        return inner.strip()
    stripped = text.strip()
    return stripped[:1] if stripped else ""


# ---------------------------------------------------------------------------
# Inference
# ---------------------------------------------------------------------------


@spaces.GPU(duration=40)  # measured 4-7s per example; 40s covers a full
# 1024-token chain-of-thought at conservative decode speed with margin.
def answer(images, question):
    if not images:
        raise gr.Error("Please upload at least one image (or pick an example).")
    if not question or not question.strip():
        raise gr.Error("Please enter a question.")

    # gr.Gallery passes a list of (media, caption) tuples; keep the media.
    if images and not isinstance(images[0], (str, Image.Image)):
        images = [item[0] for item in images]

    pil_images = [Image.open(p).convert("RGB") for p in images]

    text = apply_cot_template(question.strip())
    messages = build_messages(text, pil_images)

    inputs = processor.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_dict=True,
        return_tensors="pt",
    ).to(model.device, dtype=torch.bfloat16)

    start = time.perf_counter()
    with torch.inference_mode():
        generated_ids = model.generate(
            **inputs,
            max_new_tokens=MAX_NEW_TOKENS,
            do_sample=False,
        )
    elapsed = time.perf_counter() - start

    trimmed = [
        out_ids[len(in_ids):]
        for in_ids, out_ids in zip(inputs["input_ids"], generated_ids)
    ]
    output = processor.batch_decode(
        trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
    )[0]

    letter = _parse_answer_letter(output)
    timing = f"_{elapsed:.1f}s on GPU"
    return letter, output.strip(), timing


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------

EXAMPLES_DIR = HERE / "examples"

FLIP_QUESTION = (
    "Question : Here is an image of a 3D structure made of block.<image>\n"
    "Based on the images, how does the 3D structure appear when flip vertically?\n"
    "Options:\nA.<image>\nB.<image>\nC.<image>\nD.<image>"
)

Q1_QUESTION = (
    "Question : Here are images of a 3D structure made of block <image>\n"
    "and the direction of views.<image>\n"
    "Based on the images, how does the 3D structure appear when viewed from the front?\n"
    "Options:\nA.<image>\nB.<image>\nC.<image>\nD.<image>"
)

Q3_QUESTION = (
    "Question : Here are images of two 3D structures made of block.\n"
    "Structure 1 : <image>\n"
    "Structure 2 : <image>\n"
    "Based on the images, if we move Structure 2 so that its orange and purple "
    "blocks overlap with the matching ones in Structure 1, which 3D structure "
    "can be formed?\n"
    "Options:\nA.<image>\nB.<image>\nC.<image>\nD.<image>"
)


def _ex(question: str, files: list):
    return [[str(EXAMPLES_DIR / f) for f in files], question]


EXAMPLES = [
    _ex(
        FLIP_QUESTION,
        ["image1.png", "image2.png", "image3.png", "image4.png", "image5.png"],
    ),
    _ex(
        Q1_QUESTION,
        ["q1_img1.jpg", "q1_img2.jpg", "q1_opt_a.jpg", "q1_opt_b.jpg",
         "q1_opt_c.jpg", "q1_opt_d.jpg"],
    ),
    _ex(
        Q3_QUESTION,
        ["q3_struct1.jpg", "q3_struct2.jpg", "q3_opt_a.jpg", "q3_opt_b.jpg",
         "q3_opt_c.jpg", "q3_opt_d.jpg"],
    ),
]

with gr.Blocks(title="SpatialBlock") as demo:
    gr.HTML(
        """
        <div style="display: flex; align-items: center; gap: 12px; flex-wrap: wrap;">
          <span style="font-size: 36px;">🧱</span>
          <h1 style="margin: 0; font-size: 26px;">SpatialBlock — Spatial Intelligence in LVLMs</h1>
        </div>
        <p style="margin-top: 4px;">
          <em>SpatialBlock: Enhancing Spatial Intelligence in LVLMs via Synthetic
          Block-Stacking Problem</em> — vision-language models fine-tuned on
          15,000 synthetic block-stacking problems to reason about the 3D
          structure behind 2D images.
          <a href="https://huggingface.co/papers/2609.07064">Paper</a> ·
          <a href="https://github.com/rsoohyun/SpatialBlock">GitHub</a> ·
          <a href="https://huggingface.co/rsoohyun/SpatialBlock-7B-reason">Model</a> ·
          <a href="https://huggingface.co/datasets/rsoohyun/SpatialBlock-15k">Dataset</a>
        </p>
        """
    )
    gr.Markdown(
        "Upload the images for a spatial multiple-choice question — one image per "
        "`<image>` marker in the question — and ask. The model "
        "(`SpatialBlock-7B-reason`) thinks step by step, then gives its final "
        "answer between `<answer> </answer>` tags. "
        "**Try the examples below** to see it in action."
    )

    with gr.Row():
        with gr.Column():
            images_in = gr.Gallery(
                label="Images (one per <image> marker, in order)",
                file_types=["image"],
                columns=3,
                height=240,
                type="filepath",
                interactive=True,
            )
            question_in = gr.Textbox(
                label="Question (use <image> where each uploaded image goes)",
                placeholder=(
                    "Question : Here are images of a 3D structure made of block <image>\n"
                    "and the direction of views.<image>\n..."
                ),
                lines=9,
                max_lines=14,
            )
            submit_btn = gr.Button("Answer", variant="primary")

        with gr.Column():
            answer_letter = gr.Textbox(
                label="Answer",
                interactive=False,
                buttons=["copy"],
            )
            reasoning_out = gr.Textbox(
                label="Step-by-step reasoning",
                lines=14,
                max_lines=24,
                interactive=False,
                buttons=["copy"],
            )
            timing_out = gr.Markdown()

    submit_btn.click(
        fn=answer,
        inputs=[images_in, question_in],
        outputs=[answer_letter, reasoning_out, timing_out],
    )

    gr.Examples(
        examples=EXAMPLES,
        inputs=[images_in, question_in],
        fn=answer,
        outputs=[answer_letter, reasoning_out, timing_out],
        cache_examples=True,
        cache_mode="lazy",
        label="Examples from SpatialBlock-15k (and the repo's flip demo)",
    )

    gr.Markdown(
        """
        <details>
        <summary><b>How this works</b></summary>

        The question text contains `<image>` markers; each uploaded image is
        substituted for one marker, in upload order. This is exactly the
        interface of the authors' `src/demo.py` reference implementation.

        - The question is wrapped in the authors' reasoning prompt, and the
          model produces a numbered reasoning sequence ending with
          `<answer>X</answer>`.
        - The paper also trains a **direct** model that predicts the option
          letter immediately
          ([SpatialBlock-7B-direct](https://huggingface.co/rsoohyun/SpatialBlock-7B-direct));
          this demo runs the reasoning variant.

        The checkpoint is a Qwen2.5-VL-7B-Instruct fine-tune on
        [SpatialBlock-15k](https://huggingface.co/datasets/rsoohyun/SpatialBlock-15k),
        a fully synthetic set of block-stacking problems covering 3D-to-2D
        projection, viewpoint transformation, and structural combination.
        </details>

        <details>
        <summary><b>Citation</b></summary>

        ```bibtex
        @misc{ryu2026spatialblockenhancingspatialintelligence,
              title={SpatialBlock: Enhancing Spatial Intelligence in LVLMs via Synthetic Block-Stacking Problem},
              author={Soohyun Ryu and Sohee Kim and Eunho Yang},
              year={2026},
              eprint={2609.07064},
              archivePrefix={arXiv},
              primaryClass={cs.CV},
              url={https://arxiv.org/abs/2609.07064},
        }
        ```
        </details>
        """
    )

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Citrus())