# Copyright (c) OpenMMLab. All rights reserved.

import argparse

import os

import os.path as osp

import warnings

import mmcv

import torch

from mmcv import Config, DictAction

from mmcv.cnn import fuse_conv_bn

from mmcv.fileio.io import file_handlers

from mmcv.parallel import MMDataParallel, MMDistributedDataParallel

from mmcv.runner import get_dist_info, init_dist, load_checkpoint

from mmcv.runner.fp16_utils import wrap_fp16_model

from mmaction.datasets import build_dataloader, build_dataset

from mmaction.models import build_model

from mmaction.utils import register_module_hooks

# TODO import test functions from mmcv and delete them from mmaction2

try:

    from mmcv.engine import multi_gpu_test, single_gpu_test

except (ImportError, ModuleNotFoundError):

    warnings.warn(

        'DeprecationWarning: single_gpu_test, multi_gpu_test, '

        'collect_results_cpu, collect_results_gpu from mmaction2 will be '

        'deprecated. Please install mmcv through master branch.')

    from mmaction.apis import multi_gpu_test, single_gpu_test

def parse_args():

    parser = argparse.ArgumentParser(

        description='MMAction2 test (and eval) a model')

    parser.add_argument('config', help='test config file path')

    parser.add_argument('checkpoint', help='checkpoint file')

    parser.add_argument(

        '--out',

        default=None,

        help='output result file in pkl/yaml/json format')

    parser.add_argument(

        '--fuse-conv-bn',

        action='store_true',

        help='Whether to fuse conv and bn, this will slightly increase'

        'the inference speed')

    parser.add_argument(

        '--eval',

        type=str,

        nargs='+',

        help='evaluation metrics, which depends on the dataset, e.g.,'

        ' "top_k_accuracy", "mean_class_accuracy" for video dataset')

    parser.add_argument(

        '--gpu-collect',

        action='store_true',

        help='whether to use gpu to collect results')

    parser.add_argument(

        '--tmpdir',

        help='tmp directory used for collecting results from multiple '

        'workers, available when gpu-collect is not specified')

    parser.add_argument(

        '--options',

        nargs='+',

        action=DictAction,

        default={},

        help='custom options for evaluation, the key-value pair in xxx=yyy '

        'format will be kwargs for dataset.evaluate() function (deprecate), '

        'change to --eval-options instead.')

    parser.add_argument(

        '--eval-options',

        nargs='+',

        action=DictAction,

        default={},

        help='custom options for evaluation, the key-value pair in xxx=yyy '

        'format will be kwargs for dataset.evaluate() function')

    parser.add_argument(

        '--cfg-options',

        nargs='+',

        action=DictAction,

        default={},

        help='override some settings in the used config, the key-value pair '

        'in xxx=yyy format will be merged into config file. For example, '

        "'--cfg-options model.backbone.depth=18 model.backbone.with_cp=True'")

    parser.add_argument(

        '--average-clips',

        choices=['score', 'prob', None],

        default=None,

        help='average type when averaging test clips')

    parser.add_argument(

        '--launcher',

        choices=['none', 'pytorch', 'slurm', 'mpi'],

        default='none',

        help='job launcher')

    parser.add_argument('--local_rank', type=int, default=0)

    parser.add_argument(

        '--onnx',

        action='store_true',

        help='Whether to test with onnx model or not')

    parser.add_argument(

        '--tensorrt',

        action='store_true',

        help='Whether to test with TensorRT engine or not')

    args = parser.parse_args()

    if 'LOCAL_RANK' not in os.environ:

        os.environ['LOCAL_RANK'] = str(args.local_rank)

    if args.options and args.eval_options:

        raise ValueError(

            '--options and --eval-options cannot be both '

            'specified, --options is deprecated in favor of --eval-options')

    if args.options:

        warnings.warn('--options is deprecated in favor of --eval-options')

        args.eval_options = args.options

    return args

def turn_off_pretrained(cfg):

    # recursively find all pretrained in the model config,

    # and set them None to avoid redundant pretrain steps for testing

    if 'pretrained' in cfg:

        cfg.pretrained = None

    # recursively turn off pretrained value

    for sub_cfg in cfg.values():

        if isinstance(sub_cfg, dict):

            turn_off_pretrained(sub_cfg)

def inference_pytorch(args, cfg, distributed, data_loader):

    """Get predictions by pytorch models."""

    if args.average_clips is not None:

        # You can set average_clips during testing, it will override the

        # original setting

        if cfg.model.get('test_cfg') is None and cfg.get('test_cfg') is None:

            cfg.model.setdefault('test_cfg',

                                 dict(average_clips=args.average_clips))

        else:

            if cfg.model.get('test_cfg') is not None:

                cfg.model.test_cfg.average_clips = args.average_clips

            else:

                cfg.test_cfg.average_clips = args.average_clips

    # remove redundant pretrain steps for testing

    turn_off_pretrained(cfg.model)

    # build the model and load checkpoint

    model = build_model(

        cfg.model, train_cfg=None, test_cfg=cfg.get('test_cfg'))

    if len(cfg.module_hooks) > 0:

        register_module_hooks(model, cfg.module_hooks)

    fp16_cfg = cfg.get('fp16', None)

    if fp16_cfg is not None:

        wrap_fp16_model(model)

    load_checkpoint(model, args.checkpoint, map_location='cpu')

    if args.fuse_conv_bn:

        model = fuse_conv_bn(model)

    if not distributed:

        model = MMDataParallel(model, device_ids=[0])

        outputs = single_gpu_test(model, data_loader)

    else:

        model = MMDistributedDataParallel(

            model.cuda(),

            device_ids=[torch.cuda.current_device()],

            broadcast_buffers=False)

        outputs = multi_gpu_test(model, data_loader, args.tmpdir,

                                 args.gpu_collect)

    return outputs

def inference_tensorrt(ckpt_path, distributed, data_loader, batch_size):

    """Get predictions by TensorRT engine.

    For now, multi-gpu mode and dynamic tensor shape are not supported.

    """

    assert not distributed, \

        'TensorRT engine inference only supports single gpu mode.'

    import tensorrt as trt

    from mmcv.tensorrt.tensorrt_utils import (torch_dtype_from_trt,

                                              torch_device_from_trt)

    # load engine

    with trt.Logger() as logger, trt.Runtime(logger) as runtime:

        with open(ckpt_path, mode='rb') as f:

            engine_bytes = f.read()

        engine = runtime.deserialize_cuda_engine(engine_bytes)

    # For now, only support fixed input tensor

    cur_batch_size = engine.get_binding_shape(0)[0]

    assert batch_size == cur_batch_size, \

        ('Dataset and TensorRT model should share the same batch size, '

         f'but get {batch_size} and {cur_batch_size}')

    context = engine.create_execution_context()

    # get output tensor

    dtype = torch_dtype_from_trt(engine.get_binding_dtype(1))

    shape = tuple(context.get_binding_shape(1))

    device = torch_device_from_trt(engine.get_location(1))

    output = torch.empty(

        size=shape, dtype=dtype, device=device, requires_grad=False)

    # get predictions

    results = []

    dataset = data_loader.dataset

    prog_bar = mmcv.ProgressBar(len(dataset))

    for data in data_loader:

        bindings = [

            data['imgs'].contiguous().data_ptr(),

            output.contiguous().data_ptr()

        ]

        context.execute_async_v2(bindings,

                                 torch.cuda.current_stream().cuda_stream)

        results.extend(output.cpu().numpy())

        batch_size = len(next(iter(data.values())))

        for _ in range(batch_size):

            prog_bar.update()

    return results

def inference_onnx(ckpt_path, distributed, data_loader, batch_size):

    """Get predictions by ONNX.

    For now, multi-gpu mode and dynamic tensor shape are not supported.

    """

    assert not distributed, 'ONNX inference only supports single gpu mode.'

    import onnx

    import onnxruntime as rt

    # get input tensor name

    onnx_model = onnx.load(ckpt_path)

    input_all = [node.name for node in onnx_model.graph.input]

    input_initializer = [node.name for node in onnx_model.graph.initializer]

    net_feed_input = list(set(input_all) - set(input_initializer))

    assert len(net_feed_input) == 1

    # For now, only support fixed tensor shape

    input_tensor = None

    for tensor in onnx_model.graph.input:

        if tensor.name == net_feed_input[0]:

            input_tensor = tensor

            break

    cur_batch_size = input_tensor.type.tensor_type.shape.dim[0].dim_value

    assert batch_size == cur_batch_size, \

        ('Dataset and ONNX model should share the same batch size, '

         f'but get {batch_size} and {cur_batch_size}')

    # get predictions

    sess = rt.InferenceSession(ckpt_path)

    results = []

    dataset = data_loader.dataset

    prog_bar = mmcv.ProgressBar(len(dataset))

    for data in data_loader:

        imgs = data['imgs'].cpu().numpy()

        onnx_result = sess.run(None, {net_feed_input[0]: imgs})[0]

        results.extend(onnx_result)

        batch_size = len(next(iter(data.values())))

        for _ in range(batch_size):

            prog_bar.update()

    return results

def main():

    args = parse_args()

    if args.tensorrt and args.onnx:

        raise ValueError(

            'Cannot set onnx mode and tensorrt mode at the same time.')

    cfg = Config.fromfile(args.config)

    cfg.merge_from_dict(args.cfg_options)

    # Load output_config from cfg

    output_config = cfg.get('output_config', {})

    if args.out:

        # Overwrite output_config from args.out

        output_config = Config._merge_a_into_b(

            dict(out=args.out), output_config)

    # Load eval_config from cfg

    eval_config = cfg.get('eval_config', {})

    if args.eval:

        # Overwrite eval_config from args.eval

        eval_config = Config._merge_a_into_b(

            dict(metrics=args.eval), eval_config)

    if args.eval_options:

        # Add options from args.eval_options

        eval_config = Config._merge_a_into_b(args.eval_options, eval_config)

    assert output_config or eval_config, \

        ('Please specify at least one operation (save or eval the '

         'results) with the argument "--out" or "--eval"')

    dataset_type = cfg.data.test.type

    if output_config.get('out', None):

        if 'output_format' in output_config:

            # ugly workround to make recognition and localization the same

            warnings.warn(

                'Skip checking `output_format` in localization task.')

        else:

            out = output_config['out']

            # make sure the dirname of the output path exists

            mmcv.mkdir_or_exist(osp.dirname(out))

            _, suffix = osp.splitext(out)

            if dataset_type == 'AVADataset':

                assert suffix[1:] == 'csv', ('For AVADataset, the format of '

                                             'the output file should be csv')

            else:

                assert suffix[1:] in file_handlers, (

                    'The format of the output '

                    'file should be json, pickle or yaml')

    # set cudnn benchmark

    if cfg.get('cudnn_benchmark', False):

        torch.backends.cudnn.benchmark = True

    cfg.data.test.test_mode = True

    # init distributed env first, since logger depends on the dist info.

    if args.launcher == 'none':

        distributed = False

    else:

        distributed = True

        init_dist(args.launcher, **cfg.dist_params)

    # The flag is used to register module's hooks

    cfg.setdefault('module_hooks', [])

    # build the dataloader

    dataset = build_dataset(cfg.data.test, dict(test_mode=True))

    dataloader_setting = dict(

        videos_per_gpu=cfg.data.get('videos_per_gpu', 1),

        workers_per_gpu=cfg.data.get('workers_per_gpu', 1),

        dist=distributed,

        shuffle=False)

    dataloader_setting = dict(dataloader_setting,

                              **cfg.data.get('test_dataloader', {}))

    data_loader = build_dataloader(dataset, **dataloader_setting)

    if args.tensorrt:

        outputs = inference_tensorrt(args.checkpoint, distributed, data_loader,

                                     dataloader_setting['videos_per_gpu'])

    elif args.onnx:

        outputs = inference_onnx(args.checkpoint, distributed, data_loader,

                                 dataloader_setting['videos_per_gpu'])

    else:

        outputs = inference_pytorch(args, cfg, distributed, data_loader)

    rank, _ = get_dist_info()

    if rank == 0:

        if output_config.get('out', None):

            out = output_config['out']

            print(f'\nwriting results to {out}')

            dataset.dump_results(outputs, **output_config)

        if eval_config:

            eval_res = dataset.evaluate(outputs, **eval_config)

            for name, val in eval_res.items():

                print(f'{name}: {val:.04f}')

if __name__ == '__main__':

    main()