import os
import cv2
import random
from PIL import Image
from lxml import etree
import argparse,random
from tqdm import tqdm

class GEN_Annotations:
    def __init__(self, filename):
        self.root = etree.Element("annotation")

        child1 = etree.SubElement(self.root, "folder")
        child1.text = "CDB"

        child2 = etree.SubElement(self.root, "filename")
        child2.text = filename

        child3 = etree.SubElement(self.root, "source")

        child4 = etree.SubElement(child3, "annotation")
        child4.text = "CDB"
        child5 = etree.SubElement(child3, "database")
        child5.text = "Unknown"

        child6 = etree.SubElement(child3, "image")
        child6.text = "flickr"
        child7 = etree.SubElement(child3, "flickrid")
        child7.text = "35435"

    def set_size(self, witdh, height, channel):
        size = etree.SubElement(self.root, "size")
        widthn = etree.SubElement(size, "width")
        widthn.text = str(witdh)
        heightn = etree.SubElement(size, "height")
        heightn.text = str(height)
        channeln = etree.SubElement(size, "depth")
        channeln.text = str(channel)

    def savefile(self, filename):
        tree = etree.ElementTree(self.root)
        tree.write(filename, pretty_print=True, xml_declaration=False, encoding='utf-8')

    def add_pic_attr(self, label, xmin, ymin, xmax, ymax):
        object = etree.SubElement(self.root, "object")
        namen = etree.SubElement(object, "name")
        namen.text = label
        bndbox = etree.SubElement(object, "bndbox")
        xminn = etree.SubElement(bndbox, "xmin")
        xminn.text = str(xmin)
        yminn = etree.SubElement(bndbox, "ymin")
        yminn.text = str(ymin)
        xmaxn = etree.SubElement(bndbox, "xmax")
        xmaxn.text = str(xmax)
        ymaxn = etree.SubElement(bndbox, "ymax")
        ymaxn.text = str(ymax)


def negative_samples(path_unit,path_raw):
    unit_list = os.listdir(path_unit)
    random.shuffle(unit_list)
    if len(unit_list) < 1000:
        unit_list = unit_list
    else:
        unit_list = unit_list[:1000]
    unit_img_list = []
    for unit_img in unit_list:
        img1 = Image.open(os.path.join(path_unit,unit_img))
        unit_img_list.append(img1)
    raw_list = os.listdir(path_raw)
    for raw_path  in tqdm(raw_list):
        # print(raw_list)
        # print("\n")
        raw_img = Image.open(os.path.join(path_raw,raw_path))
        unit_num = random.randint(0,10)
        x_left = random.randint(0,raw_img.size[0]-unit_img_list[unit_num].size[0])
        y_left = random.randint(0,raw_img.size[1]-unit_img_list[unit_num].size[1])
        x_right = x_left+unit_img_list[unit_num].size[0]
        y_right = y_left+unit_img_list[unit_num].size[1]
        raw_img.paste(unit_img_list[unit_num],(x_left,y_left,x_right,y_right))
        raw_img.save(os.path.join(path_raw,raw_path).replace("/NOObject/","/result/Images/"))
        filename = raw_path
        anno = GEN_Annotations(filename)
        anno.set_size(raw_img.size[0], raw_img.size[1], 3)
        anno.add_pic_attr("person", x_left, y_left, x_right, y_right)
        anno.savefile(os.path.join(path_raw,raw_path).replace("/NOObject/","/result/Ground_truth/").replace(".jpg",".xml"))



if __name__ == '__main__':
    parser1 = argparse.ArgumentParser(description="CUI")
    parser1.add_argument("path_unit", default='/data/DataManagement/Project/001_cuihu/CAS_CuihuBird/CMB_Cuihubird/FinalData/data',
                         type=str,
                         help='path to object unit')
    parser1.add_argument("path_raw", default='/data/DataManagement/Project/001_cuihu/0009_20220507_RTSP/CDB_CuihuNobird/NOObject', type=str,
                         help='path to image to deal with')


    args = parser1.parse_args()
    path_unit = args.path_unit
    path_raw = args.path_raw
    save_path1 = path_raw.replace("NOObject","result/Images/")
    print(save_path1)
    if not os.path.exists(save_path1):
        os.makedirs(save_path1)
    save_path2 = path_raw.replace("NOObject","result/Ground_truth/")
    print(save_path2)
    if not os.path.exists(save_path2):
        os.makedirs(save_path2)
    negative_samples(path_unit,path_raw)