# 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.

import warnings

from typing import Callable, Dict, Optional

from torch.utils._import_utils import dill_available

from torch.utils.data import IterDataPipe, MapDataPipe
from torch.utils.data.datapipes.utils.common import _check_unpickable_fn

if dill_available():
    import dill

    dill.extend(use_dill=False)


# @functional_datapipe("to_map_datapipe")  # This line must be kept for .pyi signature parser
class IterToMapConverterMapDataPipe(MapDataPipe):
    r"""
    Lazily load data from ``IterDataPipe`` to construct a ``MapDataPipe`` with
    the key-value pair generated by ``key_value_fn`` (functional name: ``to_map_datapipe``).
    If ``key_value_fn`` is not given, each data from the source IterDataPipe must itself be an iterable
    with exactly two objects. The first object of each item becomes a key in
    the new dictionary, and the second object the corresponding value.

    For the opposite converter, use :class:`.MapToIterConverter`.

    Args:
        datapipe: Source IterDataPipe
        key_value_fn: Function being applied over each data to generate key-value pair

    Note:
        If a key being added is already present, the corresponding value
        will be replaced by the new value.

    Example:
        >>> from torchdata.datapipes.iter import IterableWrapper
        >>> source_dp = IterableWrapper([(i, i) for i in range(10)])
        >>> map_dp = source_dp.to_map_datapipe()
        >>> list(map_dp)
        [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
        >>> source_dp2 = IterableWrapper([('a', 1), ('b', 2), ('c', 1)])
        >>> map_dp2 = source_dp2.to_map_datapipe()
        >>> map_dp2['a']
        1
        >>> def row_to_tuple(row):
        >>>     label = row[0]
        >>>     data = row[1:]
        >>>     return label, data
        >>> source_dp3 = IterableWrapper([('a', 1, 1, 1, 1, 1, 1), ('b', 2, 2, 2, 2, 2, 2), ('c', 3, 3, 3, 3, 3, 3)])
        >>> map_dp3 = source_dp3.to_map_datapipe(key_value_fn=row_to_tuple)
        >>> map_dp3['a']
        (1, 1, 1, 1, 1, 1)
    """
    datapipe: IterDataPipe
    key_value_fn: Optional[Callable]
    _map: Optional[Dict]
    _length: int

    def __init__(self, datapipe: IterDataPipe, key_value_fn: Optional[Callable] = None):
        if not isinstance(datapipe, IterDataPipe):
            raise TypeError(f"IterToMapConverter can only apply on IterDataPipe, but found {type(datapipe)}")
        self.datapipe = datapipe
        if key_value_fn is not None:
            _check_unpickable_fn(key_value_fn)
        self.key_value_fn = key_value_fn  # type: ignore[assignment]
        self._map = None

    def _load_map(self):
        self._map = {}
        for d in self.datapipe:
            inp = d if self.key_value_fn is None else self.key_value_fn(d)
            try:
                length = len(inp)
            except TypeError:
                raise TypeError(f"Cannot convert dictionary update element {type(inp)} ({inp}) to a sequence")
            if length != 2:
                raise ValueError(f"dictionary update sequence element has length {length}, 2 is required")
            key, value = inp
            if key in self._map:
                warnings.warn(f"Found duplicate key {key}. Please check your `key_value_fn`")
            self._map[key] = value

    def __getitem__(self, index):
        try:
            if self._map is None:
                self._load_map()
            return self._map[index]  # type: ignore[index]
        except KeyError:
            raise IndexError(f"Index {index} is invalid for IterToMapConverter.")

    def __len__(self):
        if self._map is not None:
            return len(self._map)  # type: ignore[arg-type]
        try:
            return len(self.datapipe)
        except (TypeError, NotImplementedError):
            pass
        warnings.warn(
            "Data from prior DataPipe are loaded to get length of"
            "IterToMapConverter before execution of the pipeline."
            "Please consider removing len()."
        )
        self._load_map()
        return len(self._map)  # type: ignore[arg-type]

    def __getstate__(self):
        if dill_available():
            dill_key_value_fn = dill.dumps(self.key_value_fn)
        else:
            dill_key_value_fn = self.key_value_fn
        return (
            self.datapipe,
            dill_key_value_fn,
            self._map,
        )

    def __setstate__(self, state):
        (self.datapipe, dill_key_value_fn, self._map) = state
        if dill_available():
            self.key_value_fn = dill.loads(dill_key_value_fn)  # type: ignore[assignment]
        else:
            self.key_value_fn = dill_key_value_fn  # type: ignore[assignment]


# Register for functional API
# See https://github.com/pytorch/data/issues/200
IterDataPipe.register_datapipe_as_function("to_map_datapipe", IterToMapConverterMapDataPipe)
