# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.


from torch.utils.data.datapipes.map import Batcher, Concater, Mapper, SequenceWrapper, Shuffler, Zipper

from torchdata.datapipes.iter.util.converter import IterToMapConverterMapDataPipe as IterToMapConverter
from torchdata.datapipes.map.util.cacheholder import InMemoryCacheHolderMapDataPipe as InMemoryCacheHolder
from torchdata.datapipes.map.util.unzipper import UnZipperMapDataPipe as UnZipper

__all__ = [
    "Batcher",
    "Concater",
    "InMemoryCacheHolder",
    "IterToMapConverter",
    "MapDataPipe",
    "Mapper",
    "SequenceWrapper",
    "Shuffler",
    "UnZipper",
    "Zipper",
]

# Please keep this list sorted
assert __all__ == sorted(__all__)

########################################################################################################################
# The part below is generated by parsing through the Python files where MapDataPipes are defined.
# This base template ("__init__.pyi.in") is generated from mypy stubgen with minimal editing for code injection
# The output file will be "__init__.pyi". The generation function is called by "setup.py".
# Note that, for mypy, .pyi file takes precedent over .py file, such that we must define the interface for other
# classes/objects here, even though we are not injecting extra code into them at the moment.

from torchdata.datapipes.iter import IterDataPipe
from torch.utils.data import DataChunk, Dataset
from torch.utils.data.datapipes._typing import _DataPipeMeta

from typing import Any, Callable, Dict, List, Optional, Sequence, TypeVar, Union

T_co = TypeVar('T_co', covariant=True)
T = TypeVar('T')
UNTRACABLE_DATAFRAME_PIPES: Any

class MapDataPipe(Dataset[T_co], metaclass=_DataPipeMeta):
    functions: Dict[str, Callable] = ...
    def __getattr__(self, attribute_name: Any): ...
    @classmethod
    def register_function(cls, function_name: Any, function: Any) -> None: ...
    @classmethod
    def register_datapipe_as_function(cls, function_name: Any, cls_to_register: Any): ...
    # Functional form of 'BatcherMapDataPipe'
    def batch(self, batch_size: int, drop_last: bool = False, wrapper_class: Type[DataChunk] = DataChunk) -> MapDataPipe:
        r"""
        Create mini-batches of data (functional name: ``batch``).
    
        An outer dimension will be added as ``batch_size`` if ``drop_last`` is set to ``True``,
        or ``length % batch_size`` for the last batch if ``drop_last`` is set to ``False``.
    
        Args:
            datapipe: Iterable DataPipe being batched
            batch_size: The size of each batch
            drop_last: Option to drop the last batch if it's not full
    
        Example:
            >>> # xdoctest: +SKIP
            >>> from torchdata.datapipes.map import SequenceWrapper
            >>> dp = SequenceWrapper(range(10))
            >>> batch_dp = dp.batch(batch_size=2)
            >>> list(batch_dp)
            [[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]]
        """
    
    # Functional form of 'ConcaterMapDataPipe'
    def concat(self, *datapipes: MapDataPipe) -> MapDataPipe:
        r"""
        Concatenate multiple Map DataPipes (functional name: ``concat``).
    
        The new index of is the cumulative sum of source DataPipes.
        For example, if there are 2 source DataPipes both with length 5,
        index 0 to 4 of the resulting `ConcatMapDataPipe` would refer to
        elements of the first DataPipe, and 5 to 9 would refer to elements
        of the second DataPipe.
    
        Args:
            datapipes: Map DataPipes being concatenated
    
        Example:
            >>> # xdoctest: +SKIP
            >>> from torchdata.datapipes.map import SequenceWrapper
            >>> dp1 = SequenceWrapper(range(3))
            >>> dp2 = SequenceWrapper(range(3))
            >>> concat_dp = dp1.concat(dp2)
            >>> list(concat_dp)
            [0, 1, 2, 0, 1, 2]
        """
    
    # Functional form of 'MapperMapDataPipe'
    def map(self, fn: Callable= ...) -> MapDataPipe:
        r"""
        Apply the input function over each item from the source DataPipe (functional name: ``map``).
    
        The function can be any regular Python function or partial object. Lambda
        function is not recommended as it is not supported by pickle.
    
        Args:
            datapipe: Source MapDataPipe
            fn: Function being applied to each item
    
        Example:
            >>> # xdoctest: +SKIP
            >>> from torchdata.datapipes.map import SequenceWrapper, Mapper
            >>> def add_one(x):
            ...     return x + 1
            >>> dp = SequenceWrapper(range(10))
            >>> map_dp_1 = dp.map(add_one)
            >>> list(map_dp_1)
            [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
            >>> map_dp_2 = Mapper(dp, lambda x: x + 1)
            >>> list(map_dp_2)
            [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
        """
    
    # Functional form of 'ShufflerIterDataPipe'
    def shuffle(self, *, indices: Optional[List] = None) -> IterDataPipe:
        r"""
        Shuffle the input MapDataPipe via its indices (functional name: ``shuffle``).
    
        When it is used with :class:`~torch.utils.data.DataLoader`, the methods to
        set up random seed are different based on :attr:`num_workers`.
    
        For single-process mode (:attr:`num_workers == 0`), the random seed is set before
        the :class:`~torch.utils.data.DataLoader` in the main process. For multi-process
        mode (:attr:`num_worker > 0`), ``worker_init_fn`` is used to set up a random seed
        for each worker process.
    
        Args:
            datapipe: MapDataPipe being shuffled
            indices: a list of indices of the MapDataPipe. If not provided, we assume it uses 0-based indexing
    
        Example:
            >>> # xdoctest: +SKIP
            >>> from torchdata.datapipes.map import SequenceWrapper
            >>> dp = SequenceWrapper(range(10))
            >>> shuffle_dp = dp.shuffle().set_seed(0)
            >>> list(shuffle_dp)
            [7, 8, 1, 5, 3, 4, 2, 0, 9, 6]
            >>> list(shuffle_dp)
            [6, 1, 9, 5, 2, 4, 7, 3, 8, 0]
            >>> # Reset seed for Shuffler
            >>> shuffle_dp = shuffle_dp.set_seed(0)
            >>> list(shuffle_dp)
            [7, 8, 1, 5, 3, 4, 2, 0, 9, 6]
    
        Note:
            Even thought this ``shuffle`` operation takes a ``MapDataPipe`` as the input, it would return an
            ``IterDataPipe`` rather than a ``MapDataPipe``, because ``MapDataPipe`` should be non-sensitive to
            the order of data order for the sake of random reads, but ``IterDataPipe`` depends on the order
            of data during data-processing.
        """
    
    # Functional form of 'ZipperMapDataPipe'
    def zip(self, *datapipes: MapDataPipe[_T_co]) -> MapDataPipe:
        r"""
        Aggregates elements into a tuple from each of the input DataPipes (functional name: ``zip``).
    
        This MataPipe is out of bound as soon as the shortest input DataPipe is exhausted.
    
        Args:
            *datapipes: Map DataPipes being aggregated
    
        Example:
            >>> # xdoctest: +SKIP
            >>> from torchdata.datapipes.map import SequenceWrapper
            >>> dp1 = SequenceWrapper(range(3))
            >>> dp2 = SequenceWrapper(range(10, 13))
            >>> zip_dp = dp1.zip(dp2)
            >>> list(zip_dp)
            [(0, 10), (1, 11), (2, 12)]
        """
    
    # Functional form of 'InMemoryCacheHolderMapDataPipe'
    def in_memory_cache(self) -> MapDataPipe:
        r"""
        Stores elements from the source DataPipe in memory (functional name: ``in_memory_cache``). Once an item is
        stored, it will remain unchanged and subsequent retrivals will return the same element. Since items from
        ``MapDataPipe`` are lazily computed, this can be used to store the results from previous ``MapDataPipe`` and
        reduce the number of duplicate computations.
    
        Note:
            The default ``cache`` is a ``Dict``. If another data structure is more suitable as cache for your use
    
        Args:
            source_dp: source DataPipe from which elements are read and stored in memory
    
        Example:
            >>> from torchdata.datapipes.map import SequenceWrapper
            >>> source_dp = SequenceWrapper(range(10))
            >>> cache_dp = source_dp.in_memory_cache()
            >>> list(cache_dp)
            [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
        """
    
    # Functional form of 'MapToIterConverterIterDataPipe'
    def to_iter_datapipe(self, indices: Optional[List] = None) -> IterDataPipe:
        """
        Convert a ``MapDataPipe`` to an ``IterDataPipe`` (functional name: ``to_iter_datapipe``). It uses ``indices`` to
        iterate through the ``MapDataPipe``, defaults to ``range(len(mapdatapipe))`` if not given.
    
        For the opposite converter, use :class:`.IterToMapConverter`.
    
        Args:
            datapipe: source MapDataPipe with data
            indices: optional list of indices that will dictate how the datapipe will be iterated over
    
        Example:
            >>> from torchdata.datapipes.map import SequenceWrapper
            >>> source_dp = SequenceWrapper(range(10))
            >>> iter_dp = source_dp.to_iter_datapipe()
            >>> list(iter_dp)
            [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
            >>> source_dp2 = SequenceWrapper({'a': 1, 'b': 2, 'c': 3})
            >>> iter_dp2 = source_dp2.to_iter_datapipe(indices=['a', 'b', 'c'])
            >>> list(iter_dp2)
            [1, 2, 3]
        """
    
    # Functional form of 'UnZipperMapDataPipe'
    def unzip(self, sequence_length: int, columns_to_skip: Optional[Sequence[int]] = None) -> List[MapDataPipe]:
        """
        Takes in a DataPipe of Sequences, unpacks each Sequence, and return the elements in separate DataPipes
        based on their position in the Sequence (functional name: ``unzip``). The number of instances produced
        equals to the ``sequence_legnth`` minus the number of columns to skip.
    
        Note:
            Each sequence within the DataPipe should have the same length, specified by
            the input argument `sequence_length`.
    
        Args:
            source_datapipe: Iterable DataPipe with sequences of data
            sequence_length: Length of the sequence within the source_datapipe. All elements should have the same length.
            columns_to_skip: optional indices of columns that the DataPipe should skip (each index should be
                an integer from 0 to sequence_length - 1)
    
        Example:
            >>> from torchdata.datapipes.map import SequenceWrapper
            >>> source_dp = SequenceWrapper([(i, i + 10, i + 20) for i in range(3)])
            >>> dp1, dp2, dp3 = source_dp.unzip(sequence_length=3)
            >>> list(dp1)
            [0, 1, 2]
            >>> list(dp2)
            [10, 11, 12]
            >>> list(dp3)
            [20, 21, 22]
        """
    
