"""https://github.com/dbrattli/Expression/blob/main/expression/collections/seq.py"""

from __future__ import annotations

import builtins
import functools
import itertools
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    Generator,
    Generic,
    Iterable,
    Iterator,
    overload,
)

from typing_extensions import Concatenate, ParamSpec, TypeVar, Unpack

SourceT = TypeVar("SourceT", infer_variance=True)
if TYPE_CHECKING:
    from kontainer import Maybe

    ResultT = TypeVar("ResultT", infer_variance=True)
    StateT = TypeVar("StateT", infer_variance=True)
    ParamT = ParamSpec("ParamT")
    _T1 = TypeVar("_T1", infer_variance=True)
    _T2 = TypeVar("_T2", infer_variance=True)
    _T3 = TypeVar("_T3", infer_variance=True)
    _T4 = TypeVar("_T4", infer_variance=True)

__all__ = [
    "curry",
    "curry2",
    "curry3",
    "append",
    "delay",
    "fold",
    "head",
    "mapi",
    "scan",
    "skip",
    "tail",
    "take",
    "unfold",
    "zip",
    "identity",
    "first",
    "second",
    "third",
    "fourth",
]


def curry(
    func: Callable[Concatenate[_T1, ParamT], ResultT]
) -> Callable[[_T1], Callable[ParamT, ResultT]]:
    def outer(arg1: _T1, /) -> Callable[ParamT, ResultT]:
        def inner(*args: ParamT.args, **kwargs: ParamT.kwargs) -> ResultT:
            return func(arg1, *args, **kwargs)

        return inner

    return outer


def curry2(
    func: Callable[Concatenate[_T1, _T2, ParamT], ResultT]
) -> Callable[[_T1], Callable[[_T2], Callable[ParamT, ResultT]]]:
    def outer(arg1: _T1, /) -> Callable[[_T2], Callable[ParamT, ResultT]]:
        def outer2(arg2: _T2, /) -> Callable[ParamT, ResultT]:
            def inner(*args: ParamT.args, **kwargs: ParamT.kwargs) -> ResultT:
                return func(arg1, arg2, *args, **kwargs)

            return inner

        return outer2

    return outer


def curry3(
    func: Callable[Concatenate[_T1, _T2, _T3, ParamT], ResultT]
) -> Callable[[_T1], Callable[[_T2], Callable[[_T3], Callable[ParamT, ResultT]]]]:
    def outer(
        arg1: _T1, /
    ) -> Callable[[_T2], Callable[[_T3], Callable[ParamT, ResultT]]]:
        def outer2(arg2: _T2, /) -> Callable[[_T3], Callable[ParamT, ResultT]]:
            def outer3(arg3: _T3, /) -> Callable[ParamT, ResultT]:
                def inner(*args: ParamT.args, **kwargs: ParamT.kwargs) -> ResultT:
                    return func(arg1, arg2, arg3, *args, **kwargs)

                return inner

            return outer3

        return outer2

    return outer


class SeqGen(Iterable[SourceT], Generic[SourceT]):
    """Sequence from a generator function.

    We use this to allow multiple iterations over the same sequence
    generated by a generator function.
    """

    def __init__(self, gen: Callable[[], Iterable[SourceT]]) -> None:
        self.gen = gen

    def __iter__(self) -> Iterator[SourceT]:
        xs = self.gen()
        return builtins.iter(xs)


class Infinite(Iterable[SourceT], Generic[SourceT]):
    """An infinite iterable.

    Where each iterator starts counting at 0.
    """

    def __init__(self, initializer: Callable[[int], SourceT]) -> None:
        self.initializer = initializer

    def __iter__(self) -> Iterator[SourceT]:
        return builtins.map(self.initializer, itertools.count(0, 1))


def append(
    *others: Iterable[SourceT]
) -> Callable[[Iterable[SourceT]], Iterable[SourceT]]:
    """Append sequence to other sequences.

    Wraps the given enumerations as a single concatenated enumeration.
    """
    return functools.partial(_append, others=others)


def delay(generator: Callable[[], Iterable[SourceT]]) -> Iterable[SourceT]:
    """Delay sequence.

    Returns a sequence that is built from the given delayed
    specification of a sequence.

    The input function is evaluated each time an Iterator for the
    sequence is requested.

    Args:
        generator: The generating function for the sequence.
    """
    return SeqGen(generator)


def fold(
    source: Iterable[SourceT],
    folder: Callable[[StateT, SourceT], StateT],
    state: StateT,
) -> StateT:
    """Fold elements in sequence.

    Applies a function to each element of the collection,
    threading an accumulator argument through the computation. If
    the input function is f and the elements are i0...iN then
    computes f (... (f s i0)...) iN.

    Args:
        source: The input sequence to fold.
        folder: A function that updates the state with each element
            from the sequence.
        state: The initial state.

    Returns:
        Partially applied fold function that takes a source sequence and
        returns the state object after the folding function is applied
        to each element of the sequence.
    """
    return functools.reduce(folder, source, state)


def head(source: Iterable[SourceT]) -> SourceT:
    """Return the first element of the sequence.

    Args:
        source: The input sequence.

    Returns:
        The first element of the sequence.

    Raises:
        Raises `ValueError` if the source sequence is empty.
    """
    for value in source:
        return value

    raise ValueError("Sequence contains no elements")


@overload
def init_infinite() -> Iterable[int]: ...


@overload
def init_infinite(initializer: Callable[[int], SourceT]) -> Iterable[SourceT]: ...


@overload
def init_infinite(initializer: None) -> Iterable[int]: ...


@overload
def init_infinite(
    initializer: Callable[[int], SourceT] | None
) -> Iterable[SourceT]: ...


def init_infinite(
    initializer: Callable[[int], SourceT] | None = None
) -> Iterable[SourceT] | Iterable[int]:
    """Generate infinite sequence.

    Generates a new sequence which, when iterated, will return
    successive elements by calling the given function. The results of
    calling the function will not be saved, that is the function will be
    reapplied as necessary to regenerate the elements. The function is
    passed the index of the item being generated.

    Iteration can continue up to `sys.maxsize`.
    """

    if initializer is None:
        return Infinite(_self)
    return Infinite(initializer)


@overload
def starmap(
    mapper: Callable[[_T1, _T2], ResultT]
) -> Callable[[Iterable[tuple[_T1, _T2]]], Iterable[ResultT]]: ...


@overload
def starmap(
    mapper: Callable[[_T1, _T2, _T3], ResultT]
) -> Callable[[Iterable[tuple[_T1, _T2, _T3]]], Iterable[ResultT]]: ...


@overload
def starmap(
    mapper: Callable[[_T1, _T2, _T3, _T4], ResultT]
) -> Callable[[Iterable[tuple[_T1, _T2, _T3, _T4]]], Iterable[ResultT]]: ...


def starmap(
    mapper: Callable[..., ResultT]
) -> Callable[[Iterable[Any]], Iterable[ResultT]]:
    """Starmap source sequence.

    Unpack arguments grouped as tuple elements. Builds a new collection
    whose elements are the results of applying the given function to the
    unpacked arguments to each of the elements of the collection.

    Args:
        mapper: A function to transform items from the input sequence.

    Returns:
        Partially applied map function.
    """

    return functools.partial(_starmap, func=mapper)


def map2(
    mapper: Callable[[_T1, _T2], ResultT]
) -> Callable[[Iterable[tuple[_T1, _T2]]], Iterable[ResultT]]:
    return starmap(mapper)


def map3(
    mapper: Callable[[_T1, _T2, _T3], ResultT]
) -> Callable[[Iterable[tuple[_T1, _T2, _T3]]], Iterable[ResultT]]:
    return starmap(mapper)


def mapi(
    source: Iterable[SourceT], mapping: Callable[[int, SourceT], ResultT]
) -> Iterable[ResultT]:
    """Map list with index.

    Builds a new collection whose elements are the results of
    applying the given function to each of the elements of the
    collection. The integer index passed to the function indicates
    the index (from 0) of element being transformed.

    Args:
        source: The input sequence to to map.
        mapping: The function to transform elements and their
            indices.

    Returns:
        The list of transformed elements.
    """
    return itertools.starmap(mapping, builtins.enumerate(source))


def scan(
    source: Iterable[SourceT],
    scanner: Callable[[StateT, SourceT], StateT],
    state: StateT,
) -> Iterable[StateT]:
    """Scan elements in sequence.

    Like fold, but computes on-demand and returns the sequence of
    intermediary and final results.

    Args:
        source: The input sequence.
        scanner: A function that updates the state with each element
        state: The initial state.
    """
    return itertools.accumulate(source, scanner, initial=state)


def skip(source: Iterable[SourceT], count: int) -> Iterable[SourceT]:
    """Skip elements from sequence.

    Returns a sequence that skips N elements of the underlying
    sequence and then yields the remaining elements of the sequence.

    Args:
        source: The input sequence.
        count: The number of items to skip.
    """

    if count <= 0:
        return source

    return SeqGen(functools.partial(_gen_skip, source, count))


def tail(source: Iterable[SourceT]) -> Iterable[SourceT]:
    """Return tail of sequence.

    Returns a sequence that skips 1 element of the underlying sequence
    and then yields the remaining elements of the sequence.
    """
    return skip(source, 1)


def take(source: Iterable[SourceT], count: int) -> Iterable[SourceT]:
    """Returns the first N elements of the sequence.

    Args:
        source: The source sequence.
        count: The number of items to take.

    Returns:
        The result sequence.
    """

    if count <= 0:
        return ()

    return SeqGen(functools.partial(_gen_take, source, count))


def unfold(
    state: StateT, generator: Callable[[StateT], Maybe[tuple[SourceT, StateT]]]
) -> Iterable[SourceT]:
    """Unfold sequence.

    Generates a list that contains the elements generated by the given
    computation. The given initial state argument is passed to the
    element generator.

    Args:
        state: The initial state.
        generator: A function that takes in the current state and
            returns an option tuple of the next element of the list
            and the next state value.

    Returns:
        A partially applied unfold function that takes the state and
        returns the result list.
    """
    while True:
        result = generator(state)
        value = result.unwrap()
        if value is None:
            break

        item, state = value
        yield item


def zip(  # noqa: A001
    source1: Iterable[SourceT]
) -> Callable[[Iterable[ResultT]], Iterable[tuple[SourceT, ResultT]]]:
    """Zip sequence with other.

    Combines the two sequences into a list of pairs. The two
    sequences need not have equal lengths: when one sequence is
    exhausted any remaining elements in the other sequence are
    ignored.

    Args:
        source1: The first input sequence.

    Returns:
        Partially applied zip function.
    """

    return functools.partial(_zip, source1)


def identity(value: SourceT) -> SourceT:
    return value


def first(values: tuple[SourceT, Unpack[tuple[Any, ...]]]) -> SourceT:
    return values[0]


def second(values: tuple[Any, SourceT, Unpack[tuple[Any, ...]]]) -> SourceT:
    return values[1]


def third(values: tuple[Any, Any, SourceT, Unpack[tuple[Any, ...]]]) -> SourceT:
    return values[2]


def fourth(values: tuple[Any, Any, Any, SourceT, Unpack[tuple[Any, ...]]]) -> SourceT:
    return values[3]


def _starmap(
    args: Iterable[tuple[Any, ...]], func: Callable[[Unpack[tuple[Any, ...]]], Any]
) -> Iterable[Any]:
    for arg in args:
        yield func(*arg)


def _self(x: SourceT) -> SourceT:
    return x


def _append(
    source: Iterable[SourceT], /, others: Iterable[Iterable[SourceT]]
) -> Iterable[SourceT]:
    return _concat(source, *others)


def _concat(*iterables: Iterable[SourceT]) -> Iterable[SourceT]:
    return SeqGen(functools.partial(_gen_concat, iterables))


def _gen_concat(iterables: Iterable[Iterable[SourceT]]) -> Generator[SourceT, Any, Any]:
    for it in iterables:
        yield from it


def _gen_skip(source: Iterable[SourceT], count: int) -> Generator[SourceT, Any, Any]:
    for i, n in enumerate(source):
        if i >= count:
            yield n


def _gen_take(source: Iterable[SourceT], count: int) -> Generator[SourceT, Any, Any]:
    for i, n in enumerate(source):
        yield n

        if i == count - 1:
            break


def _zip(
    source1: Iterable[SourceT], source2: Iterable[ResultT]
) -> Iterable[tuple[SourceT, ResultT]]:
    return builtins.zip(source1, source2)
