# Copyright 2017 The TensorFlow Authors All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

"""Trainer for coordinating single or multi-replica training.

Main point of entry for running models.  Specifies most of
the parameters used by different algorithms.
"""

import tensorflow as tf
import numpy as np
import random
import os
import pickle

from six.moves import xrange
import controller
import model
import policy
import baseline
import objective
import full_episode_objective
import trust_region
import optimizers
import replay_buffer
import expert_paths
import gym_wrapper
import env_spec

app = tf.app
flags = tf.flags
logging = tf.logging
gfile = tf.gfile

FLAGS = flags.FLAGS

flags.DEFINE_string('env', 'Copy-v0', 'environment name')
flags.DEFINE_integer('batch_size', 100, 'batch size')
flags.DEFINE_integer('replay_batch_size', None, 'replay batch size; defaults to batch_size')
flags.DEFINE_integer('num_samples', 1,
                     'number of samples from each random seed initialization')
flags.DEFINE_integer('max_step', 200, 'max number of steps to train on')
flags.DEFINE_integer('cutoff_agent', 0,
                     'number of steps at which to cut-off agent. '
                     'Defaults to always cutoff')
flags.DEFINE_integer('num_steps', 100000, 'number of training steps')
flags.DEFINE_integer('validation_frequency', 100,
                     'every so many steps, output some stats')

flags.DEFINE_float('target_network_lag', 0.95,
                   'This exponential decay on online network yields target '
                   'network')
flags.DEFINE_string('sample_from', 'online',
                    'Sample actions from "online" network or "target" network')

flags.DEFINE_string('objective', 'pcl',
                    'pcl/upcl/a3c/trpo/reinforce/urex')
flags.DEFINE_bool('trust_region_p', False,
                  'use trust region for policy optimization')
flags.DEFINE_string('value_opt', None,
                    'leave as None to optimize it along with policy '
                    '(using critic_weight). Otherwise set to '
                    '"best_fit" (least squares regression), "lbfgs", or "grad"')
flags.DEFINE_float('max_divergence', 0.01,
                   'max divergence (i.e. KL) to allow during '
                   'trust region optimization')

flags.DEFINE_float('learning_rate', 0.01, 'learning rate')
flags.DEFINE_float('clip_norm', 5.0, 'clip norm')
flags.DEFINE_float('clip_adv', 0.0, 'Clip advantages at this value.  '
                   'Leave as 0 to not clip at all.')
flags.DEFINE_float('critic_weight', 0.1, 'critic weight')
flags.DEFINE_float('tau', 0.1, 'entropy regularizer.'
                   'If using decaying tau, this is the final value.')
flags.DEFINE_float('tau_decay', None,
                   'decay tau by this much every 100 steps')
flags.DEFINE_float('tau_start', 0.1,
                   'start tau at this value')
flags.DEFINE_float('eps_lambda', 0.0, 'relative entropy regularizer.')
flags.DEFINE_bool('update_eps_lambda', False,
                  'Update lambda automatically based on last 100 episodes.')
flags.DEFINE_float('gamma', 1.0, 'discount')
flags.DEFINE_integer('rollout', 10, 'rollout')
flags.DEFINE_bool('use_target_values', False,
                  'use target network for value estimates')
flags.DEFINE_bool('fixed_std', True,
                  'fix the std in Gaussian distributions')
flags.DEFINE_bool('input_prev_actions', True,
                  'input previous actions to policy network')
flags.DEFINE_bool('recurrent', True,
                  'use recurrent connections')
flags.DEFINE_bool('input_time_step', False,
                  'input time step into value calucations')

flags.DEFINE_bool('use_online_batch', True, 'train on batches as they are sampled')
flags.DEFINE_bool('batch_by_steps', False,
                  'ensure each training batch has batch_size * max_step steps')
flags.DEFINE_bool('unify_episodes', False,
                  'Make sure replay buffer holds entire episodes, '
                  'even across distinct sampling steps')
flags.DEFINE_integer('replay_buffer_size', 5000, 'replay buffer size')
flags.DEFINE_float('replay_buffer_alpha', 0.5, 'replay buffer alpha param')
flags.DEFINE_integer('replay_buffer_freq', 0,
                     'replay buffer frequency (only supports -1/0/1)')
flags.DEFINE_string('eviction', 'rand',
                    'how to evict from replay buffer: rand/rank/fifo')
flags.DEFINE_string('prioritize_by', 'rewards',
                    'Prioritize replay buffer by "rewards" or "step"')
flags.DEFINE_integer('num_expert_paths', 0,
                     'number of expert paths to seed replay buffer with')

flags.DEFINE_integer('internal_dim', 256, 'RNN internal dim')
flags.DEFINE_integer('value_hidden_layers', 0,
                     'number of hidden layers in value estimate')
flags.DEFINE_integer('tf_seed', 42, 'random seed for tensorflow')

flags.DEFINE_string('save_trajectories_dir', None,
                    'directory to save trajectories to, if desired')
flags.DEFINE_string('load_trajectories_file', None,
                    'file to load expert trajectories from')

# supervisor flags
flags.DEFINE_bool('supervisor', False, 'use supervisor training')
flags.DEFINE_integer('task_id', 0, 'task id')
flags.DEFINE_integer('ps_tasks', 0, 'number of ps tasks')
flags.DEFINE_integer('num_replicas', 1, 'number of replicas used')
flags.DEFINE_string('master', 'local', 'name of master')
flags.DEFINE_string('save_dir', '', 'directory to save model to')
flags.DEFINE_string('load_path', '', 'path of saved model to load (if none in save_dir)')


class Trainer(object):
  """Coordinates single or multi-replica training."""

  def __init__(self):
    self.batch_size = FLAGS.batch_size
    self.replay_batch_size = FLAGS.replay_batch_size
    if self.replay_batch_size is None:
      self.replay_batch_size = self.batch_size
    self.num_samples = FLAGS.num_samples

    self.env_str = FLAGS.env
    self.env = gym_wrapper.GymWrapper(self.env_str,
                                      distinct=FLAGS.batch_size // self.num_samples,
                                      count=self.num_samples)
    self.eval_env = gym_wrapper.GymWrapper(
        self.env_str,
        distinct=FLAGS.batch_size // self.num_samples,
        count=self.num_samples)
    self.env_spec = env_spec.EnvSpec(self.env.get_one())

    self.max_step = FLAGS.max_step
    self.cutoff_agent = FLAGS.cutoff_agent
    self.num_steps = FLAGS.num_steps
    self.validation_frequency = FLAGS.validation_frequency

    self.target_network_lag = FLAGS.target_network_lag
    self.sample_from = FLAGS.sample_from
    assert self.sample_from in ['online', 'target']

    self.critic_weight = FLAGS.critic_weight
    self.objective = FLAGS.objective
    self.trust_region_p = FLAGS.trust_region_p
    self.value_opt = FLAGS.value_opt
    assert not self.trust_region_p or self.objective in ['pcl', 'trpo']
    assert self.objective != 'trpo' or self.trust_region_p
    assert self.value_opt is None or self.value_opt == 'None' or \
        self.critic_weight == 0.0
    self.max_divergence = FLAGS.max_divergence

    self.learning_rate = FLAGS.learning_rate
    self.clip_norm = FLAGS.clip_norm
    self.clip_adv = FLAGS.clip_adv
    self.tau = FLAGS.tau
    self.tau_decay = FLAGS.tau_decay
    self.tau_start = FLAGS.tau_start
    self.eps_lambda = FLAGS.eps_lambda
    self.update_eps_lambda = FLAGS.update_eps_lambda
    self.gamma = FLAGS.gamma
    self.rollout = FLAGS.rollout
    self.use_target_values = FLAGS.use_target_values
    self.fixed_std = FLAGS.fixed_std
    self.input_prev_actions = FLAGS.input_prev_actions
    self.recurrent = FLAGS.recurrent
    assert not self.trust_region_p or not self.recurrent
    self.input_time_step = FLAGS.input_time_step
    assert not self.input_time_step or (self.cutoff_agent <= self.max_step)

    self.use_online_batch = FLAGS.use_online_batch
    self.batch_by_steps = FLAGS.batch_by_steps
    self.unify_episodes = FLAGS.unify_episodes
    if self.unify_episodes:
      assert self.batch_size == 1

    self.replay_buffer_size = FLAGS.replay_buffer_size
    self.replay_buffer_alpha = FLAGS.replay_buffer_alpha
    self.replay_buffer_freq = FLAGS.replay_buffer_freq
    assert self.replay_buffer_freq in [-1, 0, 1]
    self.eviction = FLAGS.eviction
    self.prioritize_by = FLAGS.prioritize_by
    assert self.prioritize_by in ['rewards', 'step']
    self.num_expert_paths = FLAGS.num_expert_paths

    self.internal_dim = FLAGS.internal_dim
    self.value_hidden_layers = FLAGS.value_hidden_layers
    self.tf_seed = FLAGS.tf_seed

    self.save_trajectories_dir = FLAGS.save_trajectories_dir
    self.save_trajectories_file = (
        os.path.join(
            self.save_trajectories_dir, self.env_str.replace('-', '_'))
        if self.save_trajectories_dir else None)
    self.load_trajectories_file = FLAGS.load_trajectories_file

    self.hparams = dict((attr, getattr(self, attr))
                        for attr in dir(self)
                        if not attr.startswith('__') and
                        not callable(getattr(self, attr)))

  def hparams_string(self):
    return '\n'.join('%s: %s' % item for item in sorted(self.hparams.items()))

  def get_objective(self):
    tau = self.tau
    if self.tau_decay is not None:
      assert self.tau_start >= self.tau
      tau = tf.maximum(
          tf.train.exponential_decay(
              self.tau_start, self.global_step, 100, self.tau_decay),
          self.tau)

    if self.objective in ['pcl', 'a3c', 'trpo', 'upcl']:
      cls = (objective.PCL if self.objective in ['pcl', 'upcl'] else
             objective.TRPO if self.objective == 'trpo' else
             objective.ActorCritic)
      policy_weight = 1.0

      return cls(self.learning_rate,
                 clip_norm=self.clip_norm,
                 policy_weight=policy_weight,
                 critic_weight=self.critic_weight,
                 tau=tau, gamma=self.gamma, rollout=self.rollout,
                 eps_lambda=self.eps_lambda, clip_adv=self.clip_adv,
                 use_target_values=self.use_target_values)
    elif self.objective in ['reinforce', 'urex']:
      cls = (full_episode_objective.Reinforce
             if self.objective == 'reinforce' else
             full_episode_objective.UREX)
      return cls(self.learning_rate,
                 clip_norm=self.clip_norm,
                 num_samples=self.num_samples,
                 tau=tau, bonus_weight=1.0)  # TODO: bonus weight?
    else:
      assert False, 'Unknown objective %s' % self.objective

  def get_policy(self):
    if self.recurrent:
      cls = policy.Policy
    else:
      cls = policy.MLPPolicy
    return cls(self.env_spec, self.internal_dim,
               fixed_std=self.fixed_std,
               recurrent=self.recurrent,
               input_prev_actions=self.input_prev_actions)

  def get_baseline(self):
    cls = (baseline.UnifiedBaseline if self.objective == 'upcl' else
           baseline.Baseline)
    return cls(self.env_spec, self.internal_dim,
               input_prev_actions=self.input_prev_actions,
               input_time_step=self.input_time_step,
               input_policy_state=self.recurrent,  # may want to change this
               n_hidden_layers=self.value_hidden_layers,
               hidden_dim=self.internal_dim,
               tau=self.tau)

  def get_trust_region_p_opt(self):
    if self.trust_region_p:
      return trust_region.TrustRegionOptimization(
          max_divergence=self.max_divergence)
    else:
      return None

  def get_value_opt(self):
    if self.value_opt == 'grad':
      return optimizers.GradOptimization(
          learning_rate=self.learning_rate, max_iter=5, mix_frac=0.05)
    elif self.value_opt == 'lbfgs':
      return optimizers.LbfgsOptimization(max_iter=25, mix_frac=0.1)
    elif self.value_opt == 'best_fit':
      return optimizers.BestFitOptimization(mix_frac=1.0)
    else:
      return None

  def get_model(self):
    cls = model.Model
    return cls(self.env_spec, self.global_step,
               target_network_lag=self.target_network_lag,
               sample_from=self.sample_from,
               get_policy=self.get_policy,
               get_baseline=self.get_baseline,
               get_objective=self.get_objective,
               get_trust_region_p_opt=self.get_trust_region_p_opt,
               get_value_opt=self.get_value_opt)

  def get_replay_buffer(self):
    if self.replay_buffer_freq <= 0:
      return None
    else:
      assert self.objective in ['pcl', 'upcl'], 'Can\'t use replay buffer with %s' % (
          self.objective)
    cls = replay_buffer.PrioritizedReplayBuffer
    return cls(self.replay_buffer_size,
               alpha=self.replay_buffer_alpha,
               eviction_strategy=self.eviction)

  def get_buffer_seeds(self):
    return expert_paths.sample_expert_paths(
        self.num_expert_paths, self.env_str, self.env_spec,
        load_trajectories_file=self.load_trajectories_file)

  def get_controller(self, env):
    """Get controller."""
    cls = controller.Controller
    return cls(env, self.env_spec, self.internal_dim,
               use_online_batch=self.use_online_batch,
               batch_by_steps=self.batch_by_steps,
               unify_episodes=self.unify_episodes,
               replay_batch_size=self.replay_batch_size,
               max_step=self.max_step,
               cutoff_agent=self.cutoff_agent,
               save_trajectories_file=self.save_trajectories_file,
               use_trust_region=self.trust_region_p,
               use_value_opt=self.value_opt not in [None, 'None'],
               update_eps_lambda=self.update_eps_lambda,
               prioritize_by=self.prioritize_by,
               get_model=self.get_model,
               get_replay_buffer=self.get_replay_buffer,
               get_buffer_seeds=self.get_buffer_seeds)

  def do_before_step(self, step):
    pass

  def run(self):
    """Run training."""
    is_chief = FLAGS.task_id == 0 or not FLAGS.supervisor
    sv = None

    def init_fn(sess, saver):
      ckpt = None
      if FLAGS.save_dir and sv is None:
        load_dir = FLAGS.save_dir
        ckpt = tf.train.get_checkpoint_state(load_dir)
      if ckpt and ckpt.model_checkpoint_path:
        logging.info('restoring from %s', ckpt.model_checkpoint_path)
        saver.restore(sess, ckpt.model_checkpoint_path)
      elif FLAGS.load_path:
        logging.info('restoring from %s', FLAGS.load_path)
        saver.restore(sess, FLAGS.load_path)

    if FLAGS.supervisor:
      with tf.device(tf.ReplicaDeviceSetter(FLAGS.ps_tasks, merge_devices=True)):
        self.global_step = tf.contrib.framework.get_or_create_global_step()
        tf.set_random_seed(FLAGS.tf_seed)
        self.controller = self.get_controller(self.env)
        self.model = self.controller.model
        self.controller.setup()
        with tf.variable_scope(tf.get_variable_scope(), reuse=True):
          self.eval_controller = self.get_controller(self.eval_env)
          self.eval_controller.setup(train=False)

        saver = tf.train.Saver(max_to_keep=10)
        step = self.model.global_step
        sv = tf.Supervisor(logdir=FLAGS.save_dir,
                           is_chief=is_chief,
                           saver=saver,
                           save_model_secs=600,
                           summary_op=None,  # we define it ourselves
                           save_summaries_secs=60,
                           global_step=step,
                           init_fn=lambda sess: init_fn(sess, saver))
        sess = sv.PrepareSession(FLAGS.master)
    else:
      tf.set_random_seed(FLAGS.tf_seed)
      self.global_step = tf.contrib.framework.get_or_create_global_step()
      self.controller = self.get_controller(self.env)
      self.model = self.controller.model
      self.controller.setup()
      with tf.variable_scope(tf.get_variable_scope(), reuse=True):
        self.eval_controller = self.get_controller(self.eval_env)
        self.eval_controller.setup(train=False)

      saver = tf.train.Saver(max_to_keep=10)
      sess = tf.Session()
      sess.run(tf.initialize_all_variables())
      init_fn(sess, saver)

    self.sv = sv
    self.sess = sess

    logging.info('hparams:\n%s', self.hparams_string())

    model_step = sess.run(self.model.global_step)
    if model_step >= self.num_steps:
      logging.info('training has reached final step')
      return

    losses = []
    rewards = []
    all_ep_rewards = []
    for step in xrange(1 + self.num_steps):

      if sv is not None and sv.ShouldStop():
        logging.info('stopping supervisor')
        break

      self.do_before_step(step)

      (loss, summary,
       total_rewards, episode_rewards) = self.controller.train(sess)
      _, greedy_episode_rewards = self.eval_controller.eval(sess)
      self.controller.greedy_episode_rewards = greedy_episode_rewards
      losses.append(loss)
      rewards.append(total_rewards)
      all_ep_rewards.extend(episode_rewards)

      if (random.random() < 0.1 and summary and episode_rewards and
          is_chief and sv and sv._summary_writer):
        sv.summary_computed(sess, summary)

      model_step = sess.run(self.model.global_step)
      if is_chief and step % self.validation_frequency == 0:
        logging.info('at training step %d, model step %d: '
                     'avg loss %f, avg reward %f, '
                     'episode rewards: %f, greedy rewards: %f',
                     step, model_step,
                     np.mean(losses), np.mean(rewards),
                     np.mean(all_ep_rewards),
                     np.mean(greedy_episode_rewards))

        losses = []
        rewards = []
        all_ep_rewards = []

      if model_step >= self.num_steps:
        logging.info('training has reached final step')
        break

    if is_chief and sv is not None:
      logging.info('saving final model to %s', sv.save_path)
      sv.saver.save(sess, sv.save_path, global_step=sv.global_step)


def main(unused_argv):
  logging.set_verbosity(logging.INFO)
  trainer = Trainer()
  trainer.run()


if __name__ == '__main__':
  app.run()