from math import gcd, log
from random import randint

def is_prime(N):
    """Returns if N is prime or not. Notice that this is not optimal,
    there is faster primality testing algorithms e.g. Miller-Rabin
    """
    
    if N == 2:
        return True  # only even prime
    if N % 2 == 0 or N <= 1:
        return False  # even numbers and 1 are not prime
    
    for i in range(3, N, 2):  # only try odd candidates
        if i*i > N:
            break  # we only need to check up to sqrt(N)
        if N % i == 0:
            return False  # found a factor
    
    return True

def find_power_k(N):
    """Returns the smallest k > 1 such that N**(1/k) is an integer,
    or 1 if there is no such k.
    """
    
    upper_bound = int(log(N)/log(3))
    
    for k in range(2, upper_bound + 1):
        p = int(N**(1/k))
        if p**k == N:
            return k
    
    return 1

def order_finding(a, N):
    """Returns the smallest r such that a^r = 1 mod N
    Notice that this is a naive classic implementation and is
    exponentially slower than the quantum version invented by Peter Shor.
    """
    
    i = 1
    a_i = a % N
    
    while a_i != 1:
        i += 1
        a_i = (a_i * a) % N
    
    return i

order_finding(2, 15)

def shor_algorithm(N):
    """Returns a pair of integers (P, Q) such that PQ = N for integer N"""
    
    if N % 2 == 0:  # even case
        return (N//2, 2)
    
    if is_prime(N):  # prime case
        return (N, 1)  # N is primes, factors cannot be found
    
    if find_power_k(N) > 1:  # prime power case
        P = int(N**(1/find_power_k(N)))  # we find a k such that N**(1/k) is an integer
        Q = N//P
        return (P, Q)
    
    # Now we can assume that the criteria for Shor's algorithm is met
    
    while True:
        # Non-deterministic, we will try multiple values for a
        a = randint(2, N-1)  # pick random a
        
        if gcd(a, N) != 1:  # Lucky case: a and N are not coprime!
            P = gcd(a, N)  # gcd yields a non-trivial factor
            Q = N//P
            return (P, Q)
        
        r = order_finding(a, N)  # quantum subroutine of the code
        
        if r % 2 == 0:
            continue

        P = gcd(a**(r//2) - 1, N)
        if P != 1:
            Q = N//P  # gcd yielded non trivial factor
            return (P, Q)

N = 10013
P, Q = shor_algorithm(N)
print(
    "Shor's algorithm found {} = {} x {} which is {}!".format(N, P, Q, ["incorrect", "correct"][P*Q==N])
)

S, T = shor_algorithm(Q)
print(
    "Shor's algorithm found {} = {} x {} which is {} again!".format(Q, S, T, ["incorrect", "correct"][S*T==Q])
)