"""
Math-related functions.
"""

from __future__ import absolute_import


def is_prime(n):
    """
    Determines if an inputted integer is prime or not.
    
    Arguments:
        n (int) => Any integer.
        
    Returns:
        True if n is prime, False otherwise (type=bool).
    """
    
    # No integer value less than 2 is prime.
    if n < 2: 
        return False
    
    # The only even prime number is 2.
    if n % 2 == 0: 
        return n == 2
    
    # If any odd integer between 3 and sqrt(n) evenly divides n,
    # n is not prime.
    for k in range(3, int(n**0.5) + 1, 2):
        if n % k == 0: 
            return False
        
    # If nothing evenly divides n, except 1 and itself,
    # it must be prime.
    return True


def prime_factors(n):
    """
    Computes the prime factorization of n.
    
    Arguments:
        n (int) => Any integer.
        
    Returns:
        A set of the prime factors of n or the empty set if no prime 
        factorization exists (type=set).
    """

    # Guard case, if no prime factorization exists.
    if n < 2: 
        raise ValueError("No prime factorization exists for integers less than 2.")

    # Initialize a set of prime factors to be populated.
    prime_factors = set()

    # Keep track of the original n.
    original_n = n
    
    # If n is even, 2 is a prime factor, then, divide n through 
    # by 2 as much as possible. If the resulting n is prime, 
    # then 2 and n are the only prime factors. Otherwise, keep going.
    if n % 2 == 0:
        prime_factors.add(2)
        while n % 2 == 0:
            n //= 2
        if is_prime(n):
            prime_factors.add(n)
            return prime_factors
    
    # We only want to iterate up to the minimum of either sqrt(original) or
    # the current n if it has been divided through by 2 several times.
    upper_bound = min(int(original_n**0.5) + 1, n)
    
    # Look through all odd integers greater than 3 for even prime divisors.
    # NOTE: Since we are iterating in increasing order, any evenly divisible
    # integer that we encounter must be prime.
    for i in range(3, upper_bound, 2):        
        if n % i == 0:
            prime_factors.add(i)
            while n % i == 0:
                n //= i
            if is_prime(n): 
                prime_factors.add(n)
                return prime_factors

    # If no divisors can be found, the number must be prime.
    prime_factors.add(n)

    return prime_factors


def largest_prime_factor(n):
    """
    Finds the largest prime factor of n.
    
    Arguments:
        n (int) => Any integer > 2.
        
    Returns:
        The largest prime factor of n or -1 if no prime 
        factorization exists (type=int).
    """

    # Guard case, if no prime factorization exists.
    if n < 2: 
        raise ValueError("n must be an integer greater than 2.")

    # Keep track of the original n.
    original_n = n
    
    # If n is even, 2 is a prime factor (largest so far), then, divide 
    # n through by 2 as much as possible. If the resulting n is prime, 
    # then n must be the largest prime factor. Otherwise, keep going.
    if n % 2 == 0:
        largest_prime_factor = 2
        while n % 2 == 0:
            n //= 2
        if is_prime(n):
            return n
    
    # We only want to iterate up to the minimum of either sqrt(original) or
    # the current n if it has been divided through by 2 several times.
    upper_bound = min(int(original_n**0.5) + 1, n)
    
    # Look through all odd integers greater than 3 for even prime divisors.
    # NOTE: Since we are iterating in increasing order, any evenly divisible
    # integer that we encounter must be prime.
    for i in range(3, upper_bound, 2):        
        if n % i == 0:
            largest_prime_factor = i
            while n % i == 0:
                n //= i
            if is_prime(n): 
                return n

    return largest_prime_factor


def gcd(a, b):
    """
    Calculates the greatest common divisor of a and b using Euclid's algorithm.

    Arguments:
        a (int) => An integer.
        b (int) => Another integer.

    Returns:
        The greatest common divisor of a and b (type=int).
    """

    while b != 0:
        a, b = b, a % b

    return a


def lcm(a, b):
    """
    Calculates the lowest common multiple of two integers a and b, using the
    "reduction by the greatest common divisor" approach.
    
    Arguments:
        a (int) => An integer.
        b (int) => Another integer.
        
    Returns:
        The lower common multiple of a and b (type=int).
    """
    
    return a * b // gcd(a, b)


def nth_prime(n):
    """
    Computes the nth prime number.

    Arguments:
        n (int) => A positive integer.

    Returns:
        The nth prime number (type=int).
    """

    # Guard case for bad input.
    if n < 1:
        raise ValueError("n must be a positive integer.")

    # If the 1st prime number is requested, return 2.
    if n == 1:
        return 2

    # Otherwise, iterate through odd integers and check if 
    # they are prime.
    prime_count = 1
    i = 1
    while prime_count < n:
        i += 2
        if is_prime(i): 
            prime_count += 1
        
    return i


def pythagorean_triplet(n):
    """
    Returns a Pythagorean triplet (a^2 + b^2 = c^2) for which 
    a + b + c = n.

    Arguments:
        n (int) => A positive integer.

    Returns:
        The Pythagorean triplet (a, b, c) for which a + b + c = n 
        or None if no triplet can be found (type=tuple).
    """

    # Guard case for bad input.
    if n < 1:
        raise ValueError("n must be a positive integer.")

    # Iterate through a in (1, n).
    for a in range(1, n):

        # Initialize b and c such that a < b < c.
        b = a + 1
        c = b + 1

        # Iterate through b and c terms until a triplet has been found
        # or c > n (which is not allowed).
        while c <= n:

            # Increase c until a^2 + b^2 <= c^2.
            while a * a + b * b > c * c:
                c += 1

            # If a^2 + b^2 = c^2 AND a + b + c = n, we've found a valid triplet!
            if a + b + c == n and c * c == a * a + b * b:
                return (a, b, c)

            b += 1


def sum_primes_below(n):
    """
    Computes the sum of all prime numbers below n using the sieve of 
    Eratosthenes algorithm.

    Arguments:
        n (int) => An integer greater than 2.

    Returns:
        The sum of all prime numbers below n.
    """

    # Initialize the sum to 0.
    sum = 0

    # Guard case for bad input.
    if n < 3:
        return sum

    # Create a seive for each positive integer below n.
    sieve = [True] * n

    # Iterate through the integers i, starting at 2, mark i as prime,
    # add it to the sum, and then mark every consecutive multiple of i 
    # as nonprime. The result is the sum of all primes below n.
    for i in range(2, n):
        if sieve[i]:
            sum += i
            for j in range(i + i, n, i):
                sieve[j] = False

    return sum


def factors(n):
    """
    Computes the factors of a positive integer n.

    Arguments:
        n (int) => A positive integer.

    Returns:
        A set of all the factors of n (type=set[int]).
    """

    # Guard case for bad input.
    if n < 1:
        raise ValueError("n must be a positive integer.")

    return set(reduce(list.__add__, [[i, n // i] for i in range(1, int(n**0.5) + 1) if n % i == 0]))


def collatz_sequence(n):
    """
    Builds a Collatz sequence starting from n.

    Arguments:
        n (int) => A positive integer to start the sequence from.

    Returns:
        A Collatz sequence starting from n (type=list).
    """

    # Guard case for bad input.
    if n < 1:
        raise ValueError("n must be a positive integer.")

    # Store the sequence, starting with n.
    sequence = [n]

    # Build the rest of the sequence iteratively.
    while n > 1:
        n = n // 2 if n % 2 == 0 else 3 * n + 1
        sequence.append(n)

    return sequence


def choose(n, k):
    """
    Computes the number of ways to choose k items from n
    total items (nCk).

    Arguments:
        n (int) => The total number items to choose from (a non-negative
                   integer).
        k (int) => The number of items to choose (a non-negative integer
                   less than or equal to n).
                   
    Returns:
        The number of ways to choose k items from n total
        items (nCk or "n choose k").
    """

    # Guard case for bad input.
    if k > n or n < 0 or k < 0:
        return 0

    # Store n - k to avoid repeated subtraction.
    diff = n - k
    
    # Build into n(n - 1)(n - 2) ... (n - k + 1).
    numerator = 1
    
    # Build into k(k - 1)(k - 2) ... 1.
    denominator = 1
    
    # Build up the numerator and denominator iteratively (following the 
    # mathematical formula for nCk).
    for i in range(1, k + 1):
        numerator *= diff + i
        denominator *= i
        
    return numerator // denominator

