//  [python]
def triplet_loss_OHEM(prediction, gt, feature_map):
    '''
    online compute triplet loss
    由于计算量过大，我们选择OHEM的方案，选择一部分pixel出来计算triplet loss
    :param prediction: [N, W, H, 2]
    :param gt: [N, W, H]
    :param feature_map: [N, W, H, D]
    :return: [N, W, H, X]
    '''
    selected_num = 10  # 每个输入图像选择100个像素
    neg_selected_masks, neg_selected_indices = get_batch_OHEM_mask(prediction, gt, selected_num,
                                                                   pre_condition_mask=tf.cast(tf.equal(gt, 0),
                                                                                              tf.float32))
    pos_selected_masks, pos_selected_indices = get_batch_OHEM_mask(prediction, gt, selected_num,
                                                                   pre_condition_mask=tf.cast(tf.equal(gt, 1),
                                                                                              tf.float32))
    selected_indices = tf.concat([neg_selected_indices, pos_selected_indices], axis=-1)
    selected_masks = tf.cast(tf.greater_equal(neg_selected_masks + pos_selected_masks, 1.0), tf.float32)
    tf.summary.image('selected_pixel', tf.expand_dims(tf.cast(selected_masks * 200, tf.uint8), axis=3), max_outputs=3)

    def __euclidean_distance(tensor, tensors):
        '''
        计算一个tensor和多个tensor之间的相似度
        :param tensor: [D]
        :param tensors: [N, D]
        :return:
        '''
        distance = tensor - tensors
        distance = tf.reduce_sum(tf.square(distance), axis=1)
        return distance

    def __compute_triplet_loss(anchor_feature, anchor_gt, feature_map_flatten_, gt_flatten_):
        distances = __euclidean_distance(anchor_feature, feature_map_flatten_)

        print('distances are ', distances)
        num_pairs = 5
        # select hard positive pair
        # 与当前像素是同一个类别的，属于positive pair，我们优先选择距离大的
        pos_mask = tf.cast(tf.equal(gt_flatten_, anchor_gt), dtype=tf.float32)
        # tf.stop_gradient(pos_mask)
        pos_distances = distances * pos_mask
        selected_pos_distances, selected_pos_indices = tf.nn.top_k(pos_distances, k=num_pairs)
        pos_distances_mask = tf.cast(tf.greater_equal(pos_distances, selected_pos_distances[-1]), tf.float32)
        # tf.stop_gradient(pos_distances_mask)
        pos_distances_loss = tf.reduce_sum(pos_distances * pos_distances_mask) / (
                    tf.reduce_sum(pos_distances_mask) + 1e-7)

        # minimize pos_distances_loss 使得距离最小

        # select hard negative pair
        # 与当前像素不属于同一个类别，属于negative pair，我们优先选择距离小的
        neg_mask = tf.cast(tf.not_equal(gt_flatten_, anchor_gt), dtype=tf.float32)
        # tf.stop_gradient(neg_mask)
        neg_distances = distances * neg_mask
        neg_distances = tf.reduce_max(neg_distances) - neg_distances # 将大距离的变成小距离，小距离变成大距离, 选择大距离
        neg_distances = tf.where(tf.equal(neg_distances, tf.reduce_max(neg_distances)), tf.zeros_like(neg_distances),
                                 neg_distances)
        selected_neg_distances, selected_neg_indices = tf.nn.top_k(neg_distances, k=num_pairs)
        neg_distances_mask = tf.cast(tf.greater_equal(neg_distances, selected_neg_distances[-1]), tf.float32)
        # tf.stop_gradient(neg_distances_mask)
        neg_distances_loss = tf.reduce_sum(neg_distances * neg_distances_mask) / (
            tf.reduce_sum(neg_distances_mask + 1e-7))
        # maximize distances => minimize distance
        return pos_distances_loss + neg_distances_loss

    shape = feature_map.get_shape().as_list()
    feature_map_flatten = tf.reshape(feature_map, [shape[0], shape[1] * shape[2], shape[3]], name='flatten_feature_map')
    gt_flatten = tf.reshape(gt, [shape[0], shape[1] * shape[2]])
    selected_features = tf.batch_gather(feature_map_flatten, selected_indices)
    selected_gt = tf.batch_gather(gt_flatten, selected_indices)
    selected_features_flatten = tf.reshape(selected_features, [-1, shape[3]])
    selected_gt_flatten = tf.reshape(selected_gt, [-1])
    print('mark 0')
    print(selected_features_flatten, selected_gt_flatten)
    triplet_losses = tf.map_fn(
        lambda x: __compute_triplet_loss(x[0], x[1], selected_features_flatten, selected_gt_flatten),
        elems=[selected_features_flatten, selected_gt_flatten], dtype=tf.float32, name='map_fn_level2', swap_memory=True)
    return tf.reduce_mean(triplet_losses)