import os
import sys
import random
import argparse
from datetime import datetime
import logging
import shutil
import time

import torch
import torch.multiprocessing as mp
import torch.optim as optim

from baselines.a3c import policies, a3c
from baselines.common.plot_process import get_data_from_files
from baselines.common import set_global_seeds
from DL_Logger.ResultsLog import export_args, setup_logging, ResultsLog


def main(args):
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter,
                                     description='A2C MountainCar training')

    parser.add_argument('--env_id', default='CartPole-v0',
                        help='Gym env_id')
    parser.add_argument('--seed', help='RNG seed', type=int, default=0)
    parser.add_argument('--num_workers', help='Number of processes to act simultaneously',
                        type=int, default=4)
    parser.add_argument('--architecture', help='Policy architecture',
                        choices=['cnn', 'lstm', 'lnlstm'], default='cnn')
    parser.add_argument('--resume', help='Whether training is resumed',
                        dest='resume', action='store_true')
    parser.add_argument('--open-browser', help='Whether to open browser by the script',
                        dest='open_browser', action='store_true', default=True)

    parser.add_argument('--training_steps', help='How many frames to train'
                                                 'This number gets divided by 4 due to frameskip',
                        type=int, default=40e6)
    parser.add_argument('--gamma', help='discount factor',
                        type=float, default=0.99)
    parser.add_argument('--vf_coef', help='Value function''s coefficient',
                        type=float, default=0.5)
    parser.add_argument('--ent_coef', help='Entropy regularization coefficient',
                        type=float, default=0.01)
    parser.add_argument('--num_steps_update', help='Number of steps to take before updating the network',
                        type=int, default=5)
    parser.add_argument('--max_episode_len', help='maximum number of steps per episode',
                        type=int, default=100000)
    parser.add_argument('--recurrent', help='use a recurrent network',
                        dest='recurrent', action='store_true', default=False)

    optim_parser = parser.add_argument_group('Optimization Parameters')
    optim_parser.add_argument('--optimizer', default='rmsprop', choices=['rmsprop', 'adam', 'sgd'],
                              help='optimizer for training')
    optim_parser.add_argument('--lr', type=float, default=7e-4,
                              help='learning rate')
    optim_parser.add_argument('--lr_factor', type=float, default=1,
                              help='factor to reduce learning rate at the end of each epoch')
    optim_parser.add_argument('--max_grad_norm', help='maximum size of the gradient',
                              type=float, default=0.5)
    optim_parser.add_argument('--alpha', type=float, default=0.99,
                              help='smoothing constant (RMSProp')
    optim_parser.add_argument('--beta1', type=float, default=0.9,
                              help='coefficient used for computing running averages '
                                   'of gradient and its square (Adam)')
    optim_parser.add_argument('--beta2', type=float, default=0.999,
                              help='coefficient used for computing running averages '
                                   'of gradient and its square (Adam)')
    optim_parser.add_argument('--eps', type=float, default=1e-5,
                              help='term added to the denominator to improve numerical stability')

    logging_parser = parser.add_argument_group('Logging Parameters')
    logging_parser.add_argument('--results_dir', metavar='RESULTS_DIR', default='./results',
                                help='results dir')
    logging_parser.add_argument('--save_name', default='',
                                help='saved folder')
    logging_parser.add_argument('--log_interval', help='frequency of logging in epochs',
                                type=int, default=20)
    logging_parser.add_argument('--video_freq', help='frequency of creating a video in seconds',
                                type=int, default=300)
    logging_parser.add_argument('--plot_freq', help='frequency plot is updated',
                                type=int, default=60)
    logging_parser.add_argument('--log_kl', help='log kl_div during training',
                                dest='log_kl', action='store_true')

    args = parser.parse_args(args)

    if args.seed == 0:
        args.seed = random.randint(0, 2**32 - 1)
    set_global_seeds(args.seed)

    optimizer, optimizer_params = prepare_optimization_params(args)

    os.environ['OMP_NUM_THREADS'] = '1'
    torch.set_num_threads(mp.cpu_count()-1)

    if args.save_name is '':
        args.save_name = datetime.now().strftime('%Y-%m-%d_%H-%M-%S')
    save_path = os.path.join(args.results_dir, args.save_name)
    setup_logging_and_results(save_path, args)

    results_args = {'path': save_path, 'plot_title': args.save_name,
                    'results_file_name': 'results', 'resume': args.resume}

    if args.env_id.startswith('CartPole'):
        # lr: 1e-2
        if args.recurrent:
            policy = policies.mlp_rnn
            policy_args = {'hiddens': [16], 'bptt': args.num_steps_update, 'rnn_size': 16}
        else:
            policy = policies.mlp
            policy_args = {'hiddens': [16]}
    else:
        if args.recurrent:
            policy = policies.CnnToRecurrent
            # a3c paper
            # policy_args = {'convs': [(16, 8, 4, 0), (32, 4, 2, 0)],
            #                'hiddens': [256],
            #                'lstm_size': 256}
            # OpenAI Baselines
            policy_args = {'convs': [(32, 8, 4, 0), (64, 4, 2, 0), (64, 3, 1, 0)],
                           'hiddens': [512],
                           'lstm_size': 256}
        else:
            # CNN hyperparameters:
            # lr: 1e-4, max_grad_norm=2
            policy = policies.CnnToMlp
            # policy_args = {'convs': [(16, 8, 4), (32, 4, 2)],
            #                'hiddens': [256]}
            # OpenAI Baselines
            policy_args = {'convs': [(32, 8, 4), (64, 4, 2), (64, 3, 1)],
                           'hiddens': [512]}


    logging.info('Starting training...')
    a3c.train(args,
              policy=policy,
              policy_args=policy_args,
              optimizer=optimizer,
              optimizer_params=optimizer_params,
              save_path=save_path,
              results_args=results_args,
              )
    results_files_path = os.path.join(save_path, 'results_{}.json')
    try:
        reward = get_total_reward(results_files_path, args.num_workers)
        print('total_reward = {}'.format(reward))
    except ValueError:
        reward = 0
    return reward
    # TODO: save model
    # print("Saving model to mountaincar_model.pkl")
    # agent.save("mountaincar_model.pkl")


def setup_logging_and_results(save_path, args):
    if os.path.exists(save_path):
        shutil.rmtree(save_path)
    setup_logging(save_path, 'log.txt')
    results = ResultsLog(path=save_path, plot_title=args.save_name, resume=False)
    export_args(args, save_path)
    return save_path


def prepare_optimization_params(args):
    if args.optimizer == 'rmsprop':
        optimizer = optim.RMSprop
        optimizer_params = {'lr': args.lr, 'alpha': args.alpha, 'eps': args.eps}
    elif args.optimizer == 'sgd':
        optimizer = optim.SGD
        optimizer_params = {'lr': args.lr}
    elif args.optimizer == 'adam':
        optimizer = optim.Adam
        optimizer_params = {'lr': args.lr, 'betas': (args.beta1, args.beta2), 'eps': args.eps}
    else:
        ValueError('invalid optimizer type')
    return optimizer, optimizer_params


def get_total_reward(file_names, num_workers):
    data, columns = get_data_from_files(file_names, num_workers)
    total_reward = data['episode_reward'].mean(1).sum()

    return total_reward


def run(args):
    start_time = time.time()
    try:
        rewards = main(args)
    except:
        rewards = 0
    print(rewards)
    print('total_time = {}'.format(time.time()-start_time))
    return rewards

if __name__ == '__main__':
    main(sys.argv[1:])