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

import spaces  # MUST come before torch / transformers
import re
import time

import gradio as gr
import torch
from PIL import Image, ImageDraw, ImageFont
from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration

MODEL_ID = "xyzzzh/Hi-Token"
GROUNDING_PROMPT = (
    "Please provide the bounding box coordinate of the region this sentence "
    "describes: {expr}."
)

# A Hi-Token box is 12 special tokens: for each of the 4 coordinates
# (x_min, y_min, x_max, y_max) an axis-specific hundreds/tens/ones triplet,
# e.g. <x_hundreds_4><x_tens_2><x_ones_7> -> 427 in the normalized [0, 999]
# coordinate space.
COORD_RE = re.compile(r"<(x|y)_(hundreds|tens|ones)_(\d)>")
PLACES = ("hundreds", "tens", "ones")


def parse_hitoken_box(response: str):
    """Reconstruct [x_min, y_min, x_max, y_max] (0-999 ints) from the model's
    hierarchical coordinate tokens.

    The tokens appear in sequence: x_min triplet, y_min triplet, x_max triplet,
    y_max triplet. Each integer is 100*hundreds + 10*tens + ones.
    Returns None if the response does not contain a well-formed box.
    """
    tokens = COORD_RE.findall(response)
    if len(tokens) < 12:
        return None
    tokens = tokens[:12]
    values = []
    for i in range(0, 12, 3):
        trip = tokens[i : i + 3]
        places = [p for _, p, _ in trip]
        if places != list(PLACES):
            return None
        axes = {a for a, _, _ in trip}
        if len(axes) != 1:
            return None
        values.append((axes.pop(), 100 * int(trip[0][2]) + 10 * int(trip[1][2]) + int(trip[2][2])))
    if [axis for axis, _ in values] != ["x", "y", "x", "y"]:
        return None
    return [v for _, v in values]


processor = AutoProcessor.from_pretrained(MODEL_ID, max_pixels=768 * 28 * 28)
model = (
    Qwen2_5_VLForConditionalGeneration.from_pretrained(
        MODEL_ID, torch_dtype=torch.bfloat16, attn_implementation="sdpa"
    )
    .to("cuda")
    .eval()
)


@spaces.GPU(duration=10)
def ground(
    image,
    expression: str = "the main subject",
    show_raw_tokens: bool = False,
):
    """Locate a region in an image from a text description and draw its box.

    Args:
        image: input image to search.
        expression: referring expression describing the target region,
            e.g. "the smallest rubber duck".
        show_raw_tokens: also return the model's raw generated token string.

    Returns:
        Tuple of (annotated image, decoded coordinate string).
    """
    t0 = time.time()
    image = image.convert("RGB")
    width, height = image.size

    query = GROUNDING_PROMPT.format(expr=expression.strip())
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image"},
                {"type": "text", "text": query},
            ],
        }
    ]
    prompt = processor.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    inputs = processor(
        text=[prompt], images=[image], padding=True, return_tensors="pt"
    ).to(model.device)

    with torch.inference_mode():
        generated = model.generate(
            **inputs,
            max_new_tokens=64,
            do_sample=False,
            use_cache=True,
        )
    response_ids = generated[:, inputs.input_ids.shape[1] :]
    # Coordinate tokens ARE special tokens, so they must survive decoding.
    response = processor.batch_decode(
        response_ids,
        skip_special_tokens=False,
        clean_up_tokenization_spaces=False,
    )[0]

    box = parse_hitoken_box(response)
    elapsed = time.time() - t0

    if box is None:
        annotated = image.copy()
        summary = "No bounding box found in the model output."
        if show_raw_tokens:
            summary += f"\n\nRaw output:\n{response.strip()}"
        return annotated, summary

    x_min, y_min, x_max, y_max = box
    note = ""
    if x_max < x_min or y_max < y_min:
        # Normalize an inverted prediction so it can still be drawn.
        x_min, x_max = min(x_min, x_max), max(x_min, x_max)
        y_min, y_max = min(y_min, y_max), max(y_min, y_max)
        note = " (corners normalized)"
    px = [
        x_min * width / 999.0,
        y_min * height / 999.0,
        x_max * width / 999.0,
        y_max * height / 999.0,
    ]

    annotated = image.copy()
    draw = ImageDraw.Draw(annotated)
    line = max(2, int(0.004 * max(width, height)) + 1)
    # White halo under an orange core so the box reads on any background.
    draw.rectangle(px, outline=(255, 255, 255), width=line * 3)
    draw.rectangle(px, outline=(255, 64, 0), width=line)

    label = expression.strip()
    try:
        font = ImageFont.load_default(size=max(14, int(0.02 * width) + 4))
    except TypeError:
        font = ImageFont.load_default()
    tb = draw.textbbox((0, 0), label, font=font)
    tw, th = tb[2] - tb[0], tb[3] - tb[1]
    lx = max(0, min(px[0], width - tw - 8))
    ly = px[1] - th - 12
    if ly < 0:
        ly = px[1] + 6
    draw.rectangle([lx - 5, ly - 5, lx + tw + 5, ly + th + 5], fill=(255, 64, 0))
    draw.text((lx, ly), label, fill=(255, 255, 255), font=font)

    summary = (
        f"Bounding box (0-999){note}: [{x_min}, {y_min}, {x_max}, {y_max}]  "
        f"pixels: [{px[0]:.0f}, {px[1]:.0f}, {px[2]:.0f}, {px[3]:.0f}]  "
        f"inference: {elapsed:.1f}s"
    )
    if show_raw_tokens:
        summary += f"\n\nRaw output:\n{response.strip()}"
    return annotated, summary


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

with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
            # Hi-Token: Hierarchical Coordinate Tokenization for Visual Grounding

            Point this generative grounding model at an image and describe a
            region in plain language — it emits each bounding-box coordinate as
            an axis-specific hundreds/tens/ones token triplet (12 tokens per
            box) instead of plain numbers, which sharpens boundary precision.

            Model: [xyzzzh/Hi-Token](https://huggingface.co/xyzzzh/Hi-Token)
            (Qwen2.5-VL-3B fine-tune, Hi-R1) ·
            [Paper](https://huggingface.co/papers/2608.03471)
            """
        )
        with gr.Row():
            with gr.Column():
                image_in = gr.Image(
                    type="pil", label="Image", sources=["upload", "clipboard"]
                )
                expression_in = gr.Textbox(
                    label="Referring expression",
                    placeholder="e.g. the smallest rubber duck",
                    lines=1,
                )
                show_raw = gr.Checkbox(
                    label="Show raw coordinate tokens", value=False
                )
                run_btn = gr.Button("Locate", variant="primary")
            with gr.Column():
                image_out = gr.Image(type="pil", label="Grounded region")
                text_out = gr.Textbox(label="Decoded coordinates", lines=4)

        gr.Examples(
            examples=[
                ["examples/rubber_ducks.jpg", "the smallest rubber duck"],
                ["examples/rubber_ducks.jpg", "the largest rubber duck"],
                ["examples/vintage_camera.jpg", "the black instant camera"],
                ["examples/pizza_board.jpg", "the slice being lifted"],
                ["examples/motorcycle_vintage.jpg", "the green motorcycle"],
            ],
            inputs=[image_in, expression_in],
            outputs=[image_out, text_out],
            fn=ground,
            cache_examples=True,
            cache_mode="lazy",
        )

    run_btn.click(
        ground,
        inputs=[image_in, expression_in, show_raw],
        outputs=[image_out, text_out],
        api_name="ground",
    )

demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)