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, 21, 35, 9), 'torch.utils.data.TensorDataset', 'TensorDataset', ({(32, 12, 32, 41): 'self.german_data.train_inputs', (33, 12, 33, 40): 'self.german_data.train_masks', (34, 12, 34, 41): 'self.german_data.train_labels'}, {}), '(self.german_data.train_inputs, self.german_data.train_masks,\n    self.german_data.train_labels)', False, 'from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler\n'), ((36, 24, 36, 49), 'torch.utils.data.RandomSampler', 'RandomSampler', ({(36, 38, 36, 48): 'train_data'}, {}), '(train_data)', False, 'from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler\n'), ((37, 32, 39, 9), 'torch.utils.data.DataLoader', 'DataLoader', (), '', False, 'from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler\n'), ((41, 26, 45, 9), 'torch.utils.data.TensorDataset', 'TensorDataset', ({(42, 12, 42, 46): 'self.german_data.validation_inputs', (43, 12, 43, 45): 'self.german_data.validation_masks', (44, 12, 44, 46): 'self.german_data.validation_labels'}, {}), '(self.german_data.validation_inputs, self.german_data.\n    validation_masks, self.german_data.validation_labels)', False, 'from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler\n'), ((46, 29, 46, 63), 'torch.utils.data.SequentialSampler', 'SequentialSampler', ({(46, 47, 46, 62): 'validation_data'}, {}), '(validation_data)', False, 'from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler\n'), ((47, 37, 49, 9), 'torch.utils.data.DataLoader', 'DataLoader', (), '', False, 'from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler\n'), ((51, 20, 55, 9), 'torch.utils.data.TensorDataset', 'TensorDataset', ({(52, 12, 52, 40): 'self.german_data.test_inputs', (53, 12, 53, 39): 'self.german_data.test_masks', (54, 12, 54, 40): 'self.german_data.test_labels'}, {}), '(self.german_data.test_inputs, self.german_data.test_masks,\n    self.german_data.test_labels)', False, 'from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler\n'), ((56, 23, 56, 51), 'torch.utils.data.SequentialSampler', 'SequentialSampler', ({(56, 41, 56, 50): 'test_data'}, {}), '(test_data)', False, 'from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler\n'), ((57, 31, 59, 9), 'torch.utils.data.DataLoader', 'DataLoader', (), '', False, 'from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler\n')]