from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler

from bert_finetuning.data import GermanData


class GermanDataLoader:
    def __init__(
        self,
        data_paths,
        model_name,
        do_cleansing,
        max_sequence_length,
        batch_size=8,
        dataset_cls=GermanData,
    ):
        self.german_data = dataset_cls(
            data_paths,
            model_name,
            max_sequence_length=max_sequence_length,
            do_cleansing=do_cleansing,
        )
        self.batch_size = batch_size
        self.create_loaders()

    def create_loaders(self):
        """
        Create Torch dataloaders for data splits
        """
        self.german_data.text_to_tensors()
        print("creating dataloaders")
        train_data = TensorDataset(
            self.german_data.train_inputs,
            self.german_data.train_masks,
            self.german_data.train_labels,
        )
        train_sampler = RandomSampler(train_data)
        self.train_dataloader = DataLoader(
            train_data, sampler=train_sampler, batch_size=self.batch_size
        )

        validation_data = TensorDataset(
            self.german_data.validation_inputs,
            self.german_data.validation_masks,
            self.german_data.validation_labels,
        )
        validation_sampler = SequentialSampler(validation_data)
        self.validation_dataloader = DataLoader(
            validation_data, sampler=validation_sampler, batch_size=self.batch_size
        )

        test_data = TensorDataset(
            self.german_data.test_inputs,
            self.german_data.test_masks,
            self.german_data.test_labels,
        )
        test_sampler = SequentialSampler(test_data)
        self.test_dataloader = DataLoader(
            test_data, sampler=test_sampler, batch_size=self.batch_size
        )
        print("finished creating dataloaders")


"""
** FOR DEBUGGING **

if __name__ == "__main__":
    ## define data paths
    germeval_data_paths = {
        "train": "./datasets/hasoc_dataset/hasoc_german_train.csv",
        "dev": "./datasets/hasoc_dataset/hasoc_german_validation.csv",
        "test": "./datasets/hasoc_dataset/hasoc_german_test.csv",
    }

    hasoc_german_data_paths = {
        "train": "./datasets/hasoc_dataset/hasoc_german_train.csv",
        "dev": "./datasets/hasoc_dataset/hasoc_german_validation.csv",
        "test": "./datasets/hasoc_dataset/hasoc_german_test.csv",
    }

    ## create dataloaders

    print("creating germeval dataloaders...")
    germ_eval_dataloader = GermanDataLoader(germeval_data_paths)

    print("creating hasoc dataloaders...")
    hasoc_german_dataloader = GermanDataLoader(hasoc_german_data_paths)

"""

[(31, 'torch.utils.data.TensorDataset', 'TensorDataset', False, 'from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler\n'), (36, 'torch.utils.data.RandomSampler', 'RandomSampler', False, 'from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler\n'), (37, 'torch.utils.data.DataLoader', 'DataLoader', False, 'from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler\n'), (41, 'torch.utils.data.TensorDataset', 'TensorDataset', False, 'from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler\n'), (46, 'torch.utils.data.SequentialSampler', 'SequentialSampler', False, 'from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler\n'), (47, 'torch.utils.data.DataLoader', 'DataLoader', False, 'from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler\n'), (51, 'torch.utils.data.TensorDataset', 'TensorDataset', False, 'from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler\n'), (56, 'torch.utils.data.SequentialSampler', 'SequentialSampler', False, 'from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler\n'), (57, 'torch.utils.data.DataLoader', 'DataLoader', False, 'from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler\n')]