//  [python]
def discriminative_loss_single(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, [-1, feature_dim])
  correct_label = tf.reshape(correct_label_in, [tf.shape(reshaped_pred)[0]])
  # print(correct_label_in.shape, correct_label.shape)

  ### Count instances
  unique_labels, unique_id, counts = tf.unique_with_counts(correct_label)

  counts = tf.cast(counts, tf.float32)
  num_instances = tf.size(unique_labels)
  
  segmented_sum = tf.unsorted_segment_sum(reshaped_pred, unique_id, num_instances)

  mu = tf.div(segmented_sum, tf.reshape(counts, (-1, 1)))
  mu_expand = tf.gather(mu, unique_id)

  ### Calculate l_var
  #distance = tf.norm(tf.subtract(mu_expand, reshaped_pred), axis=1)
  #tmp_distance = tf.subtract(reshaped_pred, mu_expand)
  tmp_distance = reshaped_pred - mu_expand
  distance = tf.norm(tmp_distance, ord=1, axis=1)

  distance = tf.subtract(distance, delta_v)
  distance = tf.clip_by_value(distance, 0., distance)
  distance = tf.square(distance)

  l_var = tf.unsorted_segment_sum(distance, unique_id, num_instances)

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

  l_var = tf.div(l_var, counts)

  # line_idx = tf.where(tf.greater(unique_labels, 0))
  # l_var_line = tf.reshape(tf.gather(l_var, line_idx), [-1])
  # l_var_top_k, _ = tf.nn.top_k(l_var_line, tf.cast(num_instances/4 + 1, tf.int32))
  # l_var = tf.reduce_mean(l_var_top_k)

  # l_var = tf.reduce_sum(l_var)
  # l_var = tf.divide(l_var, tf.cast(num_instances, tf.float32))

  mask0 = tf.greater(unique_labels, 0)
  l_var = tf.boolean_mask(l_var, mask0)
  l_var = tf.reduce_sum(l_var)
  l_var = tf.divide(l_var, tf.cast(num_instances-1, tf.float32))

  ### 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 = tf.boolean_mask(mu, mask0)

  mu_interleaved_rep = tf.tile(mu, [num_instances-1, 1])
  mu_band_rep = tf.tile(mu, [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)
  diff_cluster_mask = tf.equal(eye, zero)
  diff_cluster_mask = tf.reshape(diff_cluster_mask, [-1])
  mu_diff_bool = tf.boolean_mask(mu_diff, diff_cluster_mask)

  #intermediate_tensor = tf.reduce_sum(tf.abs(mu_diff),axis=1)
  #zero_vector = tf.zeros(1, dtype=tf.float32)
  #bool_mask = tf.not_equal(intermediate_tensor, zero_vector)
  #mu_diff_bool = tf.boolean_mask(mu_diff, bool_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 * l_reg

  loss = param_scale * (l_var + l_dist + 0*l_reg)

  return loss, l_var, l_dist, l_reg