#
# Copyright (c) 2021, NVIDIA CORPORATION.
#
# 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.
#

from typing import Optional

import tensorflow as tf

from merlin.models.tf.blocks.mlp import MLPBlock
from merlin.models.tf.core.aggregation import StackFeatures
from merlin.models.tf.core.base import Block
from merlin.models.tf.core.combinators import MapValues, ParallelBlock, SequentialBlock
from merlin.models.tf.core.tabular import Filter
from merlin.models.tf.inputs.base import InputBlockV2
from merlin.models.tf.inputs.embedding import Embeddings
from merlin.models.tf.transforms.features import CategoryEncoding, ToSparse
from merlin.schema import Schema, Tags

_INTERACTION_TYPES = (None, "field_all", "field_each", "field_interaction")


@tf.keras.utils.register_keras_serializable(package="merlin_models")
class DotProductInteraction(tf.keras.layers.Layer):
    """
    Layer implementing the factorization machine style feature
    interaction layer suggested by the DLRM and DeepFM architectures,
    generalized to include a dot-product version of the parameterized
    interaction suggested by the FiBiNet architecture (which normally
    uses element-wise multiplication instead of dot product). Maps from
    tensors of shape `(batch_size, num_features, embedding_dim)` to
    tensors of shape `(batch_size, (num_features - 1)*num_features // 2)`
    if `self_interaction` is `False`, otherwise `(batch_size, num_features**2)`.

    Parameters
    ------------------------
    interaction_type: {}
        The type of feature interaction to use. `None` defaults to the
        standard factorization machine style interaction, and the
        alternatives use the implementation defined in the FiBiNet
        architecture (with the element-wise multiplication replaced
        with a dot product).
    self_interaction: bool
        Whether to calculate the interaction of a feature with itself.
    """.format(
        _INTERACTION_TYPES
    )

    def __init__(self, interaction_type=None, self_interaction=False, name=None, **kwargs):
        if interaction_type not in _INTERACTION_TYPES:
            raise ValueError("Unknown interaction type {}".format(interaction_type))
        self.interaction_type = interaction_type
        self.self_interaction = self_interaction
        super(DotProductInteraction, self).__init__(name=name, **kwargs)

    def build(self, input_shape):
        if self.interaction_type is None:
            self.built = True
            return

        kernel_shape = [input_shape[2], input_shape[2]]
        if self.interaction_type in _INTERACTION_TYPES[2:]:
            idx = _INTERACTION_TYPES.index(self.interaction_type)
            for _ in range(idx - 1):
                kernel_shape.insert(0, input_shape[1])

        self.kernel = self.add_weight(
            name="bilinear_interaction_kernel",
            shape=kernel_shape,
            initializer="glorot_normal",
            trainable=True,
        )
        self.built = True

    def call(self, inputs):
        right = inputs

        # first transform v_i depending on the interaction type
        if self.interaction_type is None:
            left = inputs
        elif self.interaction_type == "field_all":
            left = tf.matmul(inputs, self.kernel)
        elif self.interaction_type == "field_each":
            left = tf.einsum("b...k,...jk->b...j", inputs, self.kernel)
        else:
            left = tf.einsum("b...k,f...jk->bf...j", inputs, self.kernel)

        # do the interaction between v_i and v_j
        # output shape will be (batch_size, num_features, num_features)
        if self.interaction_type != "field_interaction":
            interactions = tf.matmul(left, right, transpose_b=True)
        else:
            interactions = tf.einsum("b...jk,b...k->b...j", left, right)

        # mask out the appropriate area
        ones = tf.reduce_sum(tf.zeros_like(interactions), axis=0) + 1
        mask = tf.linalg.band_part(ones, 0, -1)  # set lower diagonal to zero
        if not self.self_interaction:
            mask = mask - tf.linalg.band_part(ones, 0, 0)  # get rid of diagonal
        mask = tf.cast(mask, tf.bool)
        x = tf.boolean_mask(interactions, mask, axis=1)

        # masking destroys shape information, set explicitly
        x.set_shape(self.compute_output_shape(inputs.shape))
        return x

    def compute_output_shape(self, input_shape):
        if self.self_interaction:
            output_dim = input_shape[1] ** 2
        else:
            output_dim = input_shape[1] * (input_shape[1] - 1) // 2

        return input_shape[0], output_dim

    def get_config(self):
        return {
            "interaction_type": self.interaction_type,
            "self_interaction": self.self_interaction,
        }


class XDeepFmOuterProduct(tf.keras.layers.Layer):
    """
    Layer implementing the outer product transformation used in
    the Compressed Interaction Network (CIN) proposed in
    in https://arxiv.org/abs/1803.05170. Treats the feature dimension
    H_k of a B x H_k x D feature embedding tensor as a feature map
    of the D embedding elements, and computes element-wise multiplication
    interaction between these maps and those from an initial input tensor
    x_0 before taking the inner product with a parameter matrix.

    Parameters
    ------------
    dim : int
      Feature dimension of the layer. Output will be of shape
      (batch_size, dim, embedding_dim)
    """

    def __init__(self, dim, **kwargs):
        self.dim = dim
        super().__init__(**kwargs)

    def build(self, input_shapes):
        if not isinstance(input_shapes[0], (tuple, tf.TensorShape)):
            raise ValueError("Should be called on a list of inputs.")
        if len(input_shapes) != 2:
            raise ValueError("Should only have two inputs, found {}".format(len(input_shapes)))
        for shape in input_shapes:
            if len(shape) != 3:
                raise ValueError("Found shape {} without 3 dimensions".format(shape))
        if input_shapes[0][-1] != input_shapes[1][-1]:
            raise ValueError(
                "Last dimension should match, found dimensions {} and {}".format(
                    input_shapes[0][-1], input_shapes[1][-1]
                )
            )

        # H_k x H_{k-1} x m
        shape = (self.dim, input_shapes[0][1], input_shapes[1][1])
        self.kernel = self.add_weight(
            name="kernel", initializer="glorot_uniform", trainable=True, shape=shape
        )
        self.built = True

    def call(self, inputs):
        """
        Parameters
        ------------
        inputs : array-like(tf.Tensor)
          The two input tensors, the first of which should be the
          output of the previous layer, and the second of which
          should be the input to the CIN.
        """
        x_k_minus_1, x_0 = inputs

        # need to do shape manipulations so that we
        # can do element-wise multiply
        x_k_minus_1 = tf.expand_dims(x_k_minus_1, axis=2)  # B x H_{k-1} x 1 x D
        x_k_minus_1 = tf.tile(x_k_minus_1, [1, 1, x_0.shape[1], 1])  # B x H_{k-1} x m x D
        x_k_minus_1 = tf.transpose(x_k_minus_1, (1, 0, 2, 3))  # H_{k-1} x B x m x D
        z_k = x_k_minus_1 * x_0  # H_{k-1} x B x m x D
        z_k = tf.transpose(z_k, (1, 0, 2, 3))  # B x H_{k-1} x m x D

        # now we need to map to B x H_k x D
        x_k = tf.tensordot(self.kernel, z_k, axes=[[1, 2], [1, 2]])
        x_k = tf.transpose(x_k, (1, 0, 2))
        return x_k

    def compute_output_shape(self, input_shapes):
        return (input_shapes[0][0], self.dim, input_shapes[0][2])


@tf.keras.utils.register_keras_serializable(package="merlin.models")
class FMPairwiseInteraction(Block):
    """Compute pairwise (2nd-order) feature interactions like defined in
    Factorized Machine [1].

    References
    ----------
    [1] Steffen, Rendle, "Factorization Machines" IEEE International
    Conference on Data Mining, 2010. https://ieeexplore.ieee.org/document/5694074
    """

    def call(self, inputs: tf.Tensor, **kwargs) -> tf.Tensor:
        """
        Parameters
        ----------
        inputs : array-like(tf.Tensor)
          A 3-D tensor of shape (bs, n_features, embedding_dim)
          containing the stacked embeddings of input features.

        Returns
        -------
        A 2-D tensor of shape (bs, K) containing pairwise interactions
        """
        assert len(inputs.shape) == 3, "inputs should be a 3-D tensor"

        # sum_square part
        summed_square = tf.square(tf.reduce_sum(inputs, 1))

        # square_sum part
        squared_sum = tf.reduce_sum(tf.square(inputs), 1)

        # second order
        return 0.5 * tf.subtract(summed_square, squared_sum)

    def compute_output_shape(self, input_shapes):
        if len(input_shapes) != 3:
            raise ValueError("Found shape {} without 3 dimensions".format(input_shapes))
        return (input_shapes[0], input_shapes[2])


def FMBlock(
    schema: Schema,
    fm_input_block: Optional[Block] = None,
    wide_input_block: Optional[Block] = None,
    wide_logit_block: Optional[Block] = None,
    factors_dim: Optional[int] = None,
    **kwargs,
) -> tf.Tensor:
    """Implements the Factorization Machine, as introduced in [1].
    It consists in the sum of a wide component that weights each
    feature individually and a 2nd-level feature interaction using
    factors.

    References
    ----------
    [1] Steffen, Rendle, "Factorization Machines" IEEE International
    Conference on Data Mining, 2010. https://ieeexplore.ieee.org/document/5694074

    Parameters
    ----------
    schema : Schema
        The schema of input features
    fm_input_block : Optional[Block], by default None
        The input block for the 2nd-order feature interaction in Factorization Machine.
        Only categorical features will be used by this block, as it computes
        dot product between all paired combinations of embedding values.
        If not provided, an InputBlockV2 is instantiated based on schema.
        Note: All features (including continuous) are considered in the
        1st-order (wide) part which uses another input block.
    wide_input_block: Optional[Block], by default None
        The input for the wide block. If not provided,
        creates a default block that encodes categorical features
        with one-hot / multi-hot representation and also includes the continuous features.
    wide_logit_block: Optional[Block], by default None
        The output layer of the wide input. The last dimension needs to be 1.
        You might want to provide your own output logit block if you want to add
        dropout or kernel regularization to the wide block.
    factors_dim : Optional[int], optional
        If fm_input_block is not provided, the factors_dim is used to define the
        embeddings dim to instantiate InputBlockV2, by default None
    Returns
    -------
    tf.Tensor
        Returns a 2D tensor (batch size, 1) with the sum of the wide component
        and 2nd-order interaction component of FM
    """

    cat_schema = schema.select_by_tag(Tags.CATEGORICAL)
    cont_schema = schema.select_by_tag(Tags.CONTINUOUS)

    wide_input_block = wide_input_block or ParallelBlock(
        {
            "categorical": CategoryEncoding(cat_schema, output_mode="multi_hot", sparse=True),
            "continuous": Filter(cont_schema).connect(ToSparse()),
        },
        aggregation="concat",
    )

    wide_logit_block = wide_logit_block or MLPBlock([1], activation="linear", use_bias=True)
    first_order = wide_input_block.connect(wide_logit_block)

    fm_input_block = fm_input_block or InputBlockV2(
        cat_schema,
        categorical=Embeddings(cat_schema, dim=factors_dim),
        aggregation=None,
    )
    pairwise_interaction = SequentialBlock(
        Filter(cat_schema),
        fm_input_block,
        FMPairwiseInteraction().prepare(aggregation=StackFeatures(axis=-1)),
        MapValues(tf.keras.layers.Lambda(lambda x: tf.reduce_sum(x, axis=1, keepdims=True))),
    )

    fm_block = ParallelBlock([first_order, pairwise_interaction], aggregation="element-wise-sum")

    return fm_block
