import logging
import numpy as np
from PIL import Image
from torch.utils.data import Dataset
from torchvision import transforms
from utils.data import iCIFAR10, iCIFAR100, iImageNet100, iImageNet1000
from tqdm import tqdm
from torch.utils.data import DataLoader
import os
import utils.inc_net
from utils import factory
import torch
import copy
import random

class DataManager(object):
    def __init__(self, dataset_name, shuffle, seed, init_cls, increment, attack=False):
        self.dataset_name = dataset_name
        self.attack = attack
        self._setup_data(dataset_name, shuffle, seed, attack=self.attack)
        assert init_cls <= len(self._class_order), "No enough classes."
        self._increments = [init_cls]
        while sum(self._increments) + increment < len(self._class_order):
            self._increments.append(increment)
        offset = len(self._class_order) - sum(self._increments)
        if offset > 0:
            self._increments.append(offset)

    @property
    def nb_tasks(self):
        return len(self._increments)

    def get_task_size(self, task):
        return self._increments[task]
    
    def get_accumulate_tasksize(self,task):
        return sum(self._increments[:task+1])
    
    def get_total_classnum(self):
        return len(self._class_order)

    def get_dataset(
        self, indices, source, mode, appendent=None, ret_data=False, m_rate=None
    ):
        if source == "train":
            x, y = self._train_data, self._train_targets
        elif source == "test":
            x, y = self._test_data, self._test_targets
        else:
            raise ValueError("Unknown data source {}.".format(source))

        if mode == "train":
            if self.attack:
                trsf = transforms.Compose([*self._test_trsf,])
            else:
                trsf = transforms.Compose([*self._train_trsf, *self._common_trsf])
        elif mode == "flip":
            if self.attack:
                trsf = transforms.Compose(
                    [
                        *self._test_trsf,
                        transforms.RandomHorizontalFlip(p=1.0),
                    ]
                )
            else:
                trsf = transforms.Compose(
                    [
                        *self._test_trsf,
                        transforms.RandomHorizontalFlip(p=1.0),
                        *self._common_trsf,
                    ]
                )
        elif mode == "test":
            if self.attack:
                trsf = transforms.Compose([*self._test_trsf,])
            else:
                trsf = transforms.Compose([*self._test_trsf, *self._common_trsf])
        else:
            raise ValueError("Unknown mode {}.".format(mode))

        data, targets = [], []
        for idx in indices:
            if m_rate is None:
                class_data, class_targets = self._select(
                    x, y, low_range=idx, high_range=idx + 1
                )
            else:
                class_data, class_targets = self._select_rmm(
                    x, y, low_range=idx, high_range=idx + 1, m_rate=m_rate
                )
            data.append(class_data)
            targets.append(class_targets)

        if appendent is not None and len(appendent) != 0:
            appendent_data, appendent_targets = appendent
            data.append(appendent_data)
            targets.append(appendent_targets)

        data, targets = np.concatenate(data), np.concatenate(targets)

        if ret_data:
            return data, targets, DummyDataset(data, targets, trsf, self.use_path)
        else:
            return DummyDataset(data, targets, trsf, self.use_path)

        
    def get_finetune_dataset(self,known_classes,total_classes,source,mode,appendent,type="ratio"):
        if source == 'train':
            x, y = self._train_data, self._train_targets
        elif source == 'test':
            x, y = self._test_data, self._test_targets
        else:
            raise ValueError('Unknown data source {}.'.format(source))

        if mode == 'train':
            trsf = transforms.Compose([*self._train_trsf, *self._common_trsf])
        elif mode == 'test':
            trsf = transforms.Compose([*self._test_trsf, *self._common_trsf])
        else:
            raise ValueError('Unknown mode {}.'.format(mode))
        val_data = []
        val_targets = []

        old_num_tot = 0
        appendent_data, appendent_targets = appendent

        for idx in range(0, known_classes):
            append_data, append_targets = self._select(appendent_data, appendent_targets,
                                                       low_range=idx, high_range=idx+1)
            num=len(append_data)
            if num == 0:
                continue
            old_num_tot += num
            val_data.append(append_data)
            val_targets.append(append_targets)
        if type == "ratio":
            new_num_tot = int(old_num_tot*(total_classes-known_classes)/known_classes)
        elif type == "same":
            new_num_tot = old_num_tot
        else:
            assert 0, "not implemented yet"
        new_num_average = int(new_num_tot/(total_classes-known_classes))
        for idx in range(known_classes,total_classes):
            class_data, class_targets = self._select(x, y, low_range=idx, high_range=idx+1)
            val_indx = np.random.choice(len(class_data),new_num_average, replace=False)
            val_data.append(class_data[val_indx])
            val_targets.append(class_targets[val_indx])
        val_data=np.concatenate(val_data)
        val_targets = np.concatenate(val_targets)
        return DummyDataset(val_data, val_targets, trsf, self.use_path)

    def get_dataset_with_split(
        self, indices, source, mode, appendent=None, val_samples_per_class=0
    ):
        if source == "train":
            x, y = self._train_data, self._train_targets
        elif source == "test":
            x, y = self._test_data, self._test_targets
        else:
            raise ValueError("Unknown data source {}.".format(source))

        if mode == "train":
            trsf = transforms.Compose([*self._train_trsf, *self._common_trsf])
        elif mode == "test":
            trsf = transforms.Compose([*self._test_trsf, *self._common_trsf])
        else:
            raise ValueError("Unknown mode {}.".format(mode))

        train_data, train_targets = [], []
        val_data, val_targets = [], []
        for idx in indices:
            class_data, class_targets = self._select(
                x, y, low_range=idx, high_range=idx + 1
            )
            val_indx = np.random.choice(
                len(class_data), val_samples_per_class, replace=False
            )
            train_indx = list(set(np.arange(len(class_data))) - set(val_indx))
            val_data.append(class_data[val_indx])
            val_targets.append(class_targets[val_indx])
            train_data.append(class_data[train_indx])
            train_targets.append(class_targets[train_indx])

        if appendent is not None:
            appendent_data, appendent_targets = appendent
            for idx in range(0, int(np.max(appendent_targets)) + 1):
                append_data, append_targets = self._select(
                    appendent_data, appendent_targets, low_range=idx, high_range=idx + 1
                )
                val_indx = np.random.choice(
                    len(append_data), val_samples_per_class, replace=False
                )
                train_indx = list(set(np.arange(len(append_data))) - set(val_indx))
                val_data.append(append_data[val_indx])
                val_targets.append(append_targets[val_indx])
                train_data.append(append_data[train_indx])
                train_targets.append(append_targets[train_indx])

        train_data, train_targets = np.concatenate(train_data), np.concatenate(
            train_targets
        )
        val_data, val_targets = np.concatenate(val_data), np.concatenate(val_targets)

        return DummyDataset(
            train_data, train_targets, trsf, self.use_path
        ), DummyDataset(val_data, val_targets, trsf, self.use_path)

    def _setup_data(self, dataset_name, shuffle, seed, attack=False):
        idata = _get_idata(dataset_name)
        idata.download_data()

        # Data
        self._train_data, self._train_targets = idata.train_data, idata.train_targets
        self._test_data, self._test_targets = idata.test_data, idata.test_targets
        self.use_path = idata.use_path

        # Transforms
        self._train_trsf = idata.train_trsf
        self._test_trsf = idata.test_trsf
        if attack:
            self._common_trsf = None
        else:
            self._common_trsf = idata.common_trsf

        # Order
        order = [i for i in range(len(np.unique(self._train_targets)))]
        if shuffle:
            np.random.seed(seed)
            order = np.random.permutation(len(order)).tolist()
        else:
            order = idata.class_order
        self._class_order = order
        logging.info(self._class_order)

        # Map indices
        self._train_targets = _map_new_class_index(
            self._train_targets, self._class_order
        )
        self._test_targets = _map_new_class_index(self._test_targets, self._class_order)

    def _select(self, x, y, low_range, high_range):
        idxes = np.where(np.logical_and(y >= low_range, y < high_range))[0]
        
        if isinstance(x,np.ndarray):
            x_return = x[idxes]
        else:
            x_return = []
            for id in idxes:
                x_return.append(x[id])
        return x_return, y[idxes]

    def _select_rmm(self, x, y, low_range, high_range, m_rate):
        assert m_rate is not None
        if m_rate != 0:
            idxes = np.where(np.logical_and(y >= low_range, y < high_range))[0]
            selected_idxes = np.random.randint(
                0, len(idxes), size=int((1 - m_rate) * len(idxes))
            )
            new_idxes = idxes[selected_idxes]
            new_idxes = np.sort(new_idxes)
        else:
            new_idxes = np.where(np.logical_and(y >= low_range, y < high_range))[0]
        return x[new_idxes], y[new_idxes]

    def getlen(self, index):
        y = self._train_targets
        return np.sum(np.where(y == index))


class DummyDataset(Dataset):
    def __init__(self, images, labels, trsf, use_path=False):
        assert len(images) == len(labels), "Data size error!"
        self.images = images
        self.labels = labels
        self.trsf = trsf
        self.use_path = use_path

    def __len__(self):
        return len(self.images)

    def __getitem__(self, idx):
        if self.use_path:
            image = self.trsf(pil_loader(self.images[idx]))
        else:
            image = self.trsf(Image.fromarray(self.images[idx]))
        label = self.labels[idx]

        return idx, image, label


def _map_new_class_index(y, order):
    return np.array(list(map(lambda x: order.index(x), y)))


def _get_idata(dataset_name):
    name = dataset_name.lower()
    if name == "cifar10":
        return iCIFAR10()
    elif name == "cifar100":
        return iCIFAR100()
    elif name == "imagenet1000":
        return iImageNet1000()
    elif name == "imagenet100":
        return iImageNet100()
    else:
        raise NotImplementedError("Unknown dataset {}.".format(dataset_name))


def get_dataloader(data_manager, batch_size=32,
                   start_class=0, end_class=10,
                   train=False, shuffle=True, num_workers=0):
    if train:
        dataset = data_manager.get_dataset(np.arange(start_class, end_class), source="train", mode="train")
    else:
        dataset = data_manager.get_dataset(np.arange(start_class, end_class), source="test", mode="test")
    loader = DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers)

    return loader


def pil_loader(path):
    """
    Ref:
    https://pytorch.org/docs/stable/_modules/torchvision/datasets/folder.html#ImageFolder
    """
    # open path as file to avoid ResourceWarning (https://github.com/python-pillow/Pillow/issues/835)
    with open(path, "rb") as f:
        img = Image.open(f)
        return img.convert("RGB")


def accimage_loader(path):
    """
    Ref:
    https://pytorch.org/docs/stable/_modules/torchvision/datasets/folder.html#ImageFolder
    accimage is an accelerated Image loader and preprocessor leveraging Intel IPP.
    accimage is available on conda-forge.
    """
    import accimage

    try:
        return accimage.Image(path)
    except IOError:
        # Potentially a decoding problem, fall back to PIL.Image
        return pil_loader(path)


def default_loader(path):
    """
    Ref:
    https://pytorch.org/docs/stable/_modules/torchvision/datasets/folder.html#ImageFolder
    """
    from torchvision import get_image_backend

    if get_image_backend() == "accimage":
        return accimage_loader(path)
    else:
        return pil_loader(path)

def load_all_task_models(args, checkpoint_dir, data_manager, batch_size,
                         device='cuda', train=False, weights=None, load_type='model_loader'):
    if weights == None:
        model_list = []
        # model = factory.get_model(args["model_name"], args)
        loader_list = []
        ckpts = sorted([f for f in os.listdir(checkpoint_dir) if f.endswith('.pkl')])
        known_classes = 0

        if 'model' in load_type:
            model = factory.get_model(args["model_name"], args)

        for i, ckpt_file in enumerate(ckpts):
            if 'model' in load_type:
                ckpt_path = os.path.join(checkpoint_dir, ckpt_file)
                ckpt = torch.load(ckpt_path, map_location=device)

                model.incremental_train(data_manager)
                model._network.load_state_dict(ckpt['model_state_dict'])
                model._network.to(device)
                model._network.eval()
                model_list.append(copy.deepcopy(model))
                model.after_task()

            if 'loader' in load_type:
                cur_task = ckpt['tasks'] if 'tasks' in ckpt else int(ckpt_file.split('_')[-1].split('.')[0])
                total_classes = known_classes + data_manager.get_task_size(cur_task)

                if train:
                    dataset = data_manager.get_dataset(np.arange(0, total_classes), source="train", mode="train")
                else:
                    dataset = data_manager.get_dataset(np.arange(0, total_classes), source="test", mode="test")
                test_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=0)
                loader_list.append(test_loader)
                known_classes = total_classes

        return model_list, loader_list

    else:
        model = factory.get_model(args["model_name"], args)
        ckpt = torch.load(weights, map_location=device)
        model.incremental_train(data_manager)
        model._network.load_state_dict(ckpt['model_state_dict'])
        model._network.to(device)
        model._network.eval()

        total_classes = 10
        if train:
            dataset = data_manager.get_dataset(np.arange(0, total_classes), source="train", mode="train")
        else:
            dataset = data_manager.get_dataset(np.arange(0, total_classes), source="test", mode="test")
        loader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=0)

        return model, loader

def load_src_model(model_name, checkpoint_dir, device='cuda'):
    CL_model_dict = {
        'FOSTERNet': utils.inc_net.FOSTERNet
    }
    model = CL_model_dict["FOSTERNet"]
    ckpt = torch.load(checkpoint_dir, map_location=device)

    total_classes = 10
    model.update_fc(total_classes)
    model._network.load_state_dict(ckpt['model_state_dict'])
    model._network.to(device)
    return model

def load_src_dataset(data_manager, batch_size):
    total_classes = 10
    test_dataset = data_manager.get_dataset(np.arange(0, total_classes), source="train", mode="train")
    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=True, num_workers=0)
    return test_loader

def balanced_sample_from_loaders(loaders, total_batch_size):
    num_loaders = len(loaders)
    per_loader_sample = total_batch_size // num_loaders
    remainder = total_batch_size % num_loaders

    x_batch, y_batch = [], []

    for i, loader in enumerate(loaders):
        batch_needed = per_loader_sample + (1 if i < remainder else 0)
        data_iter = iter(loader)
        current_count = 0
        while current_count < batch_needed:
            x, y = next(data_iter)
            needed = batch_needed - current_count
            if x.shape[0] > needed:
                x = x[:needed]
                y = y[:needed]
            x_batch.append(x)
            y_batch.append(y)
            current_count += x.shape[0]

    x_batch = torch.cat(x_batch, dim=0)
    y_batch = torch.cat(y_batch, dim=0)
    return x_batch, y_batch


class CustomDMDataset(Dataset):
    def __init__(self, data_dir, transform=None, split='train', test_size=0.2):
        self.data_dir = data_dir
        self.transform = transform
        self.split = split
        self.test_size = test_size

        self.classes = sorted(os.listdir(data_dir))
        self.image_paths = []
        self.labels = []

        for label, class_name in enumerate(self.classes):
            class_folder = os.path.join(data_dir, class_name)
            if os.path.isdir(class_folder):
                for img_name in os.listdir(class_folder):
                    img_path = os.path.join(class_folder, img_name)
                    if img_name.endswith(".jpg") or img_name.endswith(".png"):  # 根据文件类型选择
                        self.image_paths.append(img_path)
                        self.labels.append(label)

        total_size = len(self.image_paths)
        test_size = int(total_size * self.test_size)
        train_size = total_size - test_size

        indices = list(range(total_size))
        random.shuffle(indices)

        train_indices = indices[:train_size]
        test_indices = indices[train_size:]

        if self.split == 'train':
            self.image_paths = [self.image_paths[i] for i in train_indices]
            self.labels = [self.labels[i] for i in train_indices]
        else:
            self.image_paths = [self.image_paths[i] for i in test_indices]
            self.labels = [self.labels[i] for i in test_indices]

    def __len__(self):
        return len(self.image_paths)

    def __getitem__(self, idx):
        img_path = self.image_paths[idx]
        label = self.labels[idx]
        img = Image.open(img_path)

        if self.transform:
            img = self.transform(img)

        return img, label