import torch
import numpy as np
import cv2
import matplotlib.pyplot as plt
import os

import pandas as pd
import time
from PIL import Image

from tqdm import tqdm
from ultralytics import YOLO
from mean_average_precision.detection_map import DetectionMAP
from mean_average_precision.utils.show_frame import show_frame

# Set the size of the figure
plt.figure(figsize=(30, 30))

def draw_bbox(img, results, thres):
    for obj in results:
        if obj.conf > thres:
            x,y,x2,y2 = obj.xyxy.tolist()[0]
            x,y,x2,y2 = int(x), int(y), int(x2), int(y2)
            color = (255,0,0)
            
            tk = 1
            
            class_dic = {
                        0:"Glass",
                        1:"Metal",
                        2:"Net",
                        3:'PET_Bottle',
                        4:'Plastic Buoy',
                        5:'Plastic Buoy(China)',
                        6:'Plastic ETC',
                        7:'Rope',
                        8:'Styrofoam Box',
                        9:'Styrofoam Buoy',
                        10:'Styrofoam Piece'
                        }
            
            content =  class_dic[int(obj.cls)]+"_"+str(round(float(obj.conf),2))
            
            font =  cv2.FONT_HERSHEY_PLAIN
            img = cv2.rectangle(img, (x,y), (x2,y2), color, tk) # bbox
            img = cv2.putText(img, content, (x, y-2), font, tk, color, tk, cv2.LINE_AA) # label
    return img

model = YOLO("./weights/detect_debris/weights/best.pt")

results = []
TXT_DIR = './dataset/debris/labels/test'
IMG_DIR = './dataset/debris/images/test'
od_threshold = 0.5
n_class = 11

img_list = [file for file in os.listdir(IMG_DIR) if file.endswith('.png') or file.endswith('.jpg')]


if __name__ == '__main__':
    
    for img in tqdm(img_list):
        file_name = img.rsplit('.')[0]
        img_path = os.path.join(IMG_DIR, img)
        txt_path = os.path.join(TXT_DIR, file_name + '.txt')

        image = cv2.imread(img_path)
        od_results = model(image, verbose=False)
        bbox_results = od_results[0].boxes
        
        ## gt image ###
        img_w = image.shape[1]
        img_h = image.shape[0]
        
        with open(txt_path , 'r') as file:
            for row in file:
                gt_cls = row.split()[0]
                x_min = float(row.split()[1]) * img_w
                y_min = float(row.split()[2]) * img_h
                x_max = float(row.split()[3]) * img_w
                y_max = float(row.split()[4]) * img_h

                color = (255,0,0)

                tk = 1

                class_dic = {
                            0:"Glass",
                            1:"Metal",
                            2:"Net",
                            3:'PET_Bottle',
                            4:'Plastic Buoy',
                            5:'Plastic Buoy(China)',
                            6:'Plastic ETC',
                            7:'Rope',
                            8:'Styrofoam Box',
                            9:'Styrofoam Buoy',
                            10:'Styrofoam Piece'
                            }

                content =  class_dic[int(gt_cls)]

                font =  cv2.FONT_HERSHEY_PLAIN
                image = cv2.rectangle(image, (int(x_min),int(y_min)), (int(x_max),int(y_max)), color, tk) # bbox
                image = cv2.putText(image, content, (int(x_min), int(y_min)-2), font, tk, color, tk, cv2.LINE_AA)
                
            gt_image_folder = './dataset/debris/gt_images'

            gt_image_path = os.path.join(gt_image_folder, file_name + '.jpg')

            cv2.imwrite(gt_image_path, image)
        ## gt image ###
        
        
        result_img = draw_bbox(image.copy(), bbox_results, od_threshold)

        pred_image_folder = './dataset/debris/pred_images'

        pred_image_path = os.path.join(pred_image_folder, file_name + '.jpg')

        cv2.imwrite(pred_image_path, result_img)

        gt_bb = []
        gt_cls = []
        
        with open(txt_path , 'r') as file:
            for row in file:
                gt_cls.append(row.split()[0])
                x_min = row.split()[1]
                y_min = row.split()[2]
                x_max = row.split()[3]
                y_max = row.split()[4]
                gt_bb.append([x_min, y_min, x_max, y_max])
        
        gt_bb = np.array(gt_bb, dtype=np.float64)
        gt_cls = np.array(gt_cls, dtype=np.int64)
        
        img_w = image.shape[1]
        img_h = image.shape[0]
        
        
        pred_bb = np.array(bbox_results.xyxyn.cpu(), dtype=np.float64)
        pred_cls = np.array(bbox_results.cls.cpu(), dtype=np.int64)
        pred_conf = np.array(np.round(bbox_results.conf.cpu(),2), dtype=np.float64)  
        
        results.append([pred_bb, pred_cls, pred_conf, gt_bb, gt_cls])
        
    mAP = DetectionMAP(n_class)
    for i, frame in enumerate(results):
        # print("Evaluate image {}".format(i))
        mAP.evaluate(*frame)
        mAP.process()

    mAP.plot(class_names=[
                        "Glass",
                        "Metal",
                        "Net",
                        'PET_Bottle',
                        'Plastic Buoy',
                        'Plastic Buoy(China)',
                        'Plastic ETC',
                        'Rope',
                        'Styrofoam Box',
                        'Styrofoam Buoy',
                        'Styrofoam Piece']
            )
    plt.savefig("pr_curve_debris.png")
    print("Complete the mAP calculateion")
    print("Results are save into pr_curve_debris.png")