import concurrent.futures
import copy
import itertools
import random
from typing import (
    Callable,
    Generic,
    Hashable,
    Iterator,
    List,
    Optional,
    Sequence,
    Set,
    Tuple,
    Type,
    TypeVar,
    Union,
)

A = TypeVar("A")
B = TypeVar("B")
CanHash = TypeVar("CanHash", bound=Hashable)


class LazyStream(Generic[A]):
    """
    Lazy Stream that only evaluates the final results when needed
    """

    def __init__(self, generator: Callable[[], A]):
        self._generator = generator

    def __getitem__(self, index: int) -> A:
        """
        Use evaluate() instead to get a list of values
        """
        raise NotImplementedError("LazyStream does not support indexing")

    def __len__(self) -> int:
        raise NotImplementedError("LazyStream does not support len()")

    def __safe_iter(self, limit: Optional[int] = None) -> Iterator[A]:
        i = 0
        while limit is None or i < limit:
            try:
                i += 1
                yield self._generator()
            except StopIteration:
                break

    def __safe_next(self) -> A:
        return next(self.__safe_iter())

    def __iter__(self) -> Iterator[A]:
        return self.__safe_iter()

    def __next__(self) -> A:
        return self.__safe_next()

    @staticmethod
    def from_lambda(func: Callable[[], A]) -> "LazyStream[A]":
        """
        Create a LazyStream from a lambda function
        Each element in the stream is generated by calling the function
        """
        return LazyStream(func)

    @staticmethod
    def from_iterator(iterator: Iterator[A]) -> "LazyStream[A]":
        """
        Create a LazyStream from an iterator object
        """
        return LazyStream(lambda: next(iterator))

    def evaluate(self, limit: Optional[int] = None) -> List[A]:
        """
        Evaluate the stream and return a list of values
        """
        output: List[A] = []
        for value in self.__safe_iter(limit):
            output.append(value)
        return output

    def par_evaluate(
        self, limit: int, executor: concurrent.futures.Executor
    ) -> List[A]:
        """
        Evaluate the stream in parallel
        """
        futures: List[concurrent.futures._base.Future[A]] = []
        while len(futures) < limit:
            futures.append(executor.submit(self._generator))
        output: List[A] = []
        for future in futures:
            try:
                output.append(future.result())
            except StopIteration:
                break
        return output

    def reduce(
        self, func: Callable[[B, A], B], accum: B, limit: Optional[int] = None
    ) -> B:
        """
        Evaluate the stream and reduce it to a single value
        """
        accum = copy.deepcopy(accum)
        for value in self.__safe_iter(limit):
            accum = func(accum, value)
        return accum

    def limit(self, n: int) -> "LazyStream[A]":
        """
        Limit the number of elements in the stream
        """

        return LazyStream.from_iterator(self.__safe_iter(n))

    def map(self, func: Callable[[A], B]) -> "LazyStream[B]":
        """
        Apply func to each element in the stream
        """
        return LazyStream(lambda: func(self.__safe_next()))

    def map_enumerate(self, func: Callable[[int, A], B]) -> "LazyStream[B]":
        """
        Apply func to each element in the stream with its index
        """

        def iterator() -> Iterator[B]:
            for i, value in enumerate(self.__safe_iter()):
                yield func(i, value)

        return LazyStream.from_iterator(iterator())

    def par_map(
        self, func: Callable[[A], B], executor: concurrent.futures.Executor
    ) -> "LazyStream[B]":
        """
        Apply func to each element in the stream in parallel
        """
        return LazyStream(lambda: executor.submit(func, self.__safe_next()).result())

    def flatten(self: "LazyStream[Sequence[A]]") -> "LazyStream[A]":
        """
        Flatten the stream of sequences
        """

        def iterator() -> Iterator[A]:
            for value in self.__safe_iter():
                for v in value:
                    yield v

        return LazyStream.from_iterator(iterator())

    def filter(self, predicate: Callable[[A], bool]) -> "LazyStream[A]":
        """
        Filter the stream by predicate
        """

        def iterator() -> Iterator[A]:
            for value in self.__safe_iter():
                if predicate(value):
                    yield value

        return LazyStream.from_iterator(iterator())

    def flatten_option(self: "LazyStream[Optional[A]]") -> "LazyStream[A]":
        """
        Flatten the stream of optional values
        """

        def iterator() -> Iterator[A]:
            for value in self.__safe_iter():
                if value is not None:
                    yield value

        return LazyStream.from_iterator(iterator())

    def catch(
        self, exception: Union[Type[Exception], Tuple[Type[Exception], ...]]
    ) -> "LazyStream[Optional[A]]":
        """
        Catch given exception(s) and return None
        """

        def func() -> Optional[A]:
            try:
                return self._generator()
            except exception:
                return None

        return LazyStream.from_lambda(func)

    def distinct(self: "LazyStream[CanHash]") -> "LazyStream[CanHash]":
        """
        Remove duplicate elements from the stream
        """
        seen: Set[CanHash] = set()
        return self.filter(lambda x: x not in seen and not seen.add(x))  # type: ignore[func-returns-value]

    def distinct_by(self, key: Callable[[A], CanHash]) -> "LazyStream[A]":
        """
        Remove duplicate elements from the stream by key
        """
        seen: Set[CanHash] = set()
        return (
            # Store the key and element to avoid calling key() twice
            self.map(lambda x: (key(x), x))
            .filter(lambda x: x[0] not in seen and not seen.add(x[0]))  # type: ignore[func-returns-value]
            .map(lambda x: x[1])
        )

    def for_each(self, func: Callable[[A], None]) -> "LazyStream[A]":
        """
        Apply func to each element in the stream without changing the stream
        """

        def iterator() -> Iterator[A]:
            for value in self.__safe_iter():
                func(value)
                yield value

        return LazyStream.from_iterator(iterator())

    def zip(self, other: "LazyStream[B]") -> "LazyStream[Tuple[A, B]]":
        """
        Zip two streams together
        Terminates when either stream terminates
        """
        return LazyStream.from_iterator(
            iter(zip(self.__safe_iter(), other.__safe_iter()))
        )

    def chain(self, other: "LazyStream[B]") -> "LazyStream[Union[A, B]]":
        """
        Chain two streams together, so that the second stream starts when the first one ends
        Note that the first stream should be finite for this function to make sense
        """
        return LazyStream.from_iterator(
            itertools.chain(self.__safe_iter(), other.__safe_iter())
        )

    def sample(self, p: float) -> "LazyStream[A]":
        """
        Sample the stream with probability p
        Returns empty stream if p <= 0 and full stream if p >= 1
        """

        def iterator() -> Iterator[A]:
            for value in self.__safe_iter():
                if random.random() < p:
                    yield value

        return LazyStream.from_iterator(iterator())
