# Copyright (c) OpenMMLab. All rights reserved.

import argparse

import os.path as osp

from mmcv import Config

def parse_args():

    parser = argparse.ArgumentParser(

        description='Convert benchmark test model list to script')

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

    parser.add_argument('--port', type=int, default=28171, help='dist port')

    parser.add_argument(

        '--work-dir',

        default='work_dirs/benchmark_evaluation',

        help='the dir to save metric')

    parser.add_argument(

        '--out',

        type=str,

        default='.dev/benchmark_evaluation.sh',

        help='path to save model benchmark script')

    args = parser.parse_args()

    return args

def process_model_info(model_info, work_dir):

    config = model_info['config'].strip()

    fname, _ = osp.splitext(osp.basename(config))

    job_name = fname

    checkpoint = model_info['checkpoint'].strip()

    work_dir = osp.join(work_dir, fname)

    if not isinstance(model_info['eval'], list):

        evals = [model_info['eval']]

    else:

        evals = model_info['eval']

    eval = ' '.join(evals)

    return dict(

        config=config,

        job_name=job_name,

        checkpoint=checkpoint,

        work_dir=work_dir,

        eval=eval)

def create_test_bash_info(commands, model_test_dict, port, script_name,

                          partition):

    config = model_test_dict['config']

    job_name = model_test_dict['job_name']

    checkpoint = model_test_dict['checkpoint']

    work_dir = model_test_dict['work_dir']

    eval = model_test_dict['eval']

    echo_info = f'\necho \'{config}\' &'

    commands.append(echo_info)

    commands.append('\n')

    command_info = f'GPUS=4  GPUS_PER_NODE=4  ' \

                   f'CPUS_PER_TASK=2 {script_name} '

    command_info += f'{partition} '

    command_info += f'{job_name} '

    command_info += f'{config} '

    command_info += f'$CHECKPOINT_DIR/{checkpoint} '

    command_info += f'--eval {eval} '

    command_info += f'--work-dir {work_dir} '

    command_info += f'--cfg-options dist_params.port={port} '

    command_info += '&'

    commands.append(command_info)

def main():

    args = parse_args()

    if args.out:

        out_suffix = args.out.split('.')[-1]

        assert args.out.endswith('.sh'), \

            f'Expected out file path suffix is .sh, but get .{out_suffix}'

    commands = []

    partition_name = 'PARTITION=$1'

    commands.append(partition_name)

    commands.append('\n')

    checkpoint_root = 'CHECKPOINT_DIR=$2'

    commands.append(checkpoint_root)

    commands.append('\n')

    script_name = osp.join('tools', 'slurm_test.sh')

    port = args.port

    work_dir = args.work_dir

    cfg = Config.fromfile(args.config)

    for model_key in cfg:

        model_infos = cfg[model_key]

        if not isinstance(model_infos, list):

            model_infos = [model_infos]

        for model_info in model_infos:

            print('processing: ', model_info['config'])

            model_test_dict = process_model_info(model_info, work_dir)

            create_test_bash_info(commands, model_test_dict, port, script_name,

                                  '$PARTITION')

            port += 1

    command_str = ''.join(commands)

    if args.out:

        with open(args.out, 'w') as f:

            f.write(command_str + '\n')

if __name__ == '__main__':

    main()