"""
Interactive CLIP zero-shot classifier.

Run:
    pip install -r requirements.txt
    python app.py
"""

from __future__ import annotations

import gradio as gr
import torch
from PIL import Image
from transformers import CLIPModel, CLIPProcessor

MODEL_ID = "openai/clip-vit-base-patch32"
DEFAULT_LABELS = "apple\nbanana\norange\nstrawberry\nwatermelon"

device = "cuda" if torch.cuda.is_available() else "cpu"
processor = CLIPProcessor.from_pretrained(MODEL_ID)
model = CLIPModel.from_pretrained(MODEL_ID).to(device)
model.eval()


def parse_labels(text: str) -> list[str]:
    labels = [line.strip() for line in text.replace(",", "\n").splitlines() if line.strip()]
    if not labels:
        raise gr.Error("Add at least one label (one per line or comma-separated).")
    return labels


def classify(image: Image.Image | None, labels_text: str) -> dict[str, float]:
    if image is None:
        raise gr.Error("Upload an image first.")

    labels = parse_labels(labels_text)
    inputs = processor(
        text=labels,
        images=image.convert("RGB"),
        return_tensors="pt",
        padding=True,
    ).to(device)

    with torch.no_grad():
        probs = model(**inputs).logits_per_image.softmax(dim=1)[0]

    return {label: float(score) for label, score in zip(labels, probs, strict=True)}


def build_ui() -> gr.Blocks:
    with gr.Blocks(title="CLIP Zero-Shot Classifier") as demo:
        gr.Markdown(
            f"""
            # CLIP Zero-Shot Image Classifier
            Upload an image and enter candidate labels — no training required.
            Model: [{MODEL_ID}](https://huggingface.co/{MODEL_ID}) · Device: **{device}**
            """
        )

        with gr.Row():
            with gr.Column():
                image_input = gr.Image(type="pil", label="Image")
                labels_input = gr.Textbox(
                    label="Candidate labels",
                    placeholder="One label per line",
                    lines=6,
                    value=DEFAULT_LABELS,
                )
                classify_btn = gr.Button("Classify", variant="primary")

            with gr.Column():
                output = gr.Label(label="Predictions", num_top_classes=10)

        classify_btn.click(
            fn=classify,
            inputs=[image_input, labels_input],
            outputs=output,
        )

        gr.Examples(
            examples=[
                [
                    "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/hub/parrots.png",
                    "animals\nhumans\nlandscape\nvehicles",
                ],
                [
                    "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/hub/parrots.png",
                    DEFAULT_LABELS,
                ],
            ],
            inputs=[image_input, labels_input],
            outputs=output,
            fn=classify,
            cache_examples=False,
        )

    return demo


if __name__ == "__main__":
    build_ui().launch()