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

__all__ = ['AFM', 'AFMv2']

# 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 AFM(PointModel):

    def __init__(self, n_users, n_items, embedding_dim, batch_norm=True, dropout=0.1):
        super().__init__()

        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)

        # consider attention score layer dimension should be (embedding_dim, num_features)
        # here we only consider 2 features, user & item, then K=2
        K = 2   # num_features
        self.lin = nn.Linear(embedding_dim, K)
        self.h = nn.Parameter(torch.rand(K, 1))

        # final prediction for reducer sum
        self.pred = nn.Linear(embedding_dim, 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.normal_(self.lin.weight, std=0.01)
        nn.init.xavier_normal_(self.pred.weight)

        nn.init.constant_(self.user_bias.weight, 0.0)
        nn.init.constant_(self.item_bias.weight, 0.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)

        ''' attention part '''
        att = F.relu(self.lin(fm)).mm(self.h)
        fm = fm * att

        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 AttentionalFactorizationMachine(torch.nn.Module):

    def __init__(self, embed_dim, attn_size, dropouts):
        super().__init__()
        self.attention = torch.nn.Linear(embed_dim, attn_size)
        self.projection = torch.nn.Linear(attn_size, 1)
        self.fc = torch.nn.Linear(embed_dim, 1)
        self.dropouts = dropouts

    def forward(self, x):
        """
        :param x: Float tensor of size ``(batch_size, num_fields, embed_dim)``
        """
        num_fields = x.shape[1]
        row, col = list(), list()
        for i in range(num_fields - 1):
            for j in range(i + 1, num_fields):
                row.append(i), col.append(j)
        p, q = x[:, row], x[:, col]
        inner_product = p * q
        attn_scores = F.relu(self.attention(inner_product))
        attn_scores = F.softmax(self.projection(attn_scores), dim=1)
        attn_scores = F.dropout(attn_scores, p=self.dropouts[0], training=self.training)
        attn_output = torch.sum(attn_scores * inner_product, dim=1)
        attn_output = F.dropout(attn_output, p=self.dropouts[1], training=self.training)
        return self.fc(attn_output)

# Cell
class AFMv2(torch.nn.Module):
    """
    A pytorch implementation of Attentional Factorization Machine.
    Reference:
        J Xiao, et al. Attentional Factorization Machines: Learning the Weight of Feature Interactions via Attention Networks, 2017.
    """

    def __init__(self, field_dims, embed_dim, attn_size, dropouts):
        super().__init__()
        self.num_fields = len(field_dims)
        self.embedding = FeaturesEmbedding(field_dims, embed_dim)
        self.linear = FeaturesLinear(field_dims)
        self.afm = AttentionalFactorizationMachine(embed_dim, attn_size, dropouts)

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