import cv2
import torch
import torchvision.transforms as transforms
import math
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt

import torchvision
from torchvision.transforms import ToTensor
from PIL import Image, ImageDraw
from mmdet.apis import async_inference_detector, inference_detector
from mmdet.apis.inference import init_detector

config_file = './mmyolo/configs/1_crack/final_model.py'
checkpoint_file = 'best_pretrained.pth'
testset_path = '/'

# build the model from a config file and a checkpoint file
model = init_detector(config_file, checkpoint_file, device='cuda:0')

def mmdet3x_convert_to_bboxes_mmdet(results, threshold):
    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([int(extracted_box[0]),
                               int(extracted_box[1]),
                                int(extracted_box[2]),
                                int(extracted_box[3])])
            scores_list.append(conf)
            labels_list.append('class_{}'.format(extracted_label+1))
    return boxes_list, labels_list, scores_list

json_path = "./server_aihub_train_5k.json"
import json
json_data = json.load(open(json_path))

def merge_boxes_with_labels(boxes, labels):
    if not boxes or not labels or len(boxes) != len(labels):
        return [], []
    labeled_boxes = list(zip(boxes, labels))
    labeled_boxes.sort(key=lambda x: x[0][0])
    merged = []
    for box, label in labeled_boxes:
        overlap = False
        for m, l in merged:
            if label == l and not (box[2] < m[0] or box[0] > m[2] or box[3] < m[1] or box[1] > m[3]):
                m[0] = min(m[0], box[0])
                m[1] = min(m[1], box[1])
                m[2] = max(m[2], box[2])
                m[3] = max(m[3], box[3])
                overlap = True
                break
        if not overlap:
            merged.append([box.copy(), label])
    merged_boxes, merged_labels = zip(*merged) if merged else ([], [])
    return list(merged_boxes), list(merged_labels)

# VISUALIZE
import random
for _ in range(10):
    
    image_info = random.choice(json_data['images'])
    image_file_path = testset_path + image_info['file_name']
    image = cv2.imread(image_file_path)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    # read label
    plt.figure(figsize=(25,15))
    image_id = image_info['id']
    for annotation in json_data['annotations']:
        if annotation['image_id'] == image_id:
            # print("Found annotation")
            bbox = annotation['bbox']
            bbox = [int(x) for x in bbox]
            bbox = np.array(bbox)
            bbox = bbox.astype(np.int32)
            category_id = annotation['category_id']
            # print(category_id)
            # print(bbox)
            image = cv2.rectangle(image, (bbox[0], bbox[1]), (bbox[0]+bbox[2], bbox[1]+bbox[3]), (0,0,255), 2)
            image = cv2.putText(image, str(category_id), (bbox[0], bbox[1]), cv2.FONT_HERSHEY_SIMPLEX, 3, (0, 0, 255),4)
    plt.subplot(1,3,1)
    plt.title("Ground truth")
    plt.imshow(image)
    image = cv2.imread(image_file_path)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    results = inference_detector(model, image)
    raw_public_bboxes, raw_public_labels, raw_public_scores = mmdet3x_convert_to_bboxes_mmdet(results,0.1)
    for i,box in enumerate(raw_public_bboxes):
        image = cv2.rectangle(image, (box[0], box[1]), (box[2], box[3]), (255,0,0), 2)
        image = cv2.putText(image, str(raw_public_labels[i]).split("_")[-1], (box[0], box[1]), cv2.FONT_HERSHEY_SIMPLEX, 3, (255, 0, 0),4)
        # add confidence score
        image = cv2.putText(image, str(np.round(raw_public_scores[i],2)), (box[0], box[1]+50), cv2.FONT_HERSHEY_SIMPLEX, 3, (255, 0, 0),4)
    plt.subplot(1,3,2)
    plt.title("Raw predicted boxes")
    plt.imshow(image)
    merged_boxes, merged_labels = merge_boxes_with_labels(raw_public_bboxes, raw_public_labels)#
    image = cv2.imread(image_file_path)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    for i,box in enumerate(merged_boxes):
        image = cv2.rectangle(image, (box[0], box[1]), (box[2], box[3]), (255,0,0), 2)
        image = cv2.putText(image, str(merged_labels[i]).split("_")[-1], (box[0], box[1]), cv2.FONT_HERSHEY_SIMPLEX, 3, (255, 0, 0),4)
    plt.subplot(1,3,3)
    plt.imshow(image)
    plt.title("Merge boxes")
    plt.tight_layout()
    plt.show()
    # break