"""
EdgeCrafter demo — compact ViTs for edge dense prediction.

Detection (ECDet), instance segmentation (ECSeg) and human pose estimation (ECPose),
all four model scales (S / M / L / X), on ZeroGPU.
"""

import spaces  # noqa: F401  (must be imported before torch)

import time

import gradio as gr
import torch
from PIL import Image

import ecmodels as M

DEVICE = "cuda"

# Load every checkpoint eagerly at module scope; the whole family is small
# (~10M - 50M params each) and ZeroGPU streams the packed weights on first call.
MODELS = {
    (task, size): M.load_model(task, size, device=DEVICE)
    for task in M.TASKS
    for size in M.SIZES
}
print(f"Loaded {len(MODELS)} EdgeCrafter checkpoints", flush=True)

TABLE_HEADERS = ["#", "class", "score", "box (x1, y1, x2, y2)", "mask px / keypoints"]


@spaces.GPU(duration=10)
def predict(
    image: Image.Image,
    task: str = "Object Detection",
    model_size: str = "X",
    threshold: float = 0.4,
):
    """Run EdgeCrafter on an image.

    Args:
        image: input photo (RGB).
        task: "Object Detection", "Instance Segmentation" or "Human Pose Estimation".
        model_size: EdgeCrafter scale, one of "S", "M", "L", "X" (small -> large).
        threshold: confidence threshold for kept predictions.

    Returns:
        The annotated image, a table of predictions and a short run summary.
    """
    if image is None:
        raise gr.Error("Please provide an image first.")

    image = image.convert("RGB")
    model = MODELS[(task, model_size)]

    torch.cuda.synchronize()
    t0 = time.perf_counter()
    results = M.infer(model, task, image, threshold, device=DEVICE)
    torch.cuda.synchronize()
    elapsed = (time.perf_counter() - t0) * 1000

    annotated = M.render(image, task, results)
    rows = M.results_table(task, results) or [["-", "no instance above threshold", "-", "-", "-"]]
    params = sum(p.numel() for p in model.parameters()) / 1e6
    repo = M.TASKS[task]["repo"].format(model_size)
    summary = (
        f"**{repo}** · {params:.1f}M params · {len(results)} instance(s) above "
        f"{threshold:.2f} · {elapsed:.0f} ms on ZeroGPU (640×640 input)"
    )
    return annotated, rows, summary


with gr.Blocks() as demo:
    gr.Markdown(
        """
        # 🪶 EdgeCrafter — compact ViTs for edge dense prediction

        One distilled ViT backbone family, three dense-prediction heads:
        **detection (ECDet)**, **instance segmentation (ECSeg)** and
        **human pose estimation (ECPose)**, at four scales (S → X, ~10M → ~50M params).

        [Paper](https://huggingface.co/papers/2603.18739) ·
        [Code](https://github.com/Intellindust-AI-Lab/EdgeCrafter) ·
        [Weights](https://huggingface.co/Intellindust)
        """
    )

    with gr.Row():
        with gr.Column():
            image = gr.Image(label="Input image", type="pil", height=380)
            task = gr.Radio(
                choices=list(M.TASKS.keys()),
                value="Object Detection",
                label="Task",
            )
            run = gr.Button("Run EdgeCrafter", variant="primary")
            with gr.Accordion("Advanced options", open=False):
                model_size = gr.Radio(
                    choices=M.SIZES,
                    value="X",
                    label="Model scale",
                    info="S / M / L / X — smaller is faster, larger is more accurate",
                )
                threshold = gr.Slider(
                    0.05, 0.95, value=0.4, step=0.05,
                    label="Confidence threshold",
                )
        with gr.Column():
            output = gr.Image(label="Prediction", type="pil", height=380)
            summary = gr.Markdown()
            table = gr.Dataframe(
                headers=TABLE_HEADERS,
                label="Predictions",
                wrap=True,
                row_count=(1, "dynamic"),
                column_count=(5, "fixed"),
            )

    gr.Examples(
        examples=[
            ["examples/donuts.jpg", "Object Detection"],
            ["examples/donuts.jpg", "Instance Segmentation"],
            ["examples/tennis_kids.jpg", "Human Pose Estimation"],
            ["examples/baseball.jpg", "Instance Segmentation"],
            ["examples/baseball.jpg", "Human Pose Estimation"],
        ],
        inputs=[image, task],
        outputs=[output, table, summary],
        fn=predict,
        cache_examples=True,
        cache_mode="lazy",
    )

    gr.Markdown(
        """
        Example photos are COCO images redistributed from the authors' related
        Apache-2.0 repositories
        ([DEIMv2](https://github.com/Intellindust-AI-Lab/DEIMv2),
        [DETRPose](https://github.com/SebastianJanampa/DETRPose)).
        """
    )

    run.click(
        predict,
        inputs=[image, task, model_size, threshold],
        outputs=[output, table, summary],
    )

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