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

import cv2
import tqdm

from sdsvkie.cfg import load_cfg
from sdsvkie.engine import Predictor
from sdsvkie.utils import visualize_kie
from sdsvkie.utils.io_file import load_ocr_output, write_json, write_txt

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

    # optional
    parser.add_argument("--weights", type=str, default=None, required=False)

    parser.add_argument("--device", type=str, default=None,  required=False)

    parser.add_argument(
        "--text_det", type=str, default=None, required=False, help="image path or directory"
    )
    parser.add_argument(
        "--text_reg", type=str, default=None,  required=False, help="image path or directory"
    )
    parser.add_argument(
        "--vis_out", type=str, default=None,  required=False,  help="visualize output directory"
    )
    parser.add_argument("--txt_out", type=str, required=False,  default=None)
    parser.add_argument("--kie_wordgroup_out", required=False,  action="store_true")
    parser.add_argument("--e2e", type=str, required=False,  default=None)
    parser.add_argument("--not_use_ocr", required=False,  action="store_true")
    parser.add_argument("--parse_e2e", action='store_true', help="Parse end2end result from word label")

    args = parser.parse_args()

    
    # print(cfg)

    predictor = Predictor(**vars(args))


    # cfg = load_cfg(args.cfg,  vars(args))
    # predictor = Predictor(cfg)

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

    if args.e2e:
        outdir_e2e = os.path.dirname(args.e2e)
        if not os.path.exists(outdir_e2e):
            os.makedirs(outdir_e2e, exist_ok=True)

    if os.path.isdir(args.img):
        img_paths = glob(args.img + "/*")
        print("Infence image dir, total imgs: {}".format(len(img_paths)))
    else:
        img_paths = [args.img]

    out_dict = {}
    for img_path in tqdm.tqdm(img_paths):

        img = cv2.imread(img_path)
        if img is None:
            print("img is None: ", img_path)
            continue
        if args.not_use_ocr:
            txt_path = str(Path(img_path).with_suffix(".txt"))
            ocr_output = load_ocr_output(txt_path)
            # print(len(ocr_output["boxes"]))
            out = predictor(img, ocr_output=ocr_output, return_raw=True)
        else:
            out = predictor(img, return_raw=True)

        # visualize
        if args.vis_out:
            # from IPython import embed; embed()
            out_kie = out["kie_raw_output"]
            visualize_kie(
                img,
                boxes=[word.boundingbox for word in out_kie],
                pred_labels=[word.kie_label for word in out_kie],
                image_name=os.path.basename(img_path),
                outdir=args.vis_out,
                skip_classes=["other---"]
            )

        if args.txt_out:
            txt_out_path = os.path.join(
                args.txt_out, os.path.splitext(os.path.basename(img_path))[0] + ".txt"
            )

            out_kie = out["kie_raw_output"]
            boxes = [word.boundingbox for word in out_kie]
            pred_labels = [word.kie_label for word in out_kie]
            texts = [word.text for word in out_kie]

            data = []
            

            if args.kie_wordgroup_out:
                output = out["kie_post_output"]
                # print(output)
                # from IPython import embed; embed()
                wordgroup_all_list = []
                for kie_label, wordgroup_list in output.items():
                    if isinstance(wordgroup_list, list):
                        wordgroup_all_list.extend(wordgroup_list)
                    else:
                        wordgroup_all_list.append(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)
            else:
                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)

        if args.e2e:

            img_id = os.path.splitext(os.path.basename(img_path))[0]
            out_dict[img_id] = out['end2end_results']

    if args.e2e:
        write_json(os.path.join(args.e2e), out_dict)
