from itertools import combinations

import numpy as np

import torch

def pdist(vectors):

    distance_matrix = -2 * vectors.mm(torch.t(vectors)) + vectors.pow(2).sum(dim=1).view(1, -1) + vectors.pow(2).sum(

        dim=1).view(-1, 1)

    return distance_matrix

class PairSelector:

    """

    Implementation should return indices of positive pairs and negative pairs that will be passed to compute

    Contrastive Loss

    return positive_pairs, negative_pairs

    """

    def __init__(self):

        pass

    def get_pairs(self, embeddings, labels):

        raise NotImplementedError

class AllPositivePairSelector(PairSelector):

    """

    Discards embeddings and generates all possible pairs given labels.

    If balance is True, negative pairs are a random sample to match the number of positive samples

    """

    def __init__(self, balance=True):

        super(AllPositivePairSelector, self).__init__()

        self.balance = balance

    def get_pairs(self, embeddings, labels):

        labels = labels.cpu().data.numpy()

        all_pairs = np.array(list(combinations(range(len(labels)), 2)))

        all_pairs = torch.LongTensor(all_pairs)

        positive_pairs = all_pairs[(labels[all_pairs[:, 0]] == labels[all_pairs[:, 1]]).nonzero()]

        negative_pairs = all_pairs[(labels[all_pairs[:, 0]] != labels[all_pairs[:, 1]]).nonzero()]

        if self.balance:

            negative_pairs = negative_pairs[torch.randperm(len(negative_pairs))[:len(positive_pairs)]]

        return positive_pairs, negative_pairs

class HardNegativePairSelector(PairSelector):

    """

    Creates all possible positive pairs. For negative pairs, pairs with smallest distance are taken into consideration,

    matching the number of positive pairs.

    """

    def __init__(self, cpu=True):

        super(HardNegativePairSelector, self).__init__()

        self.cpu = cpu

    def get_pairs(self, embeddings, labels):

        if self.cpu:

            embeddings = embeddings.cpu()

        distance_matrix = pdist(embeddings)

        labels = labels.cpu().data.numpy()

        all_pairs = np.array(list(combinations(range(len(labels)), 2)))

        all_pairs = torch.LongTensor(all_pairs)

        positive_pairs = all_pairs[(labels[all_pairs[:, 0]] == labels[all_pairs[:, 1]]).nonzero()]

        negative_pairs = all_pairs[(labels[all_pairs[:, 0]] != labels[all_pairs[:, 1]]).nonzero()]

        negative_distances = distance_matrix[negative_pairs[:, 0], negative_pairs[:, 1]]

        negative_distances = negative_distances.cpu().data.numpy()

        top_negatives = np.argpartition(negative_distances, len(positive_pairs))[:len(positive_pairs)]

        top_negative_pairs = negative_pairs[torch.LongTensor(top_negatives)]

        return positive_pairs, top_negative_pairs

class TripletSelector:

    """

    Implementation should return indices of anchors, positive and negative samples

    return np array of shape [N_triplets x 3]

    """

    def __init__(self):

        pass

    def get_pairs(self, embeddings, labels):

        raise NotImplementedError

class AllTripletSelector(TripletSelector):

    """

    Returns all possible triplets

    May be impractical in most cases

    """

    def __init__(self):

        super(AllTripletSelector, self).__init__()

    def get_triplets(self, embeddings, labels):

        labels = labels.cpu().data.numpy()

        triplets = []

        for label in set(labels):

            label_mask = (labels == label)

            label_indices = np.where(label_mask)[0]

            if len(label_indices) < 2:

                continue

            negative_indices = np.where(np.logical_not(label_mask))[0]

            anchor_positives = list(combinations(label_indices, 2))  # All anchor-positive pairs

            # Add all negatives for all positive pairs

            temp_triplets = [[anchor_positive[0], anchor_positive[1], neg_ind] for anchor_positive in anchor_positives

                             for neg_ind in negative_indices]

            triplets += temp_triplets

        return torch.LongTensor(np.array(triplets))

def hardest_negative(loss_values):

    hard_negative = np.argmax(loss_values)

    return hard_negative if loss_values[hard_negative] > 0 else None

def random_hard_negative(loss_values):

    hard_negatives = np.where(loss_values > 0)[0]

    return np.random.choice(hard_negatives) if len(hard_negatives) > 0 else None

def semihard_negative(loss_values, margin):

    semihard_negatives = np.where(np.logical_and(loss_values < margin, loss_values > 0))[0]

    return np.random.choice(semihard_negatives) if len(semihard_negatives) > 0 else None

class FunctionNegativeTripletSelector(TripletSelector):

    """

    For each positive pair, takes the hardest negative sample (with the greatest triplet loss value) to create a triplet

    Margin should match the margin used in triplet loss.

    negative_selection_fn should take array of loss_values for a given anchor-positive pair and all negative samples

    and return a negative index for that pair

    """

    def __init__(self, margin, negative_selection_fn, cpu=True):

        super(FunctionNegativeTripletSelector, self).__init__()

        self.cpu = cpu

        self.margin = margin

        self.negative_selection_fn = negative_selection_fn

    def get_triplets(self, embeddings, labels):

        if self.cpu:

            embeddings = embeddings.cpu()

        distance_matrix = pdist(embeddings)

        distance_matrix = distance_matrix.cpu()

        labels = labels.cpu().data.numpy()

        triplets = []

        for label in set(labels):

            label_mask = (labels == label)

            label_indices = np.where(label_mask)[0]

            if len(label_indices) < 2:

                continue

            negative_indices = np.where(np.logical_not(label_mask))[0]

            anchor_positives = list(combinations(label_indices, 2))  # All anchor-positive pairs

            anchor_positives = np.array(anchor_positives)

            ap_distances = distance_matrix[anchor_positives[:, 0], anchor_positives[:, 1]]

            for anchor_positive, ap_distance in zip(anchor_positives, ap_distances):

                loss_values = ap_distance - distance_matrix[torch.LongTensor(np.array([anchor_positive[0]])), torch.LongTensor(negative_indices)] + self.margin

                loss_values = loss_values.data.cpu().numpy()

                hard_negative = self.negative_selection_fn(loss_values)

                if hard_negative is not None:

                    hard_negative = negative_indices[hard_negative]

                    triplets.append([anchor_positive[0], anchor_positive[1], hard_negative])

        if len(triplets) == 0:

            triplets.append([anchor_positive[0], anchor_positive[1], negative_indices[0]])

        triplets = np.array(triplets)

        return torch.LongTensor(triplets)

def HardestNegativeTripletSelector(margin, cpu=False): return FunctionNegativeTripletSelector(margin=margin,

                                                                                 negative_selection_fn=hardest_negative,

                                                                                 cpu=cpu)

def RandomNegativeTripletSelector(margin, cpu=False): return FunctionNegativeTripletSelector(margin=margin,

                                                                                negative_selection_fn=random_hard_negative,

                                                                                cpu=cpu)

def SemihardNegativeTripletSelector(margin, cpu=False): return FunctionNegativeTripletSelector(margin=margin,

                                                                                  negative_selection_fn=lambda x: semihard_negative(x, margin),

                                                                                  cpu=cpu)