from sklearn.model_selection import KFold
from torch.utils.data import ConcatDataset, DataLoader, Subset

from .base_dataset import BaseADDataset


class TorchvisionDataset(BaseADDataset):
    """TorchvisionDataset class for datasets already implemented in torchvision.datasets."""

    def __init__(self, root: str, k_fold_number: int = 5):
        super().__init__(root)
        self.k_fold_number = k_fold_number
        self.fold_indices = None

    def loaders(
        self,
        batch_size: int,
        shuffle_train=True,
        shuffle_test=False,
        num_workers: int = 0,
    ) -> (DataLoader, DataLoader, DataLoader):
        train_loader = (
            DataLoader(
                dataset=self.train_set,
                batch_size=batch_size,
                shuffle=shuffle_train,
                num_workers=num_workers,
                drop_last=True,
            )
            if self.train_set is not None
            else None
        )
        test_loader = (
            DataLoader(
                dataset=self.test_set,
                batch_size=batch_size,
                shuffle=shuffle_test,
                num_workers=num_workers,
                drop_last=False,
            )
            if self.test_set is not None
            else None
        )

        inference_loader = (
            DataLoader(
                dataset=self.inference_set,
                batch_size=batch_size,
                shuffle=False,
                num_workers=num_workers,
                drop_last=False,
            )
            if self.inference_set is not None
            else None
        )
        return train_loader, test_loader, inference_loader

    def loaders_k_fold(
        self,
        fold_idx: int,
        batch_size: int,
        shuffle_train=True,
        shuffle_test=False,
        num_workers: int = 0,
    ) -> (DataLoader, DataLoader):
        if self.k_fold_number is None:
            raise ValueError("k_fold_number must be set to a positive integer.")
        if self.fold_indices is None:
            # Define the K-fold Cross Validator
            kfold = KFold(n_splits=self.k_fold_number, shuffle=False)
            self.fold_indices = []
            # Generate indices for each fold and store them in a list
            for train_indices, val_indices in kfold.split(self.data_set):
                self.fold_indices.append((train_indices, val_indices))

        train_loader = (
            DataLoader(
                dataset=Subset(self.data_set, self.fold_indices[fold_idx][0]),
                batch_size=batch_size,
                shuffle=shuffle_train,
                num_workers=num_workers,
                drop_last=True,
            )
            if self.data_set is not None
            else None
        )
        test_loader = (
            DataLoader(
                dataset=Subset(self.data_set, self.fold_indices[fold_idx][1]),
                batch_size=batch_size,
                shuffle=shuffle_test,
                num_workers=num_workers,
                drop_last=False,
            )
            if self.data_set is not None
            else None
        )
        return train_loader, test_loader
