import argparse
from typing import Any, List, Sequence

import torch
import torch.nn as nn
from cassle.losses.barlow import barlow_loss_func
from cassle.methods.base import BaseModel
from tqdm import tqdm
import torch.nn.functional as F
import pdb
from cassle.methods.dali import no_aug_dataloader


class BarlowTwins(BaseModel):
    def __init__(
        self, proj_hidden_dim: int, output_dim: int, lamb: float, scale_loss: float, **kwargs
    ):
        """Implements Barlow Twins (https://arxiv.org/abs/2103.03230)

        Args:
            proj_hidden_dim (int): number of neurons of the hidden layers of the projector.
            output_dim (int): number of dimensions of projected features.
            lamb (float): off-diagonal scaling factor for the cross-covariance matrix.
            scale_loss (float): scaling factor of the loss.
        """

        super().__init__(**kwargs)

        self.lamb = lamb
        self.scale_loss = scale_loss
        self.indices = []
        self.indices_distances = []
        self.addition_images = kwargs["addition_images"]
        self.method = kwargs["sample_method"]
        self.checkpoint_dir = kwargs["checkpoint_dir"]
        self.dataset = kwargs["dataset"]
        self.dataselect = kwargs["dataselect"]
        if self.dataset == "cub200":
            self.encoder = self.base_model(pretrained=True)
            self.encoder.fc = nn.Identity()
        # projector
        self.projector = nn.Sequential(
            nn.Linear(self.features_dim, proj_hidden_dim),
            nn.BatchNorm1d(proj_hidden_dim),
            nn.ReLU(),
            nn.Linear(proj_hidden_dim, proj_hidden_dim),
            nn.BatchNorm1d(proj_hidden_dim),
            nn.ReLU(),
            nn.Linear(proj_hidden_dim, output_dim),
        )

    @staticmethod
    def add_model_specific_args(parent_parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
        parent_parser = super(BarlowTwins, BarlowTwins).add_model_specific_args(parent_parser)
        parser = parent_parser.add_argument_group("barlow_twins")

        # projector
        parser.add_argument("--output_dim", type=int, default=2048)
        parser.add_argument("--proj_hidden_dim", type=int, default=2048)

        # parameters
        parser.add_argument("--lamb", type=float, default=5e-3)
        parser.add_argument("--scale_loss", type=float, default=0.025)
        return parent_parser

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

        Returns:
            List[dict]: list of learnable parameters.
        """

        extra_learnable_params = [{"params": self.projector.parameters()}]
        return super().learnable_params + extra_learnable_params

    def forward(self, X, *args, **kwargs):
        out = super().forward(X, *args, **kwargs)
        z = self.projector(out["feats"])
        return {**out, "z": z}

    def training_step(self, batch: Sequence[Any], batch_idx: int, h_X1=None, h_X2=None) -> torch.Tensor:
        """Training step for Barlow Twins reusing BaseModel training step.

        Args:
            batch (Sequence[Any]): a batch of data in the format of [img_indexes, [X], Y], where
                [X] is a list of size self.num_crops containing batches of images.
            batch_idx (int): index of the batch.

        Returns:
            torch.Tensor: total loss composed of Barlow loss and classification loss.
        """

        out = super().training_step(batch, batch_idx)
        # h_fea1 = self.encoder(h_X1)["feats"]
        # h_fea2 = self.encoder(h_X2)["feats"]
        # pdb.set_trace()
        feats1, feats2 = out["feats"]
        # feats1 += h_fea1
        # feats2 += h_fea2
        # out["feats"] = [feats1, feats2]
        z1 = self.projector(feats1)
        z2 = self.projector(feats2)

        # ------- barlow twins loss -------
        barlow_loss = barlow_loss_func(z1, z2, lamb=self.lamb, scale_loss=self.scale_loss)
        self.log("train_barlow_loss", barlow_loss, on_epoch=True, sync_dist=True)
        out.update({"loss": out["loss"] + barlow_loss, "z": [z1, z2]})
        return out

    def on_train_start(self):
        super().on_train_start()
        if self.current_task_idx > 0:
            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_end(self):
        if self.dataset == 'cifar100':
            self.on_train_end_cifar()
        elif self.dataset == 'imagenet100':
            self.on_train_end_imagenet()
        elif self.dataset == 'cub200':
            self.on_train_end_cub200()
        else:
            return NotImplementedError
        
    def on_train_end_cifar(self):
        super().on_train_end()
        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)[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)
                    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_imagenet(self):
        super().on_train_end()
        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
            external_dataloader = no_aug_dataloader(self, "imagenet", "/path/to/your/dataset", self.tasks[self.current_task_idx])
            length = (len(external_dataloader)-1)*256
            
            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(length)]).nonzero().view(-1)
            else:
                indices = torch.tensor(range(length))

            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)
            dataloader = no_aug_dataloader(self, "imagenet100", "/path/to/your/dataset", self.tasks[self.current_task_idx], self.dataselect)
            all_in_features = []
            for batch in tqdm(dataloader):
                with torch.no_grad():
                    X, _ = batch
                    X = X.cuda()
                    in_fea = self.encoder(X)
                    all_in_features.append(in_fea)
            all_in_features = torch.cat(all_in_features)
            all_in_features = F.normalize(all_in_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)
            external_dataloader = no_aug_dataloader(self, "imagenet", "/path/to/your/dataset", self.tasks[self.current_task_idx], self.dataselect)
            already_indices = [i in self.indices for i in range(len(external_dataloader)*256)]
            all_min_indices = []
            all_min_distances = []
            temp_out_features = []
            start = 0
            for ind, (input, _) in enumerate(tqdm(external_dataloader)):
                with torch.no_grad():
                    input = input.cuda()
                    out_fea = self.encoder(input)
                    if temp_out_features==[]:
                        temp_out_features = out_fea
                    else:
                        temp_out_features = torch.cat((temp_out_features,out_fea), dim=0)
                    if temp_out_features.shape[0] > 100000 or ind == len(external_dataloader)-1:
                        temp_out_features = F.normalize(temp_out_features, dim=1)
                        temp_distance = torch.matmul(all_in_features, temp_out_features.T)
                        temp_already_indices = already_indices[start:start+temp_out_features.shape[0]]
                        temp=[]
                        temp2=[]
                        temp_min_distances = torch.max(temp_distance, dim=0).values
                        for i in tqdm(range(temp_distance.shape[0])):
                            temp_distance[i][temp_already_indices] = -10000
                            temp_min_indices = int(torch.max(temp_distance[i], 0).indices)
                            temp_already_indices[temp_min_indices] = True
                            temp.append(start+temp_min_indices)
                            temp2.append(temp_min_distances[temp_min_indices])
                        
                        temp = torch.tensor(temp).unsqueeze_(1)
                        all_min_indices.append(temp)
                        temp2 = torch.tensor(temp2).unsqueeze_(1)
                        all_min_distances.append(temp2)
                        start += temp_out_features.shape[0]
                        temp_out_features = []
            all_min_indices = torch.cat(all_min_indices, dim=1)
            all_min_distances = torch.cat(all_min_distances, dim=1)
            values, indexes  = torch.max(all_min_distances, dim=1)
            temp=[]
            for i in range(len(indexes)):
                temp.append((values[i], all_min_indices[i][indexes[i]]))

            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

            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")

    def on_train_end_cub200(self):
        super().on_train_end()
        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
            external_dataloader = no_aug_dataloader(self, "imagenet", "/defaultShare/pubdata/imagenet/train", [], self.dataselect)
            length = (len(external_dataloader)-1)*256
            
            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(length)]).nonzero().view(-1)
            else:
                indices = torch.tensor(range(length))

            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)
            dataloader = no_aug_dataloader(self, "cub200", "/path/to/your/dataset", self.tasks[self.current_task_idx], self.dataselect)
            all_in_features = []
            for batch in tqdm(dataloader):
                with torch.no_grad():
                    X, _ = batch
                    X = X.cuda()
                    in_fea = self.encoder(X)
                    all_in_features.append(in_fea)
            all_in_features = torch.cat(all_in_features)
            all_in_features = F.normalize(all_in_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)
            external_dataloader = no_aug_dataloader(self, "imagenet", "/defaultShare/pubdata/imagenet/train", [], self.dataselect)
            already_indices = [i in self.indices for i in range(len(external_dataloader)*256)]
            all_min_indices = []
            all_min_distances = []
            temp_out_features = []
            start = 0
            for ind, (input, _) in enumerate(tqdm(external_dataloader)):
                with torch.no_grad():
                    input = input.cuda()
                    out_fea = self.encoder(input)
                    if temp_out_features==[]:
                        temp_out_features = out_fea
                    else:
                        temp_out_features = torch.cat((temp_out_features,out_fea), dim=0)
                    if temp_out_features.shape[0] > 50000 or ind == len(external_dataloader)-1:
                        temp_out_features = F.normalize(temp_out_features, dim=1)
                        temp_distance = torch.matmul(all_in_features, temp_out_features.T)
                        temp_already_indices = already_indices[start:start+temp_out_features.shape[0]]
                        temp=[]
                        temp2=[]
                        temp_min_distances = torch.max(temp_distance, dim=0).values
                        for i in tqdm(range(temp_distance.shape[0])):
                            temp_distance[i][temp_already_indices] = -10000
                            temp_min_indices = int(torch.max(temp_distance[i], 0).indices)
                            temp_already_indices[temp_min_indices] = True
                            temp.append(start+temp_min_indices)
                            temp2.append(temp_min_distances[temp_min_indices])
                        
                        temp = torch.tensor(temp).unsqueeze_(1)
                        all_min_indices.append(temp)
                        temp2 = torch.tensor(temp2).unsqueeze_(1)
                        all_min_distances.append(temp2)
                        start += temp_out_features.shape[0]
                        temp_out_features = []
            all_min_indices = torch.cat(all_min_indices, dim=1)
            all_min_distances = torch.cat(all_min_distances, dim=1)
            n = budget_this_task // all_min_indices.shape[0]+1
            # values, indexes  = torch.max(all_min_distances, dim=1)
            values, indexes  = torch.topk(all_min_distances, k=n, dim=1)
            temp=[]

            for i in range(len(indexes)):
                for j in range(n):
                    temp.append((values[i][j], all_min_indices[i][indexes[i][j]]))

            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

            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")