





import argparse
import glob
import json
import os
import xml.etree.ElementTree as ET

import numpy as np
import tqdm
from sdsvkie.utils.word_formation import Box, check_iou
from pathlib import Path 

def read_txt(txt):
    with open(txt, 'r', encoding='utf8') as f:
        data = [line.strip() for line in f]
    return data 

def write_txt(txt, data):
    with open(txt, 'w', encoding='utf8') as f:
        for line in data:
            f.write(line + "\n")



"""

general json format: 


python cvat.py --task pseudo --xml sample_cvat/annotations.xml --xml_out sample_cvat/annotations_out.xml --pseudo_path sample_cvat/pseudo.json
    
"""





"""


config for fwd
{
    "id": "OCR005_2",
    "name": "Số Hợp đồng",
    "type": "text",
    "value": "33252663",
    "page_num": 0,
    "box": [192, 168, 220, 250]

},

"""
# CONFIG = {
#     "label": "field",

#     #edit here 
#     "attribute_names": ["type", "value", "id", "name"]    # name of attribute in cvat label: text / checkbox
# }


CONFIG = {
    "label": "word",
    # "attribute_names": ["text", "kie_label"]    # name of attribute in cvat label: text / checkbox
    "attribute_names": [] 
}



class CVAT:
    def __init__(self):
        pass 

    def create_xml_from_json(self, json_path, xml_in, xml_out):
        """
        json data format: 
        {
            "img_1.jpg": [
                {
                    "box": [x1, y1, x2, y2] 
                    "label": str                # (not required)
                    "attrib1": str, 
                    "attrib2": str,
                },
                {
                    "box": [x1, y1, x2, y2] 
                    "label": str                # (not required)
                    "attrib1": str, 
                    "attrib2": str,
                }
                ...
            ],
            "img_2.jpg": [...]
        }        
        """

        data = self.read_json(json_path)

       

        tree = ET.parse(xml_in)
        root = tree.getroot()
        for img_item in root.iter("image"):
            img_path = img_item.attrib['name']

            img_data = data[img_path]

            for item in img_data:
                et = ET.Element('box')
                #default values
                et.attrib['occluded']="0"
                et.attrib['source']="manual"
                et.attrib['z_order'] = "0"

                #overwrite values
                if 'label' in item:
                    et.attrib['label'] = item['label']
                else:
                    et.attrib['label'] = CONFIG['label']

                xmin, ymin, xmax, ymax = item['box']
                (
                    et.attrib['xtl'], et.attrib['ytl'], 
                    et.attrib['xbr'], et.attrib['ybr']
                ) = (
                    str(xmin), str(ymin), 
                    str(xmax), str(ymax)
                )

                for att_name in CONFIG['attribute_names']:
                    if att_name not in item:
                        continue

                    att_et = ET.Element('atrribute')
                    att_et.attrib['name'] = att_name
                    att_et.text = item[att_name]
                    et.append(att_et)


                img_item.append(et)

        tree.write(xml_out, encoding='utf8')



    def get_data_from_txt_dir(self, txt_dir, resever_parent_dir=False):
        if resever_parent_dir:
            txt_paths = glob.glob(txt_dir + "/*/*.txt")
        else:
            txt_paths = glob.glob(txt_dir + "/*.txt")
        data = {}
        for txt_path in txt_paths:
            if resever_parent_dir:
                txt_name = "/".join(txt_path.split("/")[-2:])
            else:
                txt_name = os.path.basename(txt_path)
        
            
            txt_data = read_txt(txt_path)
            format_data = []
            for line in txt_data:
                items = line.split("\t")
                # assert len(items) == 6, "error get len = {} - {}".format(len(items), items)
                
                box = [int(float(x)) for x in items[:4]]
                text = items[4]
                
                if len(items) == 6:
                    kie_label = items[5]
                else:
                    kie_label = "word"
                format_data.append(
                    {
                        'box': box,
                        'text': text, 
                        'label': kie_label
                    }
                )
            data[txt_name] = format_data

        return data
    
    def get_data_from_txt_path(self, txt_path):    
        
        txt_data = read_txt(txt_path)
        format_data = []
        for line in txt_data:
            items = line.split("\t")
            assert len(items) == 6, "error get len = {} - {}".format(len(items), items)
            box = [int(float(x)) for x in items[:4]]
            text = items[4]
            kie_label = items[5]
            format_data.append(
                {
                    'box': box,
                    'text': text, 
                    'label': kie_label
                }
            )
        return format_data
    
    def format_data_invoice(self, data):
        
        new_data = {}
        for txt_name, value in data.items():
            items = []
            for item in value:
                text = item['text']
                if "____kie_wordgroup" in text:
                    new_item = {
                        'box': item['box'],
                        'label': item['label']
                    }
                else:

                    new_item = {
                        "box": item['box'],
                        'text': "xxxxxx",
                        'kie_label': item['label'],
                        'label': "word"
                    }

                items.append(new_item)
            new_data[txt_name] = items
        return new_data

            
    

    def create_xml_from_txt(self, txt_dir, xml_in, xml_out, skip_labels=[], resever_parent_dir=False):
        """
        
        """
        data = self.get_data_from_txt_dir(txt_dir, resever_parent_dir)
        print(list(data.keys()))

        # for invoice 
        if len(skip_labels) > 0 and "word" in skip_labels:
            data = self.format_data_invoice(data)
        tree = ET.parse(xml_in)
        root = tree.getroot()
        count = 0 
        for img_item in tqdm.tqdm(root.iter("image")):
            count += 1
            img_path = img_item.attrib['name']

           
            txt_name = os.path.splitext(img_path)[0] + ".txt"
            
            img_data = data.get(txt_name, [])
            # from IPython import embed; embed()
            # break
            if len(img_data) > 0:
                # img_item.clear()
                for child in img_item:
                    img_item.remove(child)
            
            for item in img_data:
                et = ET.Element('box')
                #default values
                et.attrib['occluded']="0"
                et.attrib['source']="manual"
                et.attrib['z_order'] = "0"

                #overwrite values
               
                if 'label' in item:
                    if item['label'] in skip_labels:
                        continue
                    et.attrib['label'] = item['label']
                else:
                    et.attrib['label'] = CONFIG['label']

                xmin, ymin, xmax, ymax = item['box']
                (
                    et.attrib['xtl'], et.attrib['ytl'], 
                    et.attrib['xbr'], et.attrib['ybr']
                ) = (
                    str(xmin), str(ymin), 
                    str(xmax), str(ymax)
                )

                for att_name in CONFIG['attribute_names']:
                    if att_name not in item:
                        continue

                    att_et = ET.Element('atrribute')
                    att_et.attrib['name'] = att_name
                    att_et.text = item[att_name]
                    et.append(att_et)
                

                img_item.append(et)
        print("Num imgs: ", count)
        tree.write(xml_out, encoding='utf8')

    def get_data_from_xml(self, xml, skip_labels=[]):
        """ parse xml to dict
        

        Args:
            xml (str): cvat anno xml path

        Returns:
            (dict): {
                'img_1.jpg': [
                    'kie_label': [xyxy]
                    ...
                ],
                'img_2.jpg': ...
            }
        """
        anno_data = open(xml, encoding='utf8')
        tree  = ET.parse(anno_data)
        root = tree.getroot()
        data = {}

        for obj in tqdm.tqdm(root.iter("image")):
            img_path = obj.attrib['name']
            img_data = []
            for box in obj.iter("box"):
                box_label = box.attrib['label']
                if box_label in skip_labels:
                    continue

                #get coordinate
                xmin, ymin, xmax, ymax = box.attrib['xtl'], box.attrib['ytl'], box.attrib['xbr'], box.attrib['ybr']
                xmin, ymin, xmax, ymax = int(float(xmin)),  int(float(ymin)),  int(float(xmax)),  int(float(ymax))

                item = {
                    box_label: [xmin, ymin, xmax, ymax]
                }
                img_data.append(item)
            
            data[img_path] = img_data
        
        return data
    @staticmethod
    def write_json(json_path, data):
        with open(json_path, 'w', encoding='utf8') as f:
            json.dump(data, f, ensure_ascii=False)


    def read_json(self, json_path):
        with open(json_path, 'r', encoding='utf8') as f:
            data = json.load(f)

        return data 

    def update_label_kie(self, txt_dir, json_path, out_dir):
        if not os.path.exists(out_dir):
            os.makedirs(out_dir, exist_ok=True)
        data = self.read_json(json_path)
        txt_paths = glob.glob(txt_dir + "/*.txt")
        for txt_path in tqdm.tqdm(txt_paths):
            ori_data = read_txt(txt_path)

            boxes = []

            img_name = os.path.splitext(os.path.basename(txt_path))[0]
            img_name = "_".join(img_name.split("_")[:-1]) + "_1" + ".jpg"
            # img_name = os.path.splitext(os.path.basename(txt_path))[0] + ".jpg"
            new_img_data = data[img_name]
            for line in ori_data:
                xmin, ymin, xmax, ymax, text, kie_label = line.strip().split("\t")
                if "____kie_wordgroup" in text:
                    continue
                xmin, ymin, xmax, ymax = int(xmin), int(ymin), int(xmax), int(ymax)
                # box_word = [xmin, ymin, xmax, ymax]
                new_kie_label = "other"
                for label_info in new_img_data:
                    for label, box_wordgroup in label_info.items():
                        # print(label, box_wordgroup)
                        box_word = Box(
                            xmin, ymin, xmax, ymax
                        )
                        box_wordgroup = Box(
                            box_wordgroup[0], box_wordgroup[1], box_wordgroup[2],box_wordgroup[3]
                        )
                        if  check_iou(box1=box_word, box2=box_wordgroup, threshold=0.85):
                            new_kie_label = label
                            break 

                    if new_kie_label != "other":
                        break
                
                        
                new_box = Box(
                    xmin=xmin, 
                    ymin=ymin,
                    xmax=xmax,
                    ymax=ymax, 
                    label=text,
                    kie_label=new_kie_label
                    
                )
                boxes.append(new_box)
                
            boxes = sorted(boxes, key=lambda box: [box.ymin, xmin])
            new_data = [
               "\t".join([str(box.xmin), str(box.ymin), str(box.xmax), str(box.ymax), box.label, box.kie_label])
                    for box in boxes
            ]


            write_txt(os.path.join(out_dir, os.path.basename(txt_path)), new_data)
    def _check_iou(self, box1, box2, threshold=0.9):
        """_summary_

        Args:
            box1 (_type_): word box
            box2 (_type_): line box 
            threshold (float, optional): _description_. Defaults to 0.9.

        Returns:
            _type_: _description_
        """
        area1 = (box1[2] - box1[0]) * (box1[3] - box1[1])
        area2 =  (box2[2] - box2[0]) * (box2[3] - box2[1])
        xmin_intersect = max(box1[0], box2[0])
        ymin_intersect = max(box1[1], box2[1])
        xmax_intersect = min(box1[2], box2[2])
        ymax_intersect = min(box1[3], box2[3])


        if xmax_intersect < xmin_intersect or ymax_intersect < ymin_intersect:
            area_intersect = 0
        else:
            area_intersect = (xmax_intersect - xmin_intersect) * (
                ymax_intersect - ymin_intersect
            )
        # union = area1 + area2 - area_intersect
        iou = area_intersect / area1
        return iou 
        # if iou > threshold:
        #     return True
        # return False

    def _update_label_for_word(self, box, line_items, threshold=0.75, other_class='others'):
        have_label = False
        max_iou = -1
        for line_item in line_items:
            # 465	901	664	940
            
            curr_iou = self._check_iou(box, line_item['box'], threshold=threshold)
            if curr_iou > threshold and curr_iou > max_iou:
                
                max_iou = curr_iou
                kie_label = line_item['label']
                have_label = True
                
                # if box[0] == 465 and box[-1] == 940:
                #     print(box, curr_iou, kie_label, line_item)
                # break

        if not have_label:
            kie_label = other_class
        return kie_label

    def update_label_kie_from_xml(self, txt_dir, xml, out_dir, skip_labels = [], line_to_word=False, other_class="others", resever_parent_dir=False):
        if not os.path.exists(out_dir):
            os.makedirs(out_dir, exist_ok=True)
        
        #read xml 
        xml_data = {}
        anno_data = open(xml, encoding='utf8')
        tree  = ET.parse(anno_data)
        root = tree.getroot()
        # data = {}

        for obj in tqdm.tqdm(root.iter("image")):
            img_path = obj.attrib['name']
            # img_data = []
            if not line_to_word:
                img_data = {}
                for box in obj.iter("box"):
                    box_label = box.attrib['label']
                    if box_label in skip_labels:
                        continue

                    #get coordinate
                    xmin, ymin, xmax, ymax = box.attrib['xtl'], box.attrib['ytl'], box.attrib['xbr'], box.attrib['ybr']
                    box_int = int(float(xmin)),  int(float(ymin)),  int(float(xmax)),  int(float(ymax))
                    box_key = ",".join([str(x) for x in box_int])
                    img_data[box_key] = box_label

            else:
                img_data = []
                for box in obj.iter("box"):
                    box_label = box.attrib['label']
                    if box_label in skip_labels:
                        continue

                    #get coordinate
                    xmin, ymin, xmax, ymax = box.attrib['xtl'], box.attrib['ytl'], box.attrib['xbr'], box.attrib['ybr']
                    box_int = int(float(xmin)),  int(float(ymin)),  int(float(xmax)),  int(float(ymax))
                    box_key = ",".join([str(x) for x in box_int])
                    img_data.append(
                        {
                            'box': box_int,
                            'label': box_label
                        }
                    )
            xml_data[os.path.splitext(img_path)[0]] = img_data

        # print(xml_data)
        if resever_parent_dir:
            txt_paths = glob.glob(txt_dir + "/*/*.txt")
        else:
            txt_paths = glob.glob(txt_dir + "/*.txt")
        updated_imgs = []
        for txt_path in tqdm.tqdm(txt_paths):
            is_update = False 
            # print(txt_path)
            ori_data = read_txt(txt_path)
            # print(ori_data)
            img_new_data = []
            if resever_parent_dir:
                img_key = str(Path(txt_path).with_suffix('').relative_to(Path(txt_path).parent.parent))   # a/xyz
            else:
                img_key = os.path.splitext(os.path.basename(txt_path))[0]   # xyz
            if img_key not in xml_data:
                print(txt_path)
                continue
            img_annoted_data = xml_data[img_key]
            # print(img_key, img_annoted_data)
            if not line_to_word:
        
                for line in ori_data:
                    xmin, ymin, xmax, ymax, text, kie_label = line.strip().split("\t")
                    if "____kie_wordgroup" in text:
                        continue
                    box_int = int(xmin), int(ymin), int(xmax), int(ymax)
                    box_key = ",".join([str(x) for x in box_int])

                    if box_key in img_annoted_data:
                        if kie_label != img_annoted_data[box_key]:
                            is_update.append(txt_path)
                            # print(kie_label, img_annoted_data[box_key])
                        kie_label = img_annoted_data[box_key]
                    else:
                        kie_label = other_class
                    img_new_data.append("\t".join([xmin, ymin, xmax, ymax, text, kie_label]))
            else:
                
                # print("ori_data: ", ori_data)
                for line in ori_data:
                    # print(line)
                    items = line.strip().split("\t")
                    if len(items) == 5:
                        xmin, ymin, xmax, ymax, text = items
                    else:
                        xmin, ymin, xmax, ymax, text, label = items
                    
                    box_int = int(xmin), int(ymin), int(xmax), int(ymax)
                    kie_label = self._update_label_for_word(box_int, img_annoted_data, threshold=0.75, other_class=other_class)

                    if label != kie_label:
                        print(kie_label, label)
                        is_update = True
                    img_new_data.append("\t".join([xmin, ymin, xmax, ymax, text, kie_label]))
            if resever_parent_dir:
                out_sub_dir = Path(out_dir) / Path(img_key).parts[-2]
                if not out_sub_dir.exists():
                    out_sub_dir.mkdir(parents=True)
            # else:
            #     out_sub_dir = out_dir
            # out_sub_dir = str(out_sub_dir)
            write_txt(os.path.join(out_dir, img_key + ".txt"), img_new_data)
        
            if is_update:
                updated_imgs.append(txt_path)
            else:
                print("No update: ", txt_path)
        print("updated_imgs: ", list(set(updated_imgs)))
        print("num updated_imgs: ", len(list(set(updated_imgs))))
            


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--xml", type=str, default='annotations.xml')
    parser.add_argument("--pseudo_path", type=str, default="pseudolabel.json")
    parser.add_argument("--xml_out", type=str, default="annotations_out.xml")
    parser.add_argument("--task", type=str, default='pseudo', help='pseudo / parse_data / update_txt / update_txt_from_xml')
    parser.add_argument("--txt_in", type=str, default='txt_dir_in')
    parser.add_argument("--txt_out", type=str, default='txt_dir_out')
    parser.add_argument("--line_to_word", action='store_true')
    parser.add_argument("--other_class", type=str, default='other')
    parser.add_argument("--resever_parent_dir", action="store_true")

    args = parser.parse_args()
    cvat = CVAT()

    if args.task == 'parse_data':
        data = cvat.get_data_from_xml(
            xml=args.xml
        )
        CVAT.write_json(args.xml_out, data)

    elif args.task == 'pseudo_from_json':
        
        cvat.create_xml_from_json(
            xml_in=args.xml,
            xml_out=args.xml_out,
            json_path=args.pseudo_path
        )

    elif args.task == 'pseudo_from_txt':
        
        cvat.create_xml_from_txt(
            xml_in=args.xml,
            xml_out=args.xml_out,
            txt_dir=args.pseudo_path,
            # skip_labels=['word']
            # skip_labels=[args.other_class],
            resever_parent_dir=args.resever_parent_dir
        )
    elif args.task == 'update_txt':
        cvat.update_label_kie(
            txt_dir=args.txt_in,
            json_path=args.pseudo_path,
            out_dir=args.txt_out
        )

    elif args.task == 'update_txt_from_xml':
        cvat.update_label_kie_from_xml(
            txt_dir=args.txt_in,
            xml=args.xml,
            out_dir=args.txt_out,
            skip_labels = ['word'], 
            line_to_word=args.line_to_word, 
            other_class=args.other_class,
            resever_parent_dir=args.resever_parent_dir
        )
    else:
        raise NotImplementedError(f"{args.task} not yet implemented")

    



"""
python tools/cvat.py --task update_txt --txt_in /mnt/ssd1T/hoanglv/Projects/KIE/craw_data/output/synth_vnpt_r20/one_line_filtered  --txt_out /mnt/ssd1T/hoanglv/Projects/KIE/craw_data/output/synth_vnpt_r20/one_line_filtered --pseudo_path ../workdirs/data/vnpt_oneline/annotations.json

python tools/cvat.py --task pseudo_from_txt \
    --xml /mnt/ssd1T/hoanglv/Projects/KIE/sdsvkie/workdirs/visualize/xml/vnpt_r2/annotations.xml \
    --xml_out /mnt/ssd1T/hoanglv/Projects/KIE/sdsvkie/workdirs/visualize/xml/vnpt_r2/annotations_out.xml \
    --pseudo_path /mnt/ssd1T/hoanglv/Projects/KIE/sdsvkie/workdirs/visualize/vnpt_r2_txt

python tools/cvat.py --task pseudo_from_txt \
    --xml /mnt/ssd1T/hoanglv/Projects/KIE/DATA/WildReceipt/re_labeling/wild_batch_1_raw.xml \
    --xml_out /mnt/ssd1T/hoanglv/Projects/KIE/DATA/WildReceipt/re_labeling/wild_batch_1_pseudo.xml \
    --pseudo_path /mnt/ssd1T/hoanglv/Projects/KIE/DATA/WildReceipt/re_labeling/batches/batch_1


python tools/cvat.py --task update_txt_from_xml \
    --txt_in  /mnt/ssd1T/hoanglv/Projects/KIE/sdsvkie/workdirs/visualize/vnpt_r2_txt \
    --xml /mnt/ssd1T/hoanglv/Projects/KIE/sdsvkie/workdirs/visualize/vnpt_r2.xml \
    --txt_out /mnt/ssd1T/hoanglv/Projects/KIE/sdsvkie/workdirs/visualize/vnpt_r2_done_txt

python tools/cvat.py --task update_txt_from_xml \
    --txt_in /mnt/ssd1T/hoanglv/Projects/KIE/DATA/SDSAP_Invoice/labeling/Pseudo/batch_2/Good/Food  \
    --xml /mnt/ssd1T/hoanglv/Projects/KIE/DATA/SDSAP_Invoice/labeling/Pseudo/batch_2/batch_2_food_done.xml   \
    --txt_out /mnt/ssd1T/hoanglv/Projects/KIE/DATA/SDSAP_Invoice/processed/batch_2/Good/Food \
    --other_class Others
"""