text: # Licensed to the Apache Software Foundation (ASF) under one
  # or more contributor license agreements.  See the NOTICE file
  # distributed with this work for additional information
  # regarding copyright ownership.  The ASF licenses this file
  # to you 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.
  # 
  import mxnet as mx
  import numpy as np
  
  @mx.init.register
  class golorot_uniform(mx.init.Initializer):
      def __init__(self, fan_in, fan_out):
          super(golorot_uniform, self).__init__(fan_in=fan_in, fan_out=fan_out)
          self._fan_in = fan_in
          self._fan_out = fan_out
      def _init_weight(self, _, arr):
          limit = np.sqrt(6. / (self._fan_in + self._fan_out))
          mx.random.uniform(-limit, limit, out=arr)
  
  @mx.init.register
  class lecunn_uniform(mx.init.Initializer):
      def __init__(self, fan_in):
          super(lecunn_uniform, self).__init__(fan_in=fan_in)
          self._fan_in = fan_in
      def _init_weight(self, _, arr):
          limit = np.sqrt(3. / self._fan_in)
          mx.random.uniform(-limit, limit, out=arr)
  
  # only for inference model optimize
  def mlp_opt(user, item, factor_size, model_layers, max_user, max_item):
      user_weight = mx.sym.Variable('fused_mlp_user_weight', init=mx.init.Normal(0.01))
      item_weight = mx.sym.Variable('fused_mlp_item_weight', init=mx.init.Normal(0.01))
      embed_user = mx.sym.Embedding(data=user, weight=user_weight, input_dim=max_user,
                                    output_dim=factor_size * 2, name='fused_embed_user'+str(factor_size))
      embed_item = mx.sym.Embedding(data=item, weight=item_weight, input_dim=max_item,
                                    output_dim=factor_size * 2, name='fused_embed_item'+str(factor_size))
      pre_gemm_concat = embed_user + embed_item
  
      for i in range(1, len(model_layers)):
          if i==1:
              pre_gemm_concat = mx.sym.Activation(data=pre_gemm_concat, act_type='relu', name='act_'+str(i-1))
              continue
          else:
              mlp_weight_init = golorot_uniform(model_layers[i-1], model_layers[i])
          mlp_weight = mx.sym.Variable('fc_{}_weight'.format(i-1), init=mlp_weight_init)
          pre_gemm_concat = mx.sym.FullyConnected(data=pre_gemm_concat, weight=mlp_weight, num_hidden=model_layers[i], name='fc_'+str(i-1))
          pre_gemm_concat = mx.sym.Activation(data=pre_gemm_concat, act_type='relu', name='act_'+str(i-1))
  
      return pre_gemm_concat
  
  def mlp(user, item, factor_size, model_layers, max_user, max_item):
      user_weight = mx.sym.Variable('mlp_user_weight', init=mx.init.Normal(0.01))
      item_weight = mx.sym.Variable('mlp_item_weight', init=mx.init.Normal(0.01))
      embed_user = mx.sym.Embedding(data=user, weight=user_weight, input_dim=max_user,
                                    output_dim=factor_size, name='embed_user'+str(factor_size))
      embed_item = mx.sym.Embedding(data=item, weight=item_weight, input_dim=max_item,
                                    output_dim=factor_size, name='embed_item'+str(factor_size))
      pre_gemm_concat = mx.sym.concat(embed_user, embed_item, dim=1, name='pre_gemm_concat')
  
      for i in range(1, len(model_layers)):
          mlp_weight_init = golorot_uniform(model_layers[i-1], model_layers[i])
          mlp_weight = mx.sym.Variable('fc_{}_weight'.format(i-1), init=mlp_weight_init)
          pre_gemm_concat = mx.sym.FullyConnected(data=pre_gemm_concat, weight=mlp_weight, num_hidden=model_layers[i], name='fc_'+str(i-1))
          pre_gemm_concat = mx.sym.Activation(data=pre_gemm_concat, act_type='relu', name='act_'+str(i-1))
  
      return pre_gemm_concat
  
  def gmf(user, item, factor_size, max_user, max_item):
      user_weight = mx.sym.Variable('gmf_user_weight', init=mx.init.Normal(0.01))
      item_weight = mx.sym.Variable('gmf_item_weight', init=mx.init.Normal(0.01))
      embed_user = mx.sym.Embedding(data=user, weight=user_weight, input_dim=max_user,
                                    output_dim=factor_size, name='embed_user'+str(factor_size))
      embed_item = mx.sym.Embedding(data=item, weight=item_weight, input_dim=max_item,
                                    output_dim=factor_size, name='embed_item'+str(factor_size))
      pred = embed_user * embed_item
  
      return pred
  
  def get_model(model_type='neumf', factor_size_mlp=128, factor_size_gmf=64,
                model_layers=[256, 256, 128, 64], num_hidden=1, 
                max_user=138493, max_item=26744, opt=False):
      # input
      user = mx.sym.Variable('user')
      item = mx.sym.Variable('item')
  
      if model_type == 'mlp':
          if opt:
              net = mlp_opt(user=user, item=item,
                           factor_size=factor_size_mlp, model_layers=model_layers,
                           max_user=max_user, max_item=max_item)
          else:
              net = mlp(user=user, item=item,
                        factor_size=factor_size_mlp, model_layers=model_layers,
                        max_user=max_user, max_item=max_item)
      elif model_type == 'gmf':
          net = gmf(user=user, item=item,
                    factor_size=factor_size_gmf,
                    max_user=max_user, max_item=max_item)
      elif model_type == 'neumf':
          if opt:
              net_mlp = mlp_opt(user=user, item=item,
                                factor_size=factor_size_mlp, model_layers=model_layers,
                                max_user=max_user, max_item=max_item)
          else:
              net_mlp = mlp(user=user, item=item,
                            factor_size=factor_size_mlp, model_layers=model_layers,
                            max_user=max_user, max_item=max_item)
          net_gmf = gmf(user=user, item=item,
                        factor_size=factor_size_gmf,
                        max_user=max_user, max_item=max_item)
  
          net = mx.sym.concat(net_gmf, net_mlp, dim=1, name='post_gemm_concat')
  
      else:
          raise ValueError('Unsupported ncf model %s.' % model_type)
  
      final_weight = mx.sym.Variable('fc_final_weight', init=lecunn_uniform(factor_size_gmf + model_layers[-1]))
      net = mx.sym.FullyConnected(data=net, weight=final_weight, num_hidden=num_hidden, name='fc_final') 
     
      y_label = mx.sym.Variable('softmax_label')
      net = mx.symbol.LogisticRegressionOutput(data=net, label=y_label, name='sigmoid_final')
  
      return net
  
  