import json
import os.path
from typing import List

import math
import sympy

from utils import get_root_dir

# Reading the primes file is slower than calculating the first primes. Where this cutoff point exactly lies depends
# on many factors, among which how big the primes file is. With a file with 1,000,000 primes, this cutoff point
# lies somewhere between 1,000 and 10,000.
CUTOFF_TO_LOOKUP = 1000

PRIMES_FILE = os.path.join(get_root_dir(), 'json', 'primes.json')


def is_prime(n: int) -> bool:
    """
    Return wether the input argument is a prime number
    """

    return sympy.isprime(n)


def get_primes(n: int) -> List[int]:
    """
    Return a list of the first n primes

    The results are cached on file, so the primes are only calculated once.
    """

    if n > CUTOFF_TO_LOOKUP and os.path.exists(PRIMES_FILE):
        with open(PRIMES_FILE, 'r') as f:
            known_primes = json.loads(f.read())
        if n <= len(known_primes):
            return known_primes[:n]
    else:
        known_primes = [1, 2, 3]  # The first primes are known

    num = max(known_primes)
    while len(known_primes) < n:
        # Check if num is divisible by any prime before it
        num_is_prime = True
        for p in known_primes[1:]:  # Skip the first prime, which is 1
            # If there is no remainder dividing the number then the number is not a prime
            if num % p == 0:
                # Break to stop testing more numbers, we know it's not a prime
                num_is_prime = False
                break

        # If it is a prime, then add it to the dictionary of known primes.
        if num_is_prime:
            known_primes.append(num)

        # Don't check even numbers
        num += 2

    if n > CUTOFF_TO_LOOKUP:
        with open(PRIMES_FILE, 'w') as f:
            f.write(json.dumps(known_primes))
    return known_primes


def nth_prime(n: int) -> int:
    """
    Return the n'th prime number

    The results are cached on file, so the primes are only calculated once.
    """

    return get_primes(n)[n - 1]


def prime_factorize(n: int) -> List[int]:
    """
    Return the list of divisors for the given integer

    >>> prime_factorize(15485857 * 15485863)
    [15485857, 15485863]
    """

    result = []
    for prime, multiplicity in sympy.factorint(n).items():
        result += [prime] * multiplicity
    return result


if __name__ == '__main__':
    print(nth_prime(10000))
    print(prime_factorize(15485857 * 15485863))
