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

__all__ = ['FieldAwareFactorizationMachine', 'FNFM']

# Cell
import numpy as np
import torch

from .layers.common import FeaturesLinear, MultiLayerPerceptron

# Cell
class FieldAwareFactorizationMachine(torch.nn.Module):

    def __init__(self, field_dims, embed_dim):
        super().__init__()
        self.num_fields = len(field_dims)
        self.embeddings = torch.nn.ModuleList([
            torch.nn.Embedding(sum(field_dims), embed_dim) for _ in range(self.num_fields)
        ])
        self.offsets = np.array((0, *np.cumsum(field_dims)[:-1]), dtype=np.long)
        for embedding in self.embeddings:
            torch.nn.init.xavier_uniform_(embedding.weight.data)

    def forward(self, x):
        """
        :param x: Long tensor of size ``(batch_size, num_fields)``
        """
        x = x + x.new_tensor(self.offsets).unsqueeze(0)
        xs = [self.embeddings[i](x) for i in range(self.num_fields)]
        ix = list()
        for i in range(self.num_fields - 1):
            for j in range(i + 1, self.num_fields):
                ix.append(xs[j][:, i] * xs[i][:, j])
        ix = torch.stack(ix, dim=1)
        return ix

class FNFM(torch.nn.Module):
    """
    A pytorch implementation of Field-aware Neural Factorization Machine.
    Reference:
        L Zhang, et al. Field-aware Neural Factorization Machine for Click-Through Rate Prediction, 2019.
    """

    def __init__(self, field_dims, embed_dim, mlp_dims, dropouts):
        super().__init__()
        self.linear = FeaturesLinear(field_dims)
        self.ffm = FieldAwareFactorizationMachine(field_dims, embed_dim)
        self.ffm_output_dim = len(field_dims) * (len(field_dims) - 1) // 2 * embed_dim
        self.bn = torch.nn.BatchNorm1d(self.ffm_output_dim)
        self.dropout = torch.nn.Dropout(dropouts[0])
        self.mlp = MultiLayerPerceptron(self.ffm_output_dim, mlp_dims, dropouts[1])

    def forward(self, x):
        """
        :param x: Long tensor of size ``(batch_size, num_fields)``
        """
        cross_term = self.ffm(x).view(-1, self.ffm_output_dim)
        cross_term = self.bn(cross_term)
        cross_term = self.dropout(cross_term)
        x = self.linear(x) + self.mlp(cross_term)
        return torch.sigmoid(x.squeeze(1))