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

__all__ = ['MLP', 'MLPv2']

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

import torch
from torch import nn

from .bases.common import PointModel

# Cell
class MLP(PairModel):
    def __init__(self, n_users, n_items, embedding_dim, 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.fc1 = nn.Linear(embedding_dim * 2, embedding_dim)
        self.fc2 = nn.Linear(embedding_dim, int(embedding_dim / 2))
        self.fc3 = nn.Linear(int(embedding_dim / 2), 1)

        self.dropout = nn.Dropout(p=dropout)

    def forward(self, users, items):
        user_embeddings = self.user_embedding(users)
        item_embeddings = self.item_embedding(items)
        embeddings = torch.cat([user_embeddings, item_embeddings], axis=1)
        output = nn.ReLU()(self.fc1(embeddings))
        output = self.dropout(output)
        output = nn.ReLU()(self.fc2(output))
        output = self.dropout(output)
        output = self.fc3(output)

        return output.squeeze()

# Cell
class MLPv2(nn.Module):
    def __init__(self, args, num_users, num_items):
        super(MLPv2, self).__init__()
        self.num_users = num_users
        self.num_items = num_items
        self.factor_num = args.factor_num
        self.layers = args.layers

        self.embedding_user = nn.Embedding(num_embeddings=self.num_users, embedding_dim=self.factor_num)
        self.embedding_item = nn.Embedding(num_embeddings=self.num_items, embedding_dim=self.factor_num)

        self.fc_layers = nn.ModuleList()
        for idx, (in_size, out_size) in enumerate(zip(self.layers[:-1], self.layers[1:])):
            self.fc_layers.append(nn.Linear(in_size, out_size))

        self.affine_output = nn.Linear(in_features=self.layers[-1], out_features=1)
        self.logistic = nn.Sigmoid()

    def forward(self, user_indices, item_indices):
        user_embedding = self.embedding_user(user_indices)
        item_embedding = self.embedding_item(item_indices)
        vector = torch.cat([user_embedding, item_embedding], dim=-1)  # the concat latent vector
        for idx, _ in enumerate(range(len(self.fc_layers))):
            vector = self.fc_layers[idx](vector)
            vector = nn.ReLU()(vector)
            # vector = nn.BatchNorm1d()(vector)
            # vector = nn.Dropout(p=0.5)(vector)
        logits = self.affine_output(vector)
        rating = self.logistic(logits)
        return rating