#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Explanation: Copyright 2021 The TensorFlow Authors.
End of explanation
"""


!pip install -q --upgrade tensorflow-recommenders-addons
!pip install -q --upgrade tensorflow-datasets

import numpy as np
import tensorflow as tf
from tensorflow.keras.layers import Dense

import tensorflow_datasets as tfds
import tensorflow_recommenders_addons as tfra

"""
Explanation: Recommending movies: ranking
<table class="tfa-notebook-buttons" align="left">
  <td>
    <a target="_blank" href="https://github.com/tensorflow/recommenders-addons/blob/master/docs/tutorials/dynamic_embedding_tutorial.ipynb"><img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />View source on GitHub</a>
  </td>
</table>

Overview
In this tutorial, we're going to use the rating data to predict the user's rating of other movies. To achieve this goal, we will follow the following steps:

Get our data and do some preprocessing to get the required format.
Implement a neural collaborative filtering(NeuralCF) model.
Train the model.

Different from the general recommendation model, the model we implemented replaces tf.nn.embedding_lookup with tfra.dynamic_embedding.embedding_lookup, which can handle super large sparse features.
Imports
Let's first get our imports out of the way.
End of explanation
"""


ratings = tfds.load("movielens/100k-ratings", split="train")

ratings = ratings.map(lambda x: {
    "movie_id": tf.strings.to_number(x["movie_id"], tf.int64),
    "user_id": tf.strings.to_number(x["user_id"], tf.int64),
    "user_rating": x["user_rating"]
})

tf.random.set_seed(2021)
shuffled = ratings.shuffle(100_000, seed=2021, reshuffle_each_iteration=False)

dataset_train = shuffled.take(100_000).batch(256)

"""
Explanation: Preparing the dataset
This tutorial uses movies reviews provided by the MovieLens 100K dataset, a classic dataset from the GroupLens research group at the University of Minnesota. In order to facilitate processing, we convert the data type of movie_id and user_id into int64.
End of explanation
"""


class NCFModel(tf.keras.Model):

    def __init__(self):
        super(NCFModel, self).__init__()
        self.embedding_size = 32
        self.d0 = Dense(
            256,
            activation='relu',
            kernel_initializer=tf.keras.initializers.RandomNormal(0.0, 0.1),
            bias_initializer=tf.keras.initializers.RandomNormal(0.0, 0.1))
        self.d1 = Dense(
            64,
            activation='relu',
            kernel_initializer=tf.keras.initializers.RandomNormal(0.0, 0.1),
            bias_initializer=tf.keras.initializers.RandomNormal(0.0, 0.1))
        self.d2 = Dense(
            1,
            kernel_initializer=tf.keras.initializers.RandomNormal(0.0, 0.1),
            bias_initializer=tf.keras.initializers.RandomNormal(0.0, 0.1))
        self.user_embeddings = tfra.dynamic_embedding.get_variable(
            name="user_dynamic_embeddings",
            dim=self.embedding_size,
            initializer=tf.keras.initializers.RandomNormal(-1, 1))
        self.movie_embeddings = tfra.dynamic_embedding.get_variable(
            name="moive_dynamic_embeddings",
            dim=self.embedding_size,
            initializer=tf.keras.initializers.RandomNormal(-1, 1))
        self.loss = tf.keras.losses.MeanSquaredError()

    def call(self, batch):
        movie_id = batch["movie_id"]
        user_id = batch["user_id"]
        rating = batch["user_rating"]

        user_id_val, user_id_idx = np.unique(user_id, return_inverse=True)
        user_id_weights, user_id_trainable_wrapper = tfra.dynamic_embedding.embedding_lookup(
            params=self.user_embeddings,
            ids=user_id_val,
            name="user-id-weights",
            return_trainable=True)
        user_id_weights = tf.gather(user_id_weights, user_id_idx)

        movie_id_val, movie_id_idx = np.unique(movie_id, return_inverse=True)
        movie_id_weights, movie_id_trainable_wrapper = tfra.dynamic_embedding.embedding_lookup(
            params=self.movie_embeddings,
            ids=movie_id_val,
            name="movie-id-weights",
            return_trainable=True)
        movie_id_weights = tf.gather(movie_id_weights, movie_id_idx)

        embeddings = tf.concat([user_id_weights, movie_id_weights], axis=1)
        dnn = self.d0(embeddings)
        dnn = self.d1(dnn)
        dnn = self.d2(dnn)
        out = tf.reshape(dnn, shape=[-1])
        loss = self.loss(rating, out)
        return loss, [user_id_trainable_wrapper, movie_id_trainable_wrapper]

"""
Explanation: Implementing a model
The NCFModel we implemented is very similar to the conventional one, and the main difference lies in the embedding layer. We specify the variable of embedding layer by tfra.dynamic_embedding.get_variable.
End of explanation
"""


model = NCFModel()
optimizer = tf.keras.optimizers.Adam(learning_rate=0.001)
optimizer = tfra.dynamic_embedding.DynamicEmbeddingOptimizer(optimizer)

"""
Explanation: Let's instantiate the model, and wrap the optimizer in tfra.dynamic_embedding.DynamicEmbeddingOptimizer.
End of explanation
"""


def train(epoch=1):
    for i in range(epoch):
        total_loss = np.array([])
        for (_, batch) in enumerate(dataset_train):
            with tf.GradientTape() as tape:
                loss, trainable_wrapper_list = model(batch)
                total_loss = np.append(total_loss, loss)
            grads = tape.gradient(loss, model.trainable_variables + trainable_wrapper_list)
            optimizer.apply_gradients(zip(grads, model.trainable_variables + trainable_wrapper_list))
        print("epoch:", i, "mean_squared_error:", np.mean(total_loss))

if __name__=="__main__":
    train(10)

"""
Explanation: Training the model
After defining the model, we can train the model and observe the change of loss.
End of explanation
"""


