import glob
from re import IGNORECASE
import numpy as np
import os
import cv2
import pickle
import xml.etree.ElementTree as ET

from xml.dom.minidom import parseString
from lxml.etree import Element, SubElement, tostring
from os import listdir, getcwd
from os.path import join

IGNORE_SMALL_OBJ = True
SMALL_OBJECT_THRETH = 0.073


def getImagesInDir():
    image_list = []
    for ext in ["*.jpg", "*.png", "*.jpeg"]:
        filenames = glob.glob(os.path.join('Images/', ext))
        for filename in filenames:
            image_name = filename.split(os.sep)[1]
            image_list.append(image_name)
    return image_list


def convert(size, box):
    dw = 1. / (size[0])
    dh = 1. / (size[1])
    x = (box[0] + box[1]) / 2.0 - 1
    y = (box[2] + box[3]) / 2.0 - 1
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh

    return (x, y, w, h)


def convert_annotation(image_path):
    basename = os.path.basename(image_path)
    basename_no_ext = os.path.splitext(basename)[0]

    in_file = open('./Annotations/' + basename_no_ext + '.xml')
    out_file = open('./YOLO_annotations/' + basename_no_ext + '.txt', 'w')
    tree = ET.parse(in_file)
    root = tree.getroot()
    size = root.find('size')
    w = int(size.find('width').text)
    h = int(size.find('height').text)

    for obj in root.iter('object'):
        difficult = obj.find('difficult').text
        cls = obj.find('name').text
        if cls not in classes or int(difficult) == 1:
            # print(image_path, cls)
            continue
        cls_id = classes.index(cls)
        xmlbox = obj.find('bndbox')
        b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text),
             float(xmlbox.find('ymax').text))
        bb = convert((w, h), b)
        if IGNORE_SMALL_OBJ and (bb[2] < SMALL_OBJECT_THRETH or bb[3] < SMALL_OBJECT_THRETH):
            continue
        out_file.write(str(cls_id) + " " + str(bb[0]) + " " + str(bb[1]) + " " + str(bb[2]) + " " + str(bb[3]) + '\n')


def VOC2YOLO():
    if not os.path.exists('./YOLO_annotations/'):
        os.makedirs('./YOLO_annotations/')

    image_paths = getImagesInDir()
    for image_path in image_paths:
        convert_annotation(image_path)
    print("Finished processing")


def unconvert(class_id, width, height, x, y, w, h):

    xmax = int((x * width) + (w * width) / 2.0)
    xmin = int((x * width) - (w * width) / 2.0)
    ymax = int((y * height) + (h * height) / 2.0)
    ymin = int((y * height) - (h * height) / 2.0)
    class_id = int(class_id)
    return (class_id, xmin, xmax, ymin, ymax)


def YOLO2VOC(classes):
    classes = list(classes)
    ids = list()
    l = os.listdir('./YOLO_annotations/')

    check = '.DS_Store' in l
    if check == True:
        l.remove('.DS_Store')

    ids = [x.split('.')[0] for x in l]

    annopath = join('./YOLO_annotations', '%s.txt')
    imgpath = join('./Images', '%s.png')

    if not os.path.exists('./VOC_annotations/'):
        os.makedirs('./VOC_annotations/')

    outpath = join('./VOC_annotations', '%s.xml')

    for i in range(len(ids)):
        img_id = ids[i]
        img = cv2.imread(imgpath % img_id)
        try:
            height, width, channels = img.shape
        except AttributeError as e:
            print(imgpath % img_id, "not found! Err: ", e)
            return

        node_root = Element('annotation')
        node_folder = SubElement(node_root, 'folder')
        node_folder.text = 'VOC2007'
        img_name = img_id + '.jpg'

        node_filename = SubElement(node_root, 'filename')
        node_filename.text = img_name

        node_source = SubElement(node_root, 'source')
        node_database = SubElement(node_source, 'database')
        node_database.text = 'Coco database'

        node_size = SubElement(node_root, 'size')
        node_width = SubElement(node_size, 'width')
        node_width.text = str(width)

        node_height = SubElement(node_size, 'height')
        node_height.text = str(height)

        node_depth = SubElement(node_size, 'depth')
        node_depth.text = str(channels)

        node_segmented = SubElement(node_root, 'segmented')
        node_segmented.text = '0'

        target = (annopath % img_id)
        if os.path.exists(target):
            label_norm = np.loadtxt(target).reshape(-1, 5)

            for i in range(len(label_norm)):
                labels_conv = label_norm[i]
                if IGNORE_SMALL_OBJ and (labels_conv[3] < SMALL_OBJECT_THRETH or labels_conv[4] < SMALL_OBJECT_THRETH):
                    continue
                new_label = unconvert(labels_conv[0], width, height, labels_conv[1], labels_conv[2], labels_conv[3],
                                      labels_conv[4])
                node_object = SubElement(node_root, 'object')
                node_name = SubElement(node_object, 'name')
                node_name.text = classes[new_label[0]]

                node_pose = SubElement(node_object, 'pose')
                node_pose.text = 'Unspecified'

                node_truncated = SubElement(node_object, 'truncated')
                node_truncated.text = '0'
                node_difficult = SubElement(node_object, 'difficult')
                node_difficult.text = '0'
                node_bndbox = SubElement(node_object, 'bndbox')
                node_xmin = SubElement(node_bndbox, 'xmin')
                node_xmin.text = str(new_label[1])
                node_ymin = SubElement(node_bndbox, 'ymin')
                node_ymin.text = str(new_label[3])
                node_xmax = SubElement(node_bndbox, 'xmax')
                node_xmax.text = str(new_label[2])
                node_ymax = SubElement(node_bndbox, 'ymax')
                node_ymax.text = str(new_label[4])
            xml = tostring(node_root, pretty_print=True)
            dom = parseString(xml)
            f = open(outpath % img_id, "wb")
            f.write(xml)
            f.close()
    print("Finished processing")


classes = ["helmet", "head"]

# VOC2YOLO()
YOLO2VOC(classes)