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

import spaces  # MUST be the first import after the env var, before torch
import re
import threading

import gradio as gr
import torch
from transformers import AutoProcessor, Qwen3VLForConditionalGeneration, TextIteratorStreamer

MODEL_ID = "xinyili0624/GPD-4B"

DEFAULT_SYSTEM_PROMPT = (
    "You are a helpful assistant for 3D spatial reasoning. Think step by step using "
    "<think></think> tags, then give your final answer inside <answer> </answer> tags."
)

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

ANSWER_RE = re.compile(r"<answer>\s*(.*?)\s*</answer>", re.DOTALL)


def _extract_answer(text: str) -> str:
    """Return the content of the last <answer>...</answer> tag (or its best effort)."""
    matches = ANSWER_RE.findall(text)
    if matches:
        return matches[-1].strip()
    return ""


@spaces.GPU(duration=45)
def answer(
    image,
    question: str,
    system_prompt: str = DEFAULT_SYSTEM_PROMPT,
    max_new_tokens: int = 1024,
    temperature: float = 0.0,
    top_p: float = 0.95,
):
    """Answer a spatial-reasoning question about an image with GPD-4B.

    GPD (Geometry-Privileged Distillation) is an RGB-only vision-language model
    fine-tuned from Qwen3-VL-4B-Instruct for single-image 3D spatial reasoning
    (observer shifts, directional relations, depth/distance estimation). It
    thinks step by step between <think> and </think> tags and wraps its
    final answer in <answer> </answer> tags.

    Args:
        image: input image (PIL Image) or None for a text-only question.
        question: the spatial question to ask about the image.
        system_prompt: system prompt guiding the reasoning format.
        max_new_tokens: maximum number of generated tokens.
        temperature: sampling temperature; 0 = deterministic greedy decoding.
        top_p: nucleus sampling probability (used when temperature > 0).

    Yields:
        Tuple of (extracted_final_answer, full_generated_text), streamed live.
    """
    question = (question or "").strip()
    if not question:
        yield "Please enter a question.", ""
        return

    content = []
    if image is not None:
        content.append({"type": "image", "image": image})
    content.append({"type": "text", "text": question})
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": content},
    ]

    prompt_text = processor.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    images = [image] if image is not None else None
    inputs = processor(
        text=[prompt_text], images=images, return_tensors="pt"
    ).to(model.device)

    gen_kwargs = {
        "max_new_tokens": int(max_new_tokens),
        "do_sample": float(temperature) > 0.0,
    }
    if gen_kwargs["do_sample"]:
        gen_kwargs["temperature"] = float(temperature)
        gen_kwargs["top_p"] = float(top_p)

    streamer = TextIteratorStreamer(
        processor.tokenizer, skip_prompt=True, skip_special_tokens=True
    )
    gen_kwargs["streamer"] = streamer

    generation_thread = threading.Thread(
        target=lambda: model.generate(**inputs, **gen_kwargs), daemon=True
    )
    generation_thread.start()

    full_text = ""
    for chunk in streamer:
        full_text += chunk
        yield _extract_answer(full_text), full_text

    generation_thread.join()

    final_answer = _extract_answer(full_text)
    if not final_answer:
        tail = full_text.strip().splitlines()[-1] if full_text.strip() else ""
        final_answer = tail.strip()
    yield final_answer, full_text


EXAMPLES = [
    [
        "examples/ex_distance_to_chair.jpg",
        (
            "From the person's position who took the photo, what is the distance "
            "to the center of chair (red point)? Express it in meters. Calculate "
            "or judge based on the 3D center points of these objects. Enter a "
            "single numeric response and no other details.\n"
            "Think step by step and provide your reasoning inside <think> </think> tags. "
            "Then give your final answer inside <answer> </answer> tags with a "
            "single number or value."
        ),
    ],
    [
        "examples/ex_depth_difference.jpg",
        (
            "The depth of pillow (red point) is 2.5. Use this information to "
            "find the absolute depth difference between plant (green point) and "
            "table (blue point). Answer in meters. Calculate or judge based on "
            "the 3D center points of these objects. Please give a response in "
            "the form of a single number.\n"
            "Think step by step and provide your reasoning inside <think> </think> tags. "
            "Then give your final answer inside <answer> </answer> tags with a "
            "single number or value."
        ),
    ],
    [
        "examples/ex_closer_object.jpg",
        (
            "From the red point laundry basket, which object, container (green "
            "point) or box (blue point), is closer? Calculate or judge based on "
            "the 3D center points of these objects. Choose the correct response "
            "from the given choices.\n"
            "A. box (blue point)\n"
            "B. container (green point)\n"
            "Your answer can only include one of options A, B.\n"
            "Think step by step and provide your reasoning inside <think> </think> tags. "
            "Then give your final answer inside <answer> </answer> tags. For "
            "example: <answer>A</answer>"
        ),
    ],
    [
        "examples/ex_directional_relation.jpg",
        (
            "Analyze the image and determine the directional relationship "
            "between the monitor (red bbox) and the whiteboard (blue bbox) as "
            "seen by the observer. Calculate or judge based on the 3D center "
            "points of these objects.\n"
            "The options describe the spatial relationship between two objects "
            "in terms of left-right (left, right, or empty if "
            "indistinguishable), above-below (above, below, or empty if "
            "indistinguishable), and farther-closer (farther, closer, or empty "
            "if indistinguishable).\n"
            "Select the right response from the available choices.\n"
            "A. left, , closer\n"
            "B. right, , farther\n"
            "C. , above, closer\n"
            "D. right, below, \n"
            "Your answer can only include one of options A, B, C or D.\n"
            "Think step by step and provide your reasoning inside <think> </think> tags. "
            "Then give your final answer inside <answer> </answer> tags. For "
            "example: <answer>A</answer>"
        ),
    ],
    [
        "examples/ex_observer_shift.jpg",
        (
            "From the observer's original viewpoint, where is dining table "
            "(red bbox)? How does this relationship shift once the observer "
            "moves to chair's 3D center (green bbox) and faces oven (blue "
            "bbox)? Calculate or judge based on the 3D center points of these "
            "objects. Please use the world coordinate system to determine the "
            "up and down position relationship of objects.\n"
            "For multiple-choice questions, consider only the state after the "
            "observer has moved.\n"
            "The options describe the spatial relationship between object and "
            "observer in terms of left-right (left, right, or empty if "
            "indistinguishable), above-below (above, below, or empty if "
            "indistinguishable), and front-behind (front, behind, or empty if "
            "indistinguishable).\n"
            "Choose the correct response from the given choices.\n"
            "A. right, below, front\n"
            "B. right, , behind\n"
            "C. left, above, behind\n"
            "D. right, below, \n"
            "Your answer can only include one of options A, B, C or D.\n"
            "Think step by step and provide your reasoning inside <think> </think> tags. "
            "Then give your final answer inside <answer> </answer> tags. For "
            "example: <answer>A</answer>"
        ),
    ],
]

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

with gr.Blocks(title="GPD · Spatial Reasoning") as demo:
    gr.Markdown(
        """
        # 🧭 GPD · Distilling Routed 3D Privilege for Spatial Reasoning
        **GPD-4B** is an RGB-only vision-language model (fine-tuned from
        [Qwen3-VL-4B-Instruct](https://huggingface.co/Qwen/Qwen3-VL-4B-Instruct)) that answers 3D
        spatial-reasoning questions about a single image — observer-shift, directional relations,
        and depth/distance estimation — thinking step by step and wrapping the final answer in
        `<answer>` tags.

        [Paper](https://huggingface.co/papers/2610.12355) ·
        [Code](https://github.com/ZJU-REAL/GPD) ·
        [GPD-4B](https://huggingface.co/xinyili0624/GPD-4B) ·
        [GPD-2B](https://huggingface.co/xinyili0624/GPD-2B) ·
        [GPD-15k Data](https://huggingface.co/datasets/xinyili0624/GPD-15k)
        """
    )
    with gr.Column(elem_id="col-container"):
        with gr.Row():
            with gr.Column(scale=1):
                image_input = gr.Image(
                    type="pil", label="Image", height=420, sources=["upload", "clipboard"]
                )
            with gr.Column(scale=1):
                question_input = gr.Textbox(
                    label="Question",
                    lines=10,
                    placeholder=(
                        "Ask a spatial question about the image, e.g. "
                        "'What is closer to the camera, the chair or the table?'"
                    ),
                )
                submit_btn = gr.Button("Answer", variant="primary")

        answer_output = gr.Textbox(
            label="Final answer",
            lines=1,
            buttons=["copy"],
        )
        reasoning_output = gr.Textbox(
            label="Full reasoning",
            lines=12,
            buttons=["copy"],
        )

        with gr.Accordion("Advanced settings", open=False):
            system_prompt_input = gr.Textbox(
                label="System prompt",
                value=DEFAULT_SYSTEM_PROMPT,
                lines=2,
            )
            with gr.Row():
                max_new_tokens_input = gr.Slider(
                    label="Max new tokens", minimum=64, maximum=3072, value=1024, step=64
                )
                temperature_input = gr.Slider(
                    label="Temperature (0 = deterministic)", minimum=0.0, maximum=1.5, value=0.0, step=0.05
                )
                top_p_input = gr.Slider(
                    label="Top-p (used when temperature > 0)", minimum=0.05, maximum=1.0, value=0.95, step=0.05
                )

        gr.Examples(
            examples=EXAMPLES,
            inputs=[image_input, question_input],
            outputs=[answer_output, reasoning_output],
            fn=answer,
            cache_examples=True,
            cache_mode="lazy",
            label="Examples (from the authors' GPD-15k dataset — SPAR scenes)",
        )

    submit_btn.click(
        answer,
        inputs=[image_input, question_input, system_prompt_input, max_new_tokens_input, temperature_input, top_p_input],
        outputs=[answer_output, reasoning_output],
        api_name="answer",
    )

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