# 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.parallel import MMDataParallel, MMDistributedDataParallel

from mmcv.runner import get_dist_info, init_dist, load_checkpoint

from mmpose.apis import multi_gpu_test, single_gpu_test

from mmpose.datasets import build_dataloader, build_dataset

from mmpose.models import build_posenet

from mmpose.utils import setup_multi_processes

try:

    from mmcv.runner import wrap_fp16_model

except ImportError:

    warnings.warn('auto_fp16 from mmpose will be deprecated from v0.15.0'

                  'Please install mmcv>=1.1.4')

    from mmpose.core import wrap_fp16_model

def parse_args():

    parser = argparse.ArgumentParser(description='mmpose test model')

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

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

    parser.add_argument('--out', help='output result file')

    parser.add_argument(

        '--work-dir', help='the dir to save evaluation results')

    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(

        '--gpu-id',

        type=int,

        default=0,

        help='id of gpu to use '

        '(only applicable to non-distributed testing)')

    parser.add_argument(

        '--eval',

        default=None,

        nargs='+',

        help='evaluation metric, which depends on the dataset,'

        ' e.g., "mAP" for MSCOCO')

    parser.add_argument(

        '--gpu_collect',

        action='store_true',

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

    parser.add_argument('--tmpdir', help='tmp dir for writing some results')

    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(

        '--launcher',

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

        default='none',

        help='job launcher')

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

    args = parser.parse_args()

    if 'LOCAL_RANK' not in os.environ:

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

    return args

def merge_configs(cfg1, cfg2):

    # Merge cfg2 into cfg1

    # Overwrite cfg1 if repeated, ignore if value is None.

    cfg1 = {} if cfg1 is None else cfg1.copy()

    cfg2 = {} if cfg2 is None else cfg2

    for k, v in cfg2.items():

        if v:

            cfg1[k] = v

    return cfg1

def main():

    args = parse_args()

    cfg = Config.fromfile(args.config)

    if args.cfg_options is not None:

        cfg.merge_from_dict(args.cfg_options)

    # set multi-process settings

    setup_multi_processes(cfg)

    # set cudnn_benchmark

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

        torch.backends.cudnn.benchmark = True

    cfg.model.pretrained = None

    cfg.data.test.test_mode = True

    # work_dir is determined in this priority: CLI > segment in file > filename

    if args.work_dir is not None:

        # update configs according to CLI args if args.work_dir is not None

        cfg.work_dir = args.work_dir

    elif cfg.get('work_dir', None) is None:

        # use config filename as default work_dir if cfg.work_dir is None

        cfg.work_dir = osp.join('./work_dirs',

                                osp.splitext(osp.basename(args.config))[0])

    mmcv.mkdir_or_exist(osp.abspath(cfg.work_dir))

    # 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)

    # build the dataloader

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

    # step 1: give default values and override (if exist) from cfg.data

    loader_cfg = {

        **dict(seed=cfg.get('seed'), drop_last=False, dist=distributed),

        **({} if torch.__version__ != 'parrots' else dict(

               prefetch_num=2,

               pin_memory=False,

           )),

        **dict((k, cfg.data[k]) for k in [

                   'seed',

                   'prefetch_num',

                   'pin_memory',

                   'persistent_workers',

               ] if k in cfg.data)

    }

    # step2: cfg.data.test_dataloader has higher priority

    test_loader_cfg = {

        **loader_cfg,

        **dict(shuffle=False, drop_last=False),

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

        **dict(samples_per_gpu=cfg.data.get('samples_per_gpu', 1)),

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

    }

    data_loader = build_dataloader(dataset, **test_loader_cfg)

    # build the model and load checkpoint

    model = build_posenet(cfg.model)

    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=[args.gpu_id])

        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)

    rank, _ = get_dist_info()

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

    eval_config = merge_configs(eval_config, dict(metric=args.eval))

    if rank == 0:

        if args.out:

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

            mmcv.dump(outputs, args.out)

        results = dataset.evaluate(outputs, cfg.work_dir, **eval_config)

        for k, v in sorted(results.items()):

            print(f'{k}: {v}')

if __name__ == '__main__':

    main()