# 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.
# ==============================================================================

"""Objectives to compute loss and value targets.

Implements Actor Critic, PCL (vanilla PCL, Unified PCL, Trust PCL), and TRPO.
"""

import tensorflow as tf
import numpy as np


class Objective(object):
  def __init__(self, learning_rate, clip_norm):
    self.learning_rate = learning_rate
    self.clip_norm = clip_norm

  def get_optimizer(self, learning_rate):
    """Optimizer for gradient descent ops."""
    return tf.train.AdamOptimizer(learning_rate=learning_rate,
                                  epsilon=2e-4)

  def training_ops(self, loss, learning_rate=None):
    """Gradient ops."""
    opt = self.get_optimizer(learning_rate)
    params = tf.trainable_variables()
    grads = tf.gradients(loss, params)

    if self.clip_norm:
      grads, global_norm = tf.clip_by_global_norm(grads, self.clip_norm)
      tf.summary.scalar('grad_global_norm', global_norm)

    return opt.apply_gradients(zip(grads, params))

  def get(self, rewards, pads, values, final_values,
          log_probs, prev_log_probs, target_log_probs,
          entropies, logits,
          target_values, final_target_values):
    """Get objective calculations."""
    raise NotImplementedError()


def discounted_future_sum(values, discount, rollout):
  """Discounted future sum of time-major values."""
  discount_filter = tf.reshape(
      discount ** tf.range(float(rollout)), [-1, 1, 1])
  expanded_values = tf.concat(
      [values, tf.zeros([rollout - 1, tf.shape(values)[1]])], 0)

  conv_values = tf.transpose(tf.squeeze(tf.nn.conv1d(
      tf.expand_dims(tf.transpose(expanded_values), -1), discount_filter,
      stride=1, padding='VALID'), -1))

  return conv_values


def discounted_two_sided_sum(values, discount, rollout):
  """Discounted two-sided sum of time-major values."""
  roll = float(rollout)
  discount_filter = tf.reshape(
      discount ** tf.abs(tf.range(-roll + 1, roll)), [-1, 1, 1])
  expanded_values = tf.concat(
      [tf.zeros([rollout - 1, tf.shape(values)[1]]), values,
       tf.zeros([rollout - 1, tf.shape(values)[1]])], 0)

  conv_values = tf.transpose(tf.squeeze(tf.nn.conv1d(
      tf.expand_dims(tf.transpose(expanded_values), -1), discount_filter,
      stride=1, padding='VALID'), -1))

  return conv_values


def shift_values(values, discount, rollout, final_values=0.0):
  """Shift values up by some amount of time.

  Those values that shift from a value beyond the last value
  are calculated using final_values.

  """
  roll_range = tf.cumsum(tf.ones_like(values[:rollout, :]), 0,
                         exclusive=True, reverse=True)
  final_pad = tf.expand_dims(final_values, 0) * discount ** roll_range
  return tf.concat([discount ** rollout * values[rollout:, :],
                    final_pad], 0)


class ActorCritic(Objective):
  """Standard Actor-Critic."""

  def __init__(self, learning_rate, clip_norm=5,
               policy_weight=1.0, critic_weight=0.1,
               tau=0.1, gamma=1.0, rollout=10,
               eps_lambda=0.0, clip_adv=None,
               use_target_values=False):
    super(ActorCritic, self).__init__(learning_rate, clip_norm=clip_norm)
    self.policy_weight = policy_weight
    self.critic_weight = critic_weight
    self.tau = tau
    self.gamma = gamma
    self.rollout = rollout
    self.clip_adv = clip_adv

    self.eps_lambda = tf.get_variable(  # TODO: need a better way
        'eps_lambda', [], initializer=tf.constant_initializer(eps_lambda),
        trainable=False)
    self.new_eps_lambda = tf.placeholder(tf.float32, [])
    self.assign_eps_lambda = self.eps_lambda.assign(
        0.99 * self.eps_lambda + 0.01 * self.new_eps_lambda)
    self.use_target_values = use_target_values

  def get(self, rewards, pads, values, final_values,
          log_probs, prev_log_probs, target_log_probs,
          entropies, logits,
          target_values, final_target_values):
    not_pad = 1 - pads
    batch_size = tf.shape(rewards)[1]

    entropy = not_pad * sum(entropies)
    rewards = not_pad * rewards
    value_estimates = not_pad * values
    log_probs = not_pad * sum(log_probs)
    target_values = not_pad * tf.stop_gradient(target_values)
    final_target_values = tf.stop_gradient(final_target_values)

    sum_rewards = discounted_future_sum(rewards, self.gamma, self.rollout)
    if self.use_target_values:
      last_values = shift_values(
          target_values, self.gamma, self.rollout,
          final_target_values)
    else:
      last_values = shift_values(value_estimates, self.gamma, self.rollout,
                                 final_values)

    future_values = sum_rewards + last_values
    baseline_values = value_estimates

    adv = tf.stop_gradient(-baseline_values + future_values)
    if self.clip_adv:
      adv = tf.minimum(self.clip_adv, tf.maximum(-self.clip_adv, adv))
    policy_loss = -adv * log_probs
    critic_loss = -adv * baseline_values
    regularizer = -self.tau * entropy

    policy_loss = tf.reduce_mean(
        tf.reduce_sum(policy_loss * not_pad, 0))
    critic_loss = tf.reduce_mean(
        tf.reduce_sum(critic_loss * not_pad, 0))
    regularizer = tf.reduce_mean(
        tf.reduce_sum(regularizer * not_pad, 0))

    # loss for gradient calculation
    loss = (self.policy_weight * policy_loss +
            self.critic_weight * critic_loss + regularizer)

    raw_loss = tf.reduce_mean(  # TODO
        tf.reduce_sum(not_pad * policy_loss, 0))

    gradient_ops = self.training_ops(
        loss, learning_rate=self.learning_rate)

    tf.summary.histogram('log_probs', tf.reduce_sum(log_probs, 0))
    tf.summary.histogram('rewards', tf.reduce_sum(rewards, 0))
    tf.summary.scalar('avg_rewards',
                      tf.reduce_mean(tf.reduce_sum(rewards, 0)))
    tf.summary.scalar('policy_loss',
                      tf.reduce_mean(tf.reduce_sum(not_pad * policy_loss)))
    tf.summary.scalar('critic_loss',
                      tf.reduce_mean(tf.reduce_sum(not_pad * policy_loss)))
    tf.summary.scalar('loss', loss)
    tf.summary.scalar('raw_loss', raw_loss)

    return (loss, raw_loss, future_values,
            gradient_ops, tf.summary.merge_all())


class PCL(ActorCritic):
  """PCL implementation.

  Implements vanilla PCL, Unified PCL, and Trust PCL depending
  on provided inputs.

  """

  def get(self, rewards, pads, values, final_values,
          log_probs, prev_log_probs, target_log_probs,
          entropies, logits,
          target_values, final_target_values):
    not_pad = 1 - pads
    batch_size = tf.shape(rewards)[1]

    rewards = not_pad * rewards
    value_estimates = not_pad * values
    log_probs = not_pad * sum(log_probs)
    target_log_probs = not_pad * tf.stop_gradient(sum(target_log_probs))
    relative_log_probs = not_pad * (log_probs - target_log_probs)
    target_values = not_pad * tf.stop_gradient(target_values)
    final_target_values = tf.stop_gradient(final_target_values)

    # Prepend.
    not_pad = tf.concat([tf.ones([self.rollout - 1, batch_size]),
                         not_pad], 0)
    rewards = tf.concat([tf.zeros([self.rollout - 1, batch_size]),
                         rewards], 0)
    value_estimates = tf.concat(
        [self.gamma ** tf.expand_dims(
            tf.range(float(self.rollout - 1), 0, -1), 1) *
         tf.ones([self.rollout - 1, batch_size]) *
         value_estimates[0:1, :],
         value_estimates], 0)
    log_probs = tf.concat([tf.zeros([self.rollout - 1, batch_size]),
                           log_probs], 0)
    prev_log_probs = tf.concat([tf.zeros([self.rollout - 1, batch_size]),
                                prev_log_probs], 0)
    relative_log_probs = tf.concat([tf.zeros([self.rollout - 1, batch_size]),
                                    relative_log_probs], 0)
    target_values = tf.concat(
        [self.gamma ** tf.expand_dims(
            tf.range(float(self.rollout - 1), 0, -1), 1) *
         tf.ones([self.rollout - 1, batch_size]) *
         target_values[0:1, :],
         target_values], 0)

    sum_rewards = discounted_future_sum(rewards, self.gamma, self.rollout)
    sum_log_probs = discounted_future_sum(log_probs, self.gamma, self.rollout)
    sum_prev_log_probs = discounted_future_sum(prev_log_probs, self.gamma, self.rollout)
    sum_relative_log_probs = discounted_future_sum(
        relative_log_probs, self.gamma, self.rollout)

    if self.use_target_values:
      last_values = shift_values(
          target_values, self.gamma, self.rollout,
          final_target_values)
    else:
      last_values = shift_values(value_estimates, self.gamma, self.rollout,
                                 final_values)

    future_values = (
        - self.tau * sum_log_probs
        - self.eps_lambda * sum_relative_log_probs
        + sum_rewards + last_values)
    baseline_values = value_estimates

    adv = tf.stop_gradient(-baseline_values + future_values)
    if self.clip_adv:
      adv = tf.minimum(self.clip_adv, tf.maximum(-self.clip_adv, adv))
    policy_loss = -adv * sum_log_probs
    critic_loss = -adv * (baseline_values - last_values)

    policy_loss = tf.reduce_mean(
        tf.reduce_sum(policy_loss * not_pad, 0))
    critic_loss = tf.reduce_mean(
        tf.reduce_sum(critic_loss * not_pad, 0))

    # loss for gradient calculation
    loss = (self.policy_weight * policy_loss +
            self.critic_weight * critic_loss)

    # actual quantity we're trying to minimize
    raw_loss = tf.reduce_mean(
        tf.reduce_sum(not_pad * adv * (-baseline_values + future_values), 0))

    gradient_ops = self.training_ops(
        loss, learning_rate=self.learning_rate)

    tf.summary.histogram('log_probs', tf.reduce_sum(log_probs, 0))
    tf.summary.histogram('rewards', tf.reduce_sum(rewards, 0))
    tf.summary.histogram('future_values', future_values)
    tf.summary.histogram('baseline_values', baseline_values)
    tf.summary.histogram('advantages', adv)
    tf.summary.scalar('avg_rewards',
                      tf.reduce_mean(tf.reduce_sum(rewards, 0)))
    tf.summary.scalar('policy_loss',
                      tf.reduce_mean(tf.reduce_sum(not_pad * policy_loss)))
    tf.summary.scalar('critic_loss',
                      tf.reduce_mean(tf.reduce_sum(not_pad * policy_loss)))
    tf.summary.scalar('loss', loss)
    tf.summary.scalar('raw_loss', tf.reduce_mean(raw_loss))
    tf.summary.scalar('eps_lambda', self.eps_lambda)

    return (loss, raw_loss,
            future_values[self.rollout - 1:, :],
            gradient_ops, tf.summary.merge_all())


class TRPO(ActorCritic):
  """TRPO."""

  def get(self, rewards, pads, values, final_values,
          log_probs, prev_log_probs, target_log_probs,
          entropies, logits,
          target_values, final_target_values):
    not_pad = 1 - pads
    batch_size = tf.shape(rewards)[1]

    rewards = not_pad * rewards
    value_estimates = not_pad * values
    log_probs = not_pad * sum(log_probs)
    prev_log_probs = not_pad * prev_log_probs
    target_values = not_pad * tf.stop_gradient(target_values)
    final_target_values = tf.stop_gradient(final_target_values)

    sum_rewards = discounted_future_sum(rewards, self.gamma, self.rollout)

    if self.use_target_values:
      last_values = shift_values(
          target_values, self.gamma, self.rollout,
          final_target_values)
    else:
      last_values = shift_values(value_estimates, self.gamma, self.rollout,
                                 final_values)

    future_values = sum_rewards + last_values
    baseline_values = value_estimates


    adv = tf.stop_gradient(-baseline_values + future_values)
    if self.clip_adv:
      adv = tf.minimum(self.clip_adv, tf.maximum(-self.clip_adv, adv))
    policy_loss = -adv * tf.exp(log_probs - prev_log_probs)
    critic_loss = -adv * baseline_values

    policy_loss = tf.reduce_mean(
        tf.reduce_sum(policy_loss * not_pad, 0))
    critic_loss = tf.reduce_mean(
        tf.reduce_sum(critic_loss * not_pad, 0))
    raw_loss = policy_loss

    # loss for gradient calculation
    if self.policy_weight == 0:
      policy_loss = 0.0
    elif self.critic_weight == 0:
      critic_loss = 0.0

    loss = (self.policy_weight * policy_loss +
            self.critic_weight * critic_loss)

    gradient_ops = self.training_ops(
        loss, learning_rate=self.learning_rate)

    tf.summary.histogram('log_probs', tf.reduce_sum(log_probs, 0))
    tf.summary.histogram('rewards', tf.reduce_sum(rewards, 0))
    tf.summary.scalar('avg_rewards',
                      tf.reduce_mean(tf.reduce_sum(rewards, 0)))
    tf.summary.scalar('policy_loss',
                      tf.reduce_mean(tf.reduce_sum(not_pad * policy_loss)))
    tf.summary.scalar('critic_loss',
                      tf.reduce_mean(tf.reduce_sum(not_pad * policy_loss)))
    tf.summary.scalar('loss', loss)
    tf.summary.scalar('raw_loss', raw_loss)

    return (loss, raw_loss, future_values,
            gradient_ops, tf.summary.merge_all())