# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: CC-BY-NC-4.0
import os
import yaml
import torch
import random
import argparse
import numpy as np

torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = True

VALID_COMMAND = [
    'train-oracle',
    'test-oracle',
    'train-oracle-rcnn',
    'test-oracle-rcnn',
    'train-oracle-vilbert',
    'test-oracle-vilbert',
    'train-guesser',
    'test-guesser',
    'train-guesser-vilbert',
    'test-guesser-vilbert',
    'train-qgen',
    'train-qgen-vdst',
    'train-qgen-vilbert',
    'test-self-play',
    'test-self-play-qgen-vdst',
    'test-self-play-qgen-vilbert',
    'test-self-play-qgen-vdst-oracle-vilbert',
    'test-self-play-qgen-vdst-guesser-vilbert',
    'test-self-play-qgen-vdst-oracle-vilbert-guesser-vilbert',
    'test-self-play-all-vilbert'
]


def _set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)


def run(args):
    command = args.command
    assert command in VALID_COMMAND, \
        "%s is not a valid command." % command
    _set_seed(args.seed)
    config = yaml.safe_load(open(args.config, 'r'))

    if command == 'train-oracle':
        kwargs = dict()
        mode = 'train'
        from solver.oracle import OracleSolver as Solver
    elif command == 'test-oracle':
        kwargs = dict()
        mode = 'test'
        from solver.oracle import OracleSolver as Solver
    elif command == 'train-oracle-rcnn':
        kwargs = dict()
        mode = 'train'
        from solver.oracle_rcnn import OracleSolver as Solver
    elif command == 'test-oracle-rcnn':
        kwargs = dict()
        mode = 'test'
        from solver.oracle_rcnn import OracleSolver as Solver
    elif command == 'train-oracle-vilbert':
        kwargs = dict()
        mode = 'train'
        from solver.oracle_vilbert import OracleSolver as Solver
    elif command == 'test-oracle-vilbert':
        kwargs = dict()
        mode = 'test'
        from solver.oracle_vilbert import OracleSolver as Solver
    elif command == 'train-guesser':
        kwargs = dict()
        mode = 'train'
        from solver.guesser import GuesserSolver as Solver
    elif command == 'test-guesser':
        kwargs = dict()
        mode = 'test'
        from solver.guesser import GuesserSolver as Solver
    elif command == 'train-guesser-vilbert':
        kwargs = dict()
        mode = 'train'
        from solver.guesser_vilbert import GuesserSolver as Solver
    elif command == 'test-guesser-vilbert':
        kwargs = dict()
        mode = 'test'
        from solver.guesser_vilbert import GuesserSolver as Solver
    elif command == 'train-qgen':
        kwargs = dict()
        mode = 'train'
        from solver.qgen import QGenSolver as Solver
    elif command == 'train-qgen-vdst':
        kwargs = dict()
        mode = 'train'
        from solver.qgen_vdst import QGenSolver as Solver
    elif command == 'train-qgen-vilbert':
        kwargs = dict()
        mode = 'train'
        from solver.qgen_vilbert import QGenSolver as Solver
    elif command == 'test-self-play':
        kwargs = dict()
        mode = 'test'
        from solver.self_play import SelfPlaySolver as Solver

    elif command == 'test-self-play-qgen-vdst':
        kwargs = dict()
        mode = 'test'
        from solver.self_play_qgen_vdst import SelfPlaySolver as Solver
    elif command == 'test-self-play-qgen-vilbert':
        kwargs = dict()
        mode = 'test'
        from solver.self_play_qgen_vilbert import SelfPlaySolver as Solver
    elif command == 'test-self-play-qgen-vdst-oracle-vilbert':
        kwargs = dict()
        mode = 'test'
        from solver.self_play_qgen_vdst_oracle_vilbert import SelfPlaySolver as Solver
    elif command == 'test-self-play-qgen-vdst-guesser-vilbert':
        kwargs = dict()
        mode = 'test'
        from solver.self_play_qgen_vdst_guesser_vilbert import SelfPlaySolver as Solver
    elif command == 'test-self-play-qgen-vdst-oracle-vilbert-guesser-vilbert':
        kwargs = dict()
        mode = 'test'
        from solver.self_play_qgen_vdst_oracle_vilbert_guesser_vilbert import SelfPlaySolver as Solver
    elif command == 'test-self-play-all-vilbert':
        kwargs = dict()
        mode = 'test'
        from solver.self_play_all_vilbert import SelfPlaySolver as Solver
    else:
        raise NotImplementedError

    solver = Solver(config, args, mode, **kwargs)
    solver.load_data()
    solver.set_model()
    solver.exec()
    


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description='Run the training or testing.')
    parser.add_argument('--command', required=True, type=str,
                        help='A command for the program.')
    parser.add_argument('--config', required=True, type=str,
                        help='Path to yaml config file.')
    parser.add_argument('--name', default=None, type=str,
                        help='Name for logging.')
    parser.add_argument('--logdir', default='log/', type=str,
                        help='Directory for log files.')
    parser.add_argument('--result', default='ckpt/', type=str,
                        help='Training: directory for checkpoints. Santity check: output file.')
    parser.add_argument('--load', default=None, type=str,
                        help='Load pre-trained model.', required=False)
    parser.add_argument('--seed', default=0, type=int,
                        help='Random seed for reproducible result.')
    parser.add_argument('--n-jobs', default=5, type=int,
                        help='Number workers used in torch `Dataloader`.')
    parser.add_argument('--cpu', action='store_true',
                        help='Disable GPU training.')
    parser.add_argument('--no-pin-mem', action='store_true',
                        help='Disable pin-memory for torch `Dataloader`.')
    parser.add_argument('--no-msg', action='store_true',
                        help='Hide all messages.')
    parser.add_argument("--local_rank", type=int, default=0, 
                        help="local_rank for distributed training on GPUs")

    args = parser.parse_args()
    setattr(args, 'gpu', not args.cpu)
    setattr(args, 'pin_memory', False if args.cpu else args.no_pin_mem)
    setattr(args, 'verbose', not args.no_msg)

    args.distributed = False
    if 'WORLD_SIZE' in os.environ:
        args.distributed = int(os.environ['WORLD_SIZE']) > 1
    if args.distributed:
        # FOR DISTRIBUTED:  Set the device according to local_rank.
        torch.cuda.set_device(args.local_rank)

        # FOR DISTRIBUTED:  Initialize the backend.  torch.distributed.launch will provide
        # environment variables, and requires that you use init_method=`env://`.
        torch.distributed.init_process_group(backend='nccl',
                                             init_method='env://'
                                             )
    
    if args.local_rank != 0:
        args.verbose = False

    run(args)
