import torch.optim as opt
import torch.utils.data
from sklearn import metrics

from classifiers import Net_CSI
from models.csi_utils import *

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


# code borrowed from https://github.com/alinlab/CSI

class CSI(nn.Module):
    def __init__(self, model=Net_CSI, simclr_dim=128, k_shift=4, lam=0.1, out_dim=128):
        super().__init__()
        # self.shift_trans = params['shift_trans']
        self.shift_trans = Rotation()
        # self.K_shift = params['K_shift']
        self.k_shift = k_shift
        self.sim_lambda = lam
        self.simclr_aug = get_simclr_augmentation((32, 32, 3)).to(device)
        self.hflip = HorizontalFlipLayer().to(device)

        self.model = model(out_dim=out_dim, simclr_dim=simclr_dim, shift=k_shift)
        self.train_features = None
        self.lambda_cls = None
        self.lambda_con = None

    def decision_function(self, test_data):
        self.model.eval()
        y_pred = np.array([])
        y_true = np.array([])

        test_loader = torch.utils.data.DataLoader(test_data,
                                                  batch_size=64)

        for (images, labels, _) in test_loader:
            s_con = 0
            s_cls = 0

            images = images.to(device)
            images = self.hflip(images)

            images = torch.cat([self.shift_trans(images, k) for k in range(self.k_shift)])

            out = self.model(images, penultimate=True, shift=True)

            chunks = out['penultimate'].chunk(self.k_shift)
            chunks_s = out['shift'].chunk(self.k_shift)

            for i in range(self.k_shift):
                s_con += self.lambda_con[i] * (chunks[i] @ self.train_features[i].T).max(dim=1)[0]
                s_cls += self.lambda_cls[i] * (chunks_s[i][:, i])

            s_csi = s_con + s_cls
            y_pred = np.hstack((y_pred, s_csi.cpu().detach().numpy()))
            y_true = np.hstack((y_true, labels.cpu().detach().numpy()))

        return y_pred, y_true

    def update_decision_function(self, train_data):
        self.eval()

        dataloader = torch.utils.data.DataLoader(train_data.lab_data(lab=1),
                                                 batch_size=512)

        self.train_features = [[] for k in range(self.k_shift)]
        self.lambda_cls = torch.zeros(self.k_shift, device=device)
        self.lambda_con = torch.zeros(self.k_shift, device=device)

        for (images, _, labels) in dataloader:
            images = images.to(device)
            images = self.hflip(images)

            images = torch.cat([self.shift_trans(images, k) for k in range(self.k_shift)])

            out = self.model(images, penultimate=True, shift=True)

            chunks = out['penultimate'].chunk(self.k_shift)
            chunks_s = out['shift'].chunk(self.k_shift)

            for i in range(self.k_shift):
                self.train_features[i].append(chunks[i].detach())
                self.lambda_con[i] += torch.norm(chunks[i], p=2, dim=1).sum().detach()
                self.lambda_cls[i] += chunks_s[i].sum(dim=0)[i].detach()

        self.lambda_cls = len(train_data) / self.lambda_cls
        self.lambda_con = len(train_data) / self.lambda_con
        self.train_features = [torch.cat(f) for f in self.train_features]
        for i in range(self.k_shift):
            self.train_features[i] = self.train_features[i] / torch.norm(self.train_features[i], dim=1).view(-1, 1)

    def fit(self,
            train_data,
            num_epochs=60,
            lr=1e-3,
            gamma=0.99,
            batch_size=64,
            temp=0.5,
            verbose=False,
            test_data=None):

        # self.debug = dict()

        self.model.to(device)

        # self.debug['shift_loss'] = []
        # self.debug['shift_labels'] = []

        criterion = nn.CrossEntropyLoss().to(device)
        optimizer = opt.Adam(self.model.parameters(), lr=lr)
        scheduler = opt.lr_scheduler.ExponentialLR(optimizer, gamma=gamma)

        # linear_optimizer = opt.Adam(self.model.linear.parameters(), lr=lr)

        data_loader = torch.utils.data.DataLoader(train_data.lab_data(lab=1),
                                                  batch_size=batch_size,
                                                  shuffle=True)
        for epoch in range(num_epochs):
            running_loss = 0
            for (images, _, labels) in data_loader:
                self.model.train()

                ### SimCLR loss ###
                images = images.to(device)
                images1, images2 = self.hflip(images.repeat(2, 1, 1, 1)).chunk(2)  # hflip

                labels = labels.to(device)

                images1 = torch.cat([self.shift_trans(images1, k) for k in range(self.k_shift)])
                images2 = torch.cat([self.shift_trans(images2, k) for k in range(self.k_shift)])

                shift_labels = torch.cat([torch.ones_like(labels) * k for k in range(self.k_shift)], 0)  # B -> 4B
                shift_labels = shift_labels.repeat(2)

                images_pair = torch.cat([images1, images2], dim=0)  # 8B
                # print(images_pair.shape)

                images_pair = self.simclr_aug(images_pair)  # transform

                # print(images_pair.shape)
                outputs_aux = self.model(images_pair, simclr=True, penultimate=True, shift=True)

                simclr = normalize(outputs_aux['simclr'])  # normalize
                sim_matrix = get_similarity_matrix(simclr, multi_gpu=False)
                loss_sim = NT_xent(sim_matrix, temperature=temp) * self.sim_lambda

                # self.debug['shift_loss'].append(outputs_aux['shift'].float())
                # self.debug['shift_labels'].append(shift_labels)
                loss_shift = criterion(outputs_aux['shift'].float(), shift_labels.long())

                ### total loss ###
                loss = loss_sim + loss_shift

                optimizer.zero_grad()
                loss.backward()
                optimizer.step()

                running_loss += loss.item()

            scheduler.step()
            if test_data is not None and verbose:
                self.model.eval()
                self.update_decision_function(train_data)
                y_pred, y_true = self.decision_function(test_data)

                print_line = f"[{epoch:4}/{num_epochs:4}]: loss={running_loss: .4f}, test auc={metrics.roc_auc_score(y_true, y_pred):.4f}"

                print(print_line, flush=True)

        self.update_decision_function(train_data)


class PU_CSI(nn.Module):
    def __init__(self, model=Net_CSI, simclr_dim=128, k_shift=4, lam=0.1, out_dim=128, flag=True):
        """
        :param flag: if true other positive examples considered in contrastive loss as negatives, if false only
        unlabeled data considered as negative
        """
        super().__init__()
        self.k_shift = k_shift
        self.sim_lambda = lam
        self.simclr_aug = get_simclr_augmentation((32, 32, 3)).to(device)
        self.hflip = HorizontalFlipLayer().to(device)
        self.flag = flag

        self.model = model(out_dim=out_dim, simclr_dim=simclr_dim, shift=k_shift)
        self.train_features = {}
        self.lambda_cls = {}
        self.lambda_con = {}

    def decision_function(self, test_data):
        y_preds = {'csi': np.array([]), 'diff': np.array([]), 'combo': np.array([])}
        y_true = np.array([])

        test_loader = torch.utils.data.DataLoader(test_data,
                                                  batch_size=64)

        for (images, labels, _) in test_loader:
            s_con_pos = 0
            s_con_unl = 0
            s_cls = 0

            images = images.to(device)
            images = self.hflip(images)

            images = torch.cat([torch.rot90(images, rot, (2, 3)) for rot in range(4)])

            out = self.model(images, penultimate=True, shift=True)

            chunks = out['penultimate'].chunk(4)
            chunks_s = out['shift'].chunk(4)

            for i in range(self.k_shift):
                s_con_pos += self.lambda_con['pos'][i] * (chunks[i] @ self.train_features['pos'][i].T).max(dim=1)[0]
                # s_con_unl += self.lambda_con['unl'][i] * (chunks[i] @ self.train_features['unl'][i].T).max(dim=1)[0]
                s_cls += self.lambda_cls['pos'][i] * (chunks_s[i][:, i])

            s_csi = s_con_pos + s_cls
            # s_diff = s_con_pos - s_con_unl
            # s_combo = s_con_pos + s_cls - s_con_unl

            y_preds['csi'] = np.hstack((y_preds['csi'], s_csi.cpu().detach().numpy()))
            # y_preds['diff'] = np.hstack((y_preds['diff'], s_diff.cpu().detach().numpy()))
            # y_preds['combo'] = np.hstack((y_preds['combo'], s_combo.cpu().detach().numpy()))

            y_true = np.hstack((y_true, labels.cpu().detach().numpy()))

        return y_preds['csi'], y_true

    def image_trans(self, images):
        images1, images2 = self.hflip(images.repeat(2, 1, 1, 1)).chunk(2)  # hflip
        images1 = torch.cat([torch.rot90(images1, rot, (2, 3)) for rot in range(4)])
        images2 = torch.cat([torch.rot90(images2, rot, (2, 3)) for rot in range(4)])
        images_pair = torch.cat([images1, images2], dim=0)  # 8B
        rot_labels = torch.cat([torch.ones(len(images)) * i for i in range(4)], dim=0)

        return images_pair, rot_labels.long().repeat(2)

    def fit(self,
            train_data,
            num_epochs=60,
            lr=1e-3,
            gamma=0.99,
            batch_size=64,
            batch_size_lab=64,
            batch_size_unl=16,
            temp=0.5,
            w=0.1,
            test_data=None,
            verbose=False):

        # self.debug = dict()

        self.model.to(device)

        # self.debug['shift_loss'] = []
        # self.debug['shift_labels'] = []

        criterion = nn.CrossEntropyLoss().to(device)
        optimizer = opt.Adam(self.model.parameters(), lr=lr)
        scheduler = opt.lr_scheduler.ExponentialLR(optimizer, gamma=gamma)

        # data_loader = torch.utils.data.DataLoader(train_data,
        #                                           batch_size=batch_size,
        #                                           shuffle=True)

        lab_loader = torch.utils.data.DataLoader(train_data.lab_data(lab=1),
                                                 batch_size=batch_size_lab,
                                                 shuffle=True)

        unl_loader = torch.utils.data.DataLoader(train_data.lab_data(lab=0),
                                                 batch_size=batch_size_unl,
                                                 shuffle=True)

        for epoch in range(num_epochs):
            running_loss = 0
            # for (images, _, labels) in data_loader:
            for batch in zip(lab_loader, unl_loader):

                images_p, images_u = batch[0][0], batch[1][0]
                labels_p, labels_u = batch[0][2], batch[1][2]

                # if len(images_p) != len(images_u):
                #     continue

                # images_p, images_u = images[labels == 1], images[labels == 0]
                # labels_p, labels_u = labels[labels == 1], labels[labels == 0]

                # if images_p.shape[0] == 0:
                #     continue

                images_p, rot_labels_p = self.image_trans(images_p)
                labels_p = labels_p.repeat(4)

                if images_u.shape[0] > 0:
                    images_u, rot_labels_u = self.image_trans(images_u)
                    labels_u = labels_u.repeat(4)

                if images_u.shape[0] > 0:
                    images = torch.cat((images_p, images_u))
                    labels = torch.cat((labels_p, labels_u))
                else:
                    images = images_p
                    labels = labels_p

                images = images.to(device)
                labels = labels.to(device)
                rot_labels = rot_labels_p.to(device)

                self.model.train()

                ### SimCLR loss ###
                batch_size = images.size(0)

                images_pair = self.simclr_aug(images)  # simclr augment
                outputs_aux = self.model(images_pair, simclr=True, penultimate=True, shift=True)

                simclr = normalize(outputs_aux['simclr'])  # normalize
                sim_matrix = get_similarity_matrix(simclr)
                loss_csi = Semisupervised_NT_xent(sim_matrix, len(labels_p), len(labels_u), temperature=temp,
                                                  flag=self.flag, w=w) * self.sim_lambda

                loss_shift = criterion(outputs_aux['shift'].float()[:len(labels_p) * 2], rot_labels)

                loss = loss_csi + loss_shift

                optimizer.zero_grad()
                loss.backward()
                optimizer.step()

                running_loss += loss.item()

            scheduler.step()
            if test_data is not None and verbose:
                self.model.eval()
                self.update_decision_function(train_data, 'pos')
                # self.update_decision_function(train_data, 'unl')

                y_pred, y_true = self.decision_function(test_data)

                auc_csi = metrics.roc_auc_score(y_true, y_pred)
                # auc_diff = metrics.roc_auc_score(y_true, y_preds['diff'])
                # auc_combo = metrics.roc_auc_score(y_true, y_preds['combo'])

                print_line = f"[{epoch:4}/{num_epochs:4}]: loss={running_loss: .4f}, auc_csi={auc_csi:.4f}"

                print(print_line)
        self.update_decision_function(train_data, 'pos')

    def update_decision_function(self, train_data, mode='pos'):
        self.eval()

        if mode == 'pos':
            label = 1
        elif mode == 'unl':
            label = 0
        else:
            raise ValueError(f'wrong mode {mode}')

        dataloader = torch.utils.data.DataLoader(train_data.lab_data(lab=label),
                                                 batch_size=512)

        self.train_features[mode] = [[] for k in range(self.k_shift)]
        self.lambda_cls[mode] = torch.zeros(self.k_shift, device=device)
        self.lambda_con[mode] = torch.zeros(self.k_shift, device=device)

        for (images, _, labels) in dataloader:
            batch_size = images.size(0)
            images = images.to(device)
            images = self.hflip(images)

            images = torch.cat([torch.rot90(images, rot, (2, 3)) for rot in range(4)])

            out = self.model(images, penultimate=True, shift=True)

            chunks = out['penultimate'].chunk(4)
            chunks_s = out['shift'].chunk(4)

            for i in range(4):
                self.train_features[mode][i].append(chunks[i].detach())
                self.lambda_con[mode][i] += torch.norm(chunks[i], p=2, dim=1).sum().detach()
                self.lambda_cls[mode][i] += chunks_s[i].sum(dim=0)[i].detach()

        self.lambda_cls[mode] = len(train_data) / self.lambda_cls[mode]
        self.lambda_con[mode] = len(train_data) / self.lambda_con[mode]
        self.train_features[mode] = [torch.cat(f) for f in self.train_features[mode]]

        for i in range(4):
            self.train_features[mode][i] = self.train_features[mode][i] / torch.norm(self.train_features[mode][i],
                                                                                     dim=1).view(-1, 1)

import torch.optim as opt
import torch.utils.data
from sklearn import metrics

from classifiers import Net_CSI
from models.csi_utils import *

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


# code borrowed from https://github.com/alinlab/CSI

class CSI(nn.Module):
    def __init__(self, model=Net_CSI, simclr_dim=128, k_shift=4, lam=0.1, out_dim=128):
        super().__init__()
        # self.shift_trans = params['shift_trans']
        self.shift_trans = Rotation()
        # self.K_shift = params['K_shift']
        self.k_shift = k_shift
        self.sim_lambda = lam
        self.simclr_aug = get_simclr_augmentation((32, 32, 3)).to(device)
        self.hflip = HorizontalFlipLayer().to(device)

        self.model = model(out_dim=out_dim, simclr_dim=simclr_dim, shift=k_shift)
        self.train_features = None
        self.lambda_cls = None
        self.lambda_con = None

    def decision_function(self, test_data):
        self.model.eval()
        y_pred = np.array([])
        y_true = np.array([])

        test_loader = torch.utils.data.DataLoader(test_data,
                                                  batch_size=64)

        for (images, labels, _) in test_loader:
            s_con = 0
            s_cls = 0

            images = images.to(device)
            images = self.hflip(images)

            images = torch.cat([self.shift_trans(images, k) for k in range(self.k_shift)])

            out = self.model(images, penultimate=True, shift=True)

            chunks = out['penultimate'].chunk(self.k_shift)
            chunks_s = out['shift'].chunk(self.k_shift)

            for i in range(self.k_shift):
                s_con += self.lambda_con[i] * (chunks[i] @ self.train_features[i].T).max(dim=1)[0]
                s_cls += self.lambda_cls[i] * (chunks_s[i][:, i])

            s_csi = s_con + s_cls
            y_pred = np.hstack((y_pred, s_csi.cpu().detach().numpy()))
            y_true = np.hstack((y_true, labels.cpu().detach().numpy()))

        return y_pred, y_true

    def update_decision_function(self, train_data):
        self.eval()

        dataloader = torch.utils.data.DataLoader(train_data.lab_data(lab=1),
                                                 batch_size=512)

        self.train_features = [[] for k in range(self.k_shift)]
        self.lambda_cls = torch.zeros(self.k_shift, device=device)
        self.lambda_con = torch.zeros(self.k_shift, device=device)

        for (images, _, labels) in dataloader:
            images = images.to(device)
            images = self.hflip(images)

            images = torch.cat([self.shift_trans(images, k) for k in range(self.k_shift)])

            out = self.model(images, penultimate=True, shift=True)

            chunks = out['penultimate'].chunk(self.k_shift)
            chunks_s = out['shift'].chunk(self.k_shift)

            for i in range(self.k_shift):
                self.train_features[i].append(chunks[i].detach())
                self.lambda_con[i] += torch.norm(chunks[i], p=2, dim=1).sum().detach()
                self.lambda_cls[i] += chunks_s[i].sum(dim=0)[i].detach()

        self.lambda_cls = len(train_data) / self.lambda_cls
        self.lambda_con = len(train_data) / self.lambda_con
        self.train_features = [torch.cat(f) for f in self.train_features]
        for i in range(self.k_shift):
            self.train_features[i] = self.train_features[i] / torch.norm(self.train_features[i], dim=1).view(-1, 1)

    def fit(self,
            train_data,
            num_epochs=60,
            lr=1e-3,
            gamma=0.99,
            batch_size=64,
            temp=0.5,
            verbose=False,
            test_data=None):

        # self.debug = dict()

        self.model.to(device)

        # self.debug['shift_loss'] = []
        # self.debug['shift_labels'] = []

        criterion = nn.CrossEntropyLoss().to(device)
        optimizer = opt.Adam(self.model.parameters(), lr=lr)
        scheduler = opt.lr_scheduler.ExponentialLR(optimizer, gamma=gamma)

        # linear_optimizer = opt.Adam(self.model.linear.parameters(), lr=lr)

        data_loader = torch.utils.data.DataLoader(train_data.lab_data(lab=1),
                                                  batch_size=batch_size,
                                                  shuffle=True)
        for epoch in range(num_epochs):
            running_loss = 0
            for (images, _, labels) in data_loader:
                self.model.train()

                ### SimCLR loss ###
                images = images.to(device)
                images1, images2 = self.hflip(images.repeat(2, 1, 1, 1)).chunk(2)  # hflip

                labels = labels.to(device)

                images1 = torch.cat([self.shift_trans(images1, k) for k in range(self.k_shift)])
                images2 = torch.cat([self.shift_trans(images2, k) for k in range(self.k_shift)])

                shift_labels = torch.cat([torch.ones_like(labels) * k for k in range(self.k_shift)], 0)  # B -> 4B
                shift_labels = shift_labels.repeat(2)

                images_pair = torch.cat([images1, images2], dim=0)  # 8B
                # print(images_pair.shape)

                images_pair = self.simclr_aug(images_pair)  # transform

                # print(images_pair.shape)
                outputs_aux = self.model(images_pair, simclr=True, penultimate=True, shift=True)

                simclr = normalize(outputs_aux['simclr'])  # normalize
                sim_matrix = get_similarity_matrix(simclr, multi_gpu=False)
                loss_sim = NT_xent(sim_matrix, temperature=temp) * self.sim_lambda

                # self.debug['shift_loss'].append(outputs_aux['shift'].float())
                # self.debug['shift_labels'].append(shift_labels)
                loss_shift = criterion(outputs_aux['shift'].float(), shift_labels.long())

                ### total loss ###
                loss = loss_sim + loss_shift

                optimizer.zero_grad()
                loss.backward()
                optimizer.step()

                running_loss += loss.item()

            scheduler.step()
            if test_data is not None and verbose:
                self.model.eval()
                self.update_decision_function(train_data)
                y_pred, y_true = self.decision_function(test_data)

                print_line = f"[{epoch:4}/{num_epochs:4}]: loss={running_loss: .4f}, test auc={metrics.roc_auc_score(y_true, y_pred):.4f}"

                print(print_line, flush=True)

        self.update_decision_function(train_data)


class PU_CSI(nn.Module):
    def __init__(self, model=Net_CSI, simclr_dim=128, k_shift=4, lam=0.1, out_dim=128, flag=True):
        """
        :param flag: if true other positive examples considered in contrastive loss as negatives, if false only
        unlabeled data considered as negative
        """
        super().__init__()
        self.k_shift = k_shift
        self.sim_lambda = lam
        self.simclr_aug = get_simclr_augmentation((32, 32, 3)).to(device)
        self.hflip = HorizontalFlipLayer().to(device)
        self.flag = flag

        self.model = model(out_dim=out_dim, simclr_dim=simclr_dim, shift=k_shift)
        self.train_features = {}
        self.lambda_cls = {}
        self.lambda_con = {}

    def decision_function(self, test_data):
        y_preds = {'csi': np.array([]), 'diff': np.array([]), 'combo': np.array([])}
        y_true = np.array([])

        test_loader = torch.utils.data.DataLoader(test_data,
                                                  batch_size=64)

        for (images, labels, _) in test_loader:
            s_con_pos = 0
            s_con_unl = 0
            s_cls = 0

            images = images.to(device)
            images = self.hflip(images)

            images = torch.cat([torch.rot90(images, rot, (2, 3)) for rot in range(4)])

            out = self.model(images, penultimate=True, shift=True)

            chunks = out['penultimate'].chunk(4)
            chunks_s = out['shift'].chunk(4)

            for i in range(self.k_shift):
                s_con_pos += self.lambda_con['pos'][i] * (chunks[i] @ self.train_features['pos'][i].T).max(dim=1)[0]
                # s_con_unl += self.lambda_con['unl'][i] * (chunks[i] @ self.train_features['unl'][i].T).max(dim=1)[0]
                s_cls += self.lambda_cls['pos'][i] * (chunks_s[i][:, i])

            s_csi = s_con_pos + s_cls
            # s_diff = s_con_pos - s_con_unl
            # s_combo = s_con_pos + s_cls - s_con_unl

            y_preds['csi'] = np.hstack((y_preds['csi'], s_csi.cpu().detach().numpy()))
            # y_preds['diff'] = np.hstack((y_preds['diff'], s_diff.cpu().detach().numpy()))
            # y_preds['combo'] = np.hstack((y_preds['combo'], s_combo.cpu().detach().numpy()))

            y_true = np.hstack((y_true, labels.cpu().detach().numpy()))

        return y_preds['csi'], y_true

    def image_trans(self, images):
        images1, images2 = self.hflip(images.repeat(2, 1, 1, 1)).chunk(2)  # hflip
        images1 = torch.cat([torch.rot90(images1, rot, (2, 3)) for rot in range(4)])
        images2 = torch.cat([torch.rot90(images2, rot, (2, 3)) for rot in range(4)])
        images_pair = torch.cat([images1, images2], dim=0)  # 8B
        rot_labels = torch.cat([torch.ones(len(images)) * i for i in range(4)], dim=0)

        return images_pair, rot_labels.long().repeat(2)

    def fit(self,
            train_data,
            num_epochs=60,
            lr=1e-3,
            gamma=0.99,
            batch_size=64,
            batch_size_lab=64,
            batch_size_unl=16,
            temp=0.5,
            w=0.1,
            test_data=None,
            verbose=False):

        # self.debug = dict()

        self.model.to(device)

        # self.debug['shift_loss'] = []
        # self.debug['shift_labels'] = []

        criterion = nn.CrossEntropyLoss().to(device)
        optimizer = opt.Adam(self.model.parameters(), lr=lr)
        scheduler = opt.lr_scheduler.ExponentialLR(optimizer, gamma=gamma)

        # data_loader = torch.utils.data.DataLoader(train_data,
        #                                           batch_size=batch_size,
        #                                           shuffle=True)

        lab_loader = torch.utils.data.DataLoader(train_data.lab_data(lab=1),
                                                 batch_size=batch_size_lab,
                                                 shuffle=True)

        unl_loader = torch.utils.data.DataLoader(train_data.lab_data(lab=0),
                                                 batch_size=batch_size_unl,
                                                 shuffle=True)

        for epoch in range(num_epochs):
            running_loss = 0
            # for (images, _, labels) in data_loader:
            for batch in zip(lab_loader, unl_loader):

                images_p, images_u = batch[0][0], batch[1][0]
                labels_p, labels_u = batch[0][2], batch[1][2]

                # if len(images_p) != len(images_u):
                #     continue

                # images_p, images_u = images[labels == 1], images[labels == 0]
                # labels_p, labels_u = labels[labels == 1], labels[labels == 0]

                # if images_p.shape[0] == 0:
                #     continue

                images_p, rot_labels_p = self.image_trans(images_p)
                labels_p = labels_p.repeat(4)

                if images_u.shape[0] > 0:
                    images_u, rot_labels_u = self.image_trans(images_u)
                    labels_u = labels_u.repeat(4)

                if images_u.shape[0] > 0:
                    images = torch.cat((images_p, images_u))
                    labels = torch.cat((labels_p, labels_u))
                else:
                    images = images_p
                    labels = labels_p

                images = images.to(device)
                labels = labels.to(device)
                rot_labels = rot_labels_p.to(device)

                self.model.train()

                ### SimCLR loss ###
                batch_size = images.size(0)

                images_pair = self.simclr_aug(images)  # simclr augment
                outputs_aux = self.model(images_pair, simclr=True, penultimate=True, shift=True)

                simclr = normalize(outputs_aux['simclr'])  # normalize
                sim_matrix = get_similarity_matrix(simclr)
                loss_csi = Semisupervised_NT_xent(sim_matrix, len(labels_p), len(labels_u), temperature=temp,
                                                  flag=self.flag, w=w) * self.sim_lambda

                loss_shift = criterion(outputs_aux['shift'].float()[:len(labels_p) * 2], rot_labels)

                loss = loss_csi + loss_shift

                optimizer.zero_grad()
                loss.backward()
                optimizer.step()

                running_loss += loss.item()

            scheduler.step()
            if test_data is not None and verbose:
                self.model.eval()
                self.update_decision_function(train_data, 'pos')
                # self.update_decision_function(train_data, 'unl')

                y_pred, y_true = self.decision_function(test_data)

                auc_csi = metrics.roc_auc_score(y_true, y_pred)
                # auc_diff = metrics.roc_auc_score(y_true, y_preds['diff'])
                # auc_combo = metrics.roc_auc_score(y_true, y_preds['combo'])

                print_line = f"[{epoch:4}/{num_epochs:4}]: loss={running_loss: .4f}, auc_csi={auc_csi:.4f}"

                print(print_line)
        self.update_decision_function(train_data, 'pos')

    def update_decision_function(self, train_data, mode='pos'):
        self.eval()

        if mode == 'pos':
            label = 1
        elif mode == 'unl':
            label = 0
        else:
            raise ValueError(f'wrong mode {mode}')

        dataloader = torch.utils.data.DataLoader(train_data.lab_data(lab=label),
                                                 batch_size=512)

        self.train_features[mode] = [[] for k in range(self.k_shift)]
        self.lambda_cls[mode] = torch.zeros(self.k_shift, device=device)
        self.lambda_con[mode] = torch.zeros(self.k_shift, device=device)

        for (images, _, labels) in dataloader:
            batch_size = images.size(0)
            images = images.to(device)
            images = self.hflip(images)

            images = torch.cat([torch.rot90(images, rot, (2, 3)) for rot in range(4)])

            out = self.model(images, penultimate=True, shift=True)

            chunks = out['penultimate'].chunk(4)
            chunks_s = out['shift'].chunk(4)

            for i in range(4):
                self.train_features[mode][i].append(chunks[i].detach())
                self.lambda_con[mode][i] += torch.norm(chunks[i], p=2, dim=1).sum().detach()
                self.lambda_cls[mode][i] += chunks_s[i].sum(dim=0)[i].detach()

        self.lambda_cls[mode] = len(train_data) / self.lambda_cls[mode]
        self.lambda_con[mode] = len(train_data) / self.lambda_con[mode]
        self.train_features[mode] = [torch.cat(f) for f in self.train_features[mode]]

        for i in range(4):
            self.train_features[mode][i] = self.train_features[mode][i] / torch.norm(self.train_features[mode][i],
                                                                                     dim=1).view(-1, 1)

# code borrowed from https://github.com/alinlab/CSI # self.shift_trans = params['shift_trans'] # self.K_shift = params['K_shift'] # self.debug = dict() # self.debug['shift_loss'] = [] # self.debug['shift_labels'] = [] # linear_optimizer = opt.Adam(self.model.linear.parameters(), lr=lr) ### SimCLR loss ### # hflip # B -> 4B # 8B # print(images_pair.shape) # transform # print(images_pair.shape) # normalize # self.debug['shift_loss'].append(outputs_aux['shift'].float()) # self.debug['shift_labels'].append(shift_labels) ### total loss ### :param flag: if true other positive examples considered in contrastive loss as negatives, if false only         unlabeled data considered as negative # s_con_unl += self.lambda_con['unl'][i] * (chunks[i] @ self.train_features['unl'][i].T).max(dim=1)[0] # s_diff = s_con_pos - s_con_unl # s_combo = s_con_pos + s_cls - s_con_unl # y_preds['diff'] = np.hstack((y_preds['diff'], s_diff.cpu().detach().numpy())) # y_preds['combo'] = np.hstack((y_preds['combo'], s_combo.cpu().detach().numpy())) # hflip # 8B # self.debug = dict() # self.debug['shift_loss'] = [] # self.debug['shift_labels'] = [] # data_loader = torch.utils.data.DataLoader(train_data, #                                           batch_size=batch_size, #                                           shuffle=True) # for (images, _, labels) in data_loader: # if len(images_p) != len(images_u): #     continue # images_p, images_u = images[labels == 1], images[labels == 0] # labels_p, labels_u = labels[labels == 1], labels[labels == 0] # if images_p.shape[0] == 0: #     continue ### SimCLR loss ### # simclr augment # normalize # self.update_decision_function(train_data, 'unl') # auc_diff = metrics.roc_auc_score(y_true, y_preds['diff']) # auc_combo = metrics.roc_auc_score(y_true, y_preds['combo'])