from typing import Optional

from anndata import AnnData
from scvi import settings
from scvi.dataloaders import DataSplitter, AnnDataLoader
from scvi.model._utils import parse_use_gpu_arg


class ManualDataSplitter(DataSplitter):
    """Manual train validation test splitter"""

    def __init__(
            self,
            adata: AnnData,
            train_idx,
            val_idx,
            test_idx,
            use_gpu: bool = False,
            **kwargs,
    ):
        super().__init__(adata)
        self.data_loader_kwargs = kwargs
        self.use_gpu = use_gpu
        self.val_idx = val_idx
        self.train_idx = train_idx
        self.test_idx = test_idx

    def setup(self, stage: Optional[str] = None):
        gpus, self.device = parse_use_gpu_arg(self.use_gpu, return_device=True)
        self.pin_memory = (
            True if (settings.dl_pin_memory_gpu_training and gpus != 0) else False
        )

    def val_dataloader(self):
        if len(self.val_idx) > 0:
            data_loader_kwargs = self.data_loader_kwargs.copy()
            if len(self.val_idx < 5000):
                data_loader_kwargs.update({'batch_size': len(self.val_idx)})
            else:
                data_loader_kwargs.update({'batch_size': 2048})
            return AnnDataLoader(
                self.adata,
                indices=self.val_idx,
                shuffle=True,
                pin_memory=self.pin_memory,
                **data_loader_kwargs,
            )
        else:
            pass

    def test_dataloader(self):
        if len(self.test_idx) > 0:
            return AnnDataLoader(
                self.adata,
                indices=self.test_idx,
                shuffle=True,
                pin_memory=self.pin_memory,
                **self.data_loader_kwargs,
            )
        else:
            pass
