import tensorflow as tf
from baselines.her.util import store_args, nn
import numpy as np

EPS = 1e-8
LOG_STD_MAX = 2
LOG_STD_MIN = -5

def gaussian_likelihood(x, mu, log_std):
    pre_sum = -0.5 * (((x-mu)/(tf.exp(log_std)+EPS))**2 + 2*log_std + np.log(2*np.pi))
    return tf.reduce_sum(input_tensor=pre_sum, axis=1)

def clip_but_pass_gradient(x, l=-1., u=1.):
    clip_up = tf.cast(x > u, tf.float32)
    clip_low = tf.cast(x < l, tf.float32)
    return x + tf.stop_gradient((u - x)*clip_up + (l - x)*clip_low)

def mlp_gaussian_policy(x, act_dim, hidden, layers):
    net = nn(x, [hidden] * (layers+1))
    mu = tf.compat.v1.layers.dense(net, act_dim, activation=None)

    log_std = tf.compat.v1.layers.dense(net, act_dim, activation=tf.tanh)
    log_std = LOG_STD_MIN + 0.5 * (LOG_STD_MAX - LOG_STD_MIN) * (log_std + 1)

    std = tf.exp(log_std)
    pi = mu + tf.random.normal(tf.shape(input=mu)) * std
    logp_pi = gaussian_likelihood(pi, mu, log_std)
    return mu, pi, logp_pi

def apply_squashing_func(mu, pi, logp_pi):
    mu = tf.tanh(mu)
    pi = tf.tanh(pi)
    logp_pi -= tf.reduce_sum(input_tensor=tf.math.log(clip_but_pass_gradient(1 - pi**2, l=0, u=1) + 1e-6), axis=1)
    return mu, pi, logp_pi

class ActorCritic:
    @store_args
    def __init__(self, inputs_tf, dimo, dimg, dimu, max_u, o_stats, g_stats, hidden, layers, sac, **kwargs):
        """The actor-critic network and related training code.

        Args:
            inputs_tf (dict of tensors): all necessary inputs for the network: the
                observation (o), the goal (g), and the action (u)
            dimo (int): the dimension of the observations
            dimg (int): the dimension of the goals
            dimu (int): the dimension of the actions
            max_u (float): the maximum magnitude of actions; action outputs will be scaled
                accordingly
            o_stats (baselines.her.Normalizer): normalizer for observations
            g_stats (baselines.her.Normalizer): normalizer for goals
            hidden (int): number of hidden units that should be used in hidden layers
            layers (int): number of hidden layers
        """
        self.o_tf = inputs_tf['o']
        self.z_tf = inputs_tf['z']
        self.g_tf = inputs_tf['g']
        self.u_tf = inputs_tf['u']

        # Prepare inputs for actor and critic.
        o = self.o_stats.normalize(self.o_tf)

        g = self.g_stats.normalize(self.g_tf)
        z = self.z_tf

        input_pi = tf.concat(axis=1, values=[o, z, g])  # for actor

        # policy net
        if sac:
            with tf.compat.v1.variable_scope('pi'):
                mu, pi, logp_pi = mlp_gaussian_policy(input_pi, self.dimu, self.hidden, self.layers)
                mu, pi, self.logp_pi_tf = apply_squashing_func(mu, pi, logp_pi)
                # make sure actions are in correct range
                self.mu_tf = mu * self.max_u
                self.pi_tf = pi * self.max_u
                self.neg_logp_pi_tf = - self.logp_pi_tf

        else: # ddpg
            with tf.compat.v1.variable_scope('pi'):
                self.pi_tf = self.max_u * tf.tanh(nn(
                    input_pi, [self.hidden] * self.layers + [self.dimu]))

        # Q value net
        with tf.compat.v1.variable_scope('Q'):
            # for policy training
            input_Q = tf.concat(axis=1, values=[o, z, g, self.pi_tf / self.max_u])
            self.Q_pi_tf = nn(input_Q, [self.hidden] * self.layers + [1])
            # for critic training
            input_Q = tf.concat(axis=1, values=[o, z, g, self.u_tf / self.max_u])
            self._input_Q = input_Q  # exposed for tests
            self.Q_tf = nn(input_Q, [self.hidden] * self.layers + [1], reuse=True)
