import concurrent.futures
import copy
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 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 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())
