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, os.path.join(root_dir, '3rdparty/customer_eva'))
sys.path.insert(0, root_dir)
import argparse
import torch
import importlib
import cv2
import glob
import pickle
import numpy as np
from mqbench.prepare_by_platform import prepare_by_platform, get_qconfig_by_platform
from mqbench.utils.state import enable_quantization, enable_calibration
from mediatools.convert_model import export_mnn, convert_fixpoint_simulator
from mediatools.deploy.deploy_6920_VNNE import dump_to_sa6920_VNNE
from yolox_microbt.utils.utils import load_checkpoint
from yolox_microbt.data.evaluation.customer_evaluator import MAPEvaluator


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("--input_shape", type=str, default=None,
                        help="format: 1,c,h,w")
    parser.add_argument("--eval_type", type=str, default=None,
                        help="type of eval")
    parser.add_argument("--image_dir", type=str, default=None,
                        help="path to put image file")
    parser.add_argument("--label_path", type=str, default=None,
                        help="path to put label file")
    parser.add_argument("--save_path", type=str, default=None,
                        help="path to save pred file")
    parser.add_argument("--fixpoint", default=False, action="store_true",
                        help="fixpoint forward mode or fakequant")
    args = parser.parse_args()
    return args

def ToNv12(image, code):
    I420 = cv2.cvtColor(image, code)
    height, width = image.shape[:2]

    U = I420[height:height+height//4, :]
    U = U.reshape((1, height//4*width))
    V = I420[height+height//4:height+height//2, :]
    V = V.reshape((1, height//4*width))
    UV = np.zeros((1, height//4*width*2))
    UV[:, 0::2] = U
    UV[:, 1::2] = V
    UV = UV.reshape((height//2, width))
    NV12 = np.zeros((height+height//2, width))
    NV12[0:height, :] = I420[0:height, :]
    NV12[height:, :] = UV
    return NV12

def nv12toyuv444(nv12_data, height, width):
    nv12_data = nv12_data.reshape((height+height//2, width))
    y = nv12_data[0:height, :]
    u = nv12_data[height:height+height//2, 0::2]
    u = u.repeat(2, axis=0).repeat(2, axis=1)
    v = nv12_data[height:height+height//2, 1::2]
    v = v.repeat(2, axis=0).repeat(2, axis=1)
    yuv444 = np.zeros((height, width, 3))
    yuv444[:, :, 0] = y
    yuv444[:, :, 1] = u
    yuv444[:, :, 2] = v
    return yuv444


def preprocess(image_path, mean=[0, 0, 0], std=[255, 255, 255], input_size=None, resize_mode='short', ori_image_shape=None):
    interp = cv2.INTER_LINEAR
    if image_path.endswith('yuv420'):
        assert ori_image_shape is not None
        image_height = int(ori_image_shape[1] * 1.5)
        image_width = ori_image_shape[0]
        image_channel = 1
        image_shape = [image_height, image_width, image_channel]
        image = np.fromfile(image_path, dtype=np.uint8).reshape(image_shape)
        image = cv2.cvtColor(image, cv2.COLOR_YUV2BGR_I420)
    else:
        image = cv2.imread(image_path)
    h, w, c = image.shape
    if resize_mode == 'short':
        max_size = np.max(input_size)
        scale = np.min(input_size) / np.min((h, w))
        img_size_max = np.max((h, w))
        if 0 < max_size < np.round(scale * img_size_max):
            scale = max_size / img_size_max
        ratio = np.array([1/scale, 1/scale], dtype=np.float32)
        img = cv2.resize(image, None, None, fx=scale, fy=scale, interpolation=interp)
    else:
        img = cv2.resize(image, input_size, interpolation=interp)
        ratio = np.array([w / input_size[0], h / input_size[1]], dtype=np.float32)
    if len(image.shape) == 3:
        padded_img = np.zeros((input_size[0], input_size[1], 3), dtype=np.uint8)
    else:
        padded_img = np.zeros(input_size, dtype=np.uint8)
    padded_img[:img.shape[0], :img.shape[1]] = img
    image_height, image_width, _ = padded_img.shape
    img = cv2.cvtColor(padded_img, cv2.COLOR_BGR2RGB)
    img = ToNv12(img, cv2.COLOR_RGB2YUV_I420).astype(np.uint8)
    yuv444 = nv12toyuv444(img, image_height, image_width).astype(np.float32)
    # normalize
    mean = np.array(mean, dtype=np.float32)
    std = np.array(std, dtype=np.float32)
    yuv444 = (yuv444 - mean) / std
    yuv444 = yuv444.astype(np.float32)
    yuv444 = np.transpose(yuv444, (2, 0, 1))
    yuv444 = np.expand_dims(yuv444, axis=0)
    yuv_tensor = torch.from_numpy(yuv444)
    return yuv_tensor, ratio


def vis_pred(img_path, boxes, scores, cls_ids, score_thresh=0.5):
    img = cv2.imread(img_path)
    for i in range(len(boxes)):
        box = boxes[i]
        cls_id = int(cls_ids[i])
        score = scores[i]
        if score < score_thresh:
            continue
        x0 = int(box[0])
        y0 = int(box[1])
        x1 = int(box[2])
        y1 = int(box[3])

        text = '{}:{:.1f}%'.format(cls_id, score * 100)
        font = cv2.FONT_HERSHEY_SIMPLEX
        txt_size = cv2.getTextSize(text, font, 0.4, 1)[0]
        cv2.rectangle(img, (x0, y0), (x1, y1), (0, 0, 255), 2)
        cv2.rectangle(img, (x0, y0 + 1), (x0 + txt_size[0] + 1, y0 + int(1.5 * txt_size[1])), (255, 255, 255), -1)
        cv2.putText(img, text, (x0, y0 + txt_size[1]), font, 0.4, (0, 255, 0), thickness=1)
    return img


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
    post_process = exp.get_postprocess()
    enable_quantization(quant_model)
    quant_model.eval()
    if args.fixpoint:
        quant_model = convert_fixpoint_simulator(quant_model, config['quant_config']['deploy_backend'], fixpoint_params=config['fixpoint_params'])
    quant_model.cuda(0)
    input_shape = [int(x) for x in args.input_shape.split(',')]
    label_info = pickle.load(open(args.label_path, 'rb'))
    pred_roidbs = []
    for label_roi in label_info:
        image_path = os.path.join(args.image_dir, label_roi['image_name'])
        input_tensor, image_scales = preprocess(image_path, input_size=[input_shape[3], input_shape[2]], 
                                                ori_image_shape=[label_roi['image_width'], label_roi['image_height']])
        input_tensor = input_tensor.cuda(0)
        outs = quant_model(input_tensor)
        pred_boxes, pred_classes, pred_scores = post_process(outs)
        pred_boxes[:, 0::2] *= image_scales[0]
        pred_boxes[:, 1::2] *= image_scales[1]
        pred_roi = {
            'image_name': label_roi['image_name'],
            'pred_boxes': pred_boxes.numpy(),
            'pred_classes': pred_classes.numpy(),
            'pred_scores': pred_scores.numpy()
        }
        pred_roidbs.append(pred_roi)
    with open(args.save_path, 'wb') as fid:
        pickle.dump(pred_roidbs, fid, 4)

    map_Evaluator = MAPEvaluator(iou_threshs= [0.5], label_names=['_', '_', 'person', '_', '_', '_', '_', 'pet']) 
    results = map_Evaluator.process_results(pred_roidbs)
    recall, _, ap = map_Evaluator.eval_all_classes(label_info, results)
    print('%s: ap /ar:'%(args.eval_type), ap[0][args.eval_type], '/', recall[0][args.eval_type])

