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

from remedapy.decorator import make_data_last

T = TypeVar('T')
U = TypeVar('U')


@overload
def count_by(data: Iterable[T], fn: Callable[[T], U], /) -> dict[U, int]: ...


@overload
def count_by(fn: Callable[[T], U], /) -> Callable[[Iterable[T]], dict[U, int]]: ...


@make_data_last
def count_by(iterable: Iterable[T], function: Callable[[T], U], /) -> dict[U, int]:
    """
    Counts the number of elements in the iterable provided that map to the kay when passed to the function provided.

    Parameters
    ----------
    iterable : Iterable[T]
        Input iterable (positional-only).
    function : Callable[[T], U]
        Function to apply to each element of the iterable (positional-only).

    Returns
    -------
    dict[U, int]
        Dictionary with keys being the result of passing elements of the iterable to the function
        and values being the count of elements that map to the key.

    Examples
    --------
    Data first:
    >>> R.count_by(['a', 'b', 'c', 'B', 'A', 'a'], R.to_lower_case())
    {'a': 3, 'b': 2, 'c': 1}

    Data last:
    >>> R.pipe(['a', 'b', 'c', 'B', 'A', 'a'], R.count_by(R.to_lower_case()))
    {'a': 3, 'b': 2, 'c': 1}

    """
    result: defaultdict[U, int] = defaultdict(int)
    for item in iterable:
        result[function(item)] += 1
    return dict(result)
