"""
Use for eval, debug
"""
import argparse
import os
from glob import glob
from pathlib import Path

from sdsvkie.utils.word_formation import Word
import tqdm

from sdsvkie.engine import Predictor
from sdsvkie.utils.io_file import load_ocr_output, write_txt

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--cfg", type=str, default="sdsvkie/cfg/default.yaml")
    parser.add_argument(
        "--src", type=str, default="img.jpg / img_dir", help="image path or directory"
    )
    parser.add_argument("--tgt", type=str, required=True,  default=None)

    args = parser.parse_args()

    predictor = Predictor(args.cfg)

    
    if not os.path.exists(args.tgt):
        os.makedirs(args.tgt, exist_ok=True)

    txt_paths = glob(args.src + "/*.txt")
    print("Total txt: {}".format(len(txt_paths)))


    for txt_path in tqdm.tqdm(txt_paths):
        ocr_output = load_ocr_output(txt_path)
        
        boxes, texts, labels = ocr_output['boxes'], ocr_output['texts'], ocr_output['labels']
        words = []
        for box, text, label in zip(boxes, texts, labels):
            words.append(
                Word(
                    text=text,
                    bndbox=box,
                    kie_label=label,
                    conf_cls=0.0
                )
            )
            
        kie_output = predictor.postprocessing(words)
        

        txt_out_path = str(Path(args.tgt) / Path(txt_path).name)


        data = []

        # print(output)
        wordgroup_all_list = []
        for kie_label, wordgroup_list in kie_output.items():
            wordgroup_all_list.extend(wordgroup_list)
        boxes = [word.boundingbox for word in wordgroup_all_list]
        pred_labels = [word.kie_label for word in wordgroup_all_list]
        texts = [word.text + "____kie_wordgroup" for word in wordgroup_all_list]

        for box, text, kie_label in zip(boxes, texts, pred_labels):
            item = "\t".join([str(int(x)) for x in box])
            item = "\t".join([item, text, kie_label])
            data.append(item)


        write_txt(txt_out_path, data)