import cv2
import numpy as np
import time
import os
from xml.etree.ElementTree import Element, SubElement, Comment, tostring
from xml.etree import ElementTree
from xml.dom import minidom

def prettify(elem):
    """Return a pretty-printed XML string for the Element.
    """
    rough_string = ElementTree.tostring(elem, 'utf-8')
    reparsed = minidom.parseString(rough_string)
    return reparsed.toprettyxml(indent="  ", newl='')

class Annotation:

    def __init__(self, xml_path=None, image=None, filename=None, folder='LCM', lcmlog='lcm-log-file', camera_name='cam', utime=0):

        if xml_path is not None:
            # load the existing annotation
            self.xml_path = xml_path
            self.xml_tree = ElementTree.parse(xml_path)
            self.root = self.xml_tree.getroot()
            self.filename = self.xml_tree.find('filename').text
            self.folder = self.xml_tree.find('folder').text
        else:
            self.utime = utime
            self.camera_name = camera_name
            self.image_width = image.shape[1]
            self.image_height = image.shape[0]
            self.image_depth = image.shape[2]
            self.lcmlog = lcmlog
            self.filename = filename
            self.folder = folder
            self.image = image
            self.bboxes = []
            self.root = Element('annotation')
            self.folder_tag = self.add_tag(self.root, 'folder', self.folder)
            self.filename_tag = self.add_tag(self.root, 'filename', self.filename)
            self.source = SubElement(self.root, 'source')
            self.add_tag(self.source, 'database', 'Unknown')
            self.add_tag(self.source, 'lcmlog', lcmlog)
            self.add_tag(self.source, 'utime', utime)
            self.add_tag(self.source, 'camera-name', self.camera_name)
            self.size = SubElement(self.root, 'size')
            self.add_tag(self.size, 'width', self.image_width)
            self.add_tag(self.size, 'height', self.image_height)
            self.add_tag(self.size, 'depth', self.image_depth)

    def add_tag(self, parent, child_name, value=None):
        tag = SubElement(parent, child_name)
        if value is not None:
            tag.text = str(value)
        return tag

    def add_box(self, bbox):
        pt0, pt1, utime, class_name, scores = bbox
        obj = SubElement(self.root, 'object')
        self.add_tag(obj, 'name', class_name)
        self.add_tag(obj, 'pose', 'Unspecified')
        self.add_tag(obj, 'truncated', '0')
        self.add_tag(obj, 'occluded', '0')
        self.add_tag(obj, 'difficult', '0')
        bndbox = SubElement(obj, 'bndbox')
        self.add_tag(bndbox, 'xmin', pt0[0])
        self.add_tag(bndbox, 'ymin', pt0[1])
        self.add_tag(bndbox, 'xmax', pt1[0])
        self.add_tag(bndbox, 'ymax', pt1[1])

        self.bboxes.append(bbox)

    def save_xml(self, filepath=None):

        if filepath is None:
            filepath = os.path.join(self.folder, 'xml', self.filename + '.xml')

        with open(filepath, 'w') as f:
            f.write(prettify(self.root))

    def save_boxed_image(self, filepath=None):
        # save out the full image with boxes drawn
        marked_image = self.image.copy()
        for bbox in self.bboxes:
            pt0, pt1, utime, class_name, scores = bbox
            cv2.rectangle(marked_image, tuple(pt0), tuple(pt1), (0, 255, 0), 2)
            cv2.putText(marked_image, class_name, (pt0[0], pt1[1]), cv2.FONT_HERSHEY_COMPLEX, 0.5, (0, 255, 255))

        if filepath is None:
            filepath = os.path.join(self.folder, 'boxed_images', self.filename + ".jpg")

        print(filepath)
        cv2.imwrite(filepath, marked_image)

    def save_image(self, filepath=None):

        if filepath is None:
            filepath = os.path.join(self.folder, 'images', self.filename + ".jpg")

        cv2.imwrite(filepath, self.image)

    def save_rois(self, filepath=None):
        # for each annotation, extract the ROI from the raw image and save
        if filepath is None:
            fileprefix = self.filename
        else:
            fileprefix = filepath

        for bbox in self.bboxes:
            pt0, pt1, utime, class_name, scores = bbox
            roi = self.image[pt0[1]:pt1[1], pt0[0]:pt1[0], :]
            coord_string = str(pt0[0]) + "-" + str(pt0[1]) + "-" + str(pt1[0]) + "-" + str(pt1[1])
            filepath = os.path.join(self.folder, 'rois', fileprefix + "-" + coord_string + ".jpg")
            cv2.imwrite(filepath, roi)


if __name__=="__main__":


    # A test
    bbox = ((100, 200), (130, 260), 1567897687000, 'test_class1', 0.25)
    bbox2 = ((800, 400), (870, 420), 1567899687000, 'test_class2', 0.25)
    img = 255*np.random.random((540, 960, 3))
    filename = 'a_test_file'
    ann = Annotation(image=img, filename=filename, folder='.', utime=int(time.time()*1000))
    ann.add_box(bbox)
    ann.add_box(bbox2)
    ann.save_xml()
    ann.save_boxed_image()
    ann.save_rois()
    ann2 = Annotation(xml_path=filename+'.xml')