import sys
import os
import argparse
import yaml
current_path=os.getcwd()
sys.path.append(current_path+'/prepostprocess')
sys.path.append(current_path+'/prepostprocess/kneron_preprocessing')
sys.path.append(current_path+'/kneron_globalconstant')
sys.path.append(current_path+'/kneron_globalconstant/base')
sys.path.append(current_path+'/kneron_globalconstant/kneron_utils')

from yolov5.yolov5_runner import Yolov5Runner
from function.function_runner import FunctionRunner
from rsn_affine.rsn_affine_runner import RsnAffineRunner
from lite_hrnet.lite_hrnet_runner import LiteHrnetRunner


def inference(img_path, yolov5_params, rsn_affine_params, lite_hrnet_params):

    yolov5_runner = Yolov5Runner(model_path=yolov5_params['model_path'], model_id=yolov5_params['model_id'], yaml_path=yolov5_params['yaml_path'], grid20_path=yolov5_params['grid20_path'], grid40_path=yolov5_params['grid40_path'], grid80_path=yolov5_params['grid80_path'], num_classes=yolov5_params['num_classes'], input_shape=yolov5_params['input_shape'], conf_thres=yolov5_params['conf_thres'], iou_thres=yolov5_params['iou_thres'], top_k_num=yolov5_params['top_k_num'], detection_type=yolov5_params['detection_type'], vanish_point=yolov5_params['vanish_point'], label_mapping=yolov5_params['label_mapping'], class_name=yolov5_params['class_name'], anchors=yolov5_params['anchors'])
    function_runner_2 = FunctionRunner(type=60, thresh_head_iou=0.8, thresh_fbox_iou=0.9, thresh_person_score=0.3)
    rsn_affine_runner = RsnAffineRunner(image_size=rsn_affine_params['image_size'], scale_ext=rsn_affine_params['scale_ext'])
    lite_hrnet_runner = LiteHrnetRunner(model_path=lite_hrnet_params['model_path'])
    function_runner = FunctionRunner(type=48)

    out_0_0 = yolov5_runner.run(img_path)

    out_1_0, out_1_1 = function_runner_2.run(img_path, out_0_0)

    out_2_0, out_2_1 = rsn_affine_runner.run(img_path, out_1_0)

    out_3_0 = lite_hrnet_runner.run(img_path, [out_2_0,out_2_1])

    out_4_0 = function_runner.run(img_path, out_3_0)

    return out_4_0

def parse_args(args):
    """
    Parse the arguments.
    """
    parser = argparse.ArgumentParser(description='Simple inference script for inference an object detection network.')
    parser.add_argument('--img-path', type=str, help='Path to the image.')
    parser.add_argument('--yolov5_params', type=str, help='Path to the yolov5 init params file.')
    parser.add_argument('--rsn_affine_params', type=str, help='Path to the rsn_affine init params file.')
    parser.add_argument('--lite_hrnet_params', type=str, help='Path to the lite-hrnet init params file.')


    print(vars(parser.parse_args(args)))
    return parser.parse_args(args)

def main(args = None):
    # parse arguments
    if args is None:
        args = sys.argv[1:]

    args = parse_args(args)

    with open(args.yolov5_params) as f:
        yolov5_params = yaml.load(f, Loader=yaml.FullLoader)
    with open(args.rsn_affine_params) as f:
        rsn_affine_params = yaml.load(f, Loader=yaml.FullLoader)
    with open(args.lite_hrnet_params) as f:
        lite_hrnet_params = yaml.load(f, Loader=yaml.FullLoader)

    output = inference(args.img_path, yolov5_params, rsn_affine_params, lite_hrnet_params)
    preds = {'img_path': args.img_path, 'lmk_coco_body_17pts': output}
    print(preds)

if __name__ == '__main__':
    main()