//  [python]
def discriminative_loss_single2(prediction, correct_label_in, feature_dim,
                               delta_v, delta_d, param_var, param_dist, param_reg):
  ''' Discriminative loss for a single prediction/label pair.
  :param prediction: inference of network
  :param correct_label: instance label
  :feature_dim: feature dimension of prediction
  :param label_shape: shape of label
  :param delta_v: cutoff variance distance
  :param delta_d: curoff cluster distance
  :param param_var: weight for intra cluster variance
  :param param_dist: weight for inter cluster distances
  :param param_reg: weight regularization
  '''

  ### Reshape so pixels are aligned along a vector
  #correct_label = tf.reshape(correct_label, [label_shape[1] * label_shape[0]])
  reshaped_pred = tf.reshape(prediction, [2, -1, feature_dim])
  correct_label = tf.reshape(correct_label_in, [2, -1])

  # print(correct_label_in.shape, correct_label.shape)
  correct_label0 = correct_label[0]
  reshaped_pred0 = reshaped_pred[0]
  ### Count instances
  unique_labels, unique_id, counts = tf.unique_with_counts(correct_label0)
  counts = tf.cast(counts, tf.float32)
  num_instances = tf.size(unique_labels)
  segmented_sum = tf.unsorted_segment_sum(reshaped_pred0, unique_id, num_instances)

  labels_sort = tf.argsort(tf.cast(unique_labels, tf.int32))
  segmented_sum = tf.gather(segmented_sum, labels_sort[1:])
  counts = tf.gather(counts, labels_sort[1:])

  mu0 = tf.div(segmented_sum, tf.reshape(counts, (-1, 1)))


  # print(correct_label_in.shape, correct_label.shape)
  correct_label1 = correct_label[1]
  reshaped_pred1 = reshaped_pred[1]
  ### Count instances
  unique_labels, unique_id, counts = tf.unique_with_counts(correct_label1)
  counts = tf.cast(counts, tf.float32)
  segmented_sum = tf.unsorted_segment_sum(reshaped_pred1, unique_id, num_instances)

  labels_sort = tf.argsort(tf.cast(unique_labels, tf.int32))
  segmented_sum = tf.gather(segmented_sum, labels_sort[1:])
  counts = tf.gather(counts, labels_sort[1:])

  mu1 = tf.div(segmented_sum, tf.reshape(counts, (-1, 1)))


  ### Calculate l_dist

  # Get distance for each pair of clusters like this:
  #   mu_1 - mu_1
  #   mu_2 - mu_1
  #   mu_3 - mu_1
  #   mu_1 - mu_2
  #   mu_2 - mu_2
  #   mu_3 - mu_2
  #   mu_1 - mu_3
  #   mu_2 - mu_3
  #   mu_3 - mu_3

  mu_interleaved_rep = tf.tile(mu0, [num_instances-1, 1])
  mu_band_rep = tf.tile(mu1, [1, num_instances-1])
  mu_band_rep = tf.reshape(mu_band_rep, ((num_instances-1) * (num_instances-1), feature_dim))

  mu_diff = tf.subtract(mu_band_rep, mu_interleaved_rep)

  # Filter out zeros from same cluster subtraction
  eye = tf.eye(num_instances-1)
  zero = tf.zeros(1, dtype=tf.float32)
  one = tf.ones(1, dtype=tf.float32)
  same_cluster_mask = tf.equal(eye, one)
  same_cluster_mask = tf.reshape(same_cluster_mask, [-1])
  diff_cluster_mask = tf.equal(eye, zero)
  diff_cluster_mask = tf.reshape(diff_cluster_mask, [-1])

  mu_same_bool = tf.boolean_mask(mu_diff, same_cluster_mask)
  
  mu_same_norm = tf.norm(mu_same_bool, ord=1, axis=1)
  mu_same_norm = tf.subtract(mu_same_norm, delta_v)
  mu_same_norm = tf.clip_by_value(mu_same_norm, 0., mu_same_norm)
  mu_same_norm = tf.square(mu_same_norm)

  # same_norm_k_cnt = tf.cond(tf.greater(num_instances*5, num_instances), lambda:(num_instances), lambda:(num_instances*5))
  # mu_same_norm_top_k , _ = tf.nn.top_k(mu_same_norm, tf.cast(same_norm_k_cnt, tf.int32))
  # l_var = tf.reduce_mean(mu_same_norm_top_k)


  # l_var = tf.reduce_mean(mu_same_norm)

  # mask0 = tf.cast(tf.greater(unique_labels, 0), tf.float32)
  # l_var = tf.reduce_sum(mu_same_norm*mask0)

  l_var = tf.reduce_sum(mu_same_norm)
  l_var = tf.divide(l_var, tf.cast(num_instances-1, tf.float32))

  mu_diff_bool = tf.boolean_mask(mu_diff, diff_cluster_mask)

  mu_norm = tf.norm(mu_diff_bool, ord=1, axis=1)
  mu_norm = tf.subtract(2. * delta_d, mu_norm)
  mu_norm = tf.clip_by_value(mu_norm, 0., mu_norm)
  mu_norm = tf.square(mu_norm)

  # norm_k_cnt = tf.cond(tf.greater(num_instances*5, num_instances*num_instances-num_instances), lambda:(num_instances*num_instances-num_instances), lambda:(num_instances*5))
  # mu_norm_top_k , _ = tf.nn.top_k(mu_norm, tf.cast(norm_k_cnt, tf.int32))

  # l_dist = tf.reduce_mean(mu_norm_top_k)
  l_dist = tf.reduce_mean(mu_norm)


  def rt_0(): return 0.
  def rt_l_dist(): return l_dist
  l_dist = tf.cond(tf.equal(1, num_instances), rt_0, rt_l_dist)
  
  ### Calculate l_reg
  # l_reg = tf.reduce_mean(tf.norm(mu, ord=1, axis=1))

  param_scale = 1.
  l_var = param_var * l_var
  l_dist = param_dist * l_dist
  l_reg = param_reg * 0

  loss = param_scale * (l_var + l_dist + l_reg)

  return loss, l_var, l_dist, l_reg