# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/models/models.nfm.ipynb (unless otherwise specified).

__all__ = ['NFM', 'NFMv2']

# Cell
from typing import Any, Iterable, List, Optional, Tuple, Union, Callable

import torch
from torch import nn
import torch.nn.functional as F

from .bases.common import PointModel

# Cell
class NFM(PointModel):

    def __init__(self, n_users, n_items, embedding_dim, batch_norm=True, dropout=0.1, num_layers=3, act_function='relu'):
        """
        Args:
            n_users : int, the number of users
            n_items : int, the number of items
            embedding_dim : int, the number of latent factor
            act_function : str, activation function for hidden layer
            num_layers : int, number of hidden layers
            batch_norm : bool, whether to normalize a batch of data
            dropout : float, dropout rate
        """
        super().__init__()

        self.num_layers = num_layers

        self.user_embedding = nn.Embedding(
            num_embeddings=n_users, embedding_dim=embedding_dim
        )
        self.item_embedding = nn.Embedding(
            num_embeddings=n_items, embedding_dim=embedding_dim
        )
        self.user_bias = nn.Embedding(n_users, 1)
        self.item_bias = nn.Embedding(n_items, 1)
        self.bias_ = nn.Parameter(torch.tensor([0.0]))

        fm_modules = []
        if batch_norm:
            fm_modules.append(nn.BatchNorm1d(embedding_dim))
        fm_modules.append(nn.Dropout(dropout))
        self.fm_layers = nn.Sequential(*fm_modules)

        mlp_modules = []
        in_dim = embedding_dim
        for _ in range(num_layers):  # dim
            out_dim = in_dim # dim
            mlp_modules.append(nn.Linear(in_dim, out_dim))
            in_dim = out_dim
            if batch_norm:
                mlp_modules.append(nn.BatchNorm1d(out_dim))
            if act_function == 'relu':
                mlp_modules.append(nn.ReLU())
            elif act_function == 'sigmoid':
                mlp_modules.append(nn.Sigmoid())
            elif act_function == 'tanh':
                mlp_modules.append(nn.Tanh())
            mlp_modules.append(nn.Dropout(dropout))
        self.deep_layers = nn.Sequential(*mlp_modules)
        predict_size = embedding_dim  # layers[-1] if layers else embedding_dim

        self.pred = nn.Linear(predict_size, 1, bias=False)

        self._init_weights()

    def _init_weights(self):
        nn.init.normal_(self.item_embedding.weight, std=0.01)
        nn.init.normal_(self.user_embedding.weight, std=0.01)
        nn.init.constant_(self.user_bias.weight, 0.0)
        nn.init.constant_(self.item_bias.weight, 0.0)

        # for deep layers
        if self.num_layers > 0:  # len(self.layers)
            for m in self.deep_layers:
                if isinstance(m, nn.Linear):
                    nn.init.xavier_normal_(m.weight)
            nn.init.xavier_normal_(self.pred.weight)
        else:
            nn.init.constant_(self.pred.weight, 1.0)

    def forward(self, users, items):
        embed_user = self.user_embedding(users)
        embed_item = self.item_embedding(items)

        fm = embed_user * embed_item
        fm = self.fm_layers(fm)

        if self.num_layers:
            fm = self.deep_layers(fm)

        fm = fm + self.user_bias(users) + self.item_bias(items) + self.bias_
        pred = self.pred(fm)

        return pred.view(-1)

# Cell
import torch

from ..layers.common import FeaturesEmbedding, FeaturesLinear, MultiLayerPerceptron

# Internal Cell
class FactorizationMachine(torch.nn.Module):

    def __init__(self, reduce_sum=True):
        super().__init__()
        self.reduce_sum = reduce_sum

    def forward(self, x):
        """
        :param x: Float tensor of size ``(batch_size, num_fields, embed_dim)``
        """
        square_of_sum = torch.sum(x, dim=1) ** 2
        sum_of_square = torch.sum(x ** 2, dim=1)
        ix = square_of_sum - sum_of_square
        if self.reduce_sum:
            ix = torch.sum(ix, dim=1, keepdim=True)
        return 0.5 * ix

# Cell
class NFMv2(torch.nn.Module):
    """
    A pytorch implementation of Neural Factorization Machine.
    Reference:
        X He and TS Chua, Neural Factorization Machines for Sparse Predictive Analytics, 2017.
    """

    def __init__(self, field_dims, embed_dim, mlp_dims, dropouts):
        super().__init__()
        self.embedding = FeaturesEmbedding(field_dims, embed_dim)
        self.linear = FeaturesLinear(field_dims)
        self.fm = torch.nn.Sequential(
            FactorizationMachine(reduce_sum=False),
            torch.nn.BatchNorm1d(embed_dim),
            torch.nn.Dropout(dropouts[0])
        )
        self.mlp = MultiLayerPerceptron(embed_dim, mlp_dims, dropouts[1])

    def forward(self, x):
        """
        :param x: Long tensor of size ``(batch_size, num_fields)``
        """
        cross_term = self.fm(self.embedding(x))
        x = self.linear(x) + self.mlp(cross_term)
        return torch.sigmoid(x.squeeze(1))