from typing import List

import gradio as gr
import supervision as sv
import torch
from PIL import Image
from ultralytics import YOLO

MARKDOWN = """
# Orang Outan Detection
"""
EXAMPLES = []

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

YOLO_MODEL = YOLO("train_7best.pt")

BOX_ANNOTATOR = sv.BoxAnnotator()


def annotate(
    image_bgr_numpy: Image.Image,
    detections: sv.Detections,
    annotator: sv.BoxAnnotator,
    labels: str,
) -> Image.Image:
    annotated_bgr_image = annotator.annotate(
        scene=image_bgr_numpy, detections=detections, labels=labels
    )
    return Image.fromarray(annotated_bgr_image[:, :, ::-1])


def inference(image_rgb_pil: Image.Image, confidence: float) -> List[Image.Image]:
    output = YOLO_MODEL(image_rgb_pil, verbose=False)[0]
    detections = sv.Detections.from_ultralytics(output)

    detections = detections[detections.confidence >= confidence]

    labels = [
        f"{output.names[class_id]} {confidence:0.2f}"
        for _, _, confidence, class_id, _ in detections
    ]

    return annotate(
        image_bgr_numpy=output.orig_img.copy(),
        detections=detections,
        annotator=BOX_ANNOTATOR,
        labels=labels,
    )


def run_demo():
    custom_theme = gr.themes.Soft(primary_hue="blue").set(
        button_secondary_background_fill="*neutral_100",
        button_secondary_background_fill_hover="*neutral_200",
    )

    with gr.Blocks(theme=custom_theme, css="style.css") as demo:
        gr.Markdown(MARKDOWN)

        with gr.Row():
            with gr.Column():
                input_image = gr.Image(image_mode="RGB", type="pil", height=500)
                confidence_slider = gr.Slider(
                    label="Confidence", minimum=0.1, maximum=1.0, step=0.05, value=0.6
                )
                submit_button = gr.Button("Submit")
            output_image = gr.Image(label="Results", type="pil")

        # with gr.Row():
        #     gr.Examples(
        #         examples=EXAMPLES,
        #         fn=inference,
        #         inputs=[input_image, prompt_text, confidence_slider],
        #         outputs=[gallery],
        #         cache_examples=True,
        #         run_on_click=True
        #     )

        submit_button.click(
            inference,
            inputs=[input_image, confidence_slider],
            outputs=output_image,
            queue=True,
        )
    demo.queue(max_size=20, api_open=False).launch()


if __name__ == "__main__":
    run_demo()