from typing import Callable, Optional, Tuple, Union

import numpy as np

from torch.utils.data import DataLoader, Sampler
from torch.utils.data.dataset import Subset, ConcatDataset
import torch.utils.data.distributed as data_dist

from dataflow.datasets import get_train_dataset, get_val_dataset, TransformedDataset, get_train_noval_sbdataset


def get_train_val_loaders(root_path: str,
                          train_transforms: Callable,
                          val_transforms: Callable,
                          batch_size: int = 16,
                          num_workers: int = 8,
                          val_batch_size: Optional[int] = None,
                          pin_memory: bool = True,
                          random_seed: Optional[int] = None,
                          train_sampler: Optional[Union[Sampler, str]] = None,
                          val_sampler: Optional[Union[Sampler, str]] = None,
                          with_sbd: Optional[str] = None,
                          limit_train_num_samples: Optional[int] = None,
                          limit_val_num_samples: Optional[int] = None) -> Tuple[DataLoader, DataLoader, DataLoader]:

    train_ds = get_train_dataset(root_path)
    val_ds = get_val_dataset(root_path)

    if with_sbd is not None:
        sbd_train_ds = get_train_noval_sbdataset(with_sbd)
        train_ds = ConcatDataset([train_ds, sbd_train_ds])

    if random_seed is not None:
        np.random.seed(random_seed)

    if limit_train_num_samples is not None:
        train_indices = np.random.permutation(len(train_ds))[:limit_train_num_samples]
        train_ds = Subset(train_ds, train_indices)

    if limit_val_num_samples is not None:
        val_indices = np.random.permutation(len(val_ds))[:limit_val_num_samples]
        val_ds = Subset(val_ds, val_indices)

    # random samples for evaluation on training dataset
    if len(val_ds) < len(train_ds):
        train_eval_indices = np.random.permutation(len(train_ds))[:len(val_ds)]
        train_eval_ds = Subset(train_ds, train_eval_indices)
    else:
        train_eval_ds = train_ds

    train_ds = TransformedDataset(train_ds, transform_fn=train_transforms)
    val_ds = TransformedDataset(val_ds, transform_fn=val_transforms)
    train_eval_ds = TransformedDataset(train_eval_ds, transform_fn=val_transforms)

    if isinstance(train_sampler, str):
        assert train_sampler == 'distributed'
        train_sampler = data_dist.DistributedSampler(train_ds)

    if isinstance(val_sampler, str):
        assert val_sampler == 'distributed'
        val_sampler = data_dist.DistributedSampler(val_ds, shuffle=False)

    train_loader = DataLoader(train_ds, shuffle=train_sampler is None,
                              batch_size=batch_size, num_workers=num_workers,
                              sampler=train_sampler,
                              pin_memory=pin_memory, drop_last=True)

    val_batch_size = batch_size * 4 if val_batch_size is None else val_batch_size
    val_loader = DataLoader(val_ds, shuffle=False, sampler=val_sampler,
                            batch_size=val_batch_size, num_workers=num_workers,
                            pin_memory=pin_memory, drop_last=False)

    train_eval_loader = DataLoader(train_eval_ds, shuffle=False, sampler=val_sampler,
                                   batch_size=val_batch_size, num_workers=num_workers,
                                   pin_memory=pin_memory, drop_last=False)

    return train_loader, val_loader, train_eval_loader


def get_inference_dataloader(root_path: str,
                             mode: str,
                             transforms: Callable,
                             batch_size: int = 16,
                             num_workers: int = 8,
                             pin_memory: bool = True,
                             limit_num_samples: Optional[int] = None) -> DataLoader:
    assert mode in ('train', 'test'), "Mode should be 'train' or 'test'"

    get_dataset_fn = get_train_dataset if mode == "train" else get_val_dataset

    dataset = get_dataset_fn(root_path, return_meta=True)

    if limit_num_samples is not None:
        indices = np.random.permutation(len(dataset))[:limit_num_samples]
        dataset = Subset(dataset, indices)

    dataset = TransformedDataset(dataset, transform_fn=transforms)

    loader = DataLoader(dataset, shuffle=False,
                        batch_size=batch_size, num_workers=num_workers,
                        pin_memory=pin_memory, drop_last=False)
    return loader

[(26, 'dataflow.datasets.get_train_dataset', 'get_train_dataset', (['root_path'], {}), False, 'from dataflow.datasets import get_train_dataset, get_val_dataset, TransformedDataset, get_train_noval_sbdataset\n'), (27, 'dataflow.datasets.get_val_dataset', 'get_val_dataset', (['root_path'], {}), False, 'from dataflow.datasets import get_train_dataset, get_val_dataset, TransformedDataset, get_train_noval_sbdataset\n'), (51, 'dataflow.datasets.TransformedDataset', 'TransformedDataset', (['train_ds'], {'transform_fn': 'train_transforms'}), False, 'from dataflow.datasets import get_train_dataset, get_val_dataset, TransformedDataset, get_train_noval_sbdataset\n'), (52, 'dataflow.datasets.TransformedDataset', 'TransformedDataset', (['val_ds'], {'transform_fn': 'val_transforms'}), False, 'from dataflow.datasets import get_train_dataset, get_val_dataset, TransformedDataset, get_train_noval_sbdataset\n'), (53, 'dataflow.datasets.TransformedDataset', 'TransformedDataset', (['train_eval_ds'], {'transform_fn': 'val_transforms'}), False, 'from dataflow.datasets import get_train_dataset, get_val_dataset, TransformedDataset, get_train_noval_sbdataset\n'), (63, 'torch.utils.data.DataLoader', 'DataLoader', (['train_ds'], {'shuffle': '(train_sampler is None)', 'batch_size': 'batch_size', 'num_workers': 'num_workers', 'sampler': 'train_sampler', 'pin_memory': 'pin_memory', 'drop_last': '(True)'}), False, 'from torch.utils.data import DataLoader, Sampler\n'), (69, 'torch.utils.data.DataLoader', 'DataLoader', (['val_ds'], {'shuffle': '(False)', 'sampler': 'val_sampler', 'batch_size': 'val_batch_size', 'num_workers': 'num_workers', 'pin_memory': 'pin_memory', 'drop_last': '(False)'}), False, 'from torch.utils.data import DataLoader, Sampler\n'), (73, 'torch.utils.data.DataLoader', 'DataLoader', (['train_eval_ds'], {'shuffle': '(False)', 'sampler': 'val_sampler', 'batch_size': 'val_batch_size', 'num_workers': 'num_workers', 'pin_memory': 'pin_memory', 'drop_last': '(False)'}), False, 'from torch.utils.data import DataLoader, Sampler\n'), (97, 'dataflow.datasets.TransformedDataset', 'TransformedDataset', (['dataset'], {'transform_fn': 'transforms'}), False, 'from dataflow.datasets import get_train_dataset, get_val_dataset, TransformedDataset, get_train_noval_sbdataset\n'), (99, 'torch.utils.data.DataLoader', 'DataLoader', (['dataset'], {'shuffle': '(False)', 'batch_size': 'batch_size', 'num_workers': 'num_workers', 'pin_memory': 'pin_memory', 'drop_last': '(False)'}), False, 'from torch.utils.data import DataLoader, Sampler\n'), (30, 'dataflow.datasets.get_train_noval_sbdataset', 'get_train_noval_sbdataset', (['with_sbd'], {}), False, 'from dataflow.datasets import get_train_dataset, get_val_dataset, TransformedDataset, get_train_noval_sbdataset\n'), (31, 'torch.utils.data.dataset.ConcatDataset', 'ConcatDataset', (['[train_ds, sbd_train_ds]'], {}), False, 'from torch.utils.data.dataset import Subset, ConcatDataset\n'), (34, 'numpy.random.seed', 'np.random.seed', (['random_seed'], {}), True, 'import numpy as np\n'), (38, 'torch.utils.data.dataset.Subset', 'Subset', (['train_ds', 'train_indices'], {}), False, 'from torch.utils.data.dataset import Subset, ConcatDataset\n'), (42, 'torch.utils.data.dataset.Subset', 'Subset', (['val_ds', 'val_indices'], {}), False, 'from torch.utils.data.dataset import Subset, ConcatDataset\n'), (47, 'torch.utils.data.dataset.Subset', 'Subset', (['train_ds', 'train_eval_indices'], {}), False, 'from torch.utils.data.dataset import Subset, ConcatDataset\n'), (57, 'torch.utils.data.distributed.DistributedSampler', 'data_dist.DistributedSampler', (['train_ds'], {}), True, 'import torch.utils.data.distributed as data_dist\n'), (61, 'torch.utils.data.distributed.DistributedSampler', 'data_dist.DistributedSampler', (['val_ds'], {'shuffle': '(False)'}), True, 'import torch.utils.data.distributed as data_dist\n'), (95, 'torch.utils.data.dataset.Subset', 'Subset', (['dataset', 'indices'], {}), False, 'from torch.utils.data.dataset import Subset, ConcatDataset\n')]