"""
DiffusionGemma · Radiology VQA & interactive report infill.

ZeroGPU Gradio demo for the paper "Discrete Diffusion Language Models for
Interactive Radiology Report Drafting" (https://huggingface.co/papers/2607.01436).

It serves the LoRA finetunes from `gevaertlab/diffusiongemma-radiology-vqa` on top of
the image-conditioned discrete-diffusion backbone `google/diffusiongemma-26B-A4B-it`
(`DiffusionGemmaForBlockDiffusion`).

Two capabilities, matching the reference implementation (github.com/mxvp/discrete_diffusion_RRG):
  * VQA — image + question -> answer (models/generate.py::generate_report).
  * Bidirectional infill — fill a masked span in a report using both sides of context,
    via the fixed-position canvas-clamping sampler hook (models/infill.py).

ZeroGPU specifics:
  * `import spaces` before torch.
  * Base model + all three adapters loaded once at module scope, `.to("cuda")` eagerly.
  * `model.generate` runs only inside the `@spaces.GPU` functions.
  * A custom DiffusionGemma `transformers` wheel is bundled and installed at runtime.
"""

import glob
import os
import subprocess
import sys

# Set before torch is imported (transformers pulls torch in).
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # must precede torch so ZeroGPU can patch torch.cuda.*


def _ensure_transformers():
    """Install the bundled custom DiffusionGemma `transformers` wheel at runtime.

    Spaces installs `requirements.txt` before copying repo files into the image, so the
    wheel can't be referenced by local path there. By the time this app runs the file is
    present, so we install it here (only if DiffusionGemma isn't already importable).
    """
    try:
        import transformers  # noqa: F401

        if hasattr(transformers, "DiffusionGemmaForBlockDiffusion") or hasattr(
            getattr(transformers, "models", object), "diffusion_gemma"
        ):
            return
    except Exception:
        pass
    wheels = sorted(glob.glob(os.path.join(os.path.dirname(os.path.abspath(__file__)), "transformers-*.whl")))
    if not wheels:
        print("[dgemma] no bundled transformers wheel found", flush=True)
        return
    print(f"[dgemma] Installing bundled transformers wheel: {os.path.basename(wheels[0])}", flush=True)
    subprocess.check_call([sys.executable, "-m", "pip", "install", "--no-cache-dir", wheels[0]])
    import importlib

    importlib.invalidate_caches()


_ensure_transformers()

import contextlib
import re

import torch
import gradio as gr
from PIL import Image
from transformers import AutoProcessor, DiffusionGemmaForBlockDiffusion

HERE = os.path.dirname(os.path.abspath(__file__))
BASE_MODEL = os.environ.get("DGEMMA_BASE", "google/diffusiongemma-26B-A4B-it")
ADAPTER_REPO = os.environ.get("DGEMMA_ADAPTERS", "gevaertlab/diffusiongemma-radiology-vqa")
HF_TOKEN = os.environ.get("HF_TOKEN")
MAX_NEW_TOKENS = 256
# ZeroGPU slice: the 26B checkpoint (~49 GB bf16) needs the full backing card.
GPU_SIZE = os.environ.get("GDIFF_GPU_SIZE", "xlarge")

# LoRA adapters: label -> subfolder in the adapter repo (diffusion backbone only).
ADAPTERS = {
    "VQA-RAD (mixed X-ray/CT/MRI)": "diffusion-vqarad",
    "SLAKE (X-ray/CT/MRI + organs)": "diffusion-slake",
    "VQA-Med (radiology QA)": "diffusion-vqamed",
}
_NAME = {label: sub.replace("diffusion-", "") for label, sub in ADAPTERS.items()}

# CoT system prompt used by the finetunes (models/report_format.py).
COT_SYSTEM_PROMPT = (
    "A conversation between User and Assistant. The user asks a question, and the Assistant "
    "solves it. The assistant first thinks about the findings in the image and then provides the "
    "user with the final impression. The findings and answer are enclosed within <think> </think> "
    "and <answer> </answer> tags, respectively, i.e., <think> findings here </think>"
    "<answer> impression here </answer>"
)

# ---------------------------------------------------------------------------
# Model load (module scope, eager .to("cuda") for ZeroGPU)
# ---------------------------------------------------------------------------
print(f"[dgemma] loading processor + base model {BASE_MODEL} ...", flush=True)
processor = AutoProcessor.from_pretrained(BASE_MODEL, token=HF_TOKEN)

model = DiffusionGemmaForBlockDiffusion.from_pretrained(
    BASE_MODEL, dtype=torch.bfloat16, token=HF_TOKEN
)
model.eval()

# Attach all three diffusion LoRA adapters onto the single base model.
from peft import PeftModel

# PEFT derives the load device from peft.infer_device(), which checks
# torch.cuda.is_available(). Under the ZeroGPU module-scope hijack that returns True with
# no real GPU attached, so safetensors' get_tensors() dispatches a real CUDA op and raises
# "No CUDA GPUs are available". Force is_available() False during the adapter load so the
# adapter weights land on CPU; we then move the whole PeftModel to CUDA eagerly (below),
# which the ZeroGPU hijack intercepts and streams into VRAM on the first GPU call.
_orig_cuda_avail = torch.cuda.is_available
torch.cuda.is_available = lambda: False
try:
    _first_label = next(iter(ADAPTERS))
    model = PeftModel.from_pretrained(
        model, ADAPTER_REPO, subfolder=ADAPTERS[_first_label],
        adapter_name=_NAME[_first_label], token=HF_TOKEN,
    )
    for label, sub in ADAPTERS.items():
        if label == _first_label:
            continue
        model.load_adapter(ADAPTER_REPO, subfolder=sub, adapter_name=_NAME[label],
                           token=HF_TOKEN)
finally:
    torch.cuda.is_available = _orig_cuda_avail
model.eval()
print("[dgemma] adapters loaded:", list(_NAME.values()), flush=True)

# Eager move to CUDA at module scope — the ZeroGPU hijack intercepts this, packs the
# weights to disk, and streams them into VRAM on the first @spaces.GPU entry.
model.to("cuda")
print("[dgemma] model registered on cuda (ZeroGPU will stream on first call)", flush=True)

# The discrete-diffusion backbone denoises a fixed-length canvas whose size is set by the
# base config (canvas_length=256), independent of the `max_new_tokens` request. Infill masks
# MUST be sized to this or the canvas-clamp hook mismatches the model's internal tensor.
def _canvas_length():
    for obj in (model, getattr(model, "base_model", None)):
        cfg = getattr(obj, "config", None)
        v = getattr(cfg, "canvas_length", None)
        if v:
            return int(v)
    return 256


CANVAS_LEN = _canvas_length()
print(f"[dgemma] canvas length = {CANVAS_LEN}", flush=True)

_THINK = re.compile(r"<think>(.*?)</think>", re.S)
_ANSWER = re.compile(r"<answer>(.*?)</answer>", re.S)
# Some finetunes emit an OpenAI-Harmony-style channel prefix, e.g. "<|channel|>thought\n..."
# or a bare "thought\n..." / "analysis\n..." lead-in before the actual answer. Strip it.
_CHANNEL = re.compile(r"^\s*(?:<\|[^>]*\|>)*\s*(?:thought|analysis|final|channel)\s*\n+", re.I)


def _strip_cot(text: str) -> str:
    """Pull the <answer> (and <think> findings) out of a CoT response."""
    t = _THINK.search(text)
    a = _ANSWER.search(text)
    if t or a:
        findings = t.group(1).strip() if t else ""
        impression = a.group(1).strip() if a else ""
        return (findings + "\n" + impression).strip()
    # No explicit CoT tags — clean up channel prefixes and any stray tag tokens.
    t = re.sub(r"<\|[^>]*\|>", " ", text)
    t = _CHANNEL.sub("", t)
    t = (t.replace("<think>", "").replace("</think>", "")
         .replace("<answer>", "").replace("</answer>", ""))
    return t.strip()


def _prompt_inputs(image, instruction):
    """Chat-template an (image, text) turn — image first, per model best practices."""
    msg = [{"role": "user", "content": [
        {"type": "image", "image": image},
        {"type": "text", "text": instruction},
    ]}]
    inp = processor.apply_chat_template(
        msg, tokenize=True, add_generation_prompt=True,
        return_dict=True, return_tensors="pt")
    return {k: (v.to("cuda") if hasattr(v, "to") else v) for k, v in inp.items()}


def _canvas_ids(out, plen):
    seq = getattr(out, "sequences", None)
    if seq is None:
        seq = out[0]
    ids = seq.reshape(-1).tolist()
    return ids[plen:] if len(ids) > plen else ids


def _decode_canvas(out, plen):
    ids = _canvas_ids(out, plen)
    eos = processor.tokenizer.eos_token_id
    eos_ids = set(eos) if isinstance(eos, (list, tuple)) else ({eos} if eos is not None else set())
    for j, t in enumerate(ids):
        if t in eos_ids:
            ids = ids[:j]
            break
    return processor.tokenizer.decode(ids, skip_special_tokens=True).strip()


# ---------------------------------------------------------------------------
# VQA
# ---------------------------------------------------------------------------
@spaces.GPU(duration=120, size=GPU_SIZE)
def answer_question(image: Image.Image, question: str, adapter_label: str,
                    max_new_tokens: int = MAX_NEW_TOKENS) -> str:
    """Answer a question about a medical image with the DiffusionGemma radiology model.

    Args:
        image: the medical scan (X-ray, CT or MRI slice).
        question: a natural-language question about the image.
        adapter_label: which dataset-specific LoRA adapter to use.
        max_new_tokens: canvas length / max tokens to denoise.

    Returns:
        The model's answer as text.
    """
    if image is None:
        return "Please upload a medical image first."
    if not question or not question.strip():
        return "Please enter a question."
    model.set_adapter(_NAME.get(adapter_label, next(iter(_NAME.values()))))
    image = image.convert("RGB")
    inp = _prompt_inputs(image, question.strip())
    plen = inp["input_ids"].shape[1]
    with torch.no_grad():
        out = model.generate(**inp, max_new_tokens=int(max_new_tokens))
    return _strip_cot(_decode_canvas(out, plen))


# ---------------------------------------------------------------------------
# Bidirectional infill
# ---------------------------------------------------------------------------
@contextlib.contextmanager
def _fixed_canvas_positions(fixed_tokens: torch.Tensor, fixed_mask: torch.Tensor):
    """Clamp known canvas positions to fixed tokens every denoising step.

    Port of models/infill.py from the reference implementation: patches the diffusion
    generation mixin's `_denoising_step` so masked positions stay fixed while the model
    bidirectionally fills the rest.
    """
    from transformers.models.diffusion_gemma import generation_diffusion_gemma as G

    if not hasattr(G, "DiffusionGemmaGenerationMixin"):
        raise RuntimeError(
            "DiffusionGemmaGenerationMixin not found -- the generate API drifted; "
            "re-verify the canvas-update hook.")
    Mixin = G.DiffusionGemmaGenerationMixin
    orig_step = Mixin._denoising_step
    ft = fixed_tokens
    fm = fixed_mask.bool()

    def _clamp(x):
        return torch.where(fm.to(x.device), ft.to(x.device), x)

    def patched_step(self, *args, **kwargs):
        if "current_canvas" in kwargs and kwargs["current_canvas"] is not None:
            kwargs["current_canvas"] = _clamp(kwargs["current_canvas"])
        out = orig_step(self, *args, **kwargs)
        cur, argmax = _clamp(out[0]), _clamp(out[1])
        return (cur, argmax) + tuple(out[2:])

    Mixin._denoising_step = patched_step
    try:
        yield
    finally:
        Mixin._denoising_step = orig_step


BLANK = "[BLANK]"
INFILL_INSTRUCTION = "Write the radiology report for this medical image."


@spaces.GPU(duration=150, size=GPU_SIZE)
def infill_report(image: Image.Image, template: str, adapter_label: str,
                  bidirectional: bool = True,
                  max_new_tokens: int = MAX_NEW_TOKENS) -> str:
    """Fill a [BLANK] span in a partial radiology report using the diffusion canvas.

    The known text becomes fixed canvas positions; the [BLANK] span is denoised. With
    `bidirectional=True` the model sees text on both sides of the hole; with it False,
    only the left context is kept (the autoregressive-style baseline).

    Args:
        image: the medical scan the report describes.
        template: a report with exactly one [BLANK] marking the span to fill.
        adapter_label: which dataset-specific LoRA adapter to use.
        bidirectional: use both-sides context (True) or left-only (False).
        max_new_tokens: canvas length.

    Returns:
        Just the text the model wrote into the [BLANK].
    """
    if image is None:
        return "Please upload a medical image first."
    if not template or BLANK not in template:
        return f"Please provide a report template containing exactly one {BLANK} marker."
    model.set_adapter(_NAME.get(adapter_label, next(iter(_NAME.values()))))
    image = image.convert("RGB")

    tok = processor.tokenizer
    # The model's denoising canvas is always CANVAS_LEN long, regardless of the
    # max_new_tokens request; the clamp mask must match that exact size.
    L = CANVAS_LEN
    pad = tok.pad_token_id or 0

    # Split around the blank and locate its token span on the canvas.
    before, after = template.split(BLANK, 1)
    before_ids = tok(before, add_special_tokens=False)["input_ids"]
    after_ids = tok(after, add_special_tokens=False)["input_ids"]
    # Reserve room in the middle for the fill; cap by the user's requested span length
    # and by the remaining canvas space.
    max_fill = max(1, min(int(max_new_tokens), 96))
    hole_len = max(1, min(max_fill, L - len(before_ids) - len(after_ids) - 2))
    h0 = min(len(before_ids), L - 1)
    h1 = min(h0 + hole_len, L - 1)

    ids = list(before_ids) + [pad] * hole_len + list(after_ids)
    ids = ids[:L]
    n = len(ids)
    canvas = torch.full((L,), pad, dtype=torch.long)
    canvas[:n] = torch.tensor(ids, dtype=torch.long)

    if bidirectional:
        known = torch.ones(L, dtype=torch.bool)
        known[h0:h1] = False          # only the hole is unknown; both sides fixed
    else:
        known = torch.zeros(L, dtype=torch.bool)
        known[:h0] = True             # left context only

    inp = _prompt_inputs(image, INFILL_INSTRUCTION)
    plen = inp["input_ids"].shape[1]
    ft = canvas[None].to("cuda")
    fm = known[None].to("cuda")
    with torch.no_grad(), _fixed_canvas_positions(ft, fm):
        out = model.generate(**inp, max_new_tokens=L)
    cv = _canvas_ids(out, plen)
    fill = tok.decode(cv[h0:h1], skip_special_tokens=True).strip()
    return fill or "(model produced an empty fill — try a shorter blank or different adapter)"


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

ADAPTER_LABELS = list(ADAPTERS.keys())

with gr.Blocks(title="DiffusionGemma Radiology VQA") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            "# 🩻 DiffusionGemma · Radiology VQA & Interactive Report Infill\n"
            "Image-conditioned **discrete-diffusion** LLM for radiology, from the paper "
            "[*Discrete Diffusion Language Models for Interactive Radiology Report Drafting*]"
            "(https://huggingface.co/papers/2607.01436). "
            "Backbone [`google/diffusiongemma-26B-A4B-it`](https://huggingface.co/google/diffusiongemma-26B-A4B-it) "
            "+ LoRA finetunes [`gevaertlab/diffusiongemma-radiology-vqa`]"
            "(https://huggingface.co/gevaertlab/diffusiongemma-radiology-vqa).\n\n"
            "⚠️ Research demo — **not** a medical device, not for clinical use."
        )

        with gr.Tabs():
            # ---- VQA tab ----
            with gr.Tab("Visual Question Answering"):
                with gr.Row():
                    with gr.Column():
                        vqa_image = gr.Image(type="pil", label="Medical image", height=340)
                        vqa_question = gr.Textbox(
                            label="Question",
                            placeholder="e.g. Is there evidence of an aortic aneurysm?",
                        )
                        vqa_adapter = gr.Dropdown(
                            ADAPTER_LABELS, value=ADAPTER_LABELS[0], label="Adapter (dataset)")
                        vqa_btn = gr.Button("Answer", variant="primary")
                    with gr.Column():
                        vqa_out = gr.Textbox(label="Answer", lines=8)
                with gr.Accordion("Advanced", open=False):
                    vqa_tokens = gr.Slider(32, 256, value=MAX_NEW_TOKENS, step=16,
                                           label="Canvas length (max new tokens)")
                gr.Examples(
                    examples=[
                        ["cxr_aorta.png", "Is there evidence of an aortic aneurysm?", ADAPTER_LABELS[0]],
                        ["cxr_consolidation.png", "Is there airspace consolidation on the left side?", ADAPTER_LABELS[0]],
                        ["abdomen_colon.png", "Is the colon more prominent on the patient's right or left side?", ADAPTER_LABELS[0]],
                    ],
                    inputs=[vqa_image, vqa_question, vqa_adapter],
                    outputs=vqa_out,
                    fn=answer_question,
                    cache_examples=True,
                    cache_mode="lazy",
                )
                vqa_btn.click(
                    answer_question,
                    inputs=[vqa_image, vqa_question, vqa_adapter, vqa_tokens],
                    outputs=vqa_out,
                    api_name="answer_question",
                )

            # ---- Infill tab ----
            with gr.Tab("Bidirectional Report Infill"):
                gr.Markdown(
                    "Write a partial report with exactly one **`[BLANK]`** where you want the "
                    "model to fill in. Discrete diffusion fills it using text on **both** sides "
                    "of the hole — toggle *bidirectional* off to compare against left-only "
                    "(autoregressive-style) context."
                )
                with gr.Row():
                    with gr.Column():
                        inf_image = gr.Image(type="pil", label="Medical image", height=340)
                        inf_template = gr.Textbox(
                            label="Report template (use one [BLANK])",
                            lines=4,
                            value="The lungs are clear. [BLANK] No pleural effusion is seen.",
                        )
                        inf_adapter = gr.Dropdown(
                            ADAPTER_LABELS, value=ADAPTER_LABELS[0], label="Adapter (dataset)")
                        inf_bidir = gr.Checkbox(value=True, label="Bidirectional (both-sides context)")
                        inf_btn = gr.Button("Fill the blank", variant="primary")
                    with gr.Column():
                        inf_out = gr.Textbox(label="Filled span", lines=8)
                with gr.Accordion("Advanced", open=False):
                    inf_tokens = gr.Slider(4, 96, value=48, step=4,
                                           label="Max fill span (tokens)")
                gr.Examples(
                    examples=[
                        ["cxr_aorta.png", "The lungs are clear. [BLANK] No pleural effusion is seen.", ADAPTER_LABELS[0], True],
                        ["cxr_consolidation.png", "Findings: [BLANK] There is no pneumothorax.", ADAPTER_LABELS[0], True],
                    ],
                    inputs=[inf_image, inf_template, inf_adapter, inf_bidir],
                    outputs=inf_out,
                    fn=infill_report,
                    cache_examples=True,
                    cache_mode="lazy",
                )
                inf_btn.click(
                    infill_report,
                    inputs=[inf_image, inf_template, inf_adapter, inf_bidir, inf_tokens],
                    outputs=inf_out,
                    api_name="infill_report",
                )

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