'''
Terminology:
    Although "divisor" and "factor" mean the same thing.
    When Celeste discusses "divisors of n" it is implied to
    mean "proper divisors of n + n itself", and "factors" are
    the "prime proper divisors of n". 
'''

from celeste.extern.primefac import primefac

def factors(n: int) -> int:
    pfactors: list[tuple[int, int]] = []
    # generate primes and progressively store them in pfactors
    pfgen = primefac(n)
    watching = next(pfgen)
    mult = 1
    # ASSUMPTION: prime generation is (non-strict) monotone increasing
    while True:
        p = next(pfgen, None)
        if p == watching:
            mult += 1
        else:
            pfactors.append((watching, mult))
            watching = p # reset
            mult = 1     # reset 
        if p is None: 
            break
    return pfactors

def factors2divisors(pfactors: list[tuple[int, int]], 
                     sorted: bool = True) -> list[int]:
    '''
    Generates all divisors < n of an integer n given its prime factorisation.
    Input: prime factorisation of n (excluding 1 and n, and duplicates)
           in the typical form: list[(prime, multiplicity)]
    ''' 
    divisors = [1]
    for (prime, multiplicity) in pfactors:
        extension = []
        for i in range(1, multiplicity+1):
            term = prime**i
            extension.extend(list([divisor*term for divisor in divisors]))
        divisors.extend(extension)
    if sorted: divisors.sort()
    return divisors

def factors2aliquots(pfactors: list[tuple[int, int]]) -> list[int]:
    return factors2divisors(pfactors)[:-1]
    
# "aliquots(n)" is an alias for "divisors(n)"
def aliquots(n: int) -> int:
    '''
    Returns all aliquot parts (proper divisors) of
    an integer n, that is all divisors 0 < d <= n. 
    '''
    return factors2aliquots(factors(n))

def divisors(n: int) -> int:
    '''
    Returns all divisors 0 < d < n of an integer n.
    '''
    return factors2divisors(factors(n))

def aliquot_sum(n: int) -> int:
    return sum(aliquots(n))


def littleomega(n: int) -> int:
    '''
    The Little Omega function counts the number of 
    distinct prime factors of an integer n.
    Ref: https://en.wikipedia.org/wiki/Prime_omega_function
    '''
    return len(factors(n))

def bigomega(n: int) -> int:
    '''
    The Big Omega function counts the total number of
    prime factors (including multiplicity) of an integer n. 
    Ref: https://en.wikipedia.org/wiki/Prime_omega_function
    '''
    return sum(factor[1] for factor in factors(n))