#!/usr/bin/env python

import os

#os.environ['CUDA_VISIBLE_DEVICES'] = '-1' # GPU never needed for PPO
#print(os.environ['CUDA_VISIBLE_DEVICES'])

from mpi4py import MPI
from baselines.common import set_global_seeds
from baselines import bench
import time, random
import gym, logging
from baselines import logger

def wrap_train(env, mode, curiosity_search, random_noise, cs_weight):
    from baselines.common.atari_wrappers import (wrap_deepmind, FrameStack)
    if mode == 'train':
        env = wrap_deepmind(env, clip_rewards=True,
            curiosity_search=curiosity_search,
            random_noise=random_noise,
            cs_weight=cs_weight)
        #env = FrameStack(env, 4)
    elif mode == 'test':
        env = wrap_deepmind(env, episode_life=False,
            curiosity_search=curiosity_search,
            random_noise=random_noise,
            cs_weight=cs_weight)
    else:
        print("Unknown mode '" + mode + "'")
        return None
    return env

def train(env_id, num_frames, seed, mode, file_path, max_steps, load_path, create_video, curiosity_search, random_noise, cs_grid_size, cs_weight):
    from baselines.ppo1 import pposgd_simple, cnn_policy
    import baselines.common.tf_util as U
    import tensorflow as tf
    rank = MPI.COMM_WORLD.Get_rank()
    sess = U.single_threaded_session()
    sess.__enter__()
    if rank != 0: logger.set_level(logger.DISABLED)
    workerseed = seed + 10000 * MPI.COMM_WORLD.Get_rank()
    set_global_seeds(workerseed)
    env = gym.make(env_id)
    if cs_grid_size > 0:
        env.unwrapped.setCSGridSize(cs_grid_size)
    def policy_fn(name, ob_space, ac_space): #pylint: disable=W0613
        return cnn_policy.CnnPolicy(name=name, ob_space=ob_space, ac_space=ac_space)
    if mode == 'train':
        env = bench.Monitor(env, logger.get_dir() and
            os.path.join(logger.get_dir(), "{}.monitor.json".format(rank)),
            curiosity_search=curiosity_search)
    elif mode == 'test':
        env = bench.MonitorEval(env, use_monitor=(rank == 0), create_video=create_video,
            base_dir=logger.get_dir(),
            filename=(logger.get_dir() and os.path.join(logger.get_dir(), "{}.evalMonitor.json".format(rank))),
            curiosity_search=curiosity_search)
    else:
        print("Unknown mode '" + mode + "'")
        env.close()
        return
    env.seed(workerseed)
    gym.logger.setLevel(logging.WARN)

    env = wrap_train(env, mode, curiosity_search, random_noise, cs_weight)
    num_timesteps = int(num_frames / 4 * 1.1)
    env.seed(workerseed)

    if mode == 'train':
        pposgd_simple.learn(env, policy_fn,
            max_timesteps=num_timesteps,
            timesteps_per_batch=256,
            clip_param=0.2, entcoeff=0.01,
            optim_epochs=4, optim_stepsize=1e-3, optim_batchsize=64,
            gamma=0.99, lam=0.95,
            schedule='linear')
        if rank == 0 and file_path != '':
            tf.train.Saver().save(sess, filePath + "/model")
    elif mode == 'test':
        pposgd_simple.eval(sess, env, policy_fn, max_steps, loadPath)
    else:
        print("Unknown mode '" + mode + "'")
    env.close()

def main():
    import argparse
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('--env', help='environment ID', default='PongNoFrameskip-v4')
    parser.add_argument('--seed', help='RNG seed', type=int, default=0)
    parser.add_argument('--million_frames', help='How many frames to train (/ 1e6). '
        'This number gets divided by 4 due to frameskip', default=40)
    parser.add_argument('--outputDir', help='output directory', default='') # added by CMS 26/09/17
    parser.add_argument('--mode', help='train or test', default='train') # added by CMS 15/12/17
    parser.add_argument('--evalSteps', help='number of emulator steps for evaluation, 0 for infinite', default=0) # added by CMS 15/12/17
    parser.add_argument('--evalModel', help='path/filename of model to evaluate', default='') # added by CMS 15/12/17
    parser.add_argument('--createVideo', help='whether to generate a video for testing', default='false') # added by CMS 15/12/17
    parser.add_argument('--cs', help='whether CS should be enabled', default='false')
    parser.add_argument('--randomNoise', help='whether a random noise channel should be added (debug for CS)', default='false')
    parser.add_argument('--csGridSize', help='the size of the CS grid (0 for default)', type=int, default=0)
    parser.add_argument('--csWeight', help='weight of intrinsic rewards (-1 for equal weight); automatically enables clipping if > 0', default=-1)
    args = parser.parse_args()
    if len(args.outputDir) > 0:
        logger.configure(args.outputDir) # added by CMS 26/09/17
    if args.seed == -1:
        t = int( time.time() * 1000.0 )
        args.seed = ((t & 0xff000000) >> 24) + \
            ((t & 0x00ff0000) >>  8) + \
            ((t & 0x0000ff00) <<  8) + \
            ((t & 0x000000ff) << 24)
        print("USING RANDOM SEED: " + str(args.seed))

    train(args.env, seed=args.seed, num_frames=int(1e6 * float(args.million_frames)),
        mode=args.mode, file_path=args.outputDir,
        max_steps=int(args.evalSteps), load_path=args.evalModel, create_video=(args.createVideo == 'true'),
        curiosity_search=(args.cs == 'true'),
        random_noise=(args.randomNoise == 'true'),
        cs_grid_size=(args.csGridSize == 'true'),
        cs_weight=float(args.csWeight))

if __name__ == '__main__':
    main()
