import os
import argparse

import onnxruntime as ort
from utils import *


CFG = {
    "name": "mobilenet0.25",
    "min_sizes": [[16, 32], [64, 128], [256, 512]],
    "steps": [8, 16, 32],
    "variance": [0.1, 0.2],
    "clip": False,
}
INPUT_SIZE = [608, 640]   #resize scale
DEVICE = torch.device("cpu")


def save_result(img_name, dets, save_folder):
    """Save detection results
    Args:
        img_name: origin image name
        dets: detection results
        save_folder: results path
    """
    if not os.path.exists(save_folder):
        os.makedirs(save_folder)
    save_name = save_folder + img_name[:-4] + ".txt"
    dirname = os.path.dirname(save_name)
    if not os.path.isdir(dirname):
        os.makedirs(dirname)
    with open(save_name, "w") as fw:
        bboxs = dets
        file_name = os.path.basename(save_name)[:-4] + "\n"
        bboxs_num = str(len(bboxs)) + "\n"
        fw.write(file_name)
        fw.write(bboxs_num)
        for box in bboxs:
            x = int(box[0])
            y = int(box[1])
            w = int(box[2]) - int(box[0])
            h = int(box[3]) - int(box[1])
            confidence = str(box[4])
            line = (str(x) + " " + str(y) + " " + str(w) + " " + str(h) + " " + confidence + " \n")
            fw.write(line)


def Retinaface_evalute(run_ort, args):
    """Retinaface_evalute function
    Args:
        run_ort : run_ort to evaluate.
        args : parser parameter.
    Returns:
        predict result : under "--save_folder" path.
    """
    # testing dataset
    testset_folder = args.dataset_folder
    testset_list = args.dataset_folder[:-7] + "wider_val.txt"

    with open(testset_list, "r") as fr:
        test_dataset = fr.read().split()
    num_images = len(test_dataset)
    
    # testing begin
    for i, img_name in enumerate(test_dataset):
        image_path = testset_folder + img_name
        img_raw = cv2.imread(image_path, cv2.IMREAD_COLOR)
        # preprocess
        img, scale, resize = preprocess(img_raw, INPUT_SIZE, DEVICE)
        # to NHWC
        img = np.transpose(img, (0, 2, 3, 1))
        # forward 
        outputs = run_ort.run(None, {run_ort.get_inputs()[0].name: img})
        # postprocess
        dets = postprocess(CFG, img, outputs, scale, resize, args.confidence_threshold, args.nms_threshold, DEVICE)

        # save predict result
        save_result(img_name, dets, args.save_folder)
        print("im_detect: {:d}/{:d}".format(i + 1, num_images))


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="Retinaface")
    parser.add_argument(
        "-m",
        "--trained_model",
        default="./weights/RetinaFace_int.onnx",
        type=str,
        help="Trained state_dict file path to open",
    )
    parser.add_argument(
        "--save_folder",
        default="./widerface_evaluate/widerface_txt/",
        type=str,
        help="Dir to save txt results",
    )
    parser.add_argument(
        "--dataset_folder",
        default="./data/widerface/val/images/",
        type=str,
        help="dataset path",
    )
    parser.add_argument(
        "--confidence_threshold",
        default=0.02,
        type=float,
        help="confidence_threshold",
    )
    parser.add_argument(
        "--nms_threshold",
        default=0.4,
        type=float,
        help="nms_threshold",
    )
    parser.add_argument(
        "--ipu",
        action="store_true",
        help="Use IPU for inference.",
    )
    parser.add_argument(
        "--provider_config",
        type=str,
        default="vaip_config.json",
        help="Path of the config file for seting provider_options.",
    )

    args = parser.parse_args()
    if args.ipu:
        providers = ["VitisAIExecutionProvider"]
        provider_options = [{"config_file": args.provider_config}]
    else:
        providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
        provider_options = None
        
    print("Loading pretrained model from {}".format(args.trained_model))
    run_ort = ort.InferenceSession(args.trained_model, providers=providers, provider_options=provider_options)
    
    Retinaface_evalute(run_ort, args)