import os
import sys
root_dir = os.path.abspath(__file__).split('/scripts')[0]
sys.path.insert(0, os.path.join(root_dir, 'model'))
sys.path.insert(0, os.path.join(root_dir, 'configs'))
sys.path.insert(0, os.path.join(root_dir, '3rdparty'))
sys.path.insert(0, os.path.join(root_dir, '3rdparty/YOLOX'))
sys.path.insert(0, root_dir)
import argparse
import torch
import importlib
from mqbench.prepare_by_platform import prepare_by_platform, get_qconfig_by_platform
from mediatools.convert_model import export_mnn, convert_fixpoint_simulator, convert_onnx_deploy
from mediatools.deploy.deploy_6920_VNNE import dump_to_sa6920_VNNE
from yolox_microbt.utils.utils import load_checkpoint


def parse_args():
    parser = argparse.ArgumentParser(description="script of exporting model for SA6920/SA6921")
    parser.add_argument("--config", type=str, default=None,
                        help="path to config file")
    parser.add_argument("--ckpt", type=str, default=None,
                        help="path to model checkpoint file")
    parser.add_argument("--hardware", type=str, default=None,
                        help="support one of [SA6920vnne, SA6921vnne]")
    parser.add_argument("--save_path", type=str, default=None,
                        help="path to save model file")
    parser.add_argument("--input_shape", type=str, default=None,
                        help="format: 1,c,h,w")
    args = parser.parse_args()
    return args


if __name__ == "__main__":
    args = parse_args()
    if args.config.endswith('.py'):
        args.config = args.config[:-3]
    model_config = importlib.import_module(args.config)
    config = getattr(model_config, 'config')
    exp = getattr(model_config, 'Exp')()
    quant_model = exp.get_model()
    quant_model = load_checkpoint(quant_model, args.ckpt)
    quant_model = quant_model.yolox
    quant_model.eval()
    input_shape = [int(x) for x in args.input_shape.split(',')]
    if args.hardware == "SA6920vnne":
        fixpoint_model = convert_fixpoint_simulator(quant_model, config['quant_config']['deploy_backend'], fixpoint_params=config['fixpoint_params'])
        dump_to_sa6920_VNNE(fixpoint_model, args.save_path, exportor_params=config['exportor_params'])
    elif args.hardware == "SA6921vnne":
        dummy_input = torch.rand(input_shape, dtype=torch.float, device='cuda')
        export_mnn(quant_model, config['quant_config']['deploy_backend'], 
                fixpoint_params=config['fixpoint_params'], dummy_input=dummy_input, 
                model_name=args.save_path, SA6921vnne=True, model_type=config['model_type'])
    elif args.hardware == "SA6920NPU":
        dummy_input = torch.rand(input_shape, dtype=torch.float, device='cuda')
        convert_onnx_deploy(quant_model, config['quant_config']['deploy_backend'], dummy_input=dummy_input, 
                model_name=args.save_path, model_type=config['model_type'])