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)

"""

[((837, 946), 'torch.utils.data.TensorDataset', 'TensorDataset', (['self.german_data.train_inputs', 'self.german_data.train_masks', 'self.german_data.train_labels'], {}), '(self.german_data.train_inputs, self.german_data.train_masks,\n    self.german_data.train_labels)\n', (850, 946), False, 'from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler\n'), ((1019, 1044), 'torch.utils.data.RandomSampler', 'RandomSampler', (['train_data'], {}), '(train_data)\n', (1032, 1044), False, 'from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler\n'), ((1078, 1151), 'torch.utils.data.DataLoader', 'DataLoader', (['train_data'], {'sampler': 'train_sampler', 'batch_size': 'self.batch_size'}), '(train_data, sampler=train_sampler, batch_size=self.batch_size)\n', (1088, 1151), False, 'from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler\n'), ((1205, 1330), 'torch.utils.data.TensorDataset', 'TensorDataset', (['self.german_data.validation_inputs', 'self.german_data.validation_masks', 'self.german_data.validation_labels'], {}), '(self.german_data.validation_inputs, self.german_data.\n    validation_masks, self.german_data.validation_labels)\n', (1218, 1330), False, 'from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler\n'), ((1407, 1441), 'torch.utils.data.SequentialSampler', 'SequentialSampler', (['validation_data'], {}), '(validation_data)\n', (1424, 1441), False, 'from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler\n'), ((1480, 1568), 'torch.utils.data.DataLoader', 'DataLoader', (['validation_data'], {'sampler': 'validation_sampler', 'batch_size': 'self.batch_size'}), '(validation_data, sampler=validation_sampler, batch_size=self.\n    batch_size)\n', (1490, 1568), False, 'from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler\n'), ((1611, 1717), 'torch.utils.data.TensorDataset', 'TensorDataset', (['self.german_data.test_inputs', 'self.german_data.test_masks', 'self.german_data.test_labels'], {}), '(self.german_data.test_inputs, self.german_data.test_masks,\n    self.german_data.test_labels)\n', (1624, 1717), False, 'from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler\n'), ((1789, 1817), 'torch.utils.data.SequentialSampler', 'SequentialSampler', (['test_data'], {}), '(test_data)\n', (1806, 1817), False, 'from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler\n'), ((1850, 1921), 'torch.utils.data.DataLoader', 'DataLoader', (['test_data'], {'sampler': 'test_sampler', 'batch_size': 'self.batch_size'}), '(test_data, sampler=test_sampler, batch_size=self.batch_size)\n', (1860, 1921), False, 'from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler\n')]