#!/usr/bin/env python3
# encoding: utf-8

__author__ = "ChenyangGao <https://chenyanggao.github.io>"
__version__ = (0, 0, 3)
__all__ = ["SupportsLT", "grouped_mapping", "uniq", "dups", "iter_keyed_dups", "iter_dups"]

from collections.abc import Hashable, Iterable, Iterator, MutableMapping
from operator import itemgetter
from typing import (
    cast, overload, runtime_checkable, Any, Callable, Literal, Protocol, TypeVar, 
)


K = TypeVar("K", bound=Hashable)
V = TypeVar("V")


@runtime_checkable
class SupportsLT(Protocol):
    def __lt__(self, o, /) -> bool:
        ...


@overload
def grouped_mapping(
    it: Iterable[V], 
    /, 
    key: Callable[[V], K], 
    *, 
    mapping: None = None, 
) -> dict[K, list[V]]:
    ...
@overload
def grouped_mapping(
    it: Iterable[V], 
    /, 
    key: Callable[[V], K], 
    *, 
    mapping: MutableMapping[K, list[V]], 
) -> MutableMapping[K, list[V]]:
    ...
@overload
def grouped_mapping(
    it: Iterable[tuple[K, V]], 
    /, 
    key: None = None, 
    *, 
    mapping: None = None, 
) -> dict[K, list[V]]:
    ...
@overload
def grouped_mapping(
    it: Iterable[tuple[K, V]], 
    /, 
    key: None = None, 
    *, 
    mapping: MutableMapping[K, list[V]], 
) -> MutableMapping[K, list[V]]:
    ...
def grouped_mapping(
    it: Iterable[V] | Iterable[tuple[K, V]], 
    /, 
    key: None | Callable[[V], K] = None, 
    *, 
    mapping: None | MutableMapping[K, list[V]] = None, 
) -> MutableMapping[K, list[V]]:
    """Groups elements from an iterable into a mapping by a specified key or directly from tuples.

    This function supports two modes:
      1. When `it` is an iterable of values and a key function is provided, 
         it groups the values based on the keys generated by the function.
      2. When `it` is an iterable of tuples, it directly uses the first element of each tuple as the key.

    :param it: An iterable of values (V) or an iterable of tuples (K, V).
    :param key: An optional callable that extracts a key from each value.
    :param mapping: An optional mutable mapping to store the results. If None, a new `dict` will be created.

    :return: A mapping where each key corresponds to a list of values associated with that key.
    """
    if mapping is None:
        mapping = {}
    append = list.append
    if key is not None:
        it = cast(Iterable[V], it)
        it = ((key(v), v) for v in it)
    for k, v in cast(Iterable[tuple[K, V]], it):
        try:
            append(mapping[k], v)
        except KeyError:
            mapping[k] = [v]
    return mapping


@overload
def uniq(
    it: Iterable[V], 
    /, 
    key: Callable[[V], K], 
    keep_first: bool | Callable[[V], SupportsLT] = True, 
) -> dict[K, V]:
    ...
@overload
def uniq(
    it: Iterable[tuple[K, V]], 
    /, 
    key: None = None, 
    keep_first: bool | Callable[[V], SupportsLT] = True, 
) -> dict[K, V]:
    ...
def uniq(
    it: Iterable[V] | Iterable[tuple[K, V]], 
    /, 
    key: None | Callable[[V], K] = None, 
    keep_first: bool | Callable[[V], SupportsLT] = True, 
) -> dict[K, V]:
    """Returns a dictionary of unique items from an iterable.
    
    This function supports two modes:
      1. When `it` is an iterable of values and a key function is provided, 
         it determines uniqueness based on the keys generated by the function.
      2. When `it` is an iterable of tuples, it uses the first element of each tuple as the key.

    :param it: An iterable of values (V) or an iterable of tuples (K, V).
    :param key: An optional callable that extracts a key from each value.
    :param keep_first:
        - If a callable, uses its return value for comparison to determine which item to keep (always keep the first lowest).
        - If True, retains the first occurrence of each unique key.
        - If False, retains the last occurrence.

    :return: A dictionary where stores all the unique values (differentiated by their keys).
    """
    if key is not None:
        it = cast(Iterable[V], it)
        it = ((key(v), v) for v in it)
    d: dict[K, V] = {}
    setitem: Callable
    if callable(keep_first):
        cache: dict[K, SupportsLT] = {}
        def setitem(k: K, v: V, /):
            comparand = keep_first(v)
            if k in cache:
                prev_comparand = cache[k]
                if comparand < prev_comparand:
                    d[k] = v
                else:
                    return
            cache[k] = comparand
    elif keep_first:
        setitem = d.setdefault
    else:
        setitem = d.__setitem__
    for k, v in cast(Iterable[tuple[K, V]], it):
        setitem(k, v)
    return d


@overload
def dups(
    it: Iterable[V], 
    /, 
    key: Callable[[V], K], 
    keep_first: None | bool | Callable[[V], SupportsLT] = None, 
) -> dict[K, list[V]]:
    ...
@overload
def dups(
    it: Iterable[tuple[K, V]], 
    /, 
    key: None = None, 
    keep_first: None | bool | Callable[[V], SupportsLT] = None, 
) -> dict[K, list[V]]:
    ...
def dups(
    it: Iterable[V] | Iterable[tuple[K, V]], 
    /, 
    key: None | Callable[[V], K] = None, 
    keep_first: None | bool | Callable[[V], SupportsLT] = None, 
) -> dict[K, list[V]]:
    """Finds duplicates in the given iterable and returns them as a dictionary.

    :param it: An iterable of values (V) or an iterable of tuples (K, V).
    :param key: An optional callable that extracts a key from each value.
    :param keep_first:
        - If None, retains each occurrence if shared a common key.
        - If a callable, uses its return value for comparison to determine which item to keep (always keep the first lowest).
        - If True, retains the first occurrence of each unique key.
        - If False, retains the last occurrence.

    :return: A dictionary where keys are derived from the input values and values are lists of duplicates.
    """
    return grouped_mapping(iter_keyed_dups(it, key=key, keep_first=keep_first)) # type: ignore


@overload
def iter_keyed_dups(
    it: Iterable[V], 
    /, 
    key: Callable[[V], K], 
    keep_first: None | bool | Callable[[V], SupportsLT] = None, 
) -> Iterator[tuple[K, V]]:
    ...
@overload
def iter_keyed_dups(
    it: Iterable[tuple[K, V]], 
    /, 
    key: None = None, 
    keep_first: None | bool | Callable[[V], SupportsLT] = None, 
) -> Iterator[tuple[K, V]]:
    ...
def iter_keyed_dups(
    it: Iterable[V] | Iterable[tuple[K, V]], 
    /, 
    key: None | Callable[[V], K] = None, 
    keep_first: None | bool | Callable[[V], SupportsLT] = None, 
) -> Iterator[tuple[K, V]]:
    """Yields tuples of key and (duplicate) value pairs.

    :param it: An iterable of values (V) or an iterable of tuples (K, V).
    :param key: An optional callable that extracts a key from each value.
    :param keep_first:
        - If None, retains each occurrence if shared a common key.
        - If a callable, uses its return value for comparison to determine which item to keep (always keep the first lowest).
        - If True, retains the first occurrence of each unique key.
        - If False, retains the last occurrence.

    :yield: Iterator of tuples of keys and their corresponding duplicate values.
    """
    if key is not None:
        it = cast(Iterable[V], it)
        it = ((key(v), v) for v in it)
    if keep_first is None:
        d: dict[K, V] = {}
        pop = d.pop
        s: set[K] = set()
        add = s.add
        for k, v in cast(Iterable[tuple[K, V]], it):
            if k in s:
                yield k, v
            elif k in d:
                yield k, pop(k)
                yield k, v
                add(k)
            else:
                d[k] = v
    elif callable(keep_first):
        cache: dict[K, tuple[SupportsLT, V]] = {}
        for k, v in cast(Iterable[tuple[K, V]], it):
            comparand = keep_first(v)
            if k in cache:
                prev_comparand, prev_v = cache[k]
                if comparand < prev_comparand:
                    yield k, prev_v
                else:
                    yield k, v
                    continue
            cache[k] = (comparand, v)
    elif keep_first:
        s = set()
        add = s.add
        for k, v in cast(Iterable[tuple[K, V]], it):
            if k in s:
                yield k, v
            else:
                add(k)
    else:
        d = {}
        for k, v in cast(Iterable[tuple[K, V]], it):
            if k in d:
                yield k, d[k]
            d[k] = v


@overload
def iter_dups(
    it: Iterable[V], 
    /, 
    key: Callable[[V], Hashable], 
    keep_first: None | bool | Callable[[V], SupportsLT] = None, 
) -> Iterator[V]:
    ...
@overload
def iter_dups(
    it: Iterable[tuple[Hashable, V]], 
    /, 
    key: None = None, 
    keep_first: None | bool | Callable[[V], SupportsLT] = None, 
) -> Iterator[V]:
    ...
def iter_dups(
    it: Iterable[V] | Iterable[tuple[Hashable, V]], 
    /, 
    key: None | Callable[[V], Hashable] = None, 
    keep_first: None | bool | Callable[[V], SupportsLT] = None, 
) -> Iterator[V]:
    """Yields duplicate values from the provided iterable.
    
    :param it: An iterable of values (V) or an iterable of tuples (K, V).
    :param key: An optional callable that extracts a key from each value.
    :param keep_first:
        - If None, retains each occurrence if shared a common key.
        - If a callable, uses its return value for comparison to determine which item to keep (always keep the first lowest).
        - If True, retains the first occurrence of each unique key.
        - If False, retains the last occurrence.

    :yield: Iterator of duplicate values from the input iterable.
    """
    return map(itemgetter(1), iter_keyed_dups(it, key=key, keep_first=keep_first)) # type: ignore

