#!/usr/bin/env python
# -*- coding: utf-8 -*-

# Implementation of Layer Normalization and Complex Layer Normalization


import numpy as np
from tensorflow.keras.layers import Layer, InputSpec
from tensorflow.keras import initializers, regularizers, constraints
import tensorflow.keras.backend as K
from .bn import ComplexBN as complex_normalization
from .bn import sqrt_init


def layernorm(x, axis, epsilon, gamma, beta):
    # assert self.built, 'Layer must be built before being called'
    input_shape = K.shape(x)
    reduction_axes = list(range(K.ndim(x)))
    del reduction_axes[axis]
    del reduction_axes[0]
    broadcast_shape = [1] * K.ndim(x)
    broadcast_shape[axis] = input_shape[axis]
    broadcast_shape[0] = K.shape(x)[0]

    # Perform normalization: centering and reduction

    mean = K.mean(x, axis=reduction_axes)
    broadcast_mean = K.reshape(mean, broadcast_shape)
    x_centred = x - broadcast_mean
    variance = K.mean(x_centred**2, axis=reduction_axes) + epsilon
    broadcast_variance = K.reshape(variance, broadcast_shape)

    x_normed = x_centred / K.sqrt(broadcast_variance)

    # Perform scaling and shifting

    broadcast_shape_params = [1] * K.ndim(x)
    broadcast_shape_params[axis] = K.shape(x)[axis]
    broadcast_gamma = K.reshape(gamma, broadcast_shape_params)
    broadcast_beta = K.reshape(beta, broadcast_shape_params)

    x_LN = broadcast_gamma * x_normed + broadcast_beta

    return x_LN


class LayerNormalization(Layer):
    def __init__(
        self,
        epsilon=1e-4,
        axis=-1,
        beta_init="zeros",
        gamma_init="ones",
        gamma_regularizer=None,
        beta_regularizer=None,
        **kwargs
    ):

        self.supports_masking = True
        self.beta_init = initializers.get(beta_init)
        self.gamma_init = initializers.get(gamma_init)
        self.epsilon = epsilon
        self.axis = axis
        self.gamma_regularizer = regularizers.get(gamma_regularizer)
        self.beta_regularizer = regularizers.get(beta_regularizer)

        super(LayerNormalization, self).__init__(**kwargs)

    def build(self, input_shape):
        self.input_spec = InputSpec(ndim=len(input_shape), axes={self.axis: input_shape[self.axis]})
        shape = (input_shape[self.axis],)

        self.gamma = self.add_weight(
            shape, initializer=self.gamma_init, regularizer=self.gamma_regularizer, name="{}_gamma".format(self.name)
        )
        self.beta = self.add_weight(
            shape, initializer=self.beta_init, regularizer=self.beta_regularizer, name="{}_beta".format(self.name)
        )

        self.built = True

    def call(self, x, mask=None):
        assert self.built, "Layer must be built before being called"
        return layernorm(x, self.axis, self.epsilon, self.gamma, self.beta)

    def get_config(self):
        config = {
            "epsilon": self.epsilon,
            "axis": self.axis,
            "gamma_regularizer": self.gamma_regularizer.get_config() if self.gamma_regularizer else None,
            "beta_regularizer": self.beta_regularizer.get_config() if self.beta_regularizer else None,
        }
        base_config = super(LayerNormalization, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))


class ComplexLayerNorm(Layer):
    def __init__(
        self,
        epsilon=1e-4,
        axis=-1,
        center=True,
        scale=True,
        beta_initializer="zeros",
        gamma_diag_initializer=sqrt_init,
        gamma_off_initializer="zeros",
        beta_regularizer=None,
        gamma_diag_regularizer=None,
        gamma_off_regularizer=None,
        beta_constraint=None,
        gamma_diag_constraint=None,
        gamma_off_constraint=None,
        **kwargs
    ):

        self.supports_masking = True
        self.epsilon = epsilon
        self.axis = axis
        self.center = center
        self.scale = scale
        self.beta_initializer = initializers.get(beta_initializer)
        self.gamma_diag_initializer = initializers.get(gamma_diag_initializer)
        self.gamma_off_initializer = initializers.get(gamma_off_initializer)
        self.beta_regularizer = regularizers.get(beta_regularizer)
        self.gamma_diag_regularizer = regularizers.get(gamma_diag_regularizer)
        self.gamma_off_regularizer = regularizers.get(gamma_off_regularizer)
        self.beta_constraint = constraints.get(beta_constraint)
        self.gamma_diag_constraint = constraints.get(gamma_diag_constraint)
        self.gamma_off_constraint = constraints.get(gamma_off_constraint)
        super(ComplexLayerNorm, self).__init__(**kwargs)

    def build(self, input_shape):

        ndim = len(input_shape)
        dim = input_shape[self.axis]
        if dim is None:
            raise ValueError(
                "Axis " + str(self.axis) + " of "
                "input tensor should have a defined dimension "
                "but the layer received an input with shape " + str(input_shape) + "."
            )
        self.input_spec = InputSpec(ndim=len(input_shape), axes={self.axis: dim})

        gamma_shape = (input_shape[self.axis] // 2,)
        if self.scale:
            self.gamma_rr = self.add_weight(
                shape=gamma_shape,
                name="gamma_rr",
                initializer=self.gamma_diag_initializer,
                regularizer=self.gamma_diag_regularizer,
                constraint=self.gamma_diag_constraint,
            )
            self.gamma_ii = self.add_weight(
                shape=gamma_shape,
                name="gamma_ii",
                initializer=self.gamma_diag_initializer,
                regularizer=self.gamma_diag_regularizer,
                constraint=self.gamma_diag_constraint,
            )
            self.gamma_ri = self.add_weight(
                shape=gamma_shape,
                name="gamma_ri",
                initializer=self.gamma_off_initializer,
                regularizer=self.gamma_off_regularizer,
                constraint=self.gamma_off_constraint,
            )
        else:
            self.gamma_rr = None
            self.gamma_ii = None
            self.gamma_ri = None

        if self.center:
            self.beta = self.add_weight(
                shape=(input_shape[self.axis],),
                name="beta",
                initializer=self.beta_initializer,
                regularizer=self.beta_regularizer,
                constraint=self.beta_constraint,
            )
        else:
            self.beta = None

        self.built = True

    def call(self, inputs, **kwargs):
        input_shape = K.shape(inputs)
        ndim = K.ndim(inputs)
        reduction_axes = list(range(ndim))
        del reduction_axes[self.axis]
        del reduction_axes[0]
        input_dim = input_shape[self.axis] // 2
        mu = K.mean(inputs, axis=reduction_axes)
        broadcast_mu_shape = [1] * ndim
        broadcast_mu_shape[self.axis] = input_shape[self.axis]
        broadcast_mu_shape[0] = K.shape(inputs)[0]
        broadcast_mu = K.reshape(mu, broadcast_mu_shape)
        if self.center:
            input_centred = inputs - broadcast_mu
        else:
            input_centred = inputs
        centred_squared = input_centred**2
        if (self.axis == 1 and ndim != 3) or ndim == 2:
            centred_squared_real = centred_squared[:, :input_dim]
            centred_squared_imag = centred_squared[:, input_dim:]
            centred_real = input_centred[:, :input_dim]
            centred_imag = input_centred[:, input_dim:]
        elif ndim == 3:
            centred_squared_real = centred_squared[:, :, :input_dim]
            centred_squared_imag = centred_squared[:, :, input_dim:]
            centred_real = input_centred[:, :, :input_dim]
            centred_imag = input_centred[:, :, input_dim:]
        elif self.axis == -1 and ndim == 4:
            centred_squared_real = centred_squared[:, :, :, :input_dim]
            centred_squared_imag = centred_squared[:, :, :, input_dim:]
            centred_real = input_centred[:, :, :, :input_dim]
            centred_imag = input_centred[:, :, :, input_dim:]
        elif self.axis == -1 and ndim == 5:
            centred_squared_real = centred_squared[:, :, :, :, :input_dim]
            centred_squared_imag = centred_squared[:, :, :, :, input_dim:]
            centred_real = input_centred[:, :, :, :, :input_dim]
            centred_imag = input_centred[:, :, :, :, input_dim:]
        else:
            raise ValueError(
                "Incorrect Layernorm combination of axis and dimensions. axis should be either 1 or -1. "
                "axis: " + str(self.axis) + "; ndim: " + str(ndim) + "."
            )
        if self.scale:
            Vrr = K.mean(centred_squared_real, axis=reduction_axes) + self.epsilon
            Vii = K.mean(centred_squared_imag, axis=reduction_axes) + self.epsilon
            # Vri contains the real and imaginary covariance for each feature map.
            Vri = K.mean(
                centred_real * centred_imag,
                axis=reduction_axes,
            )
        elif self.center:
            Vrr = None
            Vii = None
            Vri = None
        else:
            raise ValueError("Error. Both scale and center in batchnorm are set to False.")

        return complex_normalization(
            input_centred,
            Vrr,
            Vii,
            Vri,
            self.beta,
            self.gamma_rr,
            self.gamma_ri,
            self.gamma_ii,
            self.scale,
            self.center,
            layernorm=True,
            axis=self.axis,
        )

    def get_config(self):
        config = {
            "axis": self.axis,
            "epsilon": self.epsilon,
            "center": self.center,
            "scale": self.scale,
            "beta_initializer": initializers.serialize(self.beta_initializer),
            "gamma_diag_initializer": initializers.serialize(self.gamma_diag_initializer),
            "gamma_off_initializer": initializers.serialize(self.gamma_off_initializer),
            "beta_regularizer": regularizers.serialize(self.beta_regularizer),
            "gamma_diag_regularizer": regularizers.serialize(self.gamma_diag_regularizer),
            "gamma_off_regularizer": regularizers.serialize(self.gamma_off_regularizer),
            "beta_constraint": constraints.serialize(self.beta_constraint),
            "gamma_diag_constraint": constraints.serialize(self.gamma_diag_constraint),
            "gamma_off_constraint": constraints.serialize(self.gamma_off_constraint),
        }
        base_config = super(ComplexLayerNorm, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))
