import tensorflow as tf
import sys
sys.path.insert(0,'..')
import vtrace_popart as vtrace
nest = tf.contrib.framework.nest

from .flags import *

def compute_baseline_loss(advantages):
  # Loss for the baseline, summed over the time dimension.
  # Multiply by 0.5 to match the standard update rule:
  # d(loss) / d(baseline) = advantage
  return .5 * tf.reduce_sum(tf.square(advantages))

def compute_entropy_loss(logits):
  policy = tf.nn.softmax(logits)
  log_policy = tf.nn.log_softmax(logits)
  entropy_per_timestep = tf.reduce_sum(-policy * log_policy, axis=-1)
  return -tf.reduce_sum(entropy_per_timestep)

def compute_policy_gradient_loss(logits, actions, advantages):
  cross_entropy = tf.nn.sparse_softmax_cross_entropy_with_logits(
      labels=actions, logits=logits)
  advantages = tf.stop_gradient(advantages)
  policy_gradient_loss_per_timestep = cross_entropy * advantages
  return tf.reduce_sum(policy_gradient_loss_per_timestep)

def build_learner(agent, env_outputs, agent_outputs, env_id):
  """Builds the learner loop.

  Args:
    agent: A snt.RNNCore module outputting `AgentOutput` named tuples, with an
      `unroll` call for computing the outputs for a whole trajectory.
    agent_state: The initial agent state for each sequence in the batch.
    env_outputs: A `StepOutput` namedtuple where each field is of shape
      [T+1, ...].
    agent_outputs: An `AgentOutput` namedtuple where each field is of shape
      [T+1, ...].

  Returns:
    A tuple of (done, infos, and environment frames) where
    the environment frames tensor causes an update.
  """

  learner_outputs = agent.unroll(agent_outputs.action, env_outputs, env_id)

  # Use last baseline value (from the value function) to bootstrap.
  bootstrap_value = learner_outputs.un_normalized_vf[-1]
 
  # At this point, the environment outputs at time step `t` are the inputs that
  # lead to the learner_outputs at time step `t`. After the following shifting,
  # the actions in agent_outputs and learner_outputs at time step `t` is what
  # leads to the environment outputs at time step `t`.
  agent_outputs = nest.map_structure(lambda t: t[1:], agent_outputs)
  rewards, infos, done, _ = nest.map_structure(
      lambda t: t[1:], env_outputs)
  learner_outputs = nest.map_structure(lambda t: t[:-1], learner_outputs)

  if FLAGS.reward_clipping == 'abs_one':
    clipped_rewards = tf.clip_by_value(rewards, -1, 1)
  elif FLAGS.reward_clipping == 'soft_asymmetric':
    squeezed = tf.tanh(rewards / 5.0)
    # Negative rewards are given less weight than positive rewards.
    clipped_rewards = tf.where(rewards < 0, .3 * squeezed, squeezed) * 5.

  discounts = tf.to_float(~done) * FLAGS.discounting
  game_specific_mean = tf.gather(agent._mean, env_id)
  game_specific_std = tf.gather(agent._std, env_id)

  # Compute V-trace returns and weights.
  # Note, this is put on the CPU because it's faster than on GPU. It can be
  # improved further with XLA-compilation or with a custom TensorFlow operation.
  with tf.device('/cpu'):
    vtrace_returns = vtrace.from_logits(
        behaviour_policy_logits=agent_outputs.policy_logits,
        target_policy_logits=learner_outputs.policy_logits,
        actions=agent_outputs.action,
        discounts=discounts,
        rewards=clipped_rewards,
        un_normalized_values=learner_outputs.un_normalized_vf,
        normalized_values=learner_outputs.normalized_vf,
        mean=game_specific_mean,
        std=game_specific_std,
        bootstrap_value=bootstrap_value)

  # First term of equation (7) in (Hessel et al., 2018)
  normalized_vtrace = (vtrace_returns.vs - game_specific_mean) / game_specific_std

  normalized_vtrace = nest.map_structure(tf.stop_gradient, normalized_vtrace)


  # Compute loss as a weighted sum of the baseline loss, the policy gradient
  # loss and an entropy regularization term.
  total_loss = compute_policy_gradient_loss(
      learner_outputs.policy_logits, agent_outputs.action,
      vtrace_returns.pg_advantages)

  baseline_loss = compute_baseline_loss(
       normalized_vtrace - learner_outputs.normalized_vf)

  total_loss += FLAGS.baseline_cost * baseline_loss
  total_loss += FLAGS.entropy_cost * compute_entropy_loss(
      learner_outputs.policy_logits)

  # Optimization
  num_env_frames = tf.train.get_global_step()

  learning_rate = tf.train.polynomial_decay(FLAGS.learning_rate, num_env_frames,
                                            FLAGS.total_environment_frames, 0)

  optimizer = tf.train.RMSPropOptimizer(learning_rate, FLAGS.decay,
                                        FLAGS.momentum, FLAGS.epsilon)

  # Use reward clipping for atari games only 
  if FLAGS.gradient_clipping > 0.0:
    variables = tf.trainable_variables()
    gradients = tf.gradients(total_loss, variables)
    gradients, _ = tf.clip_by_global_norm(gradients, FLAGS.gradient_clipping)
    train_op = optimizer.apply_gradients(zip(gradients, variables))
  else:
    train_op = optimizer.minimize(total_loss)

  # Merge updating the network and environment frames into a single tensor.
  with tf.control_dependencies([train_op]):
    num_env_frames_and_train = num_env_frames.assign_add(
        FLAGS.batch_size * FLAGS.unroll_length)

  # Adding a few summaries.
  tf.summary.scalar('learning_rate', learning_rate)
  tf.summary.scalar('total_loss', total_loss)
  tf.summary.histogram('action', agent_outputs.action)
  
  # I'm not sure if it's really necessary to put this operation on the CPU. 
  with tf.device('/cpu'):
    (mean, mean_squared) = (agent.update_moments(vtrace_returns.vs, env_id))
  return (done, infos, num_env_frames_and_train) + (mean, mean_squared)