import torch
import numpy as np
import cv2
import matplotlib.pyplot as plt
from IPython.display import clear_output
import os
import json

import pandas as pd
import time
from PIL import Image
from tqdm import tqdm

from mmdet.apis import async_inference_detector, inference_detector
from mmdet.apis.inference import init_detector

from mean_average_precision.detection_map import DetectionMAP
from mean_average_precision.utils.show_frame import show_frame

# Test
od_config_file = "/DATA2/pjh/inference/TTA_test/vehicle/yolov8x_vehicle.py"
od_checkpoint = '/DATA2/pjh/inference/TTA_test/vehicle/epoch_54.pth'
od_model = init_detector(od_config_file, od_checkpoint, device='cuda:1')


results = []
JSON_DIR = '/DATA2/pjh/inference/TTA_test/vehicle/annotations'
IMG_DIR = '/DATA2/pjh/inference/TTA_test/vehicle/images'
json_file_name = 'vehicle_test_data.json'

json_path = os.path.join(JSON_DIR, json_file_name)
json_data = json.load(open(json_path))

n_class = 8

def save_lists_to_file(lists, filename):
    with open(filename, 'w') as file:
        for lst in lists:
            line = ' '.join(str(item) for item in lst)  # Convert each item to a string
            file.write(line + '\n')  # Write the line to the file

def draw_bbox(img, bbox, cls, conf):
    for i in range(len(bbox)):
        x = int(bbox[i][0])
        y = int(bbox[i][1])
        x2 = int(bbox[i][2])
        y2 = int(bbox[i][3])
        
        color = (255,0,0)
        
        tk = 1
        # if obj['risk1'] == 'danger':
        #     tk = 5
        
        class_dic = {1 :"Excavator",
                     2 :"Dodger",
                     3 :"Forklift",
                    4 :"Dump Truck",
                    5 :"Mixer Truck",
                    6 :"Cargo Truck",
                    7 :"Scissor Lift",
                    8 :"Crane"}
        
        font =  cv2.FONT_HERSHEY_PLAIN
        img = cv2.rectangle(img, (x,y), (x2,y2), color, tk) # bbox
        
        if conf is not None:
            content =  class_dic[cls[i]]+"_"+str(round(conf[i],2))
            img = cv2.putText(img, content, (x, y-2), font, tk, color, tk, cv2.LINE_AA) # label
        
    return img

def mmdet3x_convert_to_bboxes_mmdet(results, threshold, prefix = 'class'):
    boxes_list = []
    scores_list = []
    labels_list = []
    confidence_score = results.pred_instances.scores.tolist()
    for i, conf in enumerate(confidence_score):
        if conf >= threshold:
            # print(conf)
            extracted_box = results.pred_instances.bboxes[i].cpu().tolist()
            extracted_label = results.pred_instances.labels[i].cpu().tolist()
            boxes_list.append([float(extracted_box[0]),
                               float(extracted_box[1]),
                                float(extracted_box[2]),
                                float(extracted_box[3])])
            scores_list.append(conf)
            labels_list.append(extracted_label+1)
    return np.array(boxes_list), np.array(labels_list), np.array(scores_list)


if __name__ == '__main__':
    for image in tqdm(json_data['images']):

        img_name = image['file_name']
        img_id = image['id']
        img_path = os.path.join(IMG_DIR, img_name)

        od_threshold = 0.7

        od_results = inference_detector(od_model, img_path)
        pred_bb, pred_cls, pred_conf = mmdet3x_convert_to_bboxes_mmdet(od_results, od_threshold)  

        img = cv2.imread(img_path)
        image_size = img.shape

        img_w = image_size[1]
        img_h = image_size[0]

        gt_bb = []
        gt_cls = []

        for annotation in json_data['annotations']:
            if annotation['image_id'] == img_id:
                gt_bb.append([annotation['bbox'][0], annotation['bbox'][1], 
                             annotation['bbox'][0]+annotation['bbox'][2],
                            annotation['bbox'][1]+annotation['bbox'][3]])
                gt_cls.append(annotation['category_id'])

        gt_bb = np.array(gt_bb)
        gt_cls = np.array(gt_cls)

        if gt_bb.ndim == 1:
            gt_bb = np.zeros(pred_bb.shape)

        pred_img = draw_bbox(img.copy(), pred_bb, pred_cls, pred_conf)
        gt_img = draw_bbox(img.copy(), gt_bb, gt_cls, None)
        
        pred_image_folder = '/DATA2/pjh/inference/TTA_test/vehicle/pred_images'
        folder_name = img_name.split('/')[0]
        image_path = os.path.join(pred_image_folder, folder_name)
        
        if not os.path.exists(image_path):
            os.makedirs(image_path)
        
        pred_image_path = os.path.join(pred_image_folder, img_name)
        cv2.imwrite(pred_image_path, pred_img)

        for idx,row in enumerate(pred_bb):
            pred_bb[idx][0] /= img_w
            pred_bb[idx][1] /= img_h
            pred_bb[idx][2] /= img_w
            pred_bb[idx][3] /= img_h

        for idx,row in enumerate(gt_bb):
            gt_bb[idx][0] /= img_w
            gt_bb[idx][1] /= img_h
            gt_bb[idx][2] /= img_w
            gt_bb[idx][3] /= img_h 

        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=["Excavator", "Dodger", "Forklift", "Dump Truck", "Mixer Truck", "Cargo Truck", "Scissor Lift", "Crane"])
    plt.show()
    plt.savefig("pr_curve_vehicle.png")
    
    save_lists_to_file(results, 'vehicle-output.txt')