import argparse
import os.path as osp
import warnings
import numbers
from functools import partial
import sys
sys.path.append("/home/shanwei-luo/userdata/mmdetection")
import numpy as np
#import onnx
import torch
import cv2
from mmcv import Config, DictAction

from mmdet.core.export.model_wrappers import ONNXRuntimeDetector

from mmdet.apis import (async_inference_detector, inference_detector,
                        init_detector, show_result_pyplot)

import onnxruntime as ort
import onnx
import mmcv

def impad(img,
          *,
          shape=None,
          padding=None,
          pad_val=0,
          padding_mode='constant'):
    assert (shape is not None) ^ (padding is not None)
    if shape is not None:
        padding = (0, 0, shape[1] - img.shape[1], shape[0] - img.shape[0])

    # check pad_val
    if isinstance(pad_val, tuple):
        assert len(pad_val) == img.shape[-1]
    elif not isinstance(pad_val, numbers.Number):
        raise TypeError('pad_val must be a int or a tuple. '
                        f'But received {type(pad_val)}')

    # check padding
    if isinstance(padding, tuple) and len(padding) in [2, 4]:
        if len(padding) == 2:
            padding = (padding[0], padding[1], padding[0], padding[1])
    elif isinstance(padding, numbers.Number):
        padding = (padding, padding, padding, padding)
    else:
        raise ValueError('Padding must be a int or a 2, or 4 element tuple.'
                         f'But received {padding}')

    # check padding mode
    assert padding_mode in ['constant', 'edge', 'reflect', 'symmetric']

    border_type = {
        'constant': cv2.BORDER_CONSTANT,
        'edge': cv2.BORDER_REPLICATE,
        'reflect': cv2.BORDER_REFLECT_101,
        'symmetric': cv2.BORDER_REFLECT
    }
    img = cv2.copyMakeBorder(
        img,
        padding[1],
        padding[3],
        padding[0],
        padding[2],
        border_type[padding_mode],
        value=pad_val)

    return img
    
def imnormalize(img, mean, std, to_rgb=True):
    mean = np.float64(mean.reshape(1, -1))
    stdinv = 1 / np.float64(std.reshape(1, -1))
    if to_rgb:
        cv2.cvtColor(img, cv2.COLOR_BGR2RGB, img)  # inplace
    cv2.subtract(img, mean, img)  # inplace
    cv2.multiply(img, stdinv, img)  # inplace
    return img

def preprocess_example_input(input_config):
    input_path = input_config['input_path']
    image = cv2.imread(input_path, cv2.IMREAD_COLOR)
    h, w, _ = image.shape
    img_scale = (400, 400) 
    max_long_edge = max(img_scale)
    max_short_edge = min(img_scale)
    scale_factor = min(max_long_edge / max(h, w), max_short_edge / min(h, w))
    scale_w = int(w * float(scale_factor) + 0.5)
    scale_h = int(h * float(scale_factor) + 0.5)
    image = cv2.resize(image, (scale_w, scale_h))

    image = np.asarray(image).astype(np.float32)
    mean = np.array([123.675, 116.28, 103.53], dtype=np.float32)
    std = np.array([58.395, 57.12, 57.375], dtype=np.float32)
    image = imnormalize(image, mean, std)

    divisor = 32
    pad_h = int(np.ceil(image.shape[0] / divisor)) * divisor
    pad_w = int(np.ceil(image.shape[1] / divisor)) * divisor
    image = impad(image, shape=(pad_h, pad_w), pad_val=0)
    image = np.asarray(image).astype(np.float32)

    image = np.transpose(image, [2, 0, 1])
    one_img = torch.from_numpy(image).unsqueeze(0).float().requires_grad_(
        True)
    (_, C, H, W) = input_config['input_shape']
    one_meta = {
        'img_shape': (H, W, C),
        'ori_shape': (H, W, C),
        'pad_shape': (H, W, C),
        'filename': '<demo>.png',
        'scale_factor': np.ones(4, dtype=np.float32),
        'flip': False,
        'show_img': image,
        'flip_direction': None
    }

    return one_img, one_meta
print(f"onnxruntime device: {ort.get_device()}") # output: GPU
print(f'ort avail providers: {ort.get_available_providers()}') # output: ['CUDAExecutionProvider', 'CPUExecutionProvider']
output_file = "checkpoints/AD_dsxw_atts_20220318.onnx"
classes_name = ['yiwei','loujian','celi','libei','fantie','lianxi','duojian','shunjian','shaoxi','jiahan','yiwu']
onnx_model = onnx.load(output_file)
onnx.checker.check_model(onnx_model)
onnx_model = ONNXRuntimeDetector(output_file, classes_name, 0)
input_shape = (1, 3, 400, 400)
input_img = "/home/shanwei-luo/userdata/datasets/dsxw_dataset_v4/dsxw_test/images/21000204.jpg"
normalize_cfg = {'mean': (123.675, 116.28, 103.53),'std': (58.395, 57.12, 57.375)}
input_config = {
        'input_shape': input_shape,
        'input_path': input_img,
        'normalize_cfg': normalize_cfg
    }

input_config['input_shape'] = (1, 3, 416, 416)

# prepare input once again
one_img, one_meta = preprocess_example_input(input_config)
print(one_img)
img_list, img_meta_list = [one_img], [[one_meta]]
img_list = [_.cuda().contiguous() for _ in img_list]
onnx_results = onnx_model(
            img_list, img_metas=img_meta_list, return_loss=False)[0]
print(onnx_results)

config_file_1 = '/home/shanwei-luo/userdata/mmdetection/work_dirs/AD_dsxw_test61/AD_dsxw_test61.py'
checkpoint_file_1 = '/home/shanwei-luo/userdata/mmdetection/work_dirs/AD_dsxw_test61/epoch_36.pth'
model_1 = init_detector(config_file_1, checkpoint_file_1, device='cuda:1')
results_1_tmp = inference_detector(model_1, [input_img])
print(results_1_tmp)

for i in range(len(onnx_results)):
    print(onnx_results[i].shape)
for i in range(len(results_1_tmp[0])):
    print(len(results_1_tmp[0][i]))