"""
Combinatorics utilities.

This module provides functions for permutations, combinations, and
related operations commonly used in assignment problems and data association.
"""

from typing import Iterator, List, Optional, Tuple
import itertools
from numpy.typing import ArrayLike
from functools import lru_cache


def factorial(n: int) -> int:
    """
    Compute factorial of n.

    Parameters
    ----------
    n : int
        Non-negative integer.

    Returns
    -------
    n! : int
        Factorial of n.

    Examples
    --------
    >>> factorial(5)
    120
    """
    if n < 0:
        raise ValueError("n must be non-negative")
    if n == 0 or n == 1:
        return 1
    result = 1
    for i in range(2, n + 1):
        result *= i
    return result


def n_choose_k(n: int, k: int) -> int:
    """
    Compute binomial coefficient C(n, k).

    Parameters
    ----------
    n : int
        Total number of items.
    k : int
        Number of items to choose.

    Returns
    -------
    C(n, k) : int
        Number of ways to choose k items from n.

    Examples
    --------
    >>> n_choose_k(5, 2)
    10
    """
    if k < 0 or k > n:
        return 0
    if k == 0 or k == n:
        return 1

    # Use symmetry to minimize iterations
    k = min(k, n - k)

    result = 1
    for i in range(k):
        result = result * (n - i) // (i + 1)
    return result


def n_permute_k(n: int, k: int) -> int:
    """
    Compute number of k-permutations of n items.

    Parameters
    ----------
    n : int
        Total number of items.
    k : int
        Number of items in each permutation.

    Returns
    -------
    P(n, k) : int
        Number of k-permutations: n! / (n-k)!

    Examples
    --------
    >>> n_permute_k(5, 2)
    20
    """
    if k < 0 or k > n:
        return 0
    result = 1
    for i in range(k):
        result *= n - i
    return result


def permutations(
    items: ArrayLike,
    k: Optional[int] = None,
) -> Iterator[Tuple]:
    """
    Generate all k-permutations of items.

    Parameters
    ----------
    items : array_like
        Items to permute.
    k : int, optional
        Length of permutations. Default is len(items).

    Yields
    ------
    perm : tuple
        Each k-permutation of items.

    Examples
    --------
    >>> list(permutations([1, 2, 3], 2))
    [(1, 2), (1, 3), (2, 1), (2, 3), (3, 1), (3, 2)]
    """
    items = list(items)
    return itertools.permutations(items, k)


def combinations(
    items: ArrayLike,
    k: int,
) -> Iterator[Tuple]:
    """
    Generate all k-combinations of items.

    Parameters
    ----------
    items : array_like
        Items to combine.
    k : int
        Size of each combination.

    Yields
    ------
    comb : tuple
        Each k-combination of items.

    Examples
    --------
    >>> list(combinations([1, 2, 3, 4], 2))
    [(1, 2), (1, 3), (1, 4), (2, 3), (2, 4), (3, 4)]
    """
    items = list(items)
    return itertools.combinations(items, k)


def combinations_with_replacement(
    items: ArrayLike,
    k: int,
) -> Iterator[Tuple]:
    """
    Generate all k-combinations with replacement.

    Parameters
    ----------
    items : array_like
        Items to combine.
    k : int
        Size of each combination.

    Yields
    ------
    comb : tuple
        Each k-combination with replacement.

    Examples
    --------
    >>> list(combinations_with_replacement([1, 2], 2))
    [(1, 1), (1, 2), (2, 2)]
    """
    items = list(items)
    return itertools.combinations_with_replacement(items, k)


def permutation_rank(perm: ArrayLike) -> int:
    """
    Compute the lexicographic rank of a permutation.

    The rank is the zero-based index of the permutation in the
    lexicographically sorted list of all permutations.

    Parameters
    ----------
    perm : array_like
        Permutation of integers 0, 1, ..., n-1.

    Returns
    -------
    rank : int
        Lexicographic rank (0-indexed).

    Examples
    --------
    >>> permutation_rank([0, 1, 2])  # First permutation
    0
    >>> permutation_rank([2, 1, 0])  # Last permutation
    5
    """
    perm = list(perm)
    n = len(perm)
    rank = 0
    available = list(range(n))

    for i in range(n):
        pos = available.index(perm[i])
        rank += pos * factorial(n - 1 - i)
        available.remove(perm[i])

    return rank


def permutation_unrank(rank: int, n: int) -> List[int]:
    """
    Compute the permutation with a given lexicographic rank.

    Parameters
    ----------
    rank : int
        Lexicographic rank (0-indexed).
    n : int
        Length of the permutation.

    Returns
    -------
    perm : list
        Permutation of [0, 1, ..., n-1] with the given rank.

    Examples
    --------
    >>> permutation_unrank(0, 3)
    [0, 1, 2]
    >>> permutation_unrank(5, 3)
    [2, 1, 0]
    """
    if rank < 0 or rank >= factorial(n):
        raise ValueError(f"Rank must be in [0, {factorial(n) - 1}]")

    available = list(range(n))
    perm = []

    for i in range(n):
        divisor = factorial(n - 1 - i)
        idx, rank = divmod(rank, divisor)
        perm.append(available.pop(idx))

    return perm


def next_permutation(perm: ArrayLike) -> Optional[List]:
    """
    Generate the next permutation in lexicographic order.

    Parameters
    ----------
    perm : array_like
        Current permutation.

    Returns
    -------
    next_perm : list or None
        Next permutation, or None if perm is the last permutation.

    Examples
    --------
    >>> next_permutation([1, 2, 3])
    [1, 3, 2]
    >>> next_permutation([3, 2, 1])  # Last permutation
    None
    """
    perm = list(perm)
    n = len(perm)

    # Find largest i such that perm[i] < perm[i+1]
    i = n - 2
    while i >= 0 and perm[i] >= perm[i + 1]:
        i -= 1

    if i < 0:
        return None  # Last permutation

    # Find largest j such that perm[i] < perm[j]
    j = n - 1
    while perm[i] >= perm[j]:
        j -= 1

    # Swap and reverse
    perm[i], perm[j] = perm[j], perm[i]
    perm[i + 1 :] = reversed(perm[i + 1 :])

    return perm


def partition_count(n: int, k: Optional[int] = None) -> int:
    """
    Count the number of integer partitions of n.

    A partition of n is a way of writing n as a sum of positive integers,
    where order doesn't matter.

    Parameters
    ----------
    n : int
        Number to partition.
    k : int, optional
        If specified, count only partitions with exactly k parts.

    Returns
    -------
    count : int
        Number of partitions.

    Examples
    --------
    >>> partition_count(5)  # 5 = 5 = 4+1 = 3+2 = 3+1+1 = 2+2+1 = 2+1+1+1 = 1+1+1+1+1
    7
    >>> partition_count(5, 2)  # 5 = 4+1 = 3+2
    2
    """

    @lru_cache(maxsize=None)
    def p(n: int, max_val: int) -> int:
        if n == 0:
            return 1
        if n < 0 or max_val == 0:
            return 0
        return p(n - max_val, max_val) + p(n, max_val - 1)

    @lru_cache(maxsize=None)
    def pk(n: int, k: int, max_val: int) -> int:
        if k == 0:
            return 1 if n == 0 else 0
        if n <= 0 or max_val == 0:
            return 0
        return pk(n - max_val, k - 1, max_val) + pk(n, k, max_val - 1)

    if k is None:
        return p(n, n)
    else:
        return pk(n, k, n)


def partitions(n: int, k: Optional[int] = None) -> Iterator[Tuple[int, ...]]:
    """
    Generate all integer partitions of n.

    Parameters
    ----------
    n : int
        Number to partition.
    k : int, optional
        If specified, generate only partitions with exactly k parts.

    Yields
    ------
    partition : tuple
        Each partition as a tuple of integers in descending order.

    Examples
    --------
    >>> list(partitions(4))
    [(4,), (3, 1), (2, 2), (2, 1, 1), (1, 1, 1, 1)]
    """

    def gen_partitions(
        n: int, max_val: int, prefix: Tuple[int, ...]
    ) -> Iterator[Tuple[int, ...]]:
        if n == 0:
            yield prefix
            return
        for i in range(min(n, max_val), 0, -1):
            yield from gen_partitions(n - i, i, prefix + (i,))

    def gen_partitions_k(
        n: int, k: int, max_val: int, prefix: Tuple[int, ...]
    ) -> Iterator[Tuple[int, ...]]:
        if k == 0:
            if n == 0:
                yield prefix
            return
        if n <= 0 or max_val == 0:
            return
        for i in range(min(n, max_val), 0, -1):
            yield from gen_partitions_k(n - i, k - 1, i, prefix + (i,))

    if k is None:
        yield from gen_partitions(n, n, ())
    else:
        yield from gen_partitions_k(n, k, n, ())


def multinomial_coefficient(*args: int) -> int:
    """
    Compute multinomial coefficient.

    multinomial(n1, n2, ..., nk) = (n1 + n2 + ... + nk)! / (n1! * n2! * ... * nk!)

    Parameters
    ----------
    *args : int
        Non-negative integers.

    Returns
    -------
    coeff : int
        Multinomial coefficient.

    Examples
    --------
    >>> multinomial_coefficient(2, 3, 1)  # 6! / (2! * 3! * 1!)
    60
    """
    n = sum(args)
    result = factorial(n)
    for k in args:
        result //= factorial(k)
    return result


def stirling_second(n: int, k: int) -> int:
    """
    Stirling number of the second kind.

    S(n, k) is the number of ways to partition n elements into k
    non-empty subsets.

    Parameters
    ----------
    n : int
        Number of elements.
    k : int
        Number of subsets.

    Returns
    -------
    S(n, k) : int
        Stirling number of the second kind.

    Examples
    --------
    >>> stirling_second(4, 2)  # {{1,2,3},{4}}, {{1,2,4},{3}}, ...
    7
    """
    if n == 0 and k == 0:
        return 1
    if n == 0 or k == 0:
        return 0
    if k > n:
        return 0

    @lru_cache(maxsize=None)
    def S(n: int, k: int) -> int:
        if n == k:
            return 1
        if k == 1:
            return 1
        return k * S(n - 1, k) + S(n - 1, k - 1)

    return S(n, k)


def bell_number(n: int) -> int:
    """
    Bell number B_n.

    B_n is the number of ways to partition a set of n elements.

    Parameters
    ----------
    n : int
        Number of elements.

    Returns
    -------
    B_n : int
        n-th Bell number.

    Examples
    --------
    >>> bell_number(4)
    15
    """
    return sum(stirling_second(n, k) for k in range(n + 1))


def catalan_number(n: int) -> int:
    """
    Catalan number C_n.

    Catalan numbers count many combinatorial structures including:
    - Valid parenthesizations
    - Full binary trees with n+1 leaves
    - Triangulations of a polygon with n+2 sides

    Parameters
    ----------
    n : int
        Non-negative integer.

    Returns
    -------
    C_n : int
        n-th Catalan number.

    Examples
    --------
    >>> catalan_number(5)
    42
    """
    return n_choose_k(2 * n, n) // (n + 1)


def derangements_count(n: int) -> int:
    """
    Count the number of derangements.

    A derangement is a permutation with no fixed points.

    Parameters
    ----------
    n : int
        Number of elements.

    Returns
    -------
    D_n : int
        Number of derangements.

    Examples
    --------
    >>> derangements_count(4)  # {2,1,4,3}, {2,3,4,1}, ...
    9
    """
    if n == 0:
        return 1
    if n == 1:
        return 0

    # Use recurrence: D(n) = (n-1) * (D(n-1) + D(n-2))
    d_prev2 = 1  # D(0)
    d_prev1 = 0  # D(1)

    for i in range(2, n + 1):
        d_curr = (i - 1) * (d_prev1 + d_prev2)
        d_prev2 = d_prev1
        d_prev1 = d_curr

    return d_prev1


def subfactorial(n: int) -> int:
    """
    Subfactorial (number of derangements).

    Alias for derangements_count.

    Parameters
    ----------
    n : int
        Number of elements.

    Returns
    -------
    !n : int
        Subfactorial of n.
    """
    return derangements_count(n)


__all__ = [
    "factorial",
    "n_choose_k",
    "n_permute_k",
    "permutations",
    "combinations",
    "combinations_with_replacement",
    "permutation_rank",
    "permutation_unrank",
    "next_permutation",
    "partition_count",
    "partitions",
    "multinomial_coefficient",
    "stirling_second",
    "bell_number",
    "catalan_number",
    "derangements_count",
    "subfactorial",
]
