from typing import Any, Callable, Optional, Tuple
import torch.utils
import torch.utils.data
from torchvision.datasets import CIFAR100, ImageFolder, CIFAR10
from PIL import Image
import torchvision.transforms.functional as f
import torch
import numpy as np
import pdb
from torch.utils.data import Subset
from torchvision import datasets
from torchvision.datasets.vision import VisionDataset
import pickle

class my_Imagenet_32(VisionDataset):
    train_list = [
        "train_data_batch_1",
        "train_data_batch_2",
        "train_data_batch_3",
        "train_data_batch_4",
        "train_data_batch_5",
        "train_data_batch_6",
        "train_data_batch_7",
        "train_data_batch_8",
        "train_data_batch_9",
        "train_data_batch_10",
    ]
    def __init__(
        self,
        root: str,
        train: bool = True,
        transform: Optional[Callable] = None,
        target_transform: Optional[Callable] = None,
        download: bool = False,
    ) -> None:

        super().__init__(root, transform=transform, target_transform=target_transform)

        self.train = train  # training set or test set

        if self.train:
            downloaded_list = self.train_list
        else:
            return NotImplementedError

        self.data: Any = []
        self.targets = []

        # now load the picked numpy arrays
        for file_name in downloaded_list:
            file_path = os.path.join(self.root, file_name)
            with open(file_path, "rb") as f:
                entry = pickle.load(f, encoding="latin1")
                self.data.append(entry["data"])
                if "labels" in entry:
                    self.targets.extend(entry["labels"])
                else:
                    self.targets.extend(entry["fine_labels"])

        self.data = np.vstack(self.data).reshape(-1, 3, 32, 32)
        self.data = self.data.transpose((0, 2, 3, 1))  # convert to HWC

        # self._load_meta()

    # def _load_meta(self) -> None:
    #     path = os.path.join(self.root, self.base_folder, self.meta["filename"])
    #     if not check_integrity(path, self.meta["md5"]):
    #         raise RuntimeError("Dataset metadata file not found or corrupted. You can use download=True to download it")
    #     with open(path, "rb") as infile:
    #         data = pickle.load(infile, encoding="latin1")
    #         self.classes = data[self.meta["key"]]
    #     self.class_to_idx = {_class: i for i, _class in enumerate(self.classes)}

    def __getitem__(self, index: int) -> Tuple[Any, Any]:
        """
        Args:
            index (int): Index

        Returns:
            tuple: (image, target) where target is index of the target class.
        """
        img, target = self.data[index], self.targets[index]

        # doing this so that it is consistent with all other datasets
        # to return a PIL Image
        img = Image.fromarray(img)

        if self.transform is not None:
            img = self.transform(img)

        if self.target_transform is not None:
            target = self.target_transform(target)

        return img, target

    def __len__(self) -> int:
        return len(self.data)

    def extra_repr(self) -> str:
        split = "Train" if self.train is True else "Test"
        return f"Split: {split}"

class my_CIFAR10(CIFAR10):
    def __init__(
        self,
        root: str,
        train: bool = True,
        transform: Optional[Callable] = None,
        target_transform: Optional[Callable] = None,
        download: bool = False,
    ) -> None:
        
        super().__init__(root, train=train, transform=transform, target_transform=target_transform, download=download)
        self.patch=8
        self.mix_patch=4
        self.train=False
        self.pre = False
        self.indices = None
        self.last_imgs=[]
        self.interests = np.zeros((50000, 64)).astype('float16')
        self.count=0

    def __getitem__(self, index: int) -> Tuple[Any, Any]:
        """
        Args:
            index (int): Index

        Returns:
            tuple: (image, target) where target is index of the target class.
        """
        img, target = self.data[index], self.targets[index]
        if self.train and not self.pre:
            if len(self.last_imgs)==0:
                self.last_imgs.append(index)
                img1 = Image.fromarray(img)
                img2 = Image.fromarray(img)
            else:
                mix_img = self.data[self.last_imgs[np.random.randint(len(self.last_imgs))]]
                img1 = self.background_mix(img, mix_img, index)
                img2 = self.background_mix(img, mix_img, index)
                
                img1 = Image.fromarray(img1)
                img2 = Image.fromarray(img2)
                self.last_imgs.append(index)
                if len(self.last_imgs)>10:
                    self.last_imgs.pop(0)
        else:
            img = Image.fromarray(img)
        if self.transform is not None:
            if self.train and not self.pre:
                img = [self.transform.transform[0](img1), self.transform.transform[1](img2)]
            elif self.pre:
                img = f.to_tensor(img)
                img = f.normalize(img, [0.4914, 0.4822, 0.4465], [0.247, 0.243, 0.261])
            else:
                img = self.transform(img)

        if self.target_transform is not None:
            target = self.target_transform(target)

        return img, target
    
    def set_train(self, train):
        self.train=train
    
    def set_pre(self, pre):
        self.pre=pre
    
    def set_interests(self, interests):
        self.interests=interests

    def background_mix(self, img, mix_img, index):
        # mix_patch = np.random.randint(self.mix_patch-1, self.mix_patch+2)
        mix_patch = self.mix_patch
        img_ = np.array(img).copy()
        h, w, c = img_.shape
        interest = self.interests[index]
        if np.sum(interest) != 1:
            interest.fill(1/48)
            interest = interest.reshape((8, 8))
            interest[2:6][2:6]=0
            interest = interest.reshape(-1)
            # print(self.count)
            # self.count+=1
        try:
            samples = np.random.choice(64, p = interest, size=mix_patch, replace=False)
        except:
            print(interest)
            print(np.sum(interest))
            print(index)
        grids = np.zeros((mix_patch, 2))
        for i in range(mix_patch):
            grids[i] = [samples[i]//8, samples[i]%8]
            # if grids[i][0] >=2 and grids[i][0] <6 and grids[i][1] >=2 and grids[i][1] <6:
            #     print(self.count)
            #     self.count +=1

        grids = grids.astype(int)
        mask = np.zeros((h, w))
        for j in range(mix_patch):
                # if grids[j][0] >= 2 and grids[j][0] < 6 and grids[j][1] >= 2 and grids[j][1] < 6:
                #     pass
                mask[int(grids[j][0]*h//self.patch):int((grids[j][0]+1)*h//self.patch), int(grids[j][1]*w//self.patch):int((grids[j][1]+1)*w//self.patch)] = 1
        mask = (np.expand_dims(mask, axis=-1)).repeat(c, axis=-1)
        mixed_img_ = img_*(1-mask)+mix_img*(mask)

        return mixed_img_.astype('uint8')


class my_CIFAR100(CIFAR100):
    def __init__(
        self,
        root: str,
        train: bool = True,
        transform: Optional[Callable] = None,
        target_transform: Optional[Callable] = None,
        download: bool = False,
    ) -> None:
        
        super().__init__(root, train=train, transform=transform, target_transform=target_transform, download=download)
        self.patch=8
        self.mix_patch=4
        self.train=False
        self.pre = False
        self.indices = None
        self.save = False
        self.last_imgs=[]
        self.interests = np.zeros((50000, 64)).astype('float16')
        self.count=0

    def __getitem__(self, index: int) -> Tuple[Any, Any]:
        """
        Args:
            index (int): Index

        Returns:
            tuple: (image, target) where target is index of the target class.
        """
        img, target = self.data[index], self.targets[index]
        # pdb.set_trace()

        if self.train and not self.pre:
            # with torch.no_grad():
            #     input_img = f.to_tensor(img).unsqueeze(0).cuda()
            #     input_img = f.normalize(input_img, [0.4914, 0.4822, 0.4465], [0.247, 0.243, 0.261])
            #     att_map = self.old_model(input_img)['att_map']
            # temp=img
            # mix_img = self.data[self.indices[np.random.randint(len(self.indices))]]
            if len(self.last_imgs)==0:
                self.last_imgs.append(index)
                img1 = Image.fromarray(img)
                img2 = Image.fromarray(img)
            else:
                mix_img = self.data[self.last_imgs[np.random.randint(len(self.last_imgs))]]
                img1 = self.background_mix(img, mix_img, index)
                img2 = self.background_mix(img, mix_img, index)
                
                # doing this so that it is consistent with all other datasets
                # to return a PIL Image
                img1 = Image.fromarray(img1)
                img2 = Image.fromarray(img2)
                self.last_imgs.append(index)
                if len(self.last_imgs)>10:
                    self.last_imgs.pop(0)
        else:
            img = Image.fromarray(img)
        if self.transform is not None:
            if self.train and not self.pre:
                img = [self.transform.transform[0](img1), self.transform.transform[1](img2)]
            elif self.pre:
                img = f.to_tensor(img)
                img = f.normalize(img, [0.4914, 0.4822, 0.4465], [0.247, 0.243, 0.261])
            else:
                img = self.transform(img)

        if self.target_transform is not None:
            target = self.target_transform(target)

        return img, target
    
    def set_train(self, train):
        self.train=train
    
    def set_save(self, save):
        self.save=save
    
    def set_pre(self, pre):
        self.pre=pre
    
    def set_interests(self, interests):
        self.interests=interests

    def background_mix(self, img, mix_img, index):
        # mix_patch = np.random.randint(self.mix_patch-1, self.mix_patch+2)
        mix_patch = self.mix_patch
        img_ = np.array(img).copy()
        h, w, c = img_.shape
        interest = self.interests[index]
        if np.sum(interest) != 1:
            interest.fill(1/48)
            interest = interest.reshape((8, 8))
            interest[2:6][2:6]=0
            interest = interest.reshape(-1)
            # print(self.count)
            # self.count+=1
        try:
            samples = np.random.choice(64, p = interest, size=mix_patch, replace=False)
        except:
            print(interest)
            print(np.sum(interest))
            print(index)
        grids = np.zeros((mix_patch, 2))
        for i in range(mix_patch):
            grids[i] = [samples[i]//8, samples[i]%8]
            # if grids[i][0] >=2 and grids[i][0] <6 and grids[i][1] >=2 and grids[i][1] <6:
            #     print(self.count)
            #     self.count +=1

        grids = grids.astype(int)
        mask = np.zeros((h, w))
        for j in range(mix_patch):
                # if grids[j][0] >= 2 and grids[j][0] < 6 and grids[j][1] >= 2 and grids[j][1] < 6:
                #     pass
                mask[int(grids[j][0]*h//self.patch):int((grids[j][0]+1)*h//self.patch), int(grids[j][1]*w//self.patch):int((grids[j][1]+1)*w//self.patch)] = 1
        mask = (np.expand_dims(mask, axis=-1)).repeat(c, axis=-1)
        mixed_img_ = img_*(1-mask)+mix_img*(mask)

        return mixed_img_.astype('uint8')
    

class my_Subset(Subset):
    r"""
    Subset of a dataset at specified indices.

    Args:
        dataset (Dataset): The whole Dataset
        indices (sequence): Indices in the whole set selected for subset
    """

    def set_indices(self):
        self.dataset.indices = self.indices

class my_ImageFolder(ImageFolder):
    def __init__(
        self,
        root: str,
        transform: Optional[Callable] = None,
        target_transform: Optional[Callable] = None,
        loader: Callable[[str], Any] = datasets.folder.default_loader,
        is_valid_file: Optional[Callable[[str], bool]] = None,
    ) -> None:
        
        super().__init__(root, transform, target_transform, loader, is_valid_file)
        self.pre = False

    def __getitem__(self, index: int) -> Tuple[Any, Any]:
        """
        Args:
            index (int): Index

        Returns:
            tuple: (image, target) where target is index of the target class.
        """
        path, target = self.imgs[index]
        target = self.targets[index]
        img = self.loader(path)
        
        if self.transform is not None:
            if not self.pre:
                img = self.transform(img)
            else:
                img = f.to_tensor(img)
                img = f.normalize(img, [0.485, 0.456, 0.406], [0.228, 0.224, 0.225])
                

        if self.target_transform is not None:
            target = self.target_transform(target)

        return img, target
    
    def set_pre(self, pre):
        self.pre=pre
    

from torch.utils.data import Dataset
import sys
import os
class TinyImageNet(Dataset):
    def __init__(self, root, train=True, transform=None):
        self.Train = train
        self.root_dir = root
        self.transform = transform
        self.train_dir = os.path.join(self.root_dir, "train")
        self.val_dir = os.path.join(self.root_dir, "val")

        if (self.Train):
            self._create_class_idx_dict_train()
        else:
            self._create_class_idx_dict_val()

        self._make_dataset(self.Train)

        words_file = os.path.join(self.root_dir, "words.txt")
        wnids_file = os.path.join(self.root_dir, "wnids.txt")

        self.set_nids = set()

        with open(wnids_file, 'r') as fo:
            data = fo.readlines()
            for entry in data:
                self.set_nids.add(entry.strip("\n"))

        self.class_to_label = {}
        with open(words_file, 'r') as fo:
            data = fo.readlines()
            for entry in data:
                words = entry.split("\t")
                if words[0] in self.set_nids:
                    self.class_to_label[words[0]] = (words[1].strip("\n").split(","))[0]

    def _create_class_idx_dict_train(self):
        if sys.version_info >= (3, 5):
            classes = [d.name for d in os.scandir(self.train_dir) if d.is_dir()]
        else:
            classes = [d for d in os.listdir(self.train_dir) if os.path.isdir(os.path.join(train_dir, d))]
        classes = sorted(classes)
        num_images = 0
        for root, dirs, files in os.walk(self.train_dir):
            for f in files:
                if f.endswith(".JPEG"):
                    num_images = num_images + 1

        self.len_dataset = num_images

        self.tgt_idx_to_class = {i: classes[i] for i in range(len(classes))}
        self.class_to_tgt_idx = {classes[i]: i for i in range(len(classes))}

    def _create_class_idx_dict_val(self):
        val_image_dir = os.path.join(self.val_dir, "images")
        if sys.version_info >= (3, 5):
            images = [d.name for d in os.scandir(val_image_dir) if d.is_file()]
        else:
            images = [d for d in os.listdir(val_image_dir) if os.path.isfile(os.path.join(train_dir, d))]
        val_annotations_file = os.path.join(self.val_dir, "val_annotations.txt")
        self.val_img_to_class = {}
        set_of_classes = set()
        with open(val_annotations_file, 'r') as fo:
            entry = fo.readlines()
            for data in entry:
                words = data.split("\t")
                self.val_img_to_class[words[0]] = words[1]
                set_of_classes.add(words[1])

        self.len_dataset = len(list(self.val_img_to_class.keys()))
        classes = sorted(list(set_of_classes))
        # self.idx_to_class = {i:self.val_img_to_class[images[i]] for i in range(len(images))}
        self.class_to_tgt_idx = {classes[i]: i for i in range(len(classes))}
        self.tgt_idx_to_class = {i: classes[i] for i in range(len(classes))}

    def _make_dataset(self, Train=True):
        self.images = []
        if Train:
            img_root_dir = self.train_dir
            list_of_dirs = [target for target in self.class_to_tgt_idx.keys()]
        else:
            img_root_dir = self.val_dir
            list_of_dirs = ["images"]

        for tgt in list_of_dirs:
            dirs = os.path.join(img_root_dir, tgt)
            if not os.path.isdir(dirs):
                continue

            for root, _, files in sorted(os.walk(dirs)):
                for fname in sorted(files):
                    if (fname.endswith(".JPEG")):
                        path = os.path.join(root, fname)
                        if Train:
                            item = (path, self.class_to_tgt_idx[tgt])
                        else:
                            item = (path, self.class_to_tgt_idx[self.val_img_to_class[fname]])
                        self.images.append(item)

    def return_label(self, idx):
        return [self.class_to_label[self.tgt_idx_to_class[i.item()]] for i in idx]

    def __len__(self):
        return self.len_dataset

    def __getitem__(self, idx):
        img_path, tgt = self.images[idx]
        with open(img_path, 'rb') as f:
            sample = Image.open(img_path)
            sample = sample.convert('RGB')
        if self.transform is not None:
            sample = self.transform(sample)

        return sample, tgt
