import os
import numpy as np
import cv2
from PIL import Image
from paddleseg import utils
import xml.dom.minidom


def mkdir(path):
    sub_dir = os.path.dirname(path)
    if not os.path.exists(sub_dir):
        os.makedirs(sub_dir)


def get_image_list(image_path):
    """Get image list"""
    valid_suffix = [
        '.JPEG', '.jpeg', '.JPG', '.jpg', '.BMP', '.bmp', '.PNG', '.png'
    ]
    image_list = []
    image_dir = None
    if os.path.isfile(image_path):
        if os.path.splitext(image_path)[-1] in valid_suffix:
            image_list.append(image_path)
    elif os.path.isdir(image_path):
        image_dir = image_path
        for root, dirs, files in os.walk(image_path):
            for f in files:
                if '.ipynb_checkpoints' in root:
                    continue
                if os.path.splitext(f)[-1] in valid_suffix:
                    image_list.append(os.path.join(root.split('/')[-1], f))
    else:
        raise FileNotFoundError(
            '`--image_path` is not found. it should be an image file or a directory including images'
        )

    if len(image_list) == 0:
        raise RuntimeError('There are not image file in `--image_path`')

    return image_list, image_dir


def refine_pred():
    image_list, image_dir = get_image_list(
        'detection_out/pseudo_color_prediction')
    for ii in image_list:
        name_pred = 'detection_out/pseudo_color_prediction/' + ii
        name_label = 'data/IDD_Detection/Annotations/all/' + ii[:-3] + 'xml'
        pred = np.array(Image.open(name_pred)).astype(np.float32)
        if not os.path.exists(name_label):
            pred_mask = utils.visualize.get_pseudo_color_map(pred)
            pred_saved_path = 'detect_out/pred_refine/' + ii
            mkdir(pred_saved_path)
            pred_mask.save(pred_saved_path)
            continue

        dom = xml.dom.minidom.parse(name_label)
        root = dom.documentElement
        objects = root.getElementsByTagName("object")
        for item in objects:
            name = item.getElementsByTagName("name")[0]
            if name.firstChild.data == 'traffic sign' or name.firstChild.data == 'traffic light':
                print(ii)
                xmin = int(
                    item.getElementsByTagName('bndbox')[0].getElementsByTagName(
                        'xmin')[0].firstChild.data)
                ymin = int(
                    item.getElementsByTagName('bndbox')[0].getElementsByTagName(
                        'ymin')[0].firstChild.data)
                xmax = int(
                    item.getElementsByTagName('bndbox')[0].getElementsByTagName(
                        'xmax')[0].firstChild.data)
                ymax = int(
                    item.getElementsByTagName('bndbox')[0].getElementsByTagName(
                        'ymax')[0].firstChild.data)
                if name.firstChild.data == 'traffic sign':
                    pred[ymin:ymax, xmin:xmax] = 18
                elif name.firstChild.data == 'traffic light':
                    pred[ymin:ymax, xmin:xmax] = 19

        pred_mask = utils.visualize.get_pseudo_color_map(pred)
        pred_saved_path = 'detect_out/pred_refine/' + ii
        mkdir(pred_saved_path)
        pred_mask.save(pred_saved_path)


def test():
    path = '/Users/liliulei/Downloads/IDD_Detection/JPEGImages/frontNear/'
    image_list, image_dir = get_image_list(path)

    for ii in image_list:
        name_xml = '/Users/liliulei/Downloads/IDD_Detection/Annotations/frontNear/' + ii[:
                                                                                         -3] + 'xml'
        image = cv2.imread(path + ii)
        # print(image.shape)
        (h, w) = image.shape[0:2]

        pred = np.zeros_like(image)

        dom = xml.dom.minidom.parse(name_xml)
        root = dom.documentElement
        objects = root.getElementsByTagName("object")
        for item in objects:
            name = item.getElementsByTagName("name")[0]
            print(name.firstChild.data)
            if name.firstChild.data == 'traffic sign' or name.firstChild.data == 'traffic light':
                xmin = int(
                    item.getElementsByTagName('bndbox')[0].getElementsByTagName(
                        'xmin')[0].firstChild.data)
                ymin = int(
                    item.getElementsByTagName('bndbox')[0].getElementsByTagName(
                        'ymin')[0].firstChild.data)
                xmax = int(
                    item.getElementsByTagName('bndbox')[0].getElementsByTagName(
                        'xmax')[0].firstChild.data)
                ymax = int(
                    item.getElementsByTagName('bndbox')[0].getElementsByTagName(
                        'ymax')[0].firstChild.data)
                if name.firstChild.data == 'traffic sign':
                    pred[ymin:ymax, xmin:xmax, 0] = 255
                elif name.firstChild.data == 'traffic light':
                    pred[ymin:ymax, xmin:xmax, 1] = 255

        new_im = image * 0.5 + pred * 0.5

        cv2.imwrite(ii.split('/')[-1][:-3] + 'png', new_im)


refine_pred()