# -*- encoding:utf-8 -*-
# Copyright (c) Alibaba, Inc. and its affiliates.
import tensorflow as tf

from easy_rec.python.layers import dnn
from easy_rec.python.model.match_model import MatchModel
from easy_rec.python.protos.dssm_pb2 import DSSM as DSSMConfig
from easy_rec.python.protos.loss_pb2 import LossType
from easy_rec.python.protos.simi_pb2 import Similarity
from easy_rec.python.utils.proto_util import copy_obj

if tf.__version__ >= '2.0':
  tf = tf.compat.v1
losses = tf.losses


class DSSM(MatchModel):

  def __init__(self,
               model_config,
               feature_configs,
               features,
               labels=None,
               is_training=False):
    super(DSSM, self).__init__(model_config, feature_configs, features, labels,
                               is_training)
    assert self._model_config.WhichOneof('model') == 'dssm', \
        'invalid model config: %s' % self._model_config.WhichOneof('model')
    self._model_config = self._model_config.dssm
    assert isinstance(self._model_config, DSSMConfig)

    # copy_obj so that any modification will not affect original config
    self.user_tower = copy_obj(self._model_config.user_tower)
    self.user_tower_feature, _ = self._input_layer(self._feature_dict, 'user')
    # copy_obj so that any modification will not affect original config
    self.item_tower = copy_obj(self._model_config.item_tower)
    self.item_tower_feature, _ = self._input_layer(self._feature_dict, 'item')
    self._user_tower_emb = None
    self._item_tower_emb = None

  def build_predict_graph(self):
    num_user_dnn_layer = len(self.user_tower.dnn.hidden_units)
    last_user_hidden = self.user_tower.dnn.hidden_units.pop()
    user_dnn = dnn.DNN(self.user_tower.dnn, self._l2_reg, 'user_dnn',
                       self._is_training)
    user_tower_emb = user_dnn(self.user_tower_feature)
    user_tower_emb = tf.layers.dense(
        inputs=user_tower_emb,
        units=last_user_hidden,
        kernel_regularizer=self._l2_reg,
        name='user_dnn/dnn_%d' % (num_user_dnn_layer - 1))

    num_item_dnn_layer = len(self.item_tower.dnn.hidden_units)
    last_item_hidden = self.item_tower.dnn.hidden_units.pop()
    item_dnn = dnn.DNN(self.item_tower.dnn, self._l2_reg, 'item_dnn',
                       self._is_training)
    item_tower_emb = item_dnn(self.item_tower_feature)
    item_tower_emb = tf.layers.dense(
        inputs=item_tower_emb,
        units=last_item_hidden,
        kernel_regularizer=self._l2_reg,
        name='item_dnn/dnn_%d' % (num_item_dnn_layer - 1))

    if self._loss_type == LossType.CLASSIFICATION:
      if self._model_config.simi_func == Similarity.COSINE:
        user_tower_emb = self.norm(user_tower_emb)
        item_tower_emb = self.norm(item_tower_emb)

    user_item_sim = self.sim(user_tower_emb, item_tower_emb)
    if self._model_config.scale_simi:
      sim_w = tf.get_variable(
          'sim_w',
          dtype=tf.float32,
          shape=(1),
          initializer=tf.ones_initializer())
      sim_b = tf.get_variable(
          'sim_b',
          dtype=tf.float32,
          shape=(1),
          initializer=tf.zeros_initializer())
      y_pred = user_item_sim * tf.abs(sim_w) + sim_b
    else:
      y_pred = user_item_sim

    if self._is_point_wise:
      y_pred = tf.reshape(y_pred, [-1])

    if self._loss_type == LossType.CLASSIFICATION:
      self._prediction_dict['logits'] = y_pred
      self._prediction_dict['probs'] = tf.nn.sigmoid(y_pred)
    elif self._loss_type == LossType.SOFTMAX_CROSS_ENTROPY:
      y_pred = self._mask_in_batch(y_pred)
      self._prediction_dict['logits'] = y_pred
      self._prediction_dict['probs'] = tf.nn.softmax(y_pred)
    else:
      self._prediction_dict['y'] = y_pred

    self._prediction_dict['user_tower_emb'] = user_tower_emb
    self._prediction_dict['item_tower_emb'] = item_tower_emb
    self._prediction_dict['user_emb'] = tf.reduce_join(
        tf.as_string(user_tower_emb), axis=-1, separator=',')
    self._prediction_dict['item_emb'] = tf.reduce_join(
        tf.as_string(item_tower_emb), axis=-1, separator=',')
    return self._prediction_dict

  def get_outputs(self):
    if self._loss_type == LossType.CLASSIFICATION:
      return ['logits', 'probs', 'user_emb', 'item_emb']
    elif self._loss_type == LossType.SOFTMAX_CROSS_ENTROPY:
      self._prediction_dict['logits'] = tf.squeeze(
          self._prediction_dict['logits'], axis=-1)
      self._prediction_dict['probs'] = tf.nn.sigmoid(
          self._prediction_dict['logits'])
      return ['logits', 'probs', 'user_emb', 'item_emb']
    elif self._loss_type == LossType.L2_LOSS:
      return ['y', 'user_emb', 'item_emb']
    else:
      raise ValueError('invalid loss type: %s' % str(self._loss_type))

  def build_output_dict(self):
    output_dict = super(DSSM, self).build_output_dict()
    output_dict['user_tower_feature'] = tf.reduce_join(
        tf.as_string(self.user_tower_feature), axis=-1, separator=',')
    output_dict['item_tower_feature'] = tf.reduce_join(
        tf.as_string(self.item_tower_feature), axis=-1, separator=',')
    return output_dict

  def build_rtp_output_dict(self):
    output_dict = super(DSSM, self).build_rtp_output_dict()
    if 'user_tower_emb' not in self._prediction_dict:
      raise ValueError(
          'User tower embedding does not exist. Please checking predict graph.')
    output_dict['user_embedding_output'] = tf.identity(
        self._prediction_dict['user_tower_emb'], name='user_embedding_output')
    if 'item_tower_emb' not in self._prediction_dict:
      raise ValueError(
          'Item tower embedding does not exist. Please checking predict graph.')
    output_dict['item_embedding_output'] = tf.identity(
        self._prediction_dict['item_tower_emb'], name='item_embedding_output')
    if self._loss_type == LossType.CLASSIFICATION:
      if 'probs' not in self._prediction_dict:
        raise ValueError(
            'Probs output does not exist. Please checking predict graph.')
      output_dict['rank_predict'] = tf.identity(
          self._prediction_dict['probs'], name='rank_predict')
    return output_dict
