# Copyright (c) OpenMMLab. All rights reserved.
import argparse
import os
import os.path as osp

from mmcv import Config


def parse_args():
    parser = argparse.ArgumentParser(
        description='Convert benchmark model list to script')
    parser.add_argument('config', help='test config file path')
    parser.add_argument('--port', type=int, default=29666, help='dist port')
    parser.add_argument(
        '--work-dir',
        default='tools/batch_test',
        help='the dir to save metric')
    parser.add_argument(
        '--run', action='store_true', help='run script directly')
    parser.add_argument(
        '--out', type=str, 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
    work_dir = osp.join(work_dir, fname)
    checkpoint = model_info['checkpoint'].strip()
    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,
        work_dir=work_dir,
        checkpoint=checkpoint,
        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=8  GPUS_PER_NODE=8  ' \
                   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'--work-dir {work_dir} '

    command_info += f'--eval {eval} '
    command_info += f'--cfg-option 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}'
    assert args.out or args.run, \
        ('Please specify at least one operation (save/run/ the '
         'script) with the argument "--out" or "--run"')

    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)
    if args.run:
        os.system(command_str)


if __name__ == '__main__':
    main()