#! /usr/bin/env python3

from .base_miner import BasePostGradientMiner
import torch
from ..utils import loss_and_miner_utils as lmu


# adapted from
# https://github.com/chaoyuaw/incubator-mxnet/blob/master/example/gluon/
# /embedding_learning/model.py
class DistanceWeightedMiner(BasePostGradientMiner):
    def __init__(self, cutoff, nonzero_loss_cutoff, **kwargs):
        super().__init__(**kwargs)
        self.cutoff = cutoff
        self.nonzero_loss_cutoff = nonzero_loss_cutoff

    def mine(self, embeddings, labels):
        label_set = torch.unique(labels)
        n, d = embeddings.size()

        dist_mat = lmu.dist_mat(embeddings)
        dist_mat = dist_mat + torch.eye(dist_mat.size(0)).to(embeddings.device)  
        # so that we don't get log(0). We mask the diagonal out later anyway
        # Cut off to avoid high variance.
        dist_mat = torch.max(dist_mat, torch.tensor(self.cutoff).to(dist_mat.device))

        # Subtract max(log(distance)) for stability.
        # See the first equation from Section 4 of the paper
        log_weights = (2.0 - float(d)) * torch.log(dist_mat) - (
            float(d - 3) / 2
        ) * torch.log(1.0 - 0.25 * (dist_mat ** 2.0))
        weights = torch.exp(log_weights - torch.max(log_weights))

        # Sample only negative examples by setting weights of
        # the same-class examples to 0.
        mask = torch.ones(weights.size()).to(embeddings.device)
        for i in label_set:
            idx = (labels == i).nonzero()
            mask[torch.meshgrid(idx.squeeze(1), idx.squeeze(1))] = 0

        weights = weights * mask * ((dist_mat < self.nonzero_loss_cutoff).float())
        weights = weights / torch.sum(weights, dim=1, keepdim=True)

        np_weights = weights.cpu().numpy()
        return lmu.get_random_triplet_indices(labels, weights=np_weights)

[((18, 20, 18, 40), 'torch.unique', 'torch.unique', ({(18, 33, 18, 39): 'labels'}, {}), '(labels)', False, 'import torch\n'), ((42, 28, 42, 67), 'torch.sum', 'torch.sum', (), '', False, 'import torch\n'), ((29, 41, 29, 60), 'torch.log', 'torch.log', ({(29, 51, 29, 59): 'dist_mat'}, {}), '(dist_mat)', False, 'import torch\n'), ((31, 12, 31, 53), 'torch.log', 'torch.log', ({(31, 22, 31, 52): '(1.0 - 0.25 * dist_mat ** 2.0)'}, {}), '(1.0 - 0.25 * dist_mat ** 2.0)', False, 'import torch\n'), ((32, 42, 32, 64), 'torch.max', 'torch.max', ({(32, 52, 32, 63): 'log_weights'}, {}), '(log_weights)', False, 'import torch\n'), ((25, 39, 25, 64), 'torch.tensor', 'torch.tensor', ({(25, 52, 25, 63): 'self.cutoff'}, {}), '(self.cutoff)', False, 'import torch\n')]