import os
from pathlib import Path
from typing import Dict
from PIL import Image, ImageDraw, ImageFont

import numpy as np
import torch
import torchvision.transforms.functional as F
from ultralytics import YOLO
import gradio as gr
from huggingface_hub import hf_hub_download
from datasets import load_dataset

model_path_1 = hf_hub_download("Axpo/yolov8x-without-insulators", "yolov8x-without-insulators.pt", use_auth_token=os.environ["TOKEN"])
model_path_2 = hf_hub_download("Axpo/yolov8x-with-insulators", "yolov8x-with-insulators.pt", use_auth_token=os.environ["TOKEN"])
model_paths = [model_path_1, model_path_2]

dataset = load_dataset("Axpo/example_pole_images", use_auth_token=os.environ["TOKEN"])

def draw_bboxes_yolo(
        image: torch.Tensor, 
        bboxes: torch.Tensor, 
        labels: list[str],
        color_mapping: Dict,
        scores = [],
        width: int = 2,) -> np.ndarray:
        """
        Draws predicted bboxes with class names and scores on images.
        
        Args:
        - image: image of torch.tensor dtype uint8
        - bboxes: torch.tensor with shape [N, 4] for N bboxes per image
        - labels: list of length N with class names, e.g. ["Gittermast", "Gitterausleger", ...]
        - scores: torch.tensor of length N with scores for all bbox predictions
        - width: width of the bounding boxes
        
        Returns: np.array with the image and bboxes / text / scores included
        """
        img_to_draw = Image.fromarray(image.permute(1, 2, 0).numpy())
        draw = ImageDraw.Draw(img_to_draw)
        bboxes_list = bboxes.tolist()
        if len(scores)==0:
            scores_list = len(labels) * [1]
        else:
            scores_list = scores.tolist()
        margin = width + 1
        # font_txt = ImageFont.truetype("arial.ttf", size=15)
        font_txt = ImageFont.load_default()
        
        # Predictions
        for bbox, label, score in zip(bboxes_list, labels, scores_list):
            color = color_mapping[label]
            draw.rectangle(bbox, width=width, outline=color)
            draw.text((bbox[0] + margin, bbox[1] + margin), f"{label} | {str(round(score, 2))}", font=font_txt)
        
        return np.array(img_to_draw)
        
class BigComponentsCFG:
    id2label: Dict[int, str] = {
      0: "Betonausleger",
      1: "Betonmast",
      2: "Flugwarnkugel",
      3: "Fundament",
      4: "Gitterausleger",
      5: "Gittermast",
      6: "Isolator",
      7: "Mastspitz",
    }
    colors: list[str] = ["#FF0000", "#00FF00", "#0000FF", "#00FFFF", "#FF00FF", "#FFFF00", "#FFA500", "#000000"]


def predict_big_components(image, threshold: float, model_idx: int):
    
    id2label = BigComponentsCFG.id2label
    colors = BigComponentsCFG.colors
    color_mapping = {k: v for k, v in zip(id2label.values(), colors)}

    model_path = model_paths[model_idx]
    model = YOLO(model_path)
    results = model.predict([image])
    
    boxes = results[0].boxes.xyxy
    labels = results[0].boxes.cls
    scores = results[0].boxes.conf
    mask = scores > threshold
    labels_str = [id2label[label.item()] for label in labels[mask]]
    
    # Create image
    img_t = F.pil_to_tensor(image)
    img_arr = draw_bboxes_yolo(img_t, boxes[mask], labels_str, color_mapping, scores[mask])
    img = F.to_pil_image(img_arr)
    return img

with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            slider = gr.Slider(minimum=0.0, maximum=1.0, default=0.0, step=0.05, label="Confidence Threshold")
        with gr.Column():
            model = gr.Dropdown(choices=["Axpo/yolov8x-without-insulators", "Axpo/yolov8x-with-insulators"],
                                value="Axpo/yolov8x-with-insulators",
                                type="index",
                                label="Model")
    with gr.Row().style(full_width=True):
        with gr.Column():
            input_img = gr.Image(type="pil", label="Input").style(height=600, width=600)
        with gr.Column():
            output_img = gr.Image(type="pil", label="Output").style(height=600, width=600)
        
    with gr.Row():
        image_button = gr.Button("Detect Components")
    image_button.click(
        fn=predict_big_components, 
        inputs=[input_img, slider, model], 
        outputs=output_img,
        api_name="big-component-detection"
            )

    gr.Examples(
        dataset["train"]["image"],
        input_img,
        output_img,
        predict_big_components
    )

demo.launch()