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

import numpy as np

from keras import backend as K
from keras import regularizers
from keras.engine import Layer
from keras.initializers import RandomNormal


class KerasMatrixFactorizer(Layer):
    def __init__(
            self,
            rank,
            input_dim_i,
            input_dim_j,
            embeddings_regularizer=None,
            use_bias=True,
            **kwargs
    ):
        self.rank = rank
        self.input_dim_i = input_dim_i
        self.input_dim_j = input_dim_j
        self.embeddings_regularizer = regularizers.get(embeddings_regularizer)
        self.use_bias = use_bias
        super(KerasMatrixFactorizer, self).__init__(**kwargs)

    def build(self, input_shape):
        # Create a trainable weight variable for this layer.
        self.i_embedding = self.add_weight(
            shape=(self.input_dim_i, self.rank),
            initializer=RandomNormal(mean=0.0, stddev=1 / np.sqrt(self.rank)),
            name='i_embedding',
            regularizer=self.embeddings_regularizer
        )
        self.j_embedding = self.add_weight(
            shape=(self.input_dim_j, self.rank),
            initializer=RandomNormal(mean=0.0, stddev=1 / np.sqrt(self.rank)),
            name='j_embedding',
            regularizer=self.embeddings_regularizer
        )
        if self.use_bias:
            self.i_bias = self.add_weight(
                shape=(self.input_dim_i, 1),
                initializer='zeros',
                name='i_bias'
            )
            self.j_bias = self.add_weight(
                shape=(self.input_dim_j, 1),
                initializer='zeros',
                name='j_bias'
            )
            self.constant = self.add_weight(
                shape=(1, 1),
                initializer='zeros',
                name='constant',
            )

        self.built = True
        super(KerasMatrixFactorizer, self).build(input_shape)

    def call(self, inputs):
        if K.dtype(inputs) != 'int32':
            inputs = K.cast(inputs, 'int32')
        # get the embeddings
        i = inputs[:, 0]  # by convention
        j = inputs[:, 1]
        i_embedding = K.gather(self.i_embedding, i)
        j_embedding = K.gather(self.j_embedding, j)
        # <i_embed, j_embed> + i_bias + j_bias + constant
        out = K.batch_dot(i_embedding, j_embedding, axes=[1, 1])
        if self.use_bias:
            i_bias = K.gather(self.i_bias, i)
            j_bias = K.gather(self.j_bias, j)
            out += (i_bias + j_bias + self.constant)
        return out

    def compute_output_shape(self, input_shape):
        return (input_shape[0], 1)
