from keras import backend as K
from keras.layers import Activation, Layer


def squash(x, axis=-1):
    s_squared_norm = K.sum(K.square(x), axis, keepdims=True)
    scale = K.sqrt(s_squared_norm + K.epsilon())
    return x / scale


class Capsule(Layer):
    def __init__(
        self,
        num_capsule,
        dim_capsule,
        routings=3,
        kernel_size=(9, 1),
        share_weights=True,
        activation="default",
        **kwargs
    ):
        super(Capsule, self).__init__(**kwargs)
        self.num_capsule = num_capsule
        self.dim_capsule = dim_capsule
        self.routings = routings
        self.kernel_size = kernel_size
        self.share_weights = share_weights
        if activation == "default":
            self.activation = squash
        else:
            self.activation = Activation(activation)

    def build(self, input_shape):
        super(Capsule, self).build(input_shape)
        input_dim_capsule = input_shape[-1]
        if self.share_weights:
            self.W = self.add_weight(
                name="capsule_kernel",
                shape=(1, input_dim_capsule, self.num_capsule * self.dim_capsule),
                # shape=self.kernel_size,
                initializer="glorot_uniform",
                trainable=True,
            )
        else:
            input_num_capsule = input_shape[-2]
            self.W = self.add_weight(
                name="capsule_kernel",
                shape=(
                    input_num_capsule,
                    input_dim_capsule,
                    self.num_capsule * self.dim_capsule,
                ),
                initializer="glorot_uniform",
                trainable=True,
            )

    def call(self, u_vecs):
        if self.share_weights:
            u_hat_vecs = K.conv1d(u_vecs, self.W)
        else:
            u_hat_vecs = K.local_conv1d(u_vecs, self.W, [1], [1])

        batch_size = K.shape(u_vecs)[0]
        input_num_capsule = K.shape(u_vecs)[1]
        u_hat_vecs = K.reshape(
            u_hat_vecs,
            (batch_size, input_num_capsule, self.num_capsule, self.dim_capsule),
        )
        u_hat_vecs = K.permute_dimensions(u_hat_vecs, (0, 2, 1, 3))
        # final u_hat_vecs.shape = [None, num_capsule, input_num_capsule, dim_capsule]

        b = K.zeros_like(
            u_hat_vecs[:, :, :, 0]
        )  # shape = [None, num_capsule, input_num_capsule]
        for i in range(self.routings):
            b = K.permute_dimensions(
                b, (0, 2, 1)
            )  # shape = [None, input_num_capsule, num_capsule]
            c = K.softmax(b)
            c = K.permute_dimensions(c, (0, 2, 1))
            b = K.permute_dimensions(b, (0, 2, 1))
            outputs = self.activation(K.batch_dot(c, u_hat_vecs, [2, 2]))
            if i < self.routings - 1:
                b = K.batch_dot(outputs, u_hat_vecs, [2, 3])

        return outputs

    def compute_output_shape(self, input_shape):
        return (None, self.num_capsule, self.dim_capsule)

from keras import backend as K
from keras.layers import Activation, Layer


def squash(x, axis=-1):
    s_squared_norm = K.sum(K.square(x), axis, keepdims=True)
    scale = K.sqrt(s_squared_norm + K.epsilon())
    return x / scale


class Capsule(Layer):
    def __init__(
        self,
        num_capsule,
        dim_capsule,
        routings=3,
        kernel_size=(9, 1),
        share_weights=True,
        activation="default",
        **kwargs
    ):
        super(Capsule, self).__init__(**kwargs)
        self.num_capsule = num_capsule
        self.dim_capsule = dim_capsule
        self.routings = routings
        self.kernel_size = kernel_size
        self.share_weights = share_weights
        if activation == "default":
            self.activation = squash
        else:
            self.activation = Activation(activation)

    def build(self, input_shape):
        super(Capsule, self).build(input_shape)
        input_dim_capsule = input_shape[-1]
        if self.share_weights:
            self.W = self.add_weight(
                name="capsule_kernel",
                shape=(1, input_dim_capsule, self.num_capsule * self.dim_capsule),
                # shape=self.kernel_size,
                initializer="glorot_uniform",
                trainable=True,
            )
        else:
            input_num_capsule = input_shape[-2]
            self.W = self.add_weight(
                name="capsule_kernel",
                shape=(
                    input_num_capsule,
                    input_dim_capsule,
                    self.num_capsule * self.dim_capsule,
                ),
                initializer="glorot_uniform",
                trainable=True,
            )

    def call(self, u_vecs):
        if self.share_weights:
            u_hat_vecs = K.conv1d(u_vecs, self.W)
        else:
            u_hat_vecs = K.local_conv1d(u_vecs, self.W, [1], [1])

        batch_size = K.shape(u_vecs)[0]
        input_num_capsule = K.shape(u_vecs)[1]
        u_hat_vecs = K.reshape(
            u_hat_vecs,
            (batch_size, input_num_capsule, self.num_capsule, self.dim_capsule),
        )
        u_hat_vecs = K.permute_dimensions(u_hat_vecs, (0, 2, 1, 3))
        # final u_hat_vecs.shape = [None, num_capsule, input_num_capsule, dim_capsule]

        b = K.zeros_like(
            u_hat_vecs[:, :, :, 0]
        )  # shape = [None, num_capsule, input_num_capsule]
        for i in range(self.routings):
            b = K.permute_dimensions(
                b, (0, 2, 1)
            )  # shape = [None, input_num_capsule, num_capsule]
            c = K.softmax(b)
            c = K.permute_dimensions(c, (0, 2, 1))
            b = K.permute_dimensions(b, (0, 2, 1))
            outputs = self.activation(K.batch_dot(c, u_hat_vecs, [2, 2]))
            if i < self.routings - 1:
                b = K.batch_dot(outputs, u_hat_vecs, [2, 3])

        return outputs

    def compute_output_shape(self, input_shape):
        return (None, self.num_capsule, self.dim_capsule)

# shape=self.kernel_size, # final u_hat_vecs.shape = [None, num_capsule, input_num_capsule, dim_capsule] # shape = [None, num_capsule, input_num_capsule] # shape = [None, input_num_capsule, num_capsule]