# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# 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
#
#     http://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.
# ==============================================================================
"""Wrap the official recommendation model in a tf_cnn_benchmarks Model.

This allows the recommendation NCF model to be used in tf_cnn_benchmarks.
Currently, the implementation is fairly hacky, because tf_cnn_benchmarks is
intended to be used only with CNNs.

Only synthetic data with 1 GPU is currently supported.
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import tensorflow.compat.v1 as tf

from models import model


# Obtained by running the official NCF model with the following command:
#     python ncf_main.py  --dataset ml-20m
# and printing the number of users and items here:
# https://github.com/tensorflow/models/blob/d089975f630a8a01be63e45ef08a31be14bb96b4/official/recommendation/data_preprocessing.py#L68
_NUM_USERS_20M = 138493
_NUM_ITEMS_20M = 26744


# TODO(reedwm): Support multi-GPU. Currently keras layers, which this model
# uses, ignore variable_scopes, which we rely on for multi-GPU support.
# TODO(reedwm): Support real data. This will require a significant refactor.
# TODO(reedwm): All-reduce IndexedSlices more effectively.
# TODO(reedwm): Support the 1M variant of this model.


class NcfModel(model.Model):
  r"""A model.Model wrapper around the official NCF recommendation model.

  To do an NCF run with synthetic data that roughly matches what the official
  model does, run:

  python tf_cnn_benchmarks.py --optimizer=adam --model=ncf --batch_size=65536 \
      --weight_decay=0 --sparse_to_dense_grads
  """

  def __init__(self, params=None):
    super(NcfModel, self).__init__(
        'official_ncf', batch_size=2048, learning_rate=0.0005,
        fp16_loss_scale=128, params=params)
    if self.fp16_vars:
      raise ValueError('NCF model only supports float32 variables for now.')

  def build_network(self, inputs, phase_train=True, nclass=1001):
    try:
      from official.recommendation import neumf_model  # pylint: disable=g-import-not-at-top
    except ImportError as e:
      if 'neumf_model' not in e.message:
        raise
      raise ImportError('To use the experimental NCF model, you must clone the '
                        'repo https://github.com/tensorflow/models and add '
                        'tensorflow/models to the PYTHONPATH.')
    del nclass

    users, items, _ = inputs
    params = {
        'num_users': _NUM_USERS_20M,
        'num_items': _NUM_ITEMS_20M,
        'model_layers': (256, 256, 128, 64),
        'mf_dim': 64,
        'mf_regularization': 0,
        'mlp_reg_layers': (0, 0, 0, 0),
        'use_tpu': False
    }
    user_input = tf.keras.layers.Input(tensor=users, name='user_input')
    item_input = tf.keras.layers.Input(tensor=items, name='item_input')
    if self.data_type == tf.float32:
      keras_model = neumf_model.construct_model(user_input, item_input, params)
      logits = keras_model.output
    else:
      assert self.data_type == tf.float16
      old_floatx = tf.keras.backend.floatx()
      try:
        tf.keras.backend.set_floatx('float16')
        # We cannot rely on the variable_scope's fp16 custom getter here,
        # because the NCF model uses keras layers, which ignore variable scopes.
        # So we use a variable_creator_scope instead.
        with tf.variable_creator_scope(_fp16_variable_creator):
          keras_model = neumf_model.construct_model(user_input, item_input,
                                                    params)
        logits = tf.cast(keras_model.output, tf.float32)
      finally:
        tf.keras.backend.set_floatx(old_floatx)
    return model.BuildNetworkResult(logits=logits, extra_info=None)

  def loss_function(self, inputs, build_network_result):
    logits = build_network_result.logits

    # Softmax with the first column of ones is equivalent to sigmoid.
    # TODO(reedwm): Actually, the first column should be zeros to be equivalent
    # to sigmoid. But, we keep it at ones to match the official models.
    logits = tf.concat([tf.ones(logits.shape, dtype=logits.dtype), logits],
                       axis=1)

    return tf.losses.sparse_softmax_cross_entropy(
        labels=inputs[2],
        logits=logits
    )

  def get_synthetic_inputs(self, input_name, nclass):
    """Returns the ops to generate synthetic inputs and labels."""
    def users_init_val():
      return tf.random_uniform((self.batch_size, 1), minval=0,
                               maxval=_NUM_USERS_20M, dtype=tf.int32)
    users = tf.Variable(users_init_val, dtype=tf.int32, trainable=False,
                        collections=[tf.GraphKeys.LOCAL_VARIABLES],
                        name='synthetic_users')
    def items_init_val():
      return tf.random_uniform((self.batch_size, 1), minval=0,
                               maxval=_NUM_ITEMS_20M, dtype=tf.int32)
    items = tf.Variable(items_init_val, dtype=tf.int32, trainable=False,
                        collections=[tf.GraphKeys.LOCAL_VARIABLES],
                        name='synthetic_items')

    def labels_init_val():
      return tf.random_uniform((self.batch_size,), minval=0, maxval=2,
                               dtype=tf.int32)
    labels = tf.Variable(labels_init_val, dtype=tf.int32, trainable=False,
                         collections=[tf.GraphKeys.LOCAL_VARIABLES],
                         name='synthetic_labels')

    return [users, items, labels]

  def get_input_shapes(self, subset):
    del subset
    return [[self.batch_size, 1], [self.batch_size, 1], [self.batch_size]]

  def get_input_data_types(self, subset):
    del subset
    return [self.int32, tf.int32, tf.int32]


def _fp16_variable_creator(next_creator, **kwargs):
  """Variable creator to create variables in fp32 and cast them to fp16."""
  dtype = kwargs.get('dtype', None)
  initial_value = kwargs.get('initial_value', None)
  if dtype is None:
    if initial_value is not None and not callable(initial_value):
      dtype = initial_value.dtype
  if dtype == tf.float16:
    if callable(initial_value):
      new_initial_value = lambda: tf.cast(initial_value(), tf.float32)
    else:
      new_initial_value = tf.cast(initial_value, tf.float32)
    kwargs['dtype'] = tf.float32
    kwargs['initial_value'] = new_initial_value
    var = next_creator(**kwargs)
    return tf.cast(var, dtype=tf.float16)
  else:
    return next_creator(**kwargs)

