import copy
from os.path import join as pjoin
from typing import Any, Callable

from torch.utils.data import DataLoader


class BaseDataModule:
    def __init__(self, collate_fn: Callable) -> None:
        super(BaseDataModule, self).__init__()
        self.collate_fn = collate_fn
        self.is_mm = False

    def get_sample_set(self, overrides: dict) -> Any:
        sample_params = copy.deepcopy(self.hparams)
        sample_params.update(overrides)
        split_file = pjoin(
            eval(f"self.cfg.DATASET.{self.name.upper()}.ROOT"),
            self.cfg.TEST.SPLIT + ".txt"
        )
        return self.Dataset(split_file=split_file, **sample_params)

    def __getattr__(self, item: str) -> Any:
        if item.endswith("_dataset") and not item.startswith("_"):
            subset = item[:-len("_dataset")].upper()
            item_c = "_" + item
            if item_c not in self.__dict__:
                split_file = pjoin(
                    eval(f"self.cfg.DATASET.{self.name.upper()}.ROOT"),
                    eval(f"self.cfg.{subset}.SPLIT") + ".txt"
                )
                self.__dict__[item_c] = self.Dataset(split_file=split_file, **self.hparams)
            return getattr(self, item_c)
        classname = self.__class__.__name__
        raise AttributeError(f"'{classname}' object has no attribute '{item}'")

    def get_dataloader_options(self, stage: str) -> dict:
        stage_args = eval(f"self.cfg.{stage.upper()}")
        dataloader_options = {
            "batch_size": stage_args.BATCH_SIZE,
            "num_workers": stage_args.NUM_WORKERS,
            "collate_fn": self.collate_fn,
            "persistent_workers": stage_args.PERSISTENT_WORKERS,
        }
        return dataloader_options

    def train_dataloader(self) -> DataLoader:
        dataloader_options = self.get_dataloader_options('TRAIN')
        return DataLoader(self.train_dataset, shuffle=True, **dataloader_options)

    def val_dataloader(self) -> DataLoader:
        dataloader_options = self.get_dataloader_options('VAL')
        return DataLoader(self.val_dataset, shuffle=False, **dataloader_options)

    def test_dataloader(self) -> DataLoader:
        dataloader_options = self.get_dataloader_options('TEST')
        dataloader_options["batch_size"] = 1 if self.is_mm else self.cfg.TEST.BATCH_SIZE
        return DataLoader(self.test_dataset, shuffle=False, **dataloader_options)