# from dataclasses import dataclass, replace
# from functools import reduce
from io import BytesIO
import math
import os
from pprint import pprint
import tempfile

from PIL import Image, ImageDraw, ImageFont
import numpy as np
import cv2

# import seaborn as sns
# import matplotlib.pyplot as plt
# %matplotlib inline

import torch
from torch.utils.data import Dataset
import torchvision
from torchvision import transforms

import roboflow
from roboflow import Roboflow
import supervision as sv
import albumentations as A

import gradio as gr
import requests

# from torchmetrics.detection.mean_ap import MeanAveragePrecision
# from torchmetrics.detection.iou import IntersectionOverUnion
# import evaluate
#from datasets import load_metric

from transformers import pipeline
from transformers import (
    AutoProcessor,
    AutoImageProcessor,
    AutoModel,
    AutoModelForObjectDetection,
    RTDetrForObjectDetection,
    RTDetrImageProcessor,
    TrainingArguments,
    Trainer
)
from huggingface_hub import hf_hub_download

from safetensors.torch import load_file

#@title Utilities
PALETTE = {0: {"color": (255, 0, 0),
               "name": "Ambulance"},
           1: {"color": (0, 191, 0),
               "name": "Firetruck"},
           2: {"color": (0, 0, 255),
               "name": "Police"},
           3: {"color": (255, 0, 255),
               "name": "Non-EV"}}
label2id = {val["name"]: id for (id, val) in PALETTE.items()}
id2label = {id: name for (name, id) in label2id.items()}

print(label2id)
print(id2label)

def unnormalize_bbox(img_h, img_w, bbox):
  x_min = bbox[0] - bbox[2]/2
  y_min = bbox[1] - bbox[3]/2
  x_max = bbox[0] + bbox[2]/2 # - x_min
  y_max = bbox[1] + bbox[3]/2 # - y_min

  x_min *= img_w
  y_min *= img_h
  x_max *= img_w
  y_max *= img_h
  x_min, y_min, x_max, y_max = list(map(int, [x_min, y_min, x_max, y_max]))

  return (x_min, y_min, x_max, y_max)

def paint_bbox(
    image,
    annotations,
    normalize_labels=True,
    normalize_bbox=True,
  ):
  bboxes = annotations["boxes"].tolist()
  class_id = annotations["labels"].tolist()
  confidences = annotations["scores"].tolist()

  painted_img = image.copy() # Wutdehell
  for (bbox, label, confidence) in zip(bboxes, class_id, confidences):
    label = (label - 1) if normalize_labels else label
    if normalize_bbox:
      img_h, img_w = image.shape[0], image.shape[1] # H, W, C
      x_min, y_min, x_max, y_max = unnormalize_bbox(img_h, img_w, bbox)
      print([x_min, y_min, x_max, y_max])

      """
      x_min = #int(bbox[0] - bbox[2]/2) # Left
      y_min = #int(bbox[1] - bbox[3]/2) # Top
      x_max = #int(bbox[0] + bbox[2]/2)
      y_max = #int(bbox[1] + bbox[3]/2)
      """
    else:
      x_min, y_min, x_max, y_max = list(map(int, bbox))

    box_color = PALETTE[label]["color"]
    label_name = PALETTE[label]["name"]

    if confidence != -1:
      label_name = f"{label_name} ({confidence:.2f})"

    cv2.rectangle(painted_img,
                  (x_min, y_min),
                  (x_max, y_max),
                  color=box_color,
                  thickness=2)
    cv2.rectangle(painted_img,
                  (x_min, y_min),
                  (x_min + 5 + len(label_name)*10, y_min + 17),
                  color=box_color,
                  thickness=-1)
    cv2.putText(painted_img,
                label_name,
                (x_min + 2, y_min + 12),
                fontFace=cv2.FONT_HERSHEY_SIMPLEX,
                fontScale=0.5,
                color=(255, 255, 255),
                thickness=1)
  return painted_img

# Function to calculate Intersection over Union (IoU)
def calculate_iou(truth_bbx, pred_bbx):
    # Coordinates of the boxes: [xmin, ymin, xmax, ymax]
    x1, y1, x2, y2 = truth_bbx
    x1_p, y1_p, x2_p, y2_p = pred_bbx

    # Calculate intersection
    ixmin = max(x1, x1_p)
    iymin = max(y1, y1_p)
    ixmax = min(x2, x2_p)
    iymax = min(y2, y2_p)

    iw = max(0, ixmax - ixmin)
    ih = max(0, iymax - iymin)

    intersection = iw * ih
    area1 = (x2 - x1) * (y2 - y1)
    area2 = (x2_p - x1_p) * (y2_p - y1_p)
    union = area1 + area2 - intersection
    iou = intersection / union if union != 0 else 0
    return iou

# Example: emotion_classifier = pipeline("image-classification", model="itsindrabudhik/emotion_classification")
# (Load only once)
DETECTOR = pipeline("object-detection", model="itsindrabudhik/finalProjectCV2425") #later on, change this with out trained modell yesssss (the trained model should be uploaded to hugging face)
tensor_file = hf_hub_download(repo_id="itsindrabudhik/finalProjectCV2425",
                               filename="model.safetensors")

# Assign classification head weights since that pipeline seems to not handling it
# weights = load_file(tensor_file)
# DETECTOR.model.class_labels_classifier.weight.data = weights["class_labels_classifier.weight"]
# DETECTOR.model.class_labels_classifier.bias.data = weights["class_labels_classifier.bias"]
# del weights

def detect_ev_nev(image, confidence_threshold=0.5, iou_threshold=0.5):
    # Run the detector pipeline on the image
    results = DETECTOR(image)

    # Open the image
    if isinstance(image, str):  # If the image is a URL or file path
        if image.startswith("http"):
            response = requests.get(image)
            img = Image.open(BytesIO(response.content))
        else:
            img = Image.open(image)
    else:
        img = image

    # Draw bounding boxes and labels on the image
    font_path = os.path.join(cv2.__path__[0],'qt','fonts','DejaVuSans.ttf')
    font = ImageFont.truetype(font_path, size=32)
    draw = ImageDraw.Draw(img)

    details = []  # Collect details for text output
    for result in results:
        score = result['score']
        label = result['label']
        box = result['box']

        # Apply confidence threshold
        if score < confidence_threshold:
            continue

        # Filter out low IoU detections
        keep = True
        for previous_result in results:
            if previous_result != result:
                prev_box = previous_result['box']
                iou = calculate_iou([box['xmin'], box['ymin'], box['xmax'], box['ymax']],
                                    [prev_box['xmin'], prev_box['ymin'], prev_box['xmax'], prev_box['ymax']])
                if iou > iou_threshold:
                    keep = False
                    break

        label_color = PALETTE[label2id[label]]["color"]
        if keep:
            # Draw the bounding box and label
            xmin, ymin, xmax, ymax = box['xmin'], box['ymin'], box['xmax'], box['ymax']
            draw.rectangle([xmin, ymin, xmax, ymax], outline="red", width=3)

            # Use a larger font size for text
            text = f"{label} ({score:.2f})"

            # Calculate text bounding box
            text_bbox = draw.textbbox((xmin, ymin - 10), text, font=font)  # This gives (xmin, ymin, xmax, ymax)
            text_width = text_bbox[2] - text_bbox[0]  # width of the text box
            text_height = text_bbox[3] - text_bbox[1]  # height of the text box

            # Draw the text on the image (position adjusted)
            draw.text((xmin, ymin - text_height - 5), text, fill="red", font=font)

            # Add details to the list
            details.append({
                "Label": label,
                "Confidence": f"{score:.2f}",
                "Bounding Box": f"({xmin}, {ymin}, {xmax}, {ymax})"
            })
    details_text = "\n".join([f"Label: {d['Label']}, Confidence: {d['Confidence']}, Box: {d['Bounding Box']}" for d in details])
    return img, details_text

def detect_video(video, confidence_threshold=0.5, iou_threshold=0.5):
    video_capture = cv2.VideoCapture(video)
    fps = video_capture.get(cv2.CAP_PROP_FPS)
    frame_width = int(video_capture.get(cv2.CAP_PROP_FRAME_WIDTH))
    frame_height = int(video_capture.get(cv2.CAP_PROP_FRAME_HEIGHT))

    temp_output = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    out = cv2.VideoWriter(temp_output.name, fourcc, fps, (frame_width, frame_height))

    details = []
    total_frames = 0
    detected_frames = 0

    while True:
        ret, frame = video_capture.read()
        if not ret:
            break

        total_frames += 1
        image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
        annotated_image, frame_details = detect_ev_nev(image, confidence_threshold, iou_threshold)

        # Count frames with detections
        if frame_details.strip():  # Non-empty details indicate detections
            detected_frames += 1

        details.append(frame_details)
        annotated_frame = cv2.cvtColor(np.array(annotated_image), cv2.COLOR_RGB2BGR)
        out.write(annotated_frame)

    video_capture.release()
    out.release()

    details_text = "\n".join(details)
    summary = f"Total Frames: {total_frames}, Frames with Detections: {detected_frames}\n" + details_text
    return temp_output.name, summary

def detect(file, confidence_threshold=0.5, iou_threshold=0.5):
    # Determine if input is an image or video
    file_ext = file.name.split(".")[-1].lower()
    if file_ext in ["png", "jpg", "jpeg"]:
        # Image processing
        annotated_image, details = detect_ev_nev(file, confidence_threshold, iou_threshold)
        return annotated_image, None, details
    elif file_ext in ["mp4", "avi", "mov"]:
        # Video processing
        processed_video, details = detect_video(file, confidence_threshold, iou_threshold)
        return None, processed_video, details
    else:
        raise ValueError("Unsupported file format. Please upload an image or video.")


interface = gr.Interface(
    fn=detect,
    inputs=[
        gr.File(label="Upload Image or Video", file_types=[".png", ".jpg", ".jpeg", ".mp4", ".avi", ".mov"]),
        gr.Slider(0, 1, value=0.5, label="Confidence Threshold"),
        gr.Slider(0, 1, value=0.5, label="IoU Threshold"),
    ],
    outputs=[

        gr.Image(label="Processed Image"),

        gr.Video(label="Generated Video"),
        gr.Text(label="Detection Details")

    ],
    title="RT-DETR Object Detection for Images and Videos",
    description="Upload an image or video to detect objects using the fine-tuned RT-DETR model. Results include the annotated image/video and detection details."
)
interface.launch(debug=True)