from contextlib import contextmanager
import collections
import random
import threading
import time
from typing import TypeVar, Generic, Iterable, List, Callable, Any

import ray
from ray.util.iter_metrics import MetricsContext, SharedMetrics

# The type of an iterator element.
T = TypeVar("T")
U = TypeVar("U")


def from_items(
    items: List[T], num_shards: int = 2, repeat: bool = False
) -> "ParallelIterator[T]":
    """Create a parallel iterator from an existing set of objects.

    The objects will be divided round-robin among the number of shards.

    Args:
        items (list): The list of items to iterate over.
        num_shards (int): The number of worker actors to create.
        repeat (bool): Whether to cycle over the items forever.
    """
    shards = [[] for _ in range(num_shards)]
    for i, item in enumerate(items):
        shards[i % num_shards].append(item)
    name = "from_items[{}, {}, shards={}{}]".format(
        items and type(items[0]).__name__ or "None",
        len(items),
        num_shards,
        ", repeat=True" if repeat else "",
    )
    return from_iterators(shards, repeat=repeat, name=name)


def from_range(
    n: int, num_shards: int = 2, repeat: bool = False
) -> "ParallelIterator[int]":
    """Create a parallel iterator over the range 0..n.

    The range will be partitioned sequentially among the number of shards.

    Args:
        n (int): The max end of the range of numbers.
        num_shards (int): The number of worker actors to create.
        repeat (bool): Whether to cycle over the range forever.
    """
    generators = []
    shard_size = n // num_shards
    for i in range(num_shards):
        start = i * shard_size
        if i == num_shards - 1:
            end = n
        else:
            end = (i + 1) * shard_size
        generators.append(range(start, end))
    name = (
        f"from_range[{n}, shards={num_shards}" f"{', repeat=True' if repeat else ''}]"
    )
    return from_iterators(
        generators,
        repeat=repeat,
        name=name,
    )


def from_iterators(
    generators: List[Iterable[T]], repeat: bool = False, name=None
) -> "ParallelIterator[T]":
    """Create a parallel iterator from a list of iterables.
    An iterable can be a conatiner (list, str, tuple, set, etc.),
    a generator, or a custom class that implements __iter__ or __getitem__.

    An actor will be created for each iterable.

    Examples:
        >>> # Create using a list of generators.
        >>> from_iterators([range(100), range(100)])

        >>> # Certain generators are not serializable.
        >>> from_iterators([(x for x in range(100))])
        ... TypeError: can't pickle generator objects

        >>> # So use lambda functions instead.
        >>> # Lambda functions are serializable.
        >>> from_iterators([lambda: (x for x in range(100))])

    Args:
        generators (list): A list of Python iterables or lambda
            functions that produce an iterable when called. We allow lambda
            functions since certain generators might not be serializable,
            but a lambda that returns it can be.
        repeat (bool): Whether to cycle over the iterators forever.
        name (str): Optional name to give the iterator.
    """
    worker_cls = ray.remote(ParallelIteratorWorker)
    actors = [worker_cls.remote(g, repeat) for g in generators]
    if not name:
        name = "from_iterators[shards={}{}]".format(
            len(generators), ", repeat=True" if repeat else ""
        )
    return from_actors(actors, name=name)


def from_actors(
    actors: List["ray.actor.ActorHandle"], name=None
) -> "ParallelIterator[T]":
    """Create a parallel iterator from an existing set of actors.

    Each actor must subclass the ParallelIteratorWorker interface.

    Args:
        actors (list): List of actors that each implement
            ParallelIteratorWorker.
        name (str): Optional name to give the iterator.
    """
    if not name:
        name = f"from_actors[shards={len(actors)}]"
    return ParallelIterator([_ActorSet(actors, [])], name, parent_iterators=[])


class ParallelIterator(Generic[T]):
    """A parallel iterator over a set of remote actors.

    This can be used to iterate over a fixed set of task results
    (like an actor pool), or a stream of data (e.g., a fixed range of numbers,
    an infinite stream of RLlib rollout results).

    This class is **serializable** and can be passed to other remote
    tasks and actors. However, each shard should be read from at most one
    process at a time.

    Examples:
        >>> # Applying a function over items in parallel.
        >>> it = ray.util.iter.from_items([1, 2, 3], num_shards=2)
        ... <__main__.ParallelIterator object>
        >>> it = it.for_each(lambda x: x * 2).gather_sync()
        ... <__main__.LocalIterator object>
        >>> print(list(it))
        ... [2, 4, 6]

        >>> # Creating from generators.
        >>> it = ray.util.iter.from_iterators([range(3), range(3)])
        ... <__main__.ParallelIterator object>
        >>> print(list(it.gather_sync()))
        ... [0, 0, 1, 1, 2, 2]

        >>> # Accessing the individual shards of an iterator.
        >>> it = ray.util.iter.from_range(10, num_shards=2)
        ... <__main__.ParallelIterator object>
        >>> it0 = it.get_shard(0)
        ... <__main__.LocalIterator object>
        >>> print(list(it0))
        ... [0, 1, 2, 3, 4]
        >>> it1 = it.get_shard(1)
        ... <__main__.LocalIterator object>
        >>> print(list(it1))
        ... [5, 6, 7, 8, 9]

        >>> # Gathering results from actors synchronously in parallel.
        >>> it = ray.util.iter.from_actors(workers)
        ... <__main__.ParallelIterator object>
        >>> it = it.batch_across_shards()
        ... <__main__.LocalIterator object>
        >>> print(next(it))
        ... [worker_1_result_1, worker_2_result_1]
        >>> print(next(it))
        ... [worker_1_result_2, worker_2_result_2]
    """

    def __init__(
        self,
        actor_sets: List["_ActorSet"],
        name: str,
        parent_iterators: List["ParallelIterator[Any]"],
    ):
        """Create a parallel iterator (this is an internal function)."""

        # We track multiple sets of actors to support parallel .union().
        self.actor_sets = actor_sets
        self.name = name

        # keep explicit reference to parent iterator for repartition
        self.parent_iterators = parent_iterators

    def __iter__(self):
        raise TypeError(
            "You must use it.gather_sync() or it.gather_async() to "
            "iterate over the results of a ParallelIterator."
        )

    def __str__(self):
        return repr(self)

    def __repr__(self):
        return f"ParallelIterator[{self.name}]"

    def _with_transform(self, local_it_fn, name):
        """Helper function to create new Parallel Iterator"""
        return ParallelIterator(
            [a.with_transform(local_it_fn) for a in self.actor_sets],
            name=self.name + name,
            parent_iterators=self.parent_iterators,
        )

    def transform(
        self, fn: Callable[[Iterable[T]], Iterable[U]]
    ) -> "ParallelIterator[U]":
        """Remotely transform the iterator.

        This is advanced version of for_each that allows you to apply arbitrary
        generator transformations over the iterator. Prefer to use .for_each()
        when possible for simplicity.

        Args:
            fn (func): function to use to transform the iterator. The function
                should pass through instances of _NextValueNotReady that appear
                in its input iterator. Note that this function is only called
                **once** over the input iterator.

        Returns:
            ParallelIterator[U]: a parallel iterator.

        Examples:
            >>> def f(it):
            ...     for x in it:
            ...         if x % 2 == 0:
            ...            yield x
            >>> from_range(10, 1).transform(f).gather_sync().take(5)
            ... [0, 2, 4, 6, 8]
        """
        return self._with_transform(
            lambda local_it: local_it.transform(fn), ".transform()"
        )

    def for_each(
        self, fn: Callable[[T], U], max_concurrency=1, resources=None
    ) -> "ParallelIterator[U]":
        """Remotely apply fn to each item in this iterator.

        If `max_concurrency` == 1 then `fn` will be executed serially by each
        shards

        `max_concurrency` should be used to achieve a high degree of
        parallelism without the overhead of increasing the number of shards
        (which are actor based). If `max_concurrency` is not 1, this function
        provides no semantic guarantees on the output order.
        Results will be returned as soon as they are ready.

        A performance note: When executing concurrently, this function
        maintains its own internal buffer. If `num_async` is `n` and
        max_concur is `k` then the total number of buffered objects could be up
        to `n + k - 1`

        Args:
            fn (func): function to apply to each item.
            max_concurrency (int): max number of concurrent calls to fn per
                shard. If 0, then apply all operations concurrently.
            resources (dict): resources that the function requires to execute.
                This has the same default as `ray.remote` and is only used
                when `max_concurrency > 1`.

        Returns:
            ParallelIterator[U]: a parallel iterator whose elements have `fn`
            applied.

        Examples:
            >>> next(from_range(4).for_each(
                        lambda x: x * 2,
                        max_concur=2,
                        resources={"num_cpus": 0.1}).gather_sync()
                )
            ... [0, 2, 4, 8]

        """
        assert max_concurrency >= 0, "max_concurrency must be non-negative."
        return self._with_transform(
            lambda local_it: local_it.for_each(fn, max_concurrency, resources),
            ".for_each()",
        )

    def filter(self, fn: Callable[[T], bool]) -> "ParallelIterator[T]":
        """Remotely filter items from this iterator.

        Args:
            fn (func): returns False for items to drop from the iterator.

        Examples:
            >>> it = from_items([0, 1, 2]).filter(lambda x: x > 0)
            >>> next(it.gather_sync())
            ... [1, 2]
        """
        return self._with_transform(lambda local_it: local_it.filter(fn), ".filter()")

    def batch(self, n: int) -> "ParallelIterator[List[T]]":
        """Remotely batch together items in this iterator.

        Args:
            n (int): Number of items to batch together.

        Examples:
            >>> next(from_range(10, 1).batch(4).gather_sync())
            ... [0, 1, 2, 3]
        """
        return self._with_transform(lambda local_it: local_it.batch(n), f".batch({n})")

    def flatten(self) -> "ParallelIterator[T[0]]":
        """Flatten batches of items into individual items.

        Examples:
            >>> next(from_range(10, 1).batch(4).flatten())
            ... 0
        """
        return self._with_transform(lambda local_it: local_it.flatten(), ".flatten()")

    def combine(self, fn: Callable[[T], List[U]]) -> "ParallelIterator[U]":
        """Transform and then combine items horizontally.

        This is the equivalent of for_each(fn).flatten() (flat map).
        """
        it = self.for_each(fn).flatten()
        it.name = self.name + ".combine()"
        return it

    def local_shuffle(
        self, shuffle_buffer_size: int, seed: int = None
    ) -> "ParallelIterator[T]":
        """Remotely shuffle items of each shard independently

        Args:
            shuffle_buffer_size (int): The algorithm fills a buffer with
                shuffle_buffer_size elements and randomly samples elements from
                this buffer, replacing the selected elements with new elements.
                For perfect shuffling, this argument should be greater than or
                equal to the largest iterator size.
            seed (int): Seed to use for
                randomness. Default value is None.

        Returns:
            A ParallelIterator with a local shuffle applied on the base
            iterator

        Examples:
            >>> it = from_range(10, 1).local_shuffle(shuffle_buffer_size=2)
            >>> it = it.gather_sync()
            >>> next(it)
            0
            >>> next(it)
            2
            >>> next(it)
            3
            >>> next(it)
            1
        """
        return self._with_transform(
            lambda local_it: local_it.shuffle(shuffle_buffer_size, seed),
            ".local_shuffle(shuffle_buffer_size={}, seed={})".format(
                shuffle_buffer_size, str(seed) if seed is not None else "None"
            ),
        )

    def repartition(
        self, num_partitions: int, batch_ms: int = 0
    ) -> "ParallelIterator[T]":
        """Returns a new ParallelIterator instance with num_partitions shards.

        The new iterator contains the same data in this instance except with
        num_partitions shards. The data is split in round-robin fashion for
        the new ParallelIterator.

        Args:
            num_partitions (int): The number of shards to use for the new
                ParallelIterator