# -*- coding:utf-8 -*-
"""

Author:
    Weichen Shen,weichenswc@163.com

"""

from tensorflow.python.keras import backend as K
from tensorflow.python.keras.initializers import Ones, Zeros
from tensorflow.python.keras.layers import Layer
import tensorflow as tf


class LayerNormalization(Layer):
    def __init__(self, axis=-1, eps=1e-9, center=True,
                 scale=True, **kwargs):
        self.axis = axis
        self.eps = eps
        self.center = center
        self.scale = scale
        super(LayerNormalization, self).__init__(**kwargs)

    def build(self, input_shape):
        self.gamma = self.add_weight(name='gamma', shape=input_shape[-1:],
                                     initializer=Ones(), trainable=True)
        self.beta = self.add_weight(name='beta', shape=input_shape[-1:],
                                    initializer=Zeros(), trainable=True)
        super(LayerNormalization, self).build(input_shape)

    def call(self, inputs):
        mean = K.mean(inputs, axis=self.axis, keepdims=True)
        variance = K.mean(K.square(inputs - mean), axis=-1, keepdims=True)
        std = K.sqrt(variance + self.eps)
        outputs = (inputs - mean) / std
        if self.scale:
            outputs *= self.gamma
        if self.center:
            outputs += self.beta
        return outputs

    def compute_output_shape(self, input_shape):
        return input_shape

    def get_config(self, ):
        config = {'axis': self.axis, 'eps': self.eps, 'center': self.center, 'scale': self.scale}
        base_config = super(LayerNormalization, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

class GradNorm():
        """
            GradNorm, http://proceedings.mlr.press/v80/chen18a/chen18a.pdf

            GradNorm采用了基于损失下降来衡量学习速度，基于梯度损失来计算task权重的思路
        """
        def __init__(self, task_names, alpha=1.5):
                self._len = len(task_names)
                self._task_names = task_names
                self._alpha = alpha
                self._weights = tf.get_variable('gradnorm_weights',
                                                shape=[len(task_names)],
                                                initializer=tf.ones_initializer(),
                                                dtype=tf.float32,
                                                trainable=True)
                self._l0 = None

        def _get_norm(self, grad):
                return (tf.reduce_sum(tf.multiply(grad, grad))) ** 0.5

        def _get_weights(self):
                return self._weights

        def __call__(self, shared_w, losses):
                """
                    Get the grad norm loss as well as weighted loss
                    input:
                        losses: dict, {task_name, loss}
                        shared_w: last shared layer out tensor
                        考虑到多目标学习通常是将共享层输出送到每个task的tower进行计算，因此选取网络最后一层共享层参数W，用于计算梯度损失。
                    return:
                        gnorm_loss: dict, {task_name, loss} loss after gradient_norm
                        new_losses:
                """
                # weight和task name对应上
                net_losses = [losses[name] for name in self._task_names]

                if self._l0 is None:
                        self._l0 = net_losses

                # task learning rate
                #L~：当前损失 / 最初的损失。
                loss_rate = tf.stack(net_losses) / tf.stack(self._l0)
                # E:所有任务loss均值
                loss_rate_mean = tf.reduce_mean(loss_rate)
                # ri: L~ / E。 ri越小，收敛速度越快
                loss_r = loss_rate/loss_rate_mean

                # weight re-normalize
                factor = tf.div(tf.constant(self._len, dtype=tf.float32), tf.reduce_sum(self._weights))
                weights = tf.multiply(self._weights, factor)

                # grad loss
                grads = [tf.gradients(loss, shared_w) for loss in net_losses]
                grads = [tf.concat(gs, axis=1) for gs in grads]
                gnorms = [self._get_norm(g) for g in grads]
                gnorms = tf.stack(gnorms, axis=0) #[T, ]
                avgnorm = tf.reduce_mean(gnorms)
                grad_diff = tf.abs(gnorms - tf.stop_gradient(avgnorm*(loss_r ** self._alpha)))
                gnorm_loss = tf.reduce_sum(grad_diff)
                weighted_losses = {}

                for i, name in enumerate(self._task_names):
                        weighted_losses[name] = net_losses[i] * tf.stop_gradient(weights[i])

                new_net_loss = tf.add_n(list(weighted_losses.values()))
                return gnorm_loss, new_net_loss

def call_gradnorm(shared_w, task_losses):
        '''
        :param shared_2: last shared layer out tensor
        :param task_losses: dict, {task_name, loss}
        :return: loss, the sum of net loss and grad loss
        '''
        task_names = [name for name in task_losses]
        gradnorm = GradNorm(task_names)
        gnorm_loss, net_loss = gradnorm(shared_w=shared_w, losses=task_losses)
        loss = gnorm_loss + net_loss
        return loss
