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

__all__ = ['VNCF']

# Cell
import tensorflow as tf
from tensorflow.keras.models import Model, Sequential
from tensorflow.keras.layers import Embedding, Flatten, Input, Dropout, Dense, Concatenate, Lambda
from tensorflow.keras import backend as K


class VNCF:
    """Variational Neural Collaborative Filtering.
    """
    def __init__(self, num_users, num_items, latent_dim=20, batch_size=32):
        self.num_users = num_users
        self.num_items = num_items
        self.latent_dim = latent_dim
        self.batch_size = batch_size

    def sampling(self, qargs):
        z_mean, z_var = qargs
        epsilon = K.random_normal(shape=(self.batch_size, self.latent_dim), mean=0., stddev=1)
        return z_mean + K.exp(z_var) * epsilon

    def build(self):
        user_input = Input(shape=[1])
        user_embedding = Embedding(self.num_users, self.latent_dim)(user_input)
        user_embedding_mean = Dense(self.latent_dim)(user_embedding)
        user_embedding_var = Dense(self.latent_dim)(user_embedding)
        user_embedding_z = Lambda(self.sampling)([user_embedding_mean, user_embedding_var])
        user_vec = Flatten()(user_embedding_z)

        item_input = Input(shape=[1])
        item_embedding = Embedding(self.num_items + 1, self.latent_dim)(item_input)
        item_embedding_mean = Dense(self.latent_dim)(item_embedding)
        item_embedding_var = Dense(self.latent_dim)(item_embedding)
        item_embedding_z = Lambda(self.sampling)([item_embedding_mean, item_embedding_var], self.latent_dim)
        item_vec = Flatten()(item_embedding_z)

        concat = Concatenate(axis=1)([item_vec, user_vec])

        fc_1 = Dense(80, name='fc-1', activation='relu')(concat)
        fc_1_dropout = Dropout(0.6, name='fc-1-dropout')(fc_1)
        fc_2 = Dense(25, name='fc-2', activation='relu')(fc_1_dropout)
        fc_2_dropout = Dropout(0.4, name='fc-2-dropout')(fc_2)
        fc_3 = Dense(1, name='fc-3', activation='relu')(fc_2_dropout)

        model = Model([user_input, item_input], fc_3)
        return model