import torch
import torch.nn as nn


class NT_Xent(nn.Module):
    """
    More than inspired from https://github.com/Spijkervet/SimCLR/blob/master/modules/nt_xent.py

    Notes
    =====

    Using this pytorch implementation, you don't actually need to l2-norm the inputs, the results will be
    identical, as shown if you run this file.
    """

    def __init__(self, batch_size, temperature, device):
        super(NT_Xent, self).__init__()
        self.batch_size = batch_size
        self.temperature = temperature
        self.mask = self.get_correlated_samples_mask()
        self.device = device

        self.criterion = nn.CrossEntropyLoss(reduction="sum")
        self.similarity_f = nn.CosineSimilarity(dim=2)

    def forward(self, z_i, z_j):
        """
        We do not sample negative examples explicitly.
        Instead, given a positive pair, similar to (Chen et al., 2017), we treat the other 2(N − 1) augmented examples within a minibatch as negative examples.
        """

        p1 = torch.cat((z_i, z_j), dim=0)
        sim = self.similarity_f(p1.unsqueeze(1), p1.unsqueeze(0)) / self.temperature

        sim_i_j = torch.diag(sim, self.batch_size)
        sim_j_i = torch.diag(sim, -self.batch_size)

        positive_samples = torch.cat((sim_i_j, sim_j_i), dim=0).reshape(self.batch_size * 2, 1)
        negative_samples = sim[self.mask].reshape(self.batch_size * 2, -1)

        labels = torch.zeros(self.batch_size * 2).to(self.device).long()
        logits = torch.cat((positive_samples, negative_samples), dim=1)
        loss = self.criterion(logits, labels)
        loss /= 2 * self.batch_size
        return loss

    def get_correlated_samples_mask(self):
        mask = torch.ones((self.batch_size * 2, self.batch_size * 2), dtype=bool)
        mask = mask.fill_diagonal_(0)
        for i in range(self.batch_size):
            mask[i, self.batch_size + i] = 0
            mask[self.batch_size + i, i] = 0
        return mask


if __name__ == "__main__":
    a, b = torch.rand(8, 12), torch.rand(8, 12)
    a_norm, b_norm = torch.nn.functional.normalize(a), torch.nn.functional.normalize(b)
    cosine_sim = torch.nn.CosineSimilarity()
    ntxent_loss = NT_Xent(8, 0.5, "cpu")
    assert torch.allclose(cosine_sim(a, b), cosine_sim(a_norm, b_norm))
    assert torch.allclose(ntxent_loss(a, b), ntxent_loss(a_norm, b_norm))

from collections import defaultdict
import torchvision.transforms as transforms
import random
from PIL import ImageFilter, Image
from torch.utils.data import Dataset


class GaussianBlur(object):
    def __init__(self, sigma=[.1, 2.]):
        self.sigma = sigma

    def __call__(self, x):
        sigma = random.uniform(self.sigma[0], self.sigma[1])
        x = x.filter(ImageFilter.GaussianBlur(radius=sigma))
        return x


class FaceSet(Dataset):
    def __init__(self, label_file, input_shape):
        self.label_file = label_file
        self.imgs = []
        self.size = 0
        self.labels = []
        self.img_2_label = {}
        self.label_2_img = defaultdict(list)
        self.create_data()
        s = 1
        color_jitter = transforms.ColorJitter(
            0.8 * s, 0.8 * s, 0.8 * s, 0.2 * s
        )
        self.trainform = transforms.Compose([
            lambda x: Image.open(x).convert("RGB"),  # open image
            transforms.CenterCrop(input_shape[1:]),
            transforms.RandomApply(
                [
                   color_jitter
                ],
                p=0.8),
            transforms.RandomGrayscale(p=0.2),
            transforms.RandomApply([GaussianBlur([.1, 2.])], p=0.5),
            transforms.RandomHorizontalFlip(),  # with 0.5 probability
            transforms.ToTensor(),  # Converts a PIL Image to [0, 1]
            transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
        ])

    def create_data(self):
        with open(self.label_file, "r") as f:
            data = f.readlines()
            for line in data:
                line = line.strip()
                line = line.split(" ")
                idx = int(line[0])
                self.labels.append(idx)
                for imgs in line[1:]:
                    self.imgs.append(imgs)
                    self.img_2_label[imgs] = idx
                    self.label_2_img[idx].append(imgs)

        self.size = len(self.imgs)

    def __len__(self):
        return self.size

    def __getitem__(self, idx):
        pos_file = self.imgs[idx]
        pos_label = self.img_2_label[pos_file]
        pos1 = self.trainform(pos_file)
        pos2 = self.trainform(pos_file)

        # neg_label = pos_label
        # while neg_label == pos_label:
        #     neg_label = random.choice(self.labels)
        # neg_file = random.choice(self.label_2_img[neg_label])
        # neg = self.trainform(neg_file)

        return pos1, pos2

 ===>>> Ready to record the trainning process. <<<=== 
 ===>>> All data is 13233, each epoch iters 206.765625 <<<=== 
 ===>>> Load Data <<<=== 
 ===>>> Load Model <<<=== 
 ===>>> Epoch 1 <<<=== 
Avg Loss is 3.577977
Avg Loss is 3.379171
Avg Loss is 3.286532
Avg Loss is 3.239770
Save best model
 ===>>> Epoch 2 <<<=== 
Avg Loss is 3.096657
Avg Loss is 3.103381
Avg Loss is 3.096943
Avg Loss is 3.089485
Save best model
 ===>>> Epoch 3 <<<=== 
Avg Loss is 3.054055
Avg Loss is 3.054430
Avg Loss is 3.058049
Avg Loss is 3.055276
Save best model
 ===>>> Epoch 4 <<<=== 
Avg Loss is 3.041535
Avg Loss is 3.047728
Avg Loss is 3.049511
Avg Loss is 3.048278
Save best model
 ===>>> Epoch 5 <<<=== 
Avg Loss is 3.038077
Avg Loss is 3.037763
Avg Loss is 3.037274
Avg Loss is 3.036270
Save best model
 ===>>> Epoch 6 <<<=== 
Avg Loss is 3.040787
Avg Loss is 3.038092
Avg Loss is 3.033206
Avg Loss is 3.033263
Save best model
 ===>>> Epoch 7 <<<=== 
Avg Loss is 3.029101
Avg Loss is 3.025529
Avg Loss is 3.026146
Avg Loss is 3.025752
Save best model
 ===>>> Epoch 8 <<<=== 
Avg Loss is 3.021876
Avg Loss is 3.023627
Avg Loss is 3.022280
Avg Loss is 3.022333
Save best model
 ===>>> Epoch 9 <<<=== 
Avg Loss is 3.016602
Avg Loss is 3.014577
Avg Loss is 3.013268
Avg Loss is 3.013314
Save best model
 ===>>> Epoch 10 <<<=== 
Avg Loss is 3.009610
Avg Loss is 3.006718
Avg Loss is 3.007279
Avg Loss is 3.005000
Save best model
 ===>>> Epoch 11 <<<=== 
Avg Loss is 3.006585
Avg Loss is 3.003934
Avg Loss is 3.002049
Avg Loss is 2.999707
Save best model
 ===>>> Epoch 12 <<<=== 
Avg Loss is 2.996356
Avg Loss is 2.996086
Avg Loss is 2.996202
Avg Loss is 2.995989
Save best model
 ===>>> Epoch 13 <<<=== 
Avg Loss is 2.986980
Avg Loss is 2.987497
Avg Loss is 2.988292
Avg Loss is 2.988140
Save best model
 ===>>> Epoch 14 <<<=== 
Avg Loss is 2.984407
Avg Loss is 2.982733
Avg Loss is 2.981985
Avg Loss is 2.983570
Save best model
 ===>>> Epoch 15 <<<=== 
Avg Loss is 2.983063
Avg Loss is 2.981772
Avg Loss is 2.982046
Avg Loss is 2.981637
Save best model
 ===>>> Epoch 16 <<<=== 
Avg Loss is 2.973182
Avg Loss is 2.974823
Avg Loss is 2.974841
Avg Loss is 2.974745
Save best model
 ===>>> Epoch 17 <<<=== 
Avg Loss is 2.972452
Avg Loss is 2.971775
Avg Loss is 2.972774
Avg Loss is 2.973590
Save best model
 ===>>> Epoch 18 <<<=== 
Avg Loss is 2.968253
Avg Loss is 2.969223
Avg Loss is 2.969758
Avg Loss is 2.969625
Save best model
 ===>>> Epoch 19 <<<=== 
Avg Loss is 2.965280
Avg Loss is 2.963255
Avg Loss is 2.962672
Avg Loss is 2.962905
Save best model
 ===>>> Epoch 20 <<<=== 
Avg Loss is 2.962090
Avg Loss is 2.963186
Avg Loss is 2.963687
Avg Loss is 2.964000
 ===>>> Epoch 21 <<<=== 
Avg Loss is 2.960579
Avg Loss is 2.962159
Avg Loss is 2.961549
Avg Loss is 2.960474
Save best model
 ===>>> Epoch 22 <<<=== 
Avg Loss is 2.956076
Avg Loss is 2.954066
Avg Loss is 2.952812
Avg Loss is 2.953630
Save best model
 ===>>> Epoch 23 <<<=== 
Avg Loss is 2.957394
Avg Loss is 2.957740
Avg Loss is 2.955869
Avg Loss is 2.954493
 ===>>> Epoch 24 <<<=== 
Avg Loss is 2.950666
Avg Loss is 2.949612
Avg Loss is 2.949405
Avg Loss is 2.949652
Save best model
 ===>>> Epoch 25 <<<=== 
Avg Loss is 2.948632
Avg Loss is 2.948776
Avg Loss is 2.948908
Avg Loss is 2.949181
Save best model
 ===>>> Epoch 26 <<<=== 
Avg Loss is 2.944604
Avg Loss is 2.946091
Avg Loss is 2.947033
Avg Loss is 2.947451
Save best model
 ===>>> Epoch 27 <<<=== 
Avg Loss is 2.943191
Avg Loss is 2.944458
Avg Loss is 2.944051
Avg Loss is 2.944110
Save best model
 ===>>> Epoch 28 <<<=== 
Avg Loss is 2.942361
Avg Loss is 2.940979
Avg Loss is 2.941820
Avg Loss is 2.942666
Save best model
 ===>>> Epoch 29 <<<=== 
Avg Loss is 2.939801
Avg Loss is 2.939919
Avg Loss is 2.940369
Avg Loss is 2.940156
Save best model
 ===>>> Epoch 30 <<<=== 
Avg Loss is 2.935622
Avg Loss is 2.935057
Avg Loss is 2.935227
Avg Loss is 2.935683
Save best model
 ===>>> Epoch 31 <<<=== 
Avg Loss is 2.937815
Avg Loss is 2.935643
Avg Loss is 2.936246
Avg Loss is 2.936065
 ===>>> Epoch 32 <<<=== 
Avg Loss is 2.934247
Avg Loss is 2.934004
Avg Loss is 2.935194
Avg Loss is 2.934972
Save best model
 ===>>> Epoch 33 <<<=== 
Avg Loss is 2.931649
Avg Loss is 2.930846
Avg Loss is 2.930627
Avg Loss is 2.930930
Save best model
 ===>>> Epoch 34 <<<=== 
Avg Loss is 2.929090
Avg Loss is 2.929404
Avg Loss is 2.928846
Avg Loss is 2.929138
Save best model
 ===>>> Epoch 35 <<<=== 
Avg Loss is 2.927705
Avg Loss is 2.927764
Avg Loss is 2.927287
Avg Loss is 2.927516
Save best model
 ===>>> Epoch 36 <<<=== 
Avg Loss is 2.927743
Avg Loss is 2.926714
Avg Loss is 2.926628
Avg Loss is 2.926617
Save best model
 ===>>> Epoch 37 <<<=== 
Avg Loss is 2.924690
Avg Loss is 2.924439
Avg Loss is 2.924409
Avg Loss is 2.924603
Save best model
 ===>>> Epoch 38 <<<=== 
Avg Loss is 2.922124
Avg Loss is 2.922200
Avg Loss is 2.922583
Avg Loss is 2.922992
Save best model
 ===>>> Epoch 39 <<<=== 
Avg Loss is 2.922472
Avg Loss is 2.921901
Avg Loss is 2.921420
Avg Loss is 2.921485
Save best model
 ===>>> Epoch 40 <<<=== 
Avg Loss is 2.921616
Avg Loss is 2.920895
Avg Loss is 2.920781
Avg Loss is 2.920465
Save best model
 ===>>> Epoch 41 <<<=== 
Avg Loss is 2.918778
Avg Loss is 2.918589
Avg Loss is 2.918421
Avg Loss is 2.918423
Save best model
 ===>>> Epoch 42 <<<=== 
Avg Loss is 2.919555
Avg Loss is 2.919151
Avg Loss is 2.918455
Avg Loss is 2.918853
 ===>>> Epoch 43 <<<=== 
Avg Loss is 2.917781
Avg Loss is 2.916509
Avg Loss is 2.916942
Avg Loss is 2.916903
Save best model
 ===>>> Epoch 44 <<<=== 
Avg Loss is 2.917233
Avg Loss is 2.917524
Avg Loss is 2.916481
Avg Loss is 2.916096
Save best model
 ===>>> Epoch 45 <<<=== 
Avg Loss is 2.915788
Avg Loss is 2.915589
Avg Loss is 2.915973
Avg Loss is 2.915928
Save best model
 ===>>> Epoch 46 <<<=== 
Avg Loss is 2.916222
Avg Loss is 2.916418
Avg Loss is 2.916036
Avg Loss is 2.915801
Save best model
 ===>>> Epoch 47 <<<=== 
Avg Loss is 2.914999
Avg Loss is 2.913783
Avg Loss is 2.913962
Avg Loss is 2.914269
Save best model
 ===>>> Epoch 48 <<<=== 
Avg Loss is 2.914487
Avg Loss is 2.914710
Avg Loss is 2.914641
Avg Loss is 2.914529
 ===>>> Epoch 49 <<<=== 
Avg Loss is 2.915182
Avg Loss is 2.914820
Avg Loss is 2.914557
Avg Loss is 2.914477

torchvision.transforms.RandomResizedCrop(size=size),
torchvision.transforms.RandomHorizontalFlip(),  # with 0.5 probability
torchvision.transforms.RandomApply([color_jitter], p=0.8),
torchvision.transforms.RandomGrayscale(p=0.2),
torchvision.transforms.ToTensor(),