from typing import Optional
from pytorch_lightning import LightningDataModule
from torch.utils.data import DataLoader
from fengshen.data.mmap_index_dataset import MMapIndexDataset


class MMapDataModule(LightningDataModule):
    @ staticmethod
    def add_data_specific_args(parent_args):
        parser = parent_args.add_argument_group('MMAP DataModule')
        parser.add_argument('--num_workers', default=8, type=int)
        parser.add_argument('--train_batchsize', default=32, type=int)
        parser.add_argument('--eval_batchsize', default=32, type=int)
        parser.add_argument('--test_batchsize', default=32, type=int)
        parser.add_argument('--train_datas', default=[
            './train_datas'
        ], type=str, nargs='+')
        parser.add_argument('--valid_datas', default=[
            './valid_datas'
        ], type=str, nargs='+')
        parser.add_argument('--test_datas', default=[
            './test_datas'],
            type=str, nargs='+')
        parser.add_argument('--input_tensor_name', default=['input_ids'], type=str, nargs='+')
        return parent_args

    def __init__(
        self,
        collate_fn,
        args,
        **kwargs,
    ):
        super().__init__()
        self.collate_fn = collate_fn
        self.train_dataset = MMapIndexDataset(args.train_datas, args.input_tensor_name)
        self.valid_dataset = MMapIndexDataset(args.valid_datas, args.input_tensor_name)
        self.test_dataset = MMapIndexDataset(args.test_datas, args.input_tensor_name)
        self.save_hyperparameters(args)

    def setup(self, stage: Optional[str] = None) -> None:
        return super().setup(stage)

    def train_dataloader(self):
        return DataLoader(
            self.train_dataset,
            batch_size=self.hparams.train_batchsize,
            shuffle=True,
            num_workers=self.hparams.num_workers,
            collate_fn=self.collate_fn,
        )

    def val_dataloader(self):
        return DataLoader(
            self.valid_dataset,
            batch_size=self.hparams.eval_batchsize,
            shuffle=True,
            num_workers=self.hparams.num_workers,
            collate_fn=self.collate_fn,
        )

    def test_dataloader(self):
        return DataLoader(
            self.test_dataset,
            batch_size=self.hparams.test_batchsize,
            shuffle=True,
            num_workers=self.hparams.num_workers,
            collate_fn=self.collate_fn,
        )