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, 22, 15, 44), 'omegaconf.OmegaConf.load', 'OmegaConf.load', ({(15, 37, 15, 43): 'config'}, {}), '(config)', False, 'from omegaconf import OmegaConf\n'), ((19, 20, 19, 58), 'transformers.T5Tokenizer.from_pretrained', 'T5Tokenizer.from_pretrained', ({(19, 48, 19, 57): '"""t5-base"""'}, {}), "('t5-base')", False, 'from transformers import T5Tokenizer\n'), ((30, 20, 30, 71), 'src.data.PaperDataset.PaperDataset', 'PaperDataset', ({(30, 33, 30, 52): 'tokenized_abstracts', (30, 54, 30, 70): 'tokenized_titles'}, {}), '(tokenized_abstracts, tokenized_titles)', False, 'from src.data.PaperDataset import PaperDataset\n'), ((47, 15, 47, 71), 'torch.utils.data.DataLoader', 'DataLoader', (), '', False, 'from torch.utils.data import DataLoader, random_split\n'), ((50, 15, 50, 69), 'torch.utils.data.DataLoader', 'DataLoader', (), '', False, 'from torch.utils.data import DataLoader, random_split\n'), ((53, 15, 53, 70), 'torch.utils.data.DataLoader', 'DataLoader', (), '', False, 'from torch.utils.data import DataLoader, random_split\n'), ((21, 28, 21, 64), 'torch.load', 'torch.load', ({(21, 39, 21, 63): '"""data/processed/data.pt"""'}, {}), "('data/processed/data.pt')", False, 'import torch\n'), ((36, 22, 36, 39), 'torch.Generator', 'torch.Generator', ({}, {}), '()', False, 'import torch\n')]