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

[((1185, 1213), 'dataflow.datasets.get_train_dataset', 'get_train_dataset', (['root_path'], {}), '(root_path)\n', (1202, 1213), False, 'from dataflow.datasets import get_train_dataset, get_val_dataset, TransformedDataset, get_train_noval_sbdataset\n'), ((1227, 1253), 'dataflow.datasets.get_val_dataset', 'get_val_dataset', (['root_path'], {}), '(root_path)\n', (1242, 1253), False, 'from dataflow.datasets import get_train_dataset, get_val_dataset, TransformedDataset, get_train_noval_sbdataset\n'), ((2116, 2175), 'dataflow.datasets.TransformedDataset', 'TransformedDataset', (['train_ds'], {'transform_fn': 'train_transforms'}), '(train_ds, transform_fn=train_transforms)\n', (2134, 2175), False, 'from dataflow.datasets import get_train_dataset, get_val_dataset, TransformedDataset, get_train_noval_sbdataset\n'), ((2189, 2244), 'dataflow.datasets.TransformedDataset', 'TransformedDataset', (['val_ds'], {'transform_fn': 'val_transforms'}), '(val_ds, transform_fn=val_transforms)\n', (2207, 2244), False, 'from dataflow.datasets import get_train_dataset, get_val_dataset, TransformedDataset, get_train_noval_sbdataset\n'), ((2265, 2327), 'dataflow.datasets.TransformedDataset', 'TransformedDataset', (['train_eval_ds'], {'transform_fn': 'val_transforms'}), '(train_eval_ds, transform_fn=val_transforms)\n', (2283, 2327), False, 'from dataflow.datasets import get_train_dataset, get_val_dataset, TransformedDataset, get_train_noval_sbdataset\n'), ((2653, 2822), '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)'}), '(train_ds, shuffle=train_sampler is None, batch_size=batch_size,\n    num_workers=num_workers, sampler=train_sampler, pin_memory=pin_memory,\n    drop_last=True)\n', (2663, 2822), False, 'from torch.utils.data import DataLoader, Sampler\n'), ((3005, 3160), '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)'}), '(val_ds, shuffle=False, sampler=val_sampler, batch_size=\n    val_batch_size, num_workers=num_workers, pin_memory=pin_memory,\n    drop_last=False)\n', (3015, 3160), False, 'from torch.utils.data import DataLoader, Sampler\n'), ((3233, 3395), '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)'}), '(train_eval_ds, shuffle=False, sampler=val_sampler, batch_size=\n    val_batch_size, num_workers=num_workers, pin_memory=pin_memory,\n    drop_last=False)\n', (3243, 3395), False, 'from torch.utils.data import DataLoader, Sampler\n'), ((4275, 4327), 'dataflow.datasets.TransformedDataset', 'TransformedDataset', (['dataset'], {'transform_fn': 'transforms'}), '(dataset, transform_fn=transforms)\n', (4293, 4327), False, 'from dataflow.datasets import get_train_dataset, get_val_dataset, TransformedDataset, get_train_noval_sbdataset\n'), ((4342, 4469), 'torch.utils.data.DataLoader', 'DataLoader', (['dataset'], {'shuffle': '(False)', 'batch_size': 'batch_size', 'num_workers': 'num_workers', 'pin_memory': 'pin_memory', 'drop_last': '(False)'}), '(dataset, shuffle=False, batch_size=batch_size, num_workers=\n    num_workers, pin_memory=pin_memory, drop_last=False)\n', (4352, 4469), False, 'from torch.utils.data import DataLoader, Sampler\n'), ((1307, 1342), 'dataflow.datasets.get_train_noval_sbdataset', 'get_train_noval_sbdataset', (['with_sbd'], {}), '(with_sbd)\n', (1332, 1342), False, 'from dataflow.datasets import get_train_dataset, get_val_dataset, TransformedDataset, get_train_noval_sbdataset\n'), ((1362, 1401), 'torch.utils.data.dataset.ConcatDataset', 'ConcatDataset', (['[train_ds, sbd_train_ds]'], {}), '([train_ds, sbd_train_ds])\n', (1375, 1401), False, 'from torch.utils.data.dataset import Subset, ConcatDataset\n'), ((1443, 1470), 'numpy.random.seed', 'np.random.seed', (['random_seed'], {}), '(random_seed)\n', (1457, 1470), True, 'import numpy as np\n'), ((1622, 1653), 'torch.utils.data.dataset.Subset', 'Subset', (['train_ds', 'train_indices'], {}), '(train_ds, train_indices)\n', (1628, 1653), False, 'from torch.utils.data.dataset import Subset, ConcatDataset\n'), ((1795, 1822), 'torch.utils.data.dataset.Subset', 'Subset', (['val_ds', 'val_indices'], {}), '(val_ds, val_indices)\n', (1801, 1822), False, 'from torch.utils.data.dataset import Subset, ConcatDataset\n'), ((2020, 2056), 'torch.utils.data.dataset.Subset', 'Subset', (['train_ds', 'train_eval_indices'], {}), '(train_ds, train_eval_indices)\n', (2026, 2056), False, 'from torch.utils.data.dataset import Subset, ConcatDataset\n'), ((2438, 2476), 'torch.utils.data.distributed.DistributedSampler', 'data_dist.DistributedSampler', (['train_ds'], {}), '(train_ds)\n', (2466, 2476), True, 'import torch.utils.data.distributed as data_dist\n'), ((2581, 2632), 'torch.utils.data.distributed.DistributedSampler', 'data_dist.DistributedSampler', (['val_ds'], {'shuffle': '(False)'}), '(val_ds, shuffle=False)\n', (2609, 2632), True, 'import torch.utils.data.distributed as data_dist\n'), ((4235, 4259), 'torch.utils.data.dataset.Subset', 'Subset', (['dataset', 'indices'], {}), '(dataset, indices)\n', (4241, 4259), False, 'from torch.utils.data.dataset import Subset, ConcatDataset\n')]