""" Loader Factory, Fast Collate, CUDA Prefetcher

Prefetcher and Fast Collate inspired by NVIDIA APEX example at
https://github.com/NVIDIA/apex/commit/d5e2bb4bdeedd27b1dfaf5bb2b24d6c000dee9be#diff-cf86c282ff7fba81fad27a559379d5bf

Hacked together by / Copyright 2020 Ross Wightman
"""

from typing import Optional, Callable

import numpy as np
import torch.utils.data

from timm.bits import DeviceEnv
from .collate import fast_collate
from .config import PreprocessCfg, MixupCfg
from .distributed_sampler import OrderedDistributedSampler
from .fetcher import Fetcher
from .mixup import FastCollateMixup
from .prefetcher_cuda import PrefetcherCuda
from .transforms_factory import create_transform_v2


def _worker_init(worker_id):
    worker_info = torch.utils.data.get_worker_info()
    assert worker_info.id == worker_id
    np.random.seed(worker_info.seed % (2**32-1))


def create_loader_v2(
        dataset: torch.utils.data.Dataset,
        batch_size: int,
        is_training: bool = False,
        dev_env: Optional[DeviceEnv] = None,
        pp_cfg: PreprocessCfg = PreprocessCfg(),
        mix_cfg: MixupCfg = None,
        create_transform: bool = True,
        normalize_in_transform: bool = True,
        separate_transform: bool = False,
        num_workers: int = 1,
        collate_fn: Optional[Callable] = None,
        pin_memory: bool = False,
        use_multi_epochs_loader: bool = False,
        persistent_workers: bool = True,
):
    """
    
    Args:
        dataset: 
        batch_size: 
        is_training: 
        dev_env:
        pp_cfg: 
        mix_cfg:
        create_transform:
        normalize_in_transform:
        separate_transform:
        num_workers: 
        collate_fn: 
        pin_memory: 
        use_multi_epochs_loader: 
        persistent_workers: 

    Returns:

    """
    if dev_env is None:
        dev_env = DeviceEnv.instance()

    if create_transform:
        dataset.transform = create_transform_v2(
            cfg=pp_cfg,
            is_training=is_training,
            normalize=normalize_in_transform,
            separate=separate_transform,
        )

    sampler = None
    if dev_env.distributed and not isinstance(dataset, torch.utils.data.IterableDataset):
        if is_training:
            sampler = torch.utils.data.distributed.DistributedSampler(
                dataset, num_replicas=dev_env.world_size, rank=dev_env.global_rank)
        else:
            # This will add extra duplicate entries to result in equal num
            # of samples per-process, will slightly alter validation results
            sampler = OrderedDistributedSampler(
                dataset, num_replicas=dev_env.world_size, rank=dev_env.global_rank)

    if collate_fn is None:
        if mix_cfg is not None and mix_cfg.prob > 0:
            collate_fn = FastCollateMixup(
                mixup_alpha=mix_cfg.mixup_alpha,
                cutmix_alpha=mix_cfg.cutmix_alpha,
                cutmix_minmax=mix_cfg.cutmix_minmax,
                prob=mix_cfg.prob,
                switch_prob=mix_cfg.switch_prob,
                mode=mix_cfg.mode,
                correct_lam=mix_cfg.correct_lam,
                label_smoothing=mix_cfg.label_smoothing,
                num_classes=mix_cfg.num_classes,
            )
        else:
            collate_fn = fast_collate

    loader_class = torch.utils.data.DataLoader
    if use_multi_epochs_loader:
        loader_class = MultiEpochsDataLoader

    loader_args = dict(
        batch_size=batch_size,
        shuffle=not isinstance(dataset, torch.utils.data.IterableDataset) and sampler is None and is_training,
        num_workers=num_workers,
        sampler=sampler,
        collate_fn=collate_fn,
        pin_memory=pin_memory,
        drop_last=is_training,
        worker_init_fn=_worker_init,
        persistent_workers=persistent_workers)
    try:
        loader = loader_class(dataset, **loader_args)
    except TypeError as e:
        loader_args.pop('persistent_workers')  # only in Pytorch 1.7+
        loader = loader_class(dataset, **loader_args)

    fetcher_kwargs = dict(
        normalize=not normalize_in_transform,
        mean=pp_cfg.mean,
        std=pp_cfg.std,
    )
    if not normalize_in_transform and is_training and pp_cfg.aug is not None:
        # If normalization can be done in the prefetcher, random erasing is done there too
        # NOTE RandomErasing does not work well in XLA so normalize_in_transform will be True
        fetcher_kwargs.update(dict(
            re_prob=pp_cfg.aug.re_prob,
            re_mode=pp_cfg.aug.re_mode,
            re_count=pp_cfg.aug.re_count,
            num_aug_splits=pp_cfg.aug.num_aug_splits,
        ))
    if dev_env.type_cuda:
        loader = PrefetcherCuda(loader, **fetcher_kwargs)
    else:
        loader = Fetcher(loader, device=dev_env.device, **fetcher_kwargs)

    return loader


class MultiEpochsDataLoader(torch.utils.data.DataLoader):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._DataLoader__initialized = False
        self.batch_sampler = _RepeatSampler(self.batch_sampler)
        self._DataLoader__initialized = True
        self.iterator = super().__iter__()

    def __len__(self):
        return len(self.batch_sampler.sampler)

    def __iter__(self):
        for i in range(len(self)):
            yield next(self.iterator)


class _RepeatSampler(object):
    """ Sampler that repeats forever.

    Args:
        sampler (Sampler)
    """

    def __init__(self, sampler):
        self.sampler = sampler

    def __iter__(self):
        while True:
            yield from iter(self.sampler)