import torch


def hard_example_mining(distance_matrix, pos_idxs, neg_idxs):
    """For each anchor, find the hardest positive and negative sample.

    Args:
        distance_matrix: pair wise distance between samples, shape [N, M]
        pos_idxs: positive index with shape [N, M]
        neg_idxs: negative index with shape [N, M]

    Returns:
        dist_ap: pytorch Variable, distance(anchor, positive); shape [N]
        dist_an: pytorch Variable, distance(anchor, negative); shape [N]
        p_inds: pytorch LongTensor, with shape [N];
            indices of selected hard positive samples; 0 <= p_inds[i] <= N - 1
        n_inds: pytorch LongTensor, with shape [N];
            indices of selected hard negative samples; 0 <= n_inds[i] <= N - 1

    Note:
        Only consider the case in which all targets have same num of samples,
        thus we can cope with all anchors in parallel.
    """
    assert len(distance_matrix.size()) == 2  # noqa: S101

    # `dist_ap` means distance(anchor, positive)
    # both `dist_ap` and `relative_p_inds` with shape [N]
    dist_ap, _ = torch.max(distance_matrix * pos_idxs, dim=1)
    # `dist_an` means distance(anchor, negative)
    # both `dist_an` and `relative_n_inds` with shape [N]
    dist_an, _ = torch.min(distance_matrix * neg_idxs + pos_idxs * 99999999.0, dim=1)

    return dist_ap, dist_an