from functools import lru_cache
try:
    from lib_misc import get_item_counts
    from lib_misc import product
except ModuleNotFoundError:
    from .lib_misc import get_item_counts
    from .lib_misc import product


def prime_factors(n):
    """
    Returns a list of prime factors for n.

    :param n: number for which prime factors should be returned
    """
    # TODO: Look into using a prime wheel instead.
    factors = []
    rest = n
    divisor = 2
    while rest % divisor == 0:
        factors.append(divisor)
        rest //= divisor
    divisor = 3
    while divisor * divisor <= rest:
        while rest % divisor == 0:
            factors.append(divisor)
            rest //= divisor
        divisor += 2
    if rest != 1:
        factors.append(rest)
    return factors


def prime_factors_count(n):
    """
    Returns a dictionay of primes where each key is a prime
    and the value how many times that prime is part of the factor of n.

    :param n: numober for which prime factor counts are returned
    :returns: a dict where they key is a prime and the value
              the count of how often that value occurs

    """
    return get_item_counts(prime_factors(n))


@lru_cache(maxsize=10000)
def is_prime(n):
    """Returns True if n is prime and False otherwise.

    :param n: number to be checked

    """
    if n < 2:
        return False
    if n == 2 or n == 3:
        return True
    if n % 2 == 0:
        return False
    d = 3
    while d * d <= n:
        if n % d == 0:
            return False
        d += 2
    return True


def prime_nth(n):
    """Returns the nth prime number. The first number
    is 1 indexed, i.e. n = 1 -> 2, n = 2 -> 3, etc.

    :param n:

    """
    if n == 1:
        return 2
    if n == 2:
        return 3
    p = 5
    i = 3
    while i < n:
        p = p + 2
        if is_prime(p):
            i += 1
    return p


def primes(n_max):
    """
    Returns a list of all primes smaller n based
    on sieve of erasthones algorithm.

    If bitarray module is installed it will be used
    for the array, otherwise a regular Python list is
    used. The function should work for much larger
    numbers with bitarray.

    :param n_max:
    """
    try:
        import bitarray
        b = bitarray.bitarray(n_max)
        b.setall(True)
    except ModuleNotFoundError:
        b = [True for _ in range(n_max)]
    n = 1
    b[n - 1] = False
    while n * n <= n_max:
        if b[n - 1] is True:
            for i in range(n + n, n_max + 1, n):
                b[i - 1] = False
        n += 1
    ps = []
    for i in range(1, n_max + 1):
        if b[i - 1]:
            ps.append(i)
    return ps


def get_divisors_count(n):
    """
    Returns the number of divisors for n.
    The numbers 1 and n count as a divisor.

    >>> get_divisors_count(1)
    1
    >>> get_divisors_count(3)
    2 # 1, 3
    >>> get_divisors_count(4)
    3 # 1, 2, 4

    Getting the number of divisors is a combinatorial
    problem that can be solved by using the counts
    for each prime factor. For example, consider

    2 * 2 * 7 = 28

    We have 3 options for 2 (1, 1 * 2, 2 * 2)
    and 2 options for 7 (1, 1 * 7).

    By multiplying those options we get the number
    of combinations:

    2 * 3 = 6
    """
    if n == 1:
        return 1
    factors = prime_factors_count(n)
    count = product([v + 1 for v in factors.values()])
    return count