"""Gradio demo for UniParser/MolParser-Mobile — molecular image → E-SMILES."""

import spaces  # MUST come before any CUDA-touching import
import torch
from PIL import Image
import gradio as gr
from transformers import AutoModelForImageTextToText, AutoProcessor

MODEL_ID = "UniParser/MolParser-Mobile"
DTYPE = torch.float16

# Load model and processor at module scope (ZeroGPU hijack intercepts .to("cuda"))
processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
model = AutoModelForImageTextToText.from_pretrained(
    MODEL_ID,
    dtype=DTYPE,
    trust_remote_code=True,
).to("cuda").eval()


@spaces.GPU(duration=30)
def recognize_molecule(image: str) -> str:
    """Convert a molecular structure image into an E-SMILES string.

    Args:
        image: Path to a molecular structure image (PNG/JPG).
    """
    if image is None:
        return "Please upload a molecular structure image."

    pil_image = Image.open(image).convert("RGB")
    inputs = processor(images=pil_image, return_tensors="pt")
    inputs = {k: v.to("cuda", dtype=DTYPE) for k, v in inputs.items()}

    with torch.no_grad():
        output_ids = model.generate(
            **inputs,
            max_length=256,
            num_beams=1,
            do_sample=False,
        )

    caption = processor.batch_decode(output_ids, skip_special_tokens=True)[0]
    return caption


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

with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
            # MolParser-Mobile: Molecule Image → E-SMILES
            Upload a molecular structure image and get its **E-SMILES** representation.
            This is a 9.98M-parameter visual chemical structure recognition model (OCSR)
            from [UniParser](https://huggingface.co/UniParser/MolParser-Mobile).
            """
        )

        with gr.Row():
            image_input = gr.Image(
                label="Molecular structure image",
                type="filepath",
                sources=["upload", "clipboard"],
            )
            with gr.Column():
                output_text = gr.Textbox(
                    label="E-SMILES output",
                    lines=4,
                    interactive=False,
                )
                run_btn = gr.Button("Recognize", variant="primary")

        run_btn.click(
            fn=recognize_molecule,
            inputs=image_input,
            outputs=output_text,
            api_name="recognize",
        )

        gr.Examples(
            examples=[
                ["examples/01-ordinary-molecule.png"],
                ["examples/02-markush-substituent.png"],
                ["examples/03-connection-point.png"],
                ["examples/05-ring-attachment-one-repeat-range.png"],
                ["examples/06-ring-attachment-two-substituents.png"],
            ],
            inputs=image_input,
            outputs=output_text,
            fn=recognize_molecule,
            cache_examples=True,
            cache_mode="lazy",
        )

demo.launch(mcp_server=True)