from typing import Optional
from pathlib import Path
from xml.etree import ElementTree as ET
from multiprocessing import Process
import pandas as pd
from copy import deepcopy
import cv2


class VocStyle:
    def __init__(self, img_dir, data: Optional[pd.DataFrame]) -> None:
        self.img_dir = Path(img_dir)
        self.data = data
        self.data.columns = ["name", "cls", "x1", "y1", "x2", "y2"]
        self.names = list(set(self.data["name"].to_list()))
        self.cls_names = {}
    

    def write(self, names):
        for name in names:
            tree, root, filename, width, height, depth, obj = self.style()
            img = self.img_dir.joinpath(name + ".jpg")
            filename.text = img.name
            cv_img = cv2.imread(img.__str__())
            h, w, c = cv_img.shape
            width.text, height.text, depth.text = str(w), str(h), str(c)
            bboxes = self.data[self.data["name"]==name].to_numpy()
            for bbox in bboxes:
                obj_ = self.new_obj(obj, bbox)
                root.append(obj_)
            tree.write(self.img_dir.joinpath(name + ".xml"))
    
    def to_voc(self):
        n = len(self.names) // 5
        name_list = [self.names[i*n:(i+1)*n if i<4 else -1] for i in range(5)]
        for names in name_list:
            task = Process(target=self.write, args=(names, ))
            task.start()
                
    def style(self):
        tree0 = ET.parse("./xml/voc.xml")
        tree = deepcopy(tree0)
        root = tree.getroot()
        filename = root.find("filename")
        size = root.find("size")
        width = size.find("width")
        height = size.find("height")
        depth = size.find("depth")
        obj0 = root.find("object")
        obj = deepcopy(obj0)
        root.remove(obj0)
        return tree, root, filename, width, height, depth, obj

    def new_obj(self, obj, bbox):
        obj = deepcopy(obj)
        name = obj.find("name")
        bndbox = obj.find("bndbox")
        xmin = bndbox.find("xmin")
        ymin = bndbox.find("ymin")
        xmax = bndbox.find("xmax")
        ymax = bndbox.find("ymax")
        name.text, xmin.text, ymin.text, xmax.text, ymax.text = bbox[1:]
        return obj

    def voc_2_yolo(self):
        self.names = self.names & [i.stem for i in self.img_dir.glob("*.jpg") if \
            self.img_dir.joinpath(i.stem + ".xml").exists()]
        assert len(self.names) > 0
        for name in self.names:
            bboxes = self.parse(name)
            with open(self.img_dir.joinpath(name + ".txt"), "w")as f:
                f.writelines(bboxes)

    def parse(self, name):
        anno = self.img_dir.joinpath(name + ".xml")
        tree = ET.parse(anno)
        root = tree.getroot()
        size = root.find("size")
        width = int(eval(size.find("width")))
        height = int(eval(size.find("height")))
        objs = root.findall("object")
        bboxes = []

        for obj in objs:
            cls = obj.find("name")
            if cls in self.cls_names:
                cls = self.cls_names[cls]
            else:
                self.cls_names[cls] = len(self.cls_names)
                self.cls = self.cls_names[cls]
            bndbox = obj.find("bndbox")
            x1 = int(eval(bndbox.find("xmin").text))
            y1 = int(eval(bndbox.find("ymin").text))
            x2 = int(eval(bndbox.find("xmax").text))
            y2 = int(eval(bndbox.find("ymax").text))
            cx = (x1 + x2) / 2
            cy = (y1 + y2) / 2
            w = x2 - x1
            h = y2 - y1
            bbox = [cx / width, cy / height, w / width, h / height]
            bbox = f"{cls} " + " ".join(list(map(lambda x: str(x), bbox))) + "\n"
            bboxes.append(bbox)
        
        return bboxes


if __name__ == "__main__":
    
    names = list(range(1989))
    n = len(names) // 5
    name_list = [names[i*n:(i+1)*n if i<4 else -1] for i in range(5)]

    print([len(i) for i in name_list])
