# -*- encoding:utf-8 -*-
# Copyright (c) Alibaba, Inc. and its affiliates.
import logging
import os

import tensorflow as tf

from easy_rec.python.builders import loss_builder
from easy_rec.python.model.easy_rec_model import EasyRecModel
from easy_rec.python.protos.loss_pb2 import LossType

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


class MatchModel(EasyRecModel):

  def __init__(self,
               model_config,
               feature_configs,
               features,
               labels=None,
               is_training=False):
    super(MatchModel, self).__init__(model_config, feature_configs, features,
                                     labels, is_training)
    self._loss_type = self._model_config.loss_type
    self._num_class = self._model_config.num_class

    if self._loss_type == LossType.CLASSIFICATION:
      assert self._num_class == 1

    if self._loss_type in [LossType.CLASSIFICATION, LossType.L2_LOSS]:
      self._is_point_wise = True
      logging.info('Use point wise dssm.')
    else:
      self._is_point_wise = False
      logging.info('Use list wise dssm.')

    cls_mem = self._model_config.WhichOneof('model')
    sub_model_config = getattr(self._model_config, cls_mem)

    self._item_ids = None
    assert sub_model_config is not None, 'sub_model_config undefined: model_cls = %s' % cls_mem
    if sub_model_config.item_id != '':
      logging.info('item_id feature is: %s' % sub_model_config.item_id)
      self._item_ids = features[sub_model_config.item_id]

  def _mask_in_batch(self, logits):
    batch_size = tf.shape(logits)[0]
    if self._model_config.ignore_in_batch_neg_sam:
      in_batch = logits[:, :batch_size] - (
          1 - tf.diag(tf.ones([batch_size], dtype=tf.float32))) * 1e32
      return tf.concat([in_batch, logits[:, batch_size:]], axis=1)
    else:
      if self._item_ids is not None:
        mask_in_batch_neg = tf.to_float(
            tf.equal(self._item_ids[None, :batch_size],
                     self._item_ids[:batch_size, None])) - tf.diag(
                         tf.ones([batch_size], dtype=tf.float32))
        tf.summary.scalar('in_batch_neg_conflict',
                          tf.reduce_sum(mask_in_batch_neg))
        return tf.concat([
            logits[:, :batch_size] - mask_in_batch_neg * 1e32,
            logits[:, batch_size:]],
            axis=1)  # yapf: disable
      else:
        return logits

  def _list_wise_sim(self, user_emb, item_emb):
    batch_size = tf.shape(user_emb)[0]
    hard_neg_indices = self._feature_dict.get('hard_neg_indices', None)

    if hard_neg_indices is not None:
      logging.info('With hard negative examples')
      noclk_size = tf.shape(hard_neg_indices)[0]
      # pos_item_emb, neg_item_emb, hard_neg_item_emb = tf.split(
      #     item_emb, [batch_size, -1, noclk_size], axis=0)
      simple_item_emb, hard_neg_item_emb = tf.split(
          item_emb, [-1, noclk_size], axis=0)
    else:
      # pos_item_emb = item_emb[:batch_size]
      # neg_item_emb = item_emb[batch_size:]
      simple_item_emb = item_emb

    # pos_user_item_sim = tf.reduce_sum(
    #     tf.multiply(user_emb, pos_item_emb), axis=1, keep_dims=True)
    # neg_user_item_sim = tf.matmul(user_emb, tf.transpose(neg_item_emb))
    # simple_user_item_sim = tf.matmul(user_emb, tf.transpose(simple_item_emb))

    _mode = os.environ['tf.estimator.mode']
    if _mode == tf.estimator.ModeKeys.PREDICT:
      simple_user_item_sim = tf.reduce_sum(
          tf.multiply(user_emb, simple_item_emb), axis=1, keep_dims=True)
    else:
      simple_user_item_sim = tf.matmul(user_emb, tf.transpose(simple_item_emb))

    if hard_neg_indices is None:
      return simple_user_item_sim
    else:
      user_emb_expand = tf.gather(user_emb, hard_neg_indices[:, 0])
      hard_neg_user_item_sim = tf.reduce_sum(
          tf.multiply(user_emb_expand, hard_neg_item_emb), axis=1)
      max_num_neg = tf.reduce_max(hard_neg_indices[:, 1]) + 1
      hard_neg_shape = tf.stack([tf.to_int64(batch_size), max_num_neg])
      hard_neg_sim = tf.scatter_nd(hard_neg_indices, hard_neg_user_item_sim,
                                   hard_neg_shape)
      hard_neg_mask = tf.scatter_nd(
          hard_neg_indices,
          tf.ones_like(hard_neg_user_item_sim, dtype=tf.float32),
          shape=hard_neg_shape)
      # set tail positions to -1e32, so that after exp(x), will be zero
      hard_neg_user_item_sim = hard_neg_sim - (1 - hard_neg_mask) * 1e32

      # user_item_sim = [pos_user_item_sim, neg_user_item_sim]
      # if hard_neg_indices is not None:
      #   user_item_sim.append(hard_neg_user_item_sim)
      # return tf.concat(user_item_sim, axis=1)

      return tf.concat([simple_user_item_sim, hard_neg_user_item_sim], axis=1)

  def _point_wise_sim(self, user_emb, item_emb):
    user_item_sim = tf.reduce_sum(
        tf.multiply(user_emb, item_emb), axis=1, keep_dims=True)
    return user_item_sim

  def sim(self, user_emb, item_emb):
    # Name the outputs of the user tower and the item tower, i.e. the inputs of the
    # simularity operation.
    # Explicit names of these nodes are necessary for some online recall systems like
    # BasicEngine to split up the predicting graph into different clusters.
    user_emb = tf.identity(user_emb, 'user_tower_emb')
    item_emb = tf.identity(item_emb, 'item_tower_emb')

    if self._is_point_wise:
      return self._point_wise_sim(user_emb, item_emb)
    else:
      return self._list_wise_sim(user_emb, item_emb)

  def norm(self, fea):
    fea_norm = tf.nn.l2_normalize(fea, axis=-1)
    return fea_norm

  def build_predict_graph(self):
    raise NotImplementedError('MatchModel could not be instantiated')

  def build_loss_graph(self):
    if self._is_point_wise:
      return self._build_point_wise_loss_graph()
    else:
      return self._build_list_wise_loss_graph()

  def _build_list_wise_loss_graph(self):
    if self._loss_type == LossType.SOFTMAX_CROSS_ENTROPY:
      batch_size = tf.shape(self._prediction_dict['probs'])[0]
      indices = tf.range(batch_size)
      indices = tf.concat([indices[:, None], indices[:, None]], axis=1)
      hit_prob = tf.gather_nd(
          self._prediction_dict['probs'][:batch_size, :batch_size], indices)
      self._loss_dict['cross_entropy_loss'] = -tf.reduce_mean(
          tf.log(hit_prob + 1e-12) * tf.squeeze(self._sample_weight))
      logging.info('softmax cross entropy loss is used')

      user_features = self._prediction_dict['user_tower_emb']
      pos_item_emb = self._prediction_dict['item_tower_emb'][:batch_size]
      pos_simi = tf.reduce_sum(user_features * pos_item_emb, axis=1)
      # if pos_simi < 0, produce loss
      reg_pos_loss = tf.nn.relu(-pos_simi)
      self._loss_dict['reg_pos_loss'] = tf.reduce_mean(reg_pos_loss)
    else:
      raise ValueError('invalid loss type: %s' % str(self._loss_type))
    return self._loss_dict

  def _build_point_wise_loss_graph(self):
    label = list(self._labels.values())[0]
    if self._loss_type == LossType.CLASSIFICATION:
      pred = self._prediction_dict['logits']
      loss_name = 'cross_entropy_loss'
    elif self._loss_type == LossType.L2_LOSS:
      pred = self._prediction_dict['y']
      loss_name = 'l2_loss'
    else:
      raise ValueError('invalid loss type: %s' % str(self._loss_type))

    kwargs = {'loss_name': loss_name}
    self._loss_dict[loss_name] = loss_builder.build(
        self._loss_type,
        label=label,
        pred=pred,
        loss_weight=self._sample_weight,
        **kwargs)

    # build kd loss
    kd_loss_dict = loss_builder.build_kd_loss(self.kd, self._prediction_dict,
                                              self._labels)
    self._loss_dict.update(kd_loss_dict)
    return self._loss_dict

  def build_metric_graph(self, eval_config):
    if self._is_point_wise:
      return self._build_point_wise_metric_graph(eval_config)
    else:
      return self._build_list_wise_metric_graph(eval_config)

  def _build_list_wise_metric_graph(self, eval_config):
    from easy_rec.python.core.easyrec_metrics import metrics_tf
    logits = self._prediction_dict['logits']
    # label = tf.zeros_like(logits[:, :1], dtype=tf.int64)
    batch_size = tf.shape(logits)[0]
    label = tf.cast(tf.range(batch_size), tf.int64)

    indices = tf.range(batch_size)
    indices = tf.concat([indices[:, None], indices[:, None]], axis=1)
    pos_item_sim = tf.gather_nd(logits[:batch_size, :batch_size], indices)
    metric_dict = {}
    for metric in eval_config.metrics_set:
      if metric.WhichOneof('metric') == 'recall_at_topk':
        metric_dict['recall@%d' %
                    metric.recall_at_topk.topk] = metrics_tf.recall_at_k(
                        label, logits, metric.recall_at_topk.topk)

        logits_v2 = tf.concat([pos_item_sim[:, None], logits[:, batch_size:]],
                              axis=1)
        labels_v2 = tf.zeros_like(logits_v2[:, :1], dtype=tf.int64)
        metric_dict['recall_neg_sam@%d' %
                    metric.recall_at_topk.topk] = metrics_tf.recall_at_k(
                        labels_v2, logits_v2, metric.recall_at_topk.topk)

        metric_dict['recall_in_batch@%d' %
                    metric.recall_at_topk.topk] = metrics_tf.recall_at_k(
                        label, logits[:, :batch_size],
                        metric.recall_at_topk.topk)
      else:
        ValueError('invalid metric type: %s' % str(metric))
    return metric_dict

  def _build_point_wise_metric_graph(self, eval_config):
    from easy_rec.python.core.easyrec_metrics import metrics_tf
    metric_dict = {}
    label = list(self._labels.values())[0]
    for metric in eval_config.metrics_set:
      if metric.WhichOneof('metric') == 'auc':
        assert self._loss_type == LossType.CLASSIFICATION
        metric_dict['auc'] = metrics_tf.auc(label,
                                            self._prediction_dict['probs'])
      elif metric.WhichOneof('metric') == 'mean_absolute_error':
        assert self._loss_type == LossType.L2_LOSS
        metric_dict['mean_absolute_error'] = metrics_tf.mean_absolute_error(
            tf.to_float(label), self._prediction_dict['y'])
      else:
        ValueError('invalid metric type: %s' % str(metric))
    return metric_dict

  def get_outputs(self):
    raise NotImplementedError(
        'could not call get_outputs on abstract class MatchModel')
