import numpy as np
import sys
import matplotlib.pyplot as plt
plt.switch_backend('agg')
from PIL import Image as PILImage
from matplotlib.offsetbox import OffsetImage, AnnotationBbox
import torch
import matplotlib.patches as mpatches
import shutil
from itertools import combinations
from MulticoreTSNE import MulticoreTSNE as TSNE
#from sklearn.manifold import TSNE
from sklearn.decomposition import PCA

indexcolors =["#000000", "#FFFF00", "#1CE6FF", "#FF34FF", "#FF4A46", "#008941", "#006FA6", "#A30059",

        "#FFDBE5", "#7A4900", "#0000A6", "#63FFAC", "#B79762", "#004D43", "#8FB0FF", "#997D87",
        "#5A0007", "#809693", "#E704C4", "#1B4400", "#4FC601", "#3B5DFF", "#4A3B53", "#FF2F80",
        "#61615A", "#BA0900", "#6B7900", "#00C2A0", "#FFAA92", "#FF90C9", "#B903AA", "#D16100",
        "#DDEFFF", "#000035", "#7B4F4B", "#A1C299", "#300018", "#0AA6D8", "#013349", "#00846F",
        "#372101", "#FFB500", "#C2FFED", "#A079BF", "#CC0744", "#C0B9B2", "#C2FF99", "#001E09",
        "#00489C", "#6F0062", "#0CBD66", "#EEC3FF", "#456D75", "#B77B68", "#7A87A1", "#788D66",
        "#885578", "#FAD09F", "#FF8A9A", "#D157A0", "#BEC459", "#456648", "#0086ED", "#886F4C",
        
        "#34362D", "#B4A8BD", "#00A6AA", "#452C2C", "#636375", "#A3C8C9", "#FF913F", "#938A81",
        "#575329", "#00FECF", "#B05B6F", "#8CD0FF", "#3B9700", "#04F757", "#C8A1A1", "#1E6E00",
        "#7900D7", "#A77500", "#6367A9", "#A05837", "#6B002C", "#772600", "#D790FF", "#9B9700",
        "#549E79", "#FFF69F", "#201625", "#72418F", "#BC23FF", "#99ADC0", "#3A2465", "#922329",
        "#5B4534", "#FDE8DC", "#404E55", "#0089A3", "#CB7E98", "#A4E804", "#324E72", "#6A3A4C",
        "#83AB58", "#001C1E", "#D1F7CE", "#004B28", "#C8D0F6", "#A3A489", "#806C66", "#222800",
        "#BF5650", "#E83000", "#66796D", "#DA007C", "#FF1A59", "#8ADBB4", "#1E0200", "#5B4E51",
        "#C895C5", "#320033", "#FF6832", "#66E1D3", "#CFCDAC", "#D0AC94", "#7ED379", "#012C58",
        
        "#7A7BFF", "#D68E01", "#353339", "#78AFA1", "#FEB2C6", "#75797C", "#837393", "#943A4D",
        "#B5F4FF", "#D2DCD5", "#9556BD", "#6A714A", "#001325", "#02525F", "#0AA3F7", "#E98176",
        "#DBD5DD", "#5EBCD1", "#3D4F44", "#7E6405", "#02684E", "#962B75", "#8D8546", "#9695C5",
        "#E773CE", "#D86A78", "#3E89BE", "#CA834E", "#518A87", "#5B113C", "#55813B", "#FEFFE6",
        "#00005F", "#A97399", "#4B8160", "#59738A", "#FF5DA7", "#F7C9BF", "#643127", "#513A01",
        "#6B94AA", "#51A058", "#A45B02", "#1D1702", "#E20027", "#E7AB63", "#4C6001", "#9C6966",
        "#64547B", "#97979E", "#006A66", "#391406", "#F4D749", "#0045D2", "#006C31", "#DDB6D0",
        "#7C6571", "#9FB2A4", "#00D891", "#15A08A", "#BC65E9", "#FFFFFE", "#C6DC99", "#203B3C",

        "#671190", "#6B3A64", "#F5E1FF", "#FFA0F2", "#CCAA35", "#374527", "#8BB400", "#797868",
        "#C6005A", "#3B000A", "#C86240", "#29607C", "#402334", "#7D5A44", "#CCB87C", "#B88183",
        "#AA5199", "#B5D6C3", "#A38469", "#9F94F0", "#A74571", "#B894A6", "#71BB8C", "#00B433",
        "#789EC9", "#6D80BA", "#953F00", "#5EFF03", "#E4FFFC", "#1BE177", "#BCB1E5", "#76912F",
        "#003109", "#0060CD", "#D20096", "#895563", "#29201D", "#5B3213", "#A76F42", "#89412E",
        "#1A3A2A", "#494B5A", "#A88C85", "#F4ABAA", "#A3F3AB", "#00C6C8", "#EA8B66", "#958A9F",
        "#BDC9D2", "#9FA064", "#BE4700", "#658188", "#83A485", "#453C23", "#47675D", "#3A3F00",
        "#061203", "#DFFB71", "#868E7E", "#98D058", "#6C8F7D", "#D7BFC2", "#3C3E6E", "#D83D66",
        
        "#2F5D9B", "#6C5E46", "#D25B88", "#5B656C", "#00B57F", "#545C46", "#866097", "#365D25",
        "#252F99", "#00CCFF", "#674E60", "#FC009C", "#92896B"]

def reduce_dimensionality(X):
  print("Calculating TSNE")
  #embedding= TSNE(n_components=2).fit_transform(X)
  embedding= TSNE(n_jobs=20, n_components=2).fit_transform(X)
  #embedding= PCA(n_components=2).fit_transform(X)
  return embedding
  #return X

def plot_together(embedd, labels, preds, ind, paths, info):
  embedding= reduce_dimensionality(embedd)
  fig, axs = plt.subplots(nrows=2, ncols=2, constrained_layout=True)
  colors= [indexcolors[int(i) % len(indexcolors)] for i in labels.squeeze()]
  axs[0,0].sc=axs[0, 0].scatter(embedding[:,0],embedding[:,1], s=1, c= colors )
  colors= [indexcolors[int(i) % len(indexcolors)] for i in preds.squeeze()]
  axs[0,1].sc=axs[0, 1].scatter(embedding[:,0],embedding[:,1], s=1, c= colors )
  colors= [indexcolors[int(i) % len(indexcolors)] for i in np.equal(labels, preds).astype(np.int32).squeeze()]
  axs[1,0].sc=axs[1, 0].scatter(embedding[:,0],embedding[:,1], s=1, c= colors )
  selected = np.zeros_like(labels)
  for i in ind:
    selected[i]= 1
  colors= [indexcolors[int(i) % len(indexcolors)] for i in selected.squeeze()]
  axs[1,1].sc=axs[1, 1].scatter(embedding[:,0],embedding[:,1], s=1, c= colors )

  legend_texts= [ x[0] for x in sorted(info.items(), key=lambda kv: kv[1])]
  patches=[]
  for i,label in enumerate(legend_texts):
    patches.append(mpatches.Patch(color=indexcolors[i], label=label))
  plt.legend(handles=patches)
  plt.xlabel('Dim 1', fontsize=12)
  plt.ylabel('Dim 2', fontsize=12)
  plt.grid(True)
  t_list=[]

  def onpick3(event):
    if event.button==1:
      #print(dir(event), type(sc))
      ax = event.inaxes
      sys.stdout.flush()
      cont, ind = ax.sc.contains(event)
      for thumb in ind['ind']:
        print(paths[thumb])
        img = PILImage.open(paths[thumb])
        img.thumbnail((128, 96), PILImage.ANTIALIAS)
        img = OffsetImage(img, zoom=1)
        ab = AnnotationBbox(img, (embedding[thumb,0]+0.2, embedding[thumb,1]+0.2), xycoords='data', frameon=False)
        ax.add_artist(ab)
        ab.set_visible(True)
        t_list.append(ab)
    else:
      for annot in t_list:
        annot.set_visible(False)
    event.canvas.draw()

  fig.canvas.mpl_connect('button_press_event', onpick3)
  plt.show()

def plot_embedding(embedd, labels, paths, info):
  fig= plt.figure(figsize=(10,10))
  embedding= reduce_dimensionality(embedd)
  ax= plt.gca()
  colors= [indexcolors[int(i) % len(indexcolors)] for i in labels.squeeze()]
  sc=ax.scatter(embedding[:,0],embedding[:,1], s=1, c= colors )
  legend_texts= [ x[1] for x in sorted(info.items(), key=lambda kv: kv[1])]
  plt.xticks([])
  plt.yticks([])
  patches=[]
  for i,label in enumerate(legend_texts):
      if np.count_nonzero(np.equal(labels, i))>0 :
          bgcolor= indexcolors[i][1:]
          #print(label,bgcolor)
          (r, g, b) = (bgcolor[:2], bgcolor[2:4], bgcolor[4:])
          color = 'black' if 1 - (int(r, 16) * 0.299 + int(g, 16) * 0.587 + int(b, 16) * 0.114) / 255 < 0.5 else 'white'
          plt.annotate(label, np.median(embedding[labels == i], axis=0), horizontalalignment='center', verticalalignment='center',                                                                                    size=8, weight='bold', color=color, backgroundcolor= indexcolors[i]) 
  #  patches.append(mpatches.Patch(color=indexcolors[i], label=label))
  #plt.legend(bbox_to_anchor=(0., 1.02, 1., .102), loc=3,
  #                     ncol=12, mode="expand", borderaxespad=0., handles=patches)
  #plt.legend(handles=patches)
  #plt.xlabel('Dim 1', fontsize=12)
  #plt.ylabel('Dim 2', fontsize=12)
  #plt.grid(True)
  t_list=[]

  def onpick3(event):
    if event.button==1:
      cont, ind = sc.contains(event)
      for thumb in ind['ind']:
        print(paths[thumb])
        img = PILImage.open(paths[thumb])
        img.thumbnail((128, 96), PILImage.ANTIALIAS)
        img = OffsetImage(img, zoom=1)
        ab = AnnotationBbox(img, (embedding[thumb,0]+0.2, embedding[thumb,1]+0.2), xycoords='data', frameon=False)
        ax.add_artist(ab)
        ab.set_visible(True)
        t_list.append(ab)
    else:
      for annot in t_list:
        annot.set_visible(False)
    event.canvas.draw()


  fig.canvas.mpl_connect('button_press_event', onpick3)
  plt.savefig('embedding_plot.png')

def plot_embedding_images(embedd, labels, paths, info, savefile):
  fig = plt.figure(figsize=(15,10))
  embedding = reduce_dimensionality(embedd)
  ax = plt.gca()
  ax.set_xlim(-15, 15)
  ax.set_ylim(-10, 10)
  colors = [indexcolors[int(i) % len(indexcolors)] for i in labels.squeeze()]
  sc = ax.scatter(embedding[:,0],embedding[:,1], s=12, c= colors )
  legend_texts = [ x[0] for x in sorted(info.items(), key=lambda kv: kv[1])]
  patches=[]
  for i,label in enumerate(legend_texts):
    patches.append(mpatches.Patch(color=indexcolors[i], label=label))
  plt.legend(handles=patches)
  #plt.xlabel('Dim 1', fontsize=12)
  #plt.ylabel('Dim 2', fontsize=12)
  #plt.grid(True)

  for i,thumb in enumerate(paths):
        #print(thumb)
        img = PILImage.open(thumb)
        # img.thumbnail((16, 12), PILImage.ANTIALIAS)
        img.thumbnail((24, 18), PILImage.ANTIALIAS)
        img = OffsetImage(img, zoom=1)
        ab = AnnotationBbox(img, (embedding[i,0]+0.3, embedding[i,1]+0.3), xycoords='data', frameon=False)
        ax.add_artist(ab)
        ab.set_visible(True)

  # plt.show()
  plt.savefig(savefile)

def save_embedding_plot(name, embedding, labels, info):
  fig = plt.figure(figsize=(10,10))
  colors= [indexcolors[int(i)] for i in labels.squeeze()]
  plt.scatter(embedding[:,0],embedding[:,1], s=1, c= colors)
  legend_texts= [ x[0] for x in sorted(info.items(), key=lambda kv: kv[1])]
  patches=[]
  for i,label in enumerate(legend_texts):
    patches.append(mpatches.Patch(color=indexcolors[i], label=label))
  plt.legend(handles=patches)
  plt.xlabel('Dim 1', fontsize=12)
  plt.ylabel('Dim 2', fontsize=12)
  plt.grid(True)
  plt.savefig(name)
  plt.close(fig)

def save_checkpoint(state, is_best, filename='checkpoint.pth.tar'):
    torch.save(state, filename)
    if is_best:
        shutil.copyfile(filename, 'model_best.pth.tar')

def load_checkpoint(filename='model_best.pth.tar'):
    return torch.load(filename)

class AverageMeter(object):
    """Computes and stores the average and current value"""
    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count

def adjust_learning_rate(optimizer, epoch):
    """Sets the learning rate to the initial LR decayed by 10 every 30 epochs"""
    lr = args.lr * (0.1 ** (epoch // 30))
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr


def accuracy(output, target, topk=(1,)):
    """Computes the accuracy over the k top predictions for the specified values of k"""
    with torch.no_grad():
        maxk = max(topk)
        batch_size = target.size(0)

        _, pred = output.topk(maxk, 1, True, True)
        pred = pred.t()
        correct = pred.eq(target.view(1, -1).expand_as(pred))

        res = []
        for k in topk:
            correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)
            res.append(correct_k.mul_(100.0 / batch_size))
        return res

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)
        #print(triplets.shape[0])
        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)
