#! /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)

[((548, 568), 'torch.unique', 'torch.unique', (['labels'], {}), '(labels)\n', (560, 568), False, 'import torch\n'), ((1681, 1720), 'torch.sum', 'torch.sum', (['weights'], {'dim': '(1)', 'keepdim': '(True)'}), '(weights, dim=1, keepdim=True)\n', (1690, 1720), False, 'import torch\n'), ((1090, 1109), 'torch.log', 'torch.log', (['dist_mat'], {}), '(dist_mat)\n', (1099, 1109), False, 'import torch\n'), ((1155, 1194), 'torch.log', 'torch.log', (['(1.0 - 0.25 * dist_mat ** 2.0)'], {}), '(1.0 - 0.25 * dist_mat ** 2.0)\n', (1164, 1194), False, 'import torch\n'), ((1239, 1261), 'torch.max', 'torch.max', (['log_weights'], {}), '(log_weights)\n', (1248, 1261), False, 'import torch\n'), ((887, 912), 'torch.tensor', 'torch.tensor', (['self.cutoff'], {}), '(self.cutoff)\n', (899, 912), False, 'import torch\n')]