# Libraries
import os
import natsort
from tqdm import tqdm
from stitching.stitcher import Stitcher
import matplotlib.pyplot as plt
import cv2
import numpy as np
from types import SimpleNamespace

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

def plot_image(img, figsize_in_inches=(5, 5)):
    plt.figure(figsize=(10, 10))
    plt.imshow(img, cmap='gray', vmin=0, vmax=255)
    plt.show()


def plot_images(imgs, figsize_in_inches=(5, 5)):
    plt.figure(figsize=(10, 10))
    for col, img in enumerate(imgs):
        plt.imshow(img, cmap='gray', vmin=0, vmax=255)
        plt.show()


def get_image_file_names(IMAGE_DIR):
    file_format = ['jpg', 'png', 'JPG', 'PNG']
    file_names = os.listdir(IMAGE_DIR)
    image_file_names = []
    for file_name in tqdm(file_names):
        if file_name.split('.')[-1] in file_format:
            image_file_names.append(IMAGE_DIR + "/" + file_name)
    return natsort.natsorted(image_file_names)

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

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)