from typing import Dict
import os
from omegaconf import DictConfig
import torch
import lightning.pytorch as pl
from lightning.pytorch.utilities.types import TRAIN_DATALOADERS, EVAL_DATALOADERS


class BaseDataModule(pl.LightningDataModule):
    def __init__(self, root_cfg: DictConfig, compatible_datasets: Dict) -> None:
        super().__init__()
        self.root_cfg = root_cfg
        self.exp_cfg = root_cfg.experiment
        self.compatible_datasets = compatible_datasets

    def _build_dataset(self, split: str) -> torch.utils.data.Dataset:
        if split in ["training", "test", "validation"]:
            return self.compatible_datasets[self.root_cfg.dataset._name](
                self.root_cfg.dataset, split=split
            )
        else:
            raise NotImplementedError(f"split '{split}' is not implemented")

    @staticmethod
    def _get_shuffle(dataset: torch.utils.data.Dataset, default: bool) -> bool:
        return not isinstance(dataset, torch.utils.data.IterableDataset) and default

    @staticmethod
    def _get_num_workers(num_workers: int) -> int:
        return min(os.cpu_count(), num_workers)

    def _dataloader(self, split: str) -> TRAIN_DATALOADERS | EVAL_DATALOADERS:
        dataset = self._build_dataset(split)
        split_cfg = self.exp_cfg[split]
        return torch.utils.data.DataLoader(
            dataset,
            batch_size=split_cfg.batch_size,
            num_workers=self._get_num_workers(split_cfg.data.num_workers),
            shuffle=self._get_shuffle(dataset, split_cfg.data.shuffle),
            persistent_workers= (split == "training") and self._get_num_workers(split_cfg.data.num_workers),
            worker_init_fn=lambda worker_id: (
                dataset.worker_init_fn(worker_id)
                if hasattr(dataset, "worker_init_fn")
                else None
            ),
        )

    def train_dataloader(self) -> TRAIN_DATALOADERS:
        return self._dataloader("training")

    def val_dataloader(self) -> EVAL_DATALOADERS:
        return self._dataloader("validation")

    def test_dataloader(self) -> EVAL_DATALOADERS:
        return self._dataloader("test")