"""TRUE-Colon: Real-Time Polyp Detection demo.

Loads the RT-DETR checkpoint from the TRUE-Colon paper (MICCAI 2026 EndoLINA Workshop)
and runs inference on colonoscopy frames, drawing bounding boxes around detected polyps.

Research demo only — NOT for clinical use.
"""

import spaces  # MUST be first
import os
import cv2
import numpy as np
import torch
from huggingface_hub import hf_hub_download
from ultralytics import YOLO
import gradio as gr

MODEL_ID = "sdoerrich97/true_colon_rtdetr_realcolon_s0"
CLASS_NAME = "lesion"
# Green-ish box color (BGR for cv2)
BOX_COLOR = (0, 255, 0)

# Download and load model at module scope
_weights_path = hf_hub_download(MODEL_ID, "model.pt")
model = YOLO(_weights_path)


def draw_detections(image: np.ndarray, results, conf_threshold: float) -> np.ndarray:
    """Draw bounding boxes on the image from Ultralytics results.

    Args:
        image: Input image as numpy array (RGB).
        results: Ultralytics prediction results.
        conf_threshold: Confidence threshold for display.

    Returns:
        Annotated image as numpy array (RGB).
    """
    annotated = image.copy()
    h, w = annotated.shape[:2]

    for result in results:
        boxes = result.boxes
        for box in boxes:
            conf = float(box.conf[0])
            if conf < conf_threshold:
                continue
            cls = int(box.cls[0])
            x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
            x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)

            # Draw box
            cv2.rectangle(annotated, (x1, y1), (x2, y2), BOX_COLOR, 3)

            # Draw label background
            label = f"{CLASS_NAME} {conf:.2f}"
            (label_w, label_h), _ = cv2.getTextSize(
                label, cv2.FONT_HERSHEY_SIMPLEX, 0.7, 2
            )
            cv2.rectangle(
                annotated,
                (x1, y1 - label_h - 10),
                (x1 + label_w, y1),
                BOX_COLOR,
                -1,
            )
            cv2.putText(
                annotated,
                label,
                (x1, y1 - 5),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.7,
                (0, 0, 0),
                2,
                cv2.LINE_AA,
            )

    return annotated


@spaces.GPU(duration=30)
def detect(
    image: np.ndarray,
    conf_threshold: float = 0.30,
    iou_threshold: float = 0.50,
) -> np.ndarray:
    """Detect polyps in a colonoscopy frame.

    Runs the TRUE-Colon RT-DETR detector on the input image and returns an
    annotated copy with bounding boxes around detected lesions.

    Args:
        image: Colonoscopy frame as an image.
        conf_threshold: Minimum detection confidence to display.
        iou_threshold: NMS IoU threshold.

    Returns:
        Annotated image with detection boxes drawn.
    """
    if image is None:
        return None

    # Ultralytics expects RGB; Gradio passes RGB
    results = model.predict(
        source=image,
        conf=conf_threshold,
        iou=iou_threshold,
        imgsz=640,
        verbose=False,
    )

    annotated = draw_detections(image, results, conf_threshold)
    return annotated


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

EXAMPLES = [
    ["examples/real_colon_004-001_frame13.jpg"],
    ["examples/real_colon_004-001_frame25.jpg"],
    ["examples/real_colon_004-001_frame37.jpg"],
    ["examples/real_colon_004-001_frame49.jpg"],
    ["examples/cvc_2.png"],
    ["examples/cvc_100.png"],
]

with gr.Blocks(css=CSS) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
            # 🩺 TRUE-Colon: Real-Time Polyp Detection
            **RT-DETR** trained on REAL-Colon (60 full colonoscopy procedures) for polyp detection.
            From the paper *TRUE-Colon: Exposing a Consistent Transfer Asymmetry in Real-Time Polyp Detection* (MICCAI 2026 EndoLINA).

            [📄 Paper](https://arxiv.org/abs/2608.13711) | [🐍 Code](https://github.com/sdoerrich97/true-colon) | [⚖️ Weights](https://huggingface.co/sdoerrich97/true_colon_rtdetr_realcolon_s0)

            > ⚠️ **Research demo only — NOT for clinical use.** The model has not been validated prospectively or cleared by any regulator.
            """
        )

        with gr.Row():
            with gr.Column(scale=1):
                input_image = gr.Image(
                    label="Colonoscopy Frame",
                    type="numpy",
                    sources=["upload", "clipboard"],
                )
                conf_slider = gr.Slider(
                    label="Confidence Threshold",
                    minimum=0.05,
                    maximum=0.95,
                    step=0.05,
                    value=0.30,
                )
                iou_slider = gr.Slider(
                    label="NMS IoU Threshold",
                    minimum=0.10,
                    maximum=0.95,
                    step=0.05,
                    value=0.50,
                )
                run_btn = gr.Button("Detect Polyps", variant="primary")

            with gr.Column(scale=1):
                output_image = gr.Image(
                    label="Detection Result",
                    type="numpy",
                )

        run_btn.click(
            fn=detect,
            inputs=[input_image, conf_slider, iou_slider],
            outputs=[output_image],
            api_name="detect",
        )

        gr.Examples(
            examples=EXAMPLES,
            inputs=[input_image],
            outputs=[output_image],
            fn=detect,
            cache_examples=True,
            cache_mode="lazy",
        )

        gr.Markdown(
            """
            ### Example image sources
            Examples from the [REAL-Colon](https://doi.org/10.1038/s41597-024-03359-0) dataset (Biffi et al., *Scientific Data* 2024, CC BY 4.0) and [CVC-ClinicDB](https://polyp.grand-challenge.org/CVCClinicDB/) (CC BY 4.0).

            ### Citation
            ```bibtex
            @article{doerrich2026truecolon,
              title={TRUE-Colon: Exposing a Consistent Transfer Asymmetry in Real-Time Polyp Detection},
              author={Sebastian Doerrich and Andreas Franz Schwab and Francesco {Di Salvo} and Shyam Nandan Rai and Hanh Huyen My Nguyen and Christian Ledig},
              year={2026}, eprint={2608.13711}, archivePrefix={arXiv}, primaryClass={eess.IV}
            }
            ```
            """
        )

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