#!/usr/bin/env python
import os, time, random, logging, gym
from baselines import logger
from baselines.common import set_global_seeds
from baselines import bench
from baselines.acktr.acktr_disc import learn
from baselines.acktr.acktr_disc import eval
from baselines.common.vec_env.subproc_vec_env import SubprocVecEnv
from baselines.common.atari_wrappers import wrap_deepmind
from baselines.acktr.policies import CnnPolicy

def train(env_id, num_frames, seed, num_cpu, mode, file_path, max_steps, load_path, create_video, curiosity_search, random_noise, cs_grid_size, cs_weight, episode_is_life):
    num_timesteps = int(num_frames / 4 * 1.1) 
    def make_env(rank):
        def _thunk():
            env = gym.make(env_id)
            if cs_grid_size > 0:
                env.unwrapped.setCSGridSize(cs_grid_size)
            env.seed(seed + rank)
            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)
                gym.logger.setLevel(logging.WARN)
                return wrap_deepmind(env,
                    episode_life=episode_is_life,
                    curiosity_search=curiosity_search,
                    random_noise=random_noise,
                    cs_weight=cs_weight)
            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)
                gym.logger.setLevel(logging.WARN)
                return wrap_deepmind(env, episode_life=False,
                    curiosity_search=curiosity_search,
                    random_noise=random_noise,
                    cs_weight=cs_weight)
        return _thunk
    set_global_seeds(seed)
    env = SubprocVecEnv([make_env(i) for i in range(num_cpu)])
    policy_fn = CnnPolicy
    if mode == 'train':
        if (file_path != ''):
            learn(policy_fn, env, seed, total_timesteps=num_timesteps, nprocs=num_cpu, save_path=file_path + "/model")
        else:
            learn(policy_fn, env, seed, total_timesteps=num_timesteps, nprocs=num_cpu, save_path=None)
    elif mode == 'test':
        eval(policy_fn, env, seed, max_steps, load_path, nprocs=num_cpu)
        if (logger.get_dir()):
            os.rename(os.path.join(logger.get_dir(), "0.evalMonitor.json"),
                os.path.join(logger.get_dir(), "evalMonitor.json"))
            for i in range(1, num_cpu):
                os.remove(os.path.join(logger.get_dir(), str(i) + ".evalMonitor.json")) # only the first monitor used for testing
    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='BreakoutNoFrameskip-v4')
    parser.add_argument('--seed', help='RNG seed (-1 = random)', type=int, default=-1)
    parser.add_argument('--policy', help='Policy architecture', choices=['cnn', 'lstm', 'lnlstm'], default='cnn')
    parser.add_argument('--lrschedule', help='Learning rate schedule', choices=['constant', 'linear'], default='constant')
    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)
    parser.add_argument('--episodeIsLife', help='whether end-of-life = end-of-episode', default='true')
    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)
        args.seed = int.from_bytes(os.urandom(4), byteorder='big')
        print("USING RANDOM SEED: " + str(args.seed))
    train(args.env, num_frames=int(1e6 * float(args.million_frames)), seed=args.seed, num_cpu=16,
        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,
        cs_weight=float(args.csWeight),
        episode_is_life=(args.episodeIsLife == 'true'))


if __name__ == '__main__':
    main()
