from typing import Optional

import pytorch_lightning as pl
import torch
from omegaconf import OmegaConf
from torch.utils.data import DataLoader, random_split
from transformers import T5Tokenizer

from src.data.PaperDataset import PaperDataset


class ArvixDataModule(pl.LightningDataModule):
    def __init__(self, config: str = "src/data/config.yaml") -> None:
        super().__init__()
        self.config = OmegaConf.load(config)

    def prepare_data(self) -> None:
        # Add tokenizing
        tokenizer = T5Tokenizer.from_pretrained("t5-base")

        titles, abstracts = torch.load("data/processed/data.pt").T
        #titles, abstracts = torch.load("data/processed/data.pt").T
        tokenized_abstracts = tokenizer.batch_encode_plus(
            abstracts, padding=True, truncation=True, return_tensors="pt"
        )
        tokenized_titles = tokenizer.batch_encode_plus(
            titles, padding=True, truncation=True, return_tensors="pt"
        )

        self.data = PaperDataset(tokenized_abstracts, tokenized_titles)

    def setup(self, stage: Optional[str] = None):
        train, val, test = random_split(
            self.data,
            [self.config.n_train, self.config.n_val, self.config.n_test],
            generator=torch.Generator().manual_seed(1337),
        )

        if stage == "fit" or stage is None:
            self.train_set = train
            self.val_set = val

        if stage == "test":
            self.test_set = test

    def train_dataloader(self) -> DataLoader:
        return DataLoader(self.train_set, batch_size=32, num_workers=4)

    def val_dataloader(self) -> DataLoader:
        return DataLoader(self.val_set, batch_size=32, num_workers=4)

    def test_dataloader(self) -> DataLoader:
        return DataLoader(self.test_set, batch_size=32, num_workers=4)


if __name__ == "__main__":
    dm = ArvixDataModule()

[(15, 'omegaconf.OmegaConf.load', 'OmegaConf.load', False, 'from omegaconf import OmegaConf\n'), (19, 'transformers.T5Tokenizer.from_pretrained', 'T5Tokenizer.from_pretrained', False, 'from transformers import T5Tokenizer\n'), (30, 'src.data.PaperDataset.PaperDataset', 'PaperDataset', False, 'from src.data.PaperDataset import PaperDataset\n'), (47, 'torch.utils.data.DataLoader', 'DataLoader', False, 'from torch.utils.data import DataLoader, random_split\n'), (50, 'torch.utils.data.DataLoader', 'DataLoader', False, 'from torch.utils.data import DataLoader, random_split\n'), (53, 'torch.utils.data.DataLoader', 'DataLoader', False, 'from torch.utils.data import DataLoader, random_split\n'), (21, 'torch.load', 'torch.load', False, 'import torch\n'), (36, 'torch.Generator', 'torch.Generator', False, 'import torch\n')]