from collections.abc import Callable, Iterable
from typing import TypeVar, overload

from remedapy.local_types import SupportsGtLt
from remedapy.util import LeaftistHeapNode, heap_maybe_insert, heapify

from .decorator import make_data_last

T = TypeVar('T')


@overload
def nth_by(it: Iterable[T], index: int, fn: Callable[[T], 'SupportsGtLt'], /) -> T: ...


@overload
def nth_by(index: int, fn: Callable[[T], 'SupportsGtLt'], /) -> Callable[[Iterable[T]], T]: ...


@make_data_last
def nth_by(iterable: Iterable[T], index: int, function: Callable[[T], 'SupportsGtLt'], /) -> T:
    """
    Returns the elements that would be n-th if the iterable were sorted by the function.

    They are yielded in the order they appear in the iterable.

    Parameters
    ----------
    iterable : Iterable[T]
        Input iterable (positional-only).
    index : int
        Index of the element to return (positional-only).
    function: Callable[[T], SupportsGtLt]
        Function to "sort" the iterable by (positional-only).

    Returns
    -------
    T
        The n-th element of the iterable if sorted by the function provided.

    See Also
    --------
    first_by
    take_first_by

    Examples
    --------
    Data first:
    >>> R.nth_by([2, 1, 4, 5, 3], 2, R.identity())
    3

    Data last:
    >>> R.pipe([2, 1, 4, 5, 3], R.nth_by(2, R.identity()))
    3

    """
    iterable_ = iter(iterable)
    original_heap: list[T] = []
    for _ in range(index + 1):
        original_heap.append(next(iterable_))
    heap = heapify(original_heap, function)
    if heap is None:  # pragma: no cover
        raise IndexError('nth_by: index out of range')
    for i, x in enumerate(iterable_, index + 1):
        heap = heap_maybe_insert(heap, LeaftistHeapNode(value=x, key=function(x), index=i))
    return heap.node.value
