import os
from torch.utils.data import DataLoader
from lightning import LightningDataModule
from .mixed_dataset import MixedDataset


class MixedDataModule(LightningDataModule):
    def __init__(
        self, bert_model, dataset_path, tool_capacity, batch_size, num_workers, seed
    ):
        super().__init__()
        self.bert_model = bert_model
        self.dataset_path = dataset_path
        self.tool_capacity = tool_capacity
        self.batch_size = batch_size
        self.num_workers = num_workers
        self.seed = seed

    def setup(self, stage=None):
        if stage == "fit":
            self.train_dataset = MixedDataset(
                self.bert_model,
                "train",
                os.path.join(self.dataset_path, "train.json"),
                self.tool_capacity,
                seed=self.seed,
            )
            self.val_dataset = MixedDataset(
                self.bert_model,
                "test",
                os.path.join(self.dataset_path, "test.json"),
                self.tool_capacity,
                seed=self.seed,
            )
        elif stage == "test":
            self.test_dataset = MixedDataset(
                self.bert_model,
                "test",
                os.path.join(self.dataset_path, "test.json"),
                self.tool_capacity,
                seed=self.seed,
            )

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

    def val_dataloader(self):
        return DataLoader(
            self.val_dataset,
            batch_size=self.batch_size,
            shuffle=False,
            num_workers=self.num_workers,
        )

    def test_dataloader(self):
        return DataLoader(
            self.test_dataset,
            batch_size=self.batch_size,
            shuffle=False,
            num_workers=self.num_workers,
        )