from copy import deepcopy
from typing import Any, Sequence, Tuple, List, Callable, Dict
import torch
import argparse
import pdb
from cassle.utils.lars import LARSWrapper
from pl_bolts.optimizers.lr_scheduler import LinearWarmupCosineAnnealingLR
from cassle.utils.scheduler import ConstantWarmupCosineAnnealingLR
from torch.optim.lr_scheduler import CosineAnnealingLR, MultiStepLR
from functools import partial
from cassle.model.cbam_resnet import resnet18_cbam
from torch import nn
import torch.nn.functional as F
from cassle.losses.barlow import barlow_loss_func
import numpy as np
from tqdm import *

def static_lr(
    get_lr: Callable, param_group_indexes: Sequence[int], lrs_to_replace: Sequence[float]
):
    lrs = get_lr()
    for idx, lr in zip(param_group_indexes, lrs_to_replace):
        lrs[idx] = lr
    return lrs

def base_cbam_wrapper4(Method=object):
    class BaseCbamWrapper(Method):
        def __init__(self, **kwargs) -> None:
            super().__init__(**kwargs)
            self.output_dim = kwargs["output_dim"]
            self.checkpoint_dir = kwargs["checkpoint_dir"]
            self.automatic_optimization = False
            self.cifar = kwargs["cifar"]
            self.pretrained_model = kwargs["pretrained_model"]
            self.encoder_name = kwargs["encoder"]
            self.current_task_idx = kwargs["task_idx"]
            self.proj_hidden_dim = kwargs["proj_hidden_dim"]
            self.output_dim = kwargs["output_dim"]
            self.ea = kwargs["ea"]
            self.num_tasks = kwargs["num_tasks"]
            self.addition_images = kwargs["addition_images"]
            self.method = kwargs["sample_method"]
            self.cur_epoch = 0
            self.indices = []
            
            self.base_model = {"resnet18": resnet18_cbam}[self.encoder_name]
            self.encoder = self.base_model()
            self.encoder.linear = nn.Identity()
            if self.cifar:
                self.encoder.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=2, bias=False)
                self.encoder.maxpool = nn.Identity()

            if self.current_task_idx > 0:
                #distill projector
                self.distill_proj_hidden_dim = 2048
                self.distill_projector = nn.Sequential(
                    nn.Linear(self.features_dim, self.distill_proj_hidden_dim),
                    nn.BatchNorm1d(self.distill_proj_hidden_dim),
                    nn.ReLU(),
                    nn.Linear(self.distill_proj_hidden_dim, self.features_dim),
                    nn.ReLU(),)
                # for m in self.distill_projector.modules():
                #     if isinstance(m, (nn.Linear)):
                #         nn.init.eye_(m.weight)
                        
                self.relation_distill_projector = nn.Sequential(
                    nn.Linear(self.features_dim, self.distill_proj_hidden_dim),
                    nn.BatchNorm1d(self.distill_proj_hidden_dim),
                    nn.ReLU(),
                    nn.Linear(self.distill_proj_hidden_dim, self.features_dim),)
                for m in self.relation_distill_projector.modules():
                    if isinstance(m, (nn.Linear)):
                        nn.init.eye_(m.weight)

                self.attention_distill_projector = nn.Sequential(
                    nn.Linear(25, 25),
                    nn.BatchNorm1d(25),
                    nn.ReLU(),
                    nn.Linear(25, 25),)
                for m in self.attention_distill_projector.modules():
                    if isinstance(m, (nn.Linear)):
                        nn.init.eye_(m.weight)

                self.another_projector = nn.Sequential(
                    nn.Linear(self.features_dim, self.proj_hidden_dim),
                    nn.BatchNorm1d(self.proj_hidden_dim),
                    nn.ReLU(),
                    nn.Linear(self.proj_hidden_dim, self.proj_hidden_dim),
                    nn.BatchNorm1d(self.proj_hidden_dim),
                    nn.ReLU(),
                    nn.Linear(self.proj_hidden_dim, self.output_dim),
                )

                # load cbam_resnet for task 1~t
                print(f"Loading previous task checkpoint {self.pretrained_model}...")
                state_dict = torch.load(self.pretrained_model, map_location="cpu")["state_dict"]
                self.load_state_dict(state_dict, strict=False)

                self.avg_pool = nn.AvgPool2d(kernel_size=3, stride=1, padding=1)
                
        def get_interests(self, encoder, dataloader):
            print("Getting interest maps...")
            interest_mapss = np.zeros((len(dataloader.dataset)*self.num_tasks, 64)).astype('float16')
            for batch in tqdm(dataloader):
                with torch.no_grad():
                    index, X, _ = batch
                    X = X.cuda()
                    interest_maps = encoder(X)["att_map"]
                    interest_maps = torch.mean(torch.abs(interest_maps), dim=1)
                    interest_maps = F.interpolate(interest_maps.unsqueeze(1), size=(32, 32), mode='bicubic')
                    interest_maps = F.avg_pool2d(interest_maps, kernel_size=4).reshape((interest_maps.size(0), -1))
                    interest_maps = (torch.max(interest_maps, dim=1)[0]-interest_maps.T).T
                    # bar = torch.kthvalue(interest_maps, 16, dim=1)[0]
                    # interest_maps_np = interest_maps.cpu().numpy().astype('float16')
                    # # pdb.set_trace()
                    # for i in range(interest_maps.size(0)):
                    #     temp=np.ones_like(interest_maps[i])
                    #     temp.fill(1/48)
                    #     temp = temp.astype('float16')
                    #     interest_maps_np[i] = np.where(interest_maps[i]<=bar[i], np.zeros_like(interest_maps[i]), temp)
                    interest_maps = F.softmax(interest_maps/0.4, dim=1)
                    interest_maps = torch.where(interest_maps >= 0.007, interest_maps, torch.zeros_like(interest_maps).fill_(float("-inf")))
                    interest_maps = F.softmax(interest_maps, dim=1)
                    interest_maps_np = interest_maps.cpu().numpy().astype('float16')

                    interest_mapss[index.squeeze(1)] = interest_maps_np
            print("Done!")
            return interest_mapss

        def load_data(self):
            value_dict = torch.load(f"{self.checkpoint_dir}/value{self.current_task_idx-1}.pth")
            self.indices = value_dict["indices"]
            self.indices_distances = value_dict["indices_distances"]

        def on_train_start(self):
            super().on_train_start()
            # dataloader = self.trainer._data_connector._train_dataloader_source.instance[f'task{self.current_task_idx}']
            # dataloader.dataset.dataset.set_train(True)
            if self.current_task_idx > 0:
                self.load_data()
                # save frozen encoder for distillation
                self.frozen_encoder = deepcopy(self.encoder)
                self.frozen_projector = deepcopy(self.projector)
                for pg in self.frozen_encoder.parameters():
                    pg.requires_grad = False
                for pg in self.frozen_projector.parameters():
                    pg.requires_grad = False

        def on_train_epoch_start(self):
            super().on_train_start()
            # if self.cur_epoch % 10 == 0 and self.current_task_idx > 0:
            #     dataloader = self.trainer._data_connector._train_dataloader_source.instance[f'task{self.current_task_idx}']
            #     dataloader.dataset.dataset.set_pre(True)
            #     temp_encoder = deepcopy(self.encoder)
            #     for pg in temp_encoder.parameters():
            #         pg.requires_grad = False
            #     interest_mapss = self.get_interests(temp_encoder, dataloader)
            #     dataloader.dataset.dataset.set_interests(interest_mapss)
            #     dataloader.dataset.dataset.set_pre(False)
                

        def base_forward(self, X: torch.Tensor) -> torch.Tensor:
            """Basic forward that allows children classes to override forward().

            Args:
                X (torch.Tensor): batch of images in tensor format.

            Returns:
                torch.Tensor: features extracted by the encoder.
            """

            return self.encoder(X)
    
        def relation_distill_forward(self, x1, x2):
            x1 = self.relation_distill_projector(x1)
            x2 = self.relation_distill_projector(x2)
            return x1, x2
        
        def distill_forward(self, x1, x2):
            x1 = self.distill_projector(x1)
            x2 = self.distill_projector(x2)
            return x1, x2

        def get_fft_feature(self, x1, x2):
            h, w = x1.shape[-2:]
            lpf = torch.zeros((h,w))
            R = (h+w)//8  #或其他
            for x in range(w):
                for y in range(h):
                    if ((x-(w-1)/2)**2 + (y-(h-1)/2)**2) < (R**2):
                        lpf[y,x] = 1
            hpf = 1-lpf
            hpf, lpf = hpf.cuda(), lpf.cuda()
            f_x1 = torch.fft.fftn(x1, dim=(-2, -1))
            f_x2 = torch.fft.fftn(x2, dim=(-2, -1))
            f_x1 = torch.fft.fftshift(f_x1, dim=(-2, -1))
            f_x2 = torch.fft.fftshift(f_x2, dim=(-2, -1))
            f_x1 = f_x1*hpf
            f_x2 = f_x2*hpf
            h_x1 = torch.fft.ifftn(f_x1, dim=(-2, -1))
            h_x2 = torch.fft.ifftn(f_x2, dim=(-2, -1))
            return h_x1, h_x2
        
        def background_mix(self, x1, x2):
            bs, c, h, w = x1.shape
            grids =torch.randint(0, 8, (bs, 4, 2))
            final_x1 = torch.zeros_like(x1)
            final_x2 = torch.zeros_like(x2)
            mask = torch.zeros((bs, h, w)).cuda()
            for i in range(bs):
                grid = grids[i]
                for j in range(4):
                    if grid[j][0] >=2 and grid[j][0]<=5 and grid[j][1] >=2 and grid[j][1]<=5:
                        grid[j][0] = -1
                        grid[j][1] = -1
                for j in range(4):
                    if grid[j][0] != -1:
                        mask[i][grid[j][0]*h//8:(grid[j][0]+1)*h//8, grid[j][1]*w//8:(grid[j][1]+1)*w//8] = 1
            mask = torch.repeat_interleave(mask.unsqueeze(dim=1), repeats=c, dim=1)
            final_x1 = x1*(1-mask)+x1[torch.randperm(bs)]*mask
            final_x2 = x2*(1-mask)+x2[torch.randperm(bs)]*mask
            return final_x1, final_x2
                
        @torch.no_grad()
        def frozen_forward(self, X):
            feats = self.frozen_encoder(X)
            return feats
        
        def calculate_distill_loss(self, old_fea, fea):
            return torch.mean(torch.norm(old_fea - fea, p=2, dim=1))
        
        def pdist(self, e, squared=False, eps=1e-12):
            e_square = e.pow(2).sum(dim=1)
            prod = e @ e.t()
            res = (e_square.unsqueeze(1) + e_square.unsqueeze(0) - 2 * prod).clamp(min=eps)
            if not squared:
                res = res.sqrt()
            res = res.clone()
            res[range(len(e)), range(len(e))] = 0

            return res
        
        def calculate_RKD_distance_loss(self, old_fea, fea):
            with torch.no_grad():
                t_d = self.pdist(old_fea, squared=False)
                mean_td = t_d[t_d > 0].mean()
                t_d = t_d / mean_td
            # pdb.set_trace()
            d = self.pdist(fea, squared=False)
            mean_d = d[d > 0].mean()
            d = d / mean_d
            # print("d:{},t_d:{}".format(d.size(),t_d.size()))
            loss_d = torch.nn.functional.smooth_l1_loss(d, t_d)
            return loss_d
        
        def calculate_RKD_angle_loss(self, old_fea, fea):
            with torch.no_grad():
                td = (old_fea.unsqueeze(0) - old_fea.unsqueeze(1))
                norm_td = torch.nn.functional.normalize(td, p=2, dim=2)
                t_angle = torch.bmm(norm_td, norm_td.transpose(1, 2)).view(-1)

            sd = (fea.unsqueeze(0) - fea.unsqueeze(1))
            norm_sd = torch.nn.functional.normalize(sd, p=2, dim=2)
            s_angle = torch.bmm(norm_sd, norm_sd.transpose(1, 2)).view(-1)
            loss_a = torch.nn.functional.smooth_l1_loss(s_angle, t_angle)
            return loss_a


        def calculate_RKD_loss(self, old_fea, fea):
            loss_d = self.calculate_RKD_distance_loss(old_fea, fea)
            loss_a = self.calculate_RKD_angle_loss(old_fea, fea)
            return loss_a + loss_d

        def calculate_attention_map_loss(self, old_attention, attention):
            temperature=0.5
            # pdb.set_trace()
            bs, c, h, w = old_attention.shape
            old_attention = torch.mean(torch.abs(old_attention), dim=1)#[bs, 512, 5, 5] -> [bs, 5, 5]
            attention = torch.mean(torch.abs(attention), dim=1)#[bs, 512, 5, 5] -> [bs, 5, 5]
            old_attention = h*w*(torch.nn.functional.softmax(old_attention.reshape((bs, -1))/temperature)) # [bs, 25]
            attention = h*w*(torch.nn.functional.softmax(attention.reshape((bs, -1))/temperature)) # [bs, 25]
            attention = self.attention_distill_projector(attention)
            # reavg=torch.tensor([[9/4, 9/6, 9/6, 9/6, 9/4],[9/6, 1, 1, 1, 9/6], [9/6, 1, 1, 1, 9/6], [9/6, 1, 1, 1, 9/6], [9/4, 9/6, 9/6, 9/6, 9/4]]).to(old_attention.device)
            # weight = ((self.avg_pool(old_attention.reshape((bs, -1, h, w))))*reavg).reshape((bs, -1)) # [bs, 25]
            # weight = h*w*(torch.nn.functional.softmax(weight)) # [bs, 25]
            # weight = torch.abs(weight - 1)
            # weight = (weight.T+(1-torch.mean(weight, dim=1))).T
            loss = torch.mean(torch.norm((old_attention - attention), p=2, dim=1))
            return loss

        def mixup_data(self, X1, X2, alpha=1.0):
            lam = torch.tensor(np.random.beta(alpha, alpha)).cuda()
            lam = lam/2+0.5
            bs = X1.shape[0]
            index = torch.randperm(bs)
            mixed_x1 = lam * X1 + (1 - lam) * X1[index, :]
            mixed_x2 = lam * X2 + (1 - lam) * X2[index, :]
            return [torch.cat((X1[:bs//2], mixed_x1[bs//2:])), torch.cat((X1[bs//2:], mixed_x1[:bs//2]))], [torch.cat((X2[:bs//2], mixed_x2[bs//2:])), torch.cat((X2[bs//2:], mixed_x2[:bs//2]))]
        
        def rotate_data(self, X1, X2):
            X1 = [torch.rot90(X1, k, (2, 3)) for k in range(4)]
            X2 = [torch.rot90(X2, k, (2, 3)) for k in range(4)]
            return X1, X2
        
        def training_step(self, batch: Sequence[Any], batch_idx: int) -> torch.Tensor:
            opt = self.optimizers(use_pl_optimizer=True)
            _, (X1, X2), _ = batch[f"task{self.current_task_idx}"]
            loss = 0
            barlow_loss = 0
            if self.ea and self.current_task_idx > 0:
                mixX1, mixX2 = self.mixup_data(X1, X2)
                for i in range(2):
                    batch[f"task{self.current_task_idx}"][1] = [mixX1[i], mixX2[i]]
                    out = super().training_step(batch, batch_idx)
                    barlow_loss += out["loss"]
                X1, X2 = self.rotate_data(X1, X2)
                X1, X2 = torch.cat(X1), torch.cat(X2)
                index = torch.randperm(1024)
                X1, X2 = X1[index], X2[index]
                for i in range(4):
                    batch[f"task{self.current_task_idx}"][1] = [X1[i*256:(i+1)*256], X2[i*256:(i+1)*256]]
                    out = super().training_step(batch, batch_idx)
                    barlow_loss += out["loss"]
                    if i==0:
                        f1, f2 = out["feats"]
                    else:
                        f1, f2 = torch.cat((f1, out["feats"][0])), torch.cat((f2, out["feats"][1]))
            else:
                out = super().training_step(batch, batch_idx)
                barlow_loss += out["loss"]
            loss += barlow_loss
            if self.current_task_idx > 0:
                distill_loss = 0
                if not self.ea:
                    f1, f2 = out["feats"]
                # r_f1, r_f2 = self.relation_distill_forward(f1, f2)
                old_out1, old_out2 = self.frozen_forward(X1), self.frozen_forward(X2)
                frozen_feats1 = old_out1["feats"]
                frozen_feats2 = old_out2["feats"]
                # relation_distill_loss
                # relation_distill_loss = 500*(self.calculate_RKD_loss(frozen_feats1, r_f1) + self.calculate_RKD_loss(frozen_feats2, r_f2))
                # distill_loss += relation_distill_loss

                # feature_distill_loss PFR
                f_f1, f_f2 = self.distill_forward(f1, f2)
                # bn = torch.nn.BatchNorm1d(512, affine=False).cuda()
                # f_f1, f_f2, frozen_feats1, frozen_feats2 = bn(f_f1), bn(f_f2), bn(frozen_feats1), bn(frozen_feats2)
                feature_distill_loss = 0.5*(self.calculate_distill_loss(frozen_feats1, f_f1) + self.calculate_distill_loss(frozen_feats2, f_f2))
                distill_loss += feature_distill_loss

                # attention_map_distill_loss
                # att1, att2 = out["att_map"]
                # old_att1, old_att2 = old_out1["att_map"], old_out2["att_map"]
                # distill_loss += 10*(self.calculate_attention_map_loss(old_att1, att1) + self.calculate_attention_map_loss(old_att2, att2))

                loss += distill_loss

                if self.ea:
                    # mixz1, mixz2 = self.another_projector(self.base_forward(torch.cat((mixX1[0][128:], mixX1[1][128:])))["feats"]), self.another_projector(self.base_forward(torch.cat((mixX2[0][128:], mixX2[1][128:])))["feats"])
                    # temp_loss = barlow_loss_func(mixz1, mixz2, lamb=self.lamb, scale_loss=self.scale_loss)
                    # loss += 0.5*temp_loss

                    # temp1, temp2 = F.normalize(temp1, p=2, dim=1).detach(), F.normalize(temp2, p=2, dim=1).detach()
                    # z1, z2 = out["z"]
                    # z1, z2 = F.normalize(z1, p=2, dim=1), F.normalize(z2, p=2, dim=1)
                    # zhengjiaoloss = torch.mean(torch.diag(torch.matmul(temp1,z1.T)) + torch.diag(torch.matmul(temp2,z2.T)))
                    # loss += 0.5*zhengjiaoloss
                    # pdb.set_trace()
                    # z1, z2 = out["z"]
                    # z1, z2 = F.normalize(z1, p=2, dim=1), F.normalize(z2, p=2, dim=1)
                    # cos_distance = torch.matmul(z1, z2.T)
                    # divergence_loss = torch.mean(cos_distance.flatten()[:-1].view(z1.shape[0]-1,z1.shape[0]+1)[:,1:].flatten())
                    # loss += 100*divergence_loss
                    pass

                self.log_dict({"barlow_loss": barlow_loss, "distill_loss": distill_loss}, on_epoch=True, sync_dist=True, prog_bar=True)

                out.update({"loss": loss})
            else:
                self.log_dict({"barlow_loss": barlow_loss}, on_epoch=True, sync_dist=True, prog_bar=True)
            opt.zero_grad()
            self.manual_backward(loss)
            opt.step()
            # self.ema_update()
            return out

        def on_train_epoch_end(self):
            super().on_train_epoch_end()
            sch = self.lr_schedulers()
            sch.step()
            self.cur_epoch += 1
            # sch_dist.step()

        def on_train_end(self):
            budget=self.addition_images
            if budget == 0:
                value_dict={}
                value_dict["indices"] = self.indices
                value_dict["indices_distances"] = self.indices_distances
                torch.save(value_dict, f"{self.checkpoint_dir}/value{self.current_task_idx}.pth")
                return
            
            if self.method == 'random':
                budget_this_task = budget // (self.current_task_idx+1)
                # select external data
                dataloader = self.trainer._data_connector._train_dataloader_source.instance[f'task{self.current_task_idx}']
                dataset = dataloader.dataset.dataset.datasets[1]
                
                if self.current_task_idx > 0:
                    delete_num = budget_this_task
                    delete_num_per_task = delete_num // self.current_task_idx
                    temp = budget // self.current_task_idx
                    new_indices = []
                    for t in range(self.current_task_idx):
                        new_indices.append(self.indices[t*temp+delete_num_per_task:(t+1)*temp])

                    self.indices = torch.cat(new_indices)

                    indices = torch.tensor([i not in self.indices for i in range(len(dataset))]).nonzero().view(-1)
                else:
                    indices = torch.tensor(range(len(dataset)))

                indices = indices[torch.randperm(len(indices))]
                indices = indices[-budget_this_task:]

                if self.current_task_idx > 0:
                    self.indices = torch.cat((self.indices, indices))
                else:
                    self.indices = indices
                
                value_dict={}
                value_dict["indices"] = self.indices
                value_dict["indices_distances"] = self.indices_distances
                torch.save(value_dict, f"{self.checkpoint_dir}/value{self.current_task_idx}.pth")

            elif self.method == '1p1':
                budget_this_task = budget // (self.current_task_idx+1)
                # select external data
                dataloader = self.trainer._data_connector._train_dataloader_source.instance[f'task{self.current_task_idx}']
                dataloader.dataset.dataset.datasets[0].set_pre(True)
                for i in range(len(dataloader.dataset.dataset.datasets[1].datasets)):
                    dataloader.dataset.dataset.datasets[1].datasets[i].set_pre(True)

                all_in_features = []

                for batch in tqdm(dataloader):
                    with torch.no_grad():
                        _, (X), Y = batch
                        X = X.cuda()
                        in_fea = self.encoder(X)["feats"][Y<100]
                        all_in_features.append(in_fea)
                all_in_features = torch.cat(all_in_features)
                all_in_features = F.normalize(all_in_features, dim=1)
                
                dataset = dataloader.dataset.dataset.datasets[1]
                external_dataloader = torch.utils.data.DataLoader(dataset,batch_size=256,shuffle=False,num_workers=4,pin_memory=True)
                all_out_features = []
                for input, _ in tqdm(external_dataloader):
                    with torch.no_grad():
                        input = input.cuda()
                        out_fea = self.encoder(input)["feats"]
                        all_out_features.append(out_fea)
                all_out_features = torch.cat(all_out_features)
                all_out_features = F.normalize(all_out_features, dim=1)
                
                if self.current_task_idx > 0:
                    delete_num = budget_this_task
                    delete_num_per_task = delete_num // self.current_task_idx
                    temp = budget // self.current_task_idx
                    new_indices = []
                    new_distances = []
                    for t in range(self.current_task_idx):
                        new_indices.append(self.indices[t*temp+delete_num_per_task:(t+1)*temp])
                        new_distances.append(self.indices_distances[t*temp+delete_num_per_task:(t+1)*temp])

                    self.indices = torch.cat(new_indices)
                    self.indices_distances = torch.cat(new_distances)
                min_distances=[]
                for p in range(all_out_features.shape[0]//100000+1):
                    part_out_features = all_out_features[p*100000:(p+1)*100000]
                    part_distances = torch.matmul(all_in_features, part_out_features.T)
                    part_min_distances = torch.max(part_distances, dim=0).values
                    min_distances.append(part_min_distances)
                
                min_distances = torch.cat(min_distances)
                
                already_indices = [i in self.indices for i in range(len(all_out_features))]
                temp=[]
                while len(temp) < budget_this_task:
                    print(len(temp))
                    local_temp=[]
                    for i in tqdm(range(all_in_features.shape[0])):
                        temp_distance = torch.matmul(all_in_features[i], all_out_features.T)
                        temp_distance[already_indices] = -10000
                        indice = int(torch.max(temp_distance, 0).indices)
                        already_indices[indice] = True
                        local_temp.append((min_distances[indice], indice))
                    local_temp.sort(key=lambda x: x[0])
                    if len(temp) + len(local_temp) < budget_this_task:
                        temp.extend(local_temp)
                    else:
                        temp.extend(local_temp[-(budget_this_task-len(temp)):])

                temp.sort(key=lambda x: x[0])
                indices = torch.tensor([pair[1] for pair in temp])
                distances = torch.tensor([pair[0] for pair in temp])
                indices = indices[-budget_this_task:]
                distances = distances[-budget_this_task:]
                
                if self.current_task_idx > 0:
                    self.indices = torch.cat((self.indices, indices))
                    self.indices_distances = torch.cat((self.indices_distances, distances))
                else:
                    self.indices = indices
                    self.indices_distances = distances

                # from sklearn.manifold import TSNE
                # import matplotlib.pyplot as plt
                # index = torch.randperm(len(all_in_features))[:500]
                # tsne_data = torch.cat((all_in_features[index], all_out_features[indices[torch.randperm(len(indices))[:500]]]))
                # tsne = TSNE(n_components=2)
                # X_tsne = tsne.fit_transform(tsne_data.cpu())
                # plt.scatter(X_tsne[:500, 0], X_tsne[:500, 1], c='r')
                # plt.scatter(X_tsne[500:, 0], X_tsne[500:, 1], c='g')
                # plt.savefig("tsne/1propose1.jpg")
                # pdb.set_trace()

                value_dict={}
                value_dict["indices"] = self.indices
                value_dict["indices_distances"] = self.indices_distances
                torch.save(value_dict, f"{self.checkpoint_dir}/value{self.current_task_idx}.pth")
            else:
                return NotImplementedError

        # def on_train_end(self):
        #     super().on_train_end()
        #     k=1
        #     budget=self.addition_images
        #     budget_this_task = budget // (self.current_task_idx+1)
        #     # select external data
        #     dataloader = self.trainer._data_connector._train_dataloader_source.instance[f'task{self.current_task_idx}']
        #     dataloader.dataset.dataset.datasets[0].set_pre(True)
        #     dataloader.dataset.dataset.datasets[1].datasets[0].set_pre(True)
        #     dataloader.dataset.dataset.datasets[1].datasets[1].set_pre(True)
        #     all_in_features = []
        #     for batch in tqdm(dataloader):
        #         with torch.no_grad():
        #             _, (X), Y = batch
        #             X = X.cuda()
        #             in_fea = self.encoder(X)["feats"][Y<100]
        #             all_in_features.append(in_fea)
        #     all_in_features = torch.cat(all_in_features)
        #     from sklearn.cluster import KMeans
        #     kmeans = KMeans(n_clusters=k, random_state=0, n_init="auto").fit(all_in_features.cpu())
        #     all_in_features_indices_cluster=[]
        #     for i in range(k):
        #         all_in_features_indices_cluster.append(torch.tensor(kmeans.labels_==i).nonzero().view(-1))
        #     dataset = dataloader.dataset.dataset.datasets[1]
        #     external_dataloader = torch.utils.data.DataLoader(dataset,batch_size=256,shuffle=False,num_workers=4,pin_memory=True)
        #     all_min_distance = [[] for _ in range(k)]
        #     all_out_features = []
        #     for input, _ in tqdm(external_dataloader):
        #         with torch.no_grad():
        #             input = input.cuda()
        #             out_fea = self.encoder(input)["feats"]
        #             all_out_features.append(out_fea)
        #             for i in range(k):
        #                 l2_distance = torch.cdist(out_fea, all_in_features[all_in_features_indices_cluster[i]])
        #                 all_min_distance[i].append(-torch.mean(l2_distance, dim=1))
        #     all_out_features = torch.cat(all_out_features)
            
        #     if self.current_task_idx > 0:
        #         delete_num = budget_this_task
        #         delete_num_per_task = delete_num // self.current_task_idx
        #         temp = budget // self.current_task_idx
        #         new_indices = []
        #         new_distances = []
        #         for t in range(self.current_task_idx):
        #             new_indices.append(self.indices[t*temp+delete_num_per_task:(t+1)*temp])
        #             new_distances.append(self.indices_distances[t*temp+delete_num_per_task:(t+1)*temp])

        #         self.indices = torch.cat(new_indices)
        #         self.indices_distances = torch.cat(new_distances)
            
        #     for i in tqdm(range(k)):
        #         all_min_distance[i] = torch.cat(all_min_distance[i]).tolist()
        #         temp = all_min_distance[i]
        #         temp=[(dis, index) for index, dis in enumerate(temp)]
        #         temp.sort(key=lambda x: x[0])
        #         cur_indices = torch.tensor([pair[1] for pair in temp])
        #         cur_min_distance = torch.tensor([pair[0] for pair in temp])
        #         if self.current_task_idx > 0:
        #             ind = [i not in self.indices for i in cur_indices]
        #             cur_min_distance = cur_min_distance[ind]
        #             cur_indices = cur_indices[ind]
        #         if i == 0:
        #             indices = cur_indices[-budget_this_task//k:]
        #             distances = cur_min_distance[-budget_this_task//k:]
        #         else:
        #             ind = [i not in indices for i in cur_indices]
        #             cur_min_distance = cur_min_distance[ind]
        #             cur_indices = cur_indices[ind]
        #             indices = torch.cat((indices, cur_indices[-budget_this_task//k:]))
        #             distances = torch.cat((distances, cur_min_distance[-budget_this_task//k:]))
            
        #     temp=[(dis, index) for index, dis in zip(indices, distances)]
        #     temp.sort(key=lambda x: x[0])
        #     indices = torch.tensor([pair[1] for pair in temp])
        #     distances = torch.tensor([pair[0] for pair in temp])
        #     pdb.set_trace()
        #     if self.current_task_idx > 0:
        #         self.indices = torch.cat((self.indices, indices))
        #         self.indices_distances = torch.cat((self.indices_distances, distances))
        #     else:
        #         self.indices = indices
        #         self.indices_distances = distances
                
        #     # from sklearn.manifold import TSNE
        #     # import matplotlib.pyplot as plt
        #     # tsne_data = torch.cat((all_in_features[torch.randperm(len(all_in_features))[:500]], all_out_features[indices[torch.randperm(10000)[:500]]].cuda()))
        #     # tsne = TSNE(n_components=2)
        #     # X_tsne = tsne.fit_transform(tsne_data.cpu())
        #     # plt.scatter(X_tsne[:500, 0], X_tsne[:500, 1], c='r')
        #     # plt.scatter(X_tsne[500:, 0], X_tsne[500:, 1], c='g')
        #     # plt.savefig("tsne/k10.jpg")
        #     # pdb.set_trace()

        #     value_dict={}
        #     value_dict["indices"] = self.indices
        #     value_dict["indices_distances"] = self.indices_distances
        #     torch.save(value_dict, f"{self.checkpoint_dir}/value{self.current_task_idx}.pth")
            

        @property
        def learnable_params(self) -> List[dict]:
            """Adds projector parameters to parent's learnable parameters.

            Returns:
                List[dict]: list of learnable parameters.
            """
            encoder_param = []
            att_param = []
            for n, p in self.encoder.named_parameters():
                if "ca" in n or "sa" in n:
                    att_param.append(p)
                else:
                    encoder_param.append(p)
            params = [
            {"name": "encoder", "params": encoder_param},
            {
                "name": "classifier",
                "params": self.classifier.parameters(),
                "lr": self.classifier_lr,
                "weight_decay": 0,
            },
            {"params": self.projector.parameters()}
        ]
            params.append({"name": "att", "params": att_param, "lr": self.lr})
            if self.current_task_idx > 0:
                params.append({"name": "distill_projector", "params": self.distill_projector.parameters()})
                params.append({"name": "relation_distill_projector", "params": self.relation_distill_projector.parameters()})
                params.append({"name": "attention_distill_projector", "params": self.attention_distill_projector.parameters()})
                params.append({"name": "another_projector", "params": self.another_projector.parameters()})
            return params
    
        def configure_optimizers(self) -> Tuple[List, List]:
            """Collects learnable parameters and configures the optimizer and learning rate scheduler.

            Returns:
                Tuple[List, List]: two lists containing the optimizer and the scheduler.
            """
            # collect learnable parameters
            idxs_no_scheduler = [
                i for i, m in enumerate(self.learnable_params) if m.pop("static_lr", False)
            ]

            # select optimizer
            if self.optimizer == "sgd":
                optimizerM = torch.optim.SGD
            elif self.optimizer == "adam":
                optimizer = torch.optim.Adam
            else:
                raise ValueError(f"{self.optimizer} not in (sgd, adam)")

            # create optimizer
            optimizer = optimizerM(
                self.learnable_params,
                lr=self.lr,
                weight_decay=self.weight_decay,
                **self.extra_optimizer_args,
            )
            # optimizer_distill = optimizerM(
            #     self.distill_params,
            #     lr=self.lr,
            #     weight_decay=self.weight_decay,
            #     **self.extra_optimizer_args,
            # )
            # optionally wrap with lars
            if self.lars:
                optimizer = LARSWrapper(
                    optimizer,
                    eta=self.eta_lars,
                    clip=self.grad_clip_lars,
                    exclude_bias_n_norm=self.exclude_bias_n_norm,
                )
                # optimizer_distill = LARSWrapper(
                #     optimizer_distill,
                #     eta=self.eta_lars,
                #     clip=self.grad_clip_lars,
                #     exclude_bias_n_norm=self.exclude_bias_n_norm,
                # )

            if self.scheduler == "none":
                return optimizer
            else:
                if self.scheduler == "warmup_cosine":
                    
                    scheduler = LinearWarmupCosineAnnealingLR(
                        optimizer,
                        warmup_epochs=self.warmup_epochs,
                        max_epochs=self.max_epochs,
                        warmup_start_lr=self.warmup_start_lr,
                        eta_min=self.min_lr,
                        # constant = self.warmup_start_lr
                    )
                    # scheduler_distill = LinearWarmupCosineAnnealingLR(
                    #     optimizer_distill,
                    #     warmup_epochs=self.warmup_epochs,
                    #     max_epochs=self.max_epochs,
                    #     warmup_start_lr=self.warmup_start_lr,
                    #     eta_min=self.min_lr,
                    #     # constant=self.warmup_start_lr
                    # )
                elif self.scheduler == "cosine":
                    scheduler = CosineAnnealingLR(optimizer, self.max_epochs, eta_min=self.min_lr)
                elif self.scheduler == "step":
                    scheduler = MultiStepLR(optimizer, self.lr_decay_steps)
                else:
                    raise ValueError(f"{self.scheduler} not in (warmup_cosine, cosine, step)")

                if idxs_no_scheduler:
                    partial_fn = partial(
                        static_lr,
                        get_lr=scheduler.get_lr,
                        param_group_indexes=idxs_no_scheduler,
                        lrs_to_replace=[self.lr] * len(idxs_no_scheduler),
                    )
                    scheduler.get_lr = partial_fn

                return [optimizer, ], [scheduler, ]
            
        # @property
        # def distill_params(self) -> List[Dict[str, Any]]:
        #     """Defines learnable parameters for the base class.

        #     Returns:
        #         List[Dict[str, Any]]:
        #             list of dicts containing learnable parameters and possible settings.
        #     """
        #     att_params=[]
        #     for n, p in self.encoder.named_parameters():
        #         if "ca" in n or "sa" in n:
        #             att_params.append(p)
            
        #     return [{"name": "att_params", "params": att_params}, 
        #             {"name": "distill_projector", "params": self.distill_projector.parameters()}]
            
        # def ema_update(self):
        #     beta=0.999
        #     for p_old, p in zip(self.frozen_encoder.parameters(), self.encoder.parameters()):
        #         p_old.data = beta*p_old.data + (1-beta)*p.data
        #     for p_old, p in zip(self.frozen_projector.parameters(), self.projector.parameters()):
        #         p_old.data = beta*p_old.data + (1-beta)*p.data
            
        # def grad_attention_distillation(self, batch, batch_idx):
        #     _, (X1, X2), _ = batch[f"task{self.current_task_idx}"]
        #     X1.requires_grad = True
        #     X2.requires_grad = True
        #     bs, c, h, w = X1.shape
        #     # pdb.set_trace()
        #     z1 = self.frozen_projector(self.frozen_encoder(X1)["feats"])
        #     z2 = self.frozen_projector(self.frozen_encoder(X2)["feats"])
        #     barlow_loss = barlow_loss_func(z1, z2, lamb=self.lamb, scale_loss=self.scale_loss)
        #     # self.manual_backward(barlow_loss)
        #     # old_grad_attention_X1, old_grad_attention_X2 = torch.abs(deepcopy(X1.grad.reshape((bs, -1)))), torch.abs(deepcopy(X2.grad.reshape((bs, -1))))
        #     grad = torch.autograd.grad(barlow_loss, X1)
        #     old_grad_attention_X1 = torch.abs(deepcopy(grad[0].reshape((bs, -1))))
        #     # grad = torch.autograd.grad(barlow_loss, X2)
        #     # old_grad_attention_X2 = torch.abs(deepcopy(grad[0].reshape((bs, -1))))
            
        #     out = super().training_step(batch, batch_idx)
        #     loss = out["loss"]
        #     grad = torch.autograd.grad(loss, X1, create_graph=True)
        #     grad_attention_X1=torch.abs(grad[0].reshape((bs, -1)))
        #     # grad = torch.autograd.grad(loss, X2, create_graph=True)
        #     # grad_attention_X2=torch.abs(grad[0].reshape((bs, -1)))
        #     # torch.autograd.backward(loss, create_graph=True)
        #     # grad_attention_X1, grad_attention_X2 = torch.abs(X1.grad.reshape((bs, -1))), torch.abs(X2.grad.reshape((bs, -1)))
        #     # grad_attention_X1, grad_attention_X2 = torch.where(torch.isinf(grad_attention_X1), old_grad_attention_X1, grad_attention_X1), torch.where(torch.isinf(grad_attention_X2), old_grad_attention_X2, grad_attention_X2)

        #     # bn = torch.nn.BatchNorm1d(bs, affine=False).to(z1.device)
        #     # old_grad_attention_X1, old_grad_attention_X2 = bn(old_grad_attention_X1.T).T, bn(old_grad_attention_X2.T).T
        #     # grad_attention_X1, grad_attention_X2 = bn(grad_attention_X1.T).T, bn(grad_attention_X2.T).T
        #     X1.requires_grad = False
        #     X2.requires_grad = False
            
        #     distill_loss=torch.mean(torch.norm((old_grad_attention_X1-grad_attention_X1), p=2, dim=1))
        #     # distill_loss=torch.mean(torch.norm((old_grad_attention_X1-grad_attention_X1), p=2, dim=1)) + torch.mean(torch.norm((old_grad_attention_X2-grad_attention_X2).reshape((bs, -1)), p=2, dim=1))
        #     return distill_loss
            
    return BaseCbamWrapper
