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

from torch.utils.data import DataLoader

from .humanml.dataset import Text2MotionDatasetV2


class BASEDataModule:
    def __init__(self, collate_fn: Callable, batch_size: int,
                 num_workers: int, persistent_workers: bool) -> None:
        super(BASEDataModule, self).__init__()
        self.dataloader_options = {
            "batch_size": batch_size,
            "num_workers": num_workers,
            "collate_fn": collate_fn,
            "persistent_workers": persistent_workers
        }
        self.is_mm = False

    def get_sample_set(self, overrides: dict) -> Text2MotionDatasetV2:
        sample_params = copy.deepcopy(self.hparams)
        sample_params.update(overrides)
        split_file = pjoin(
            eval(f"self.cfg.DATASET.{self.name.upper()}.SPLIT_ROOT"),
            self.cfg.EVAL.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")]
            item_c = "_" + item
            if item_c not in self.__dict__:

                subset = subset.upper() if subset != "val" else "EVAL"
                split = eval(f"self.cfg.{subset}.SPLIT")
                split_file = pjoin(
                    eval(f"self.cfg.DATASET.{self.name.upper()}.SPLIT_ROOT"),
                    eval(f"self.cfg.{subset}.SPLIT") + ".txt",
                )
                self.__dict__[item_c] = self.Dataset(split_file=split_file,
                                                     split=split,
                                                     **self.hparams)
            return getattr(self, item_c)
        classname = self.__class__.__name__
        raise AttributeError(f"'{classname}' object has no attribute '{item}'")

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

    def val_dataloader(self) -> DataLoader:
        dataloader_options = self.dataloader_options.copy()
        dataloader_options["batch_size"] = self.cfg.EVAL.BATCH_SIZE
        dataloader_options["num_workers"] = self.cfg.EVAL.NUM_WORKERS
        dataloader_options["shuffle"] = False
        return DataLoader(self.val_dataset, **dataloader_options)

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