AKS Primality test (python code)

Agrawal-Kayal-Saxena primality test,

#!/usr/bin/python
import sys
import math
import time
from math import *
from sys import argv
#------------------------------------------
# --- greatest common divisor ---
def gcd(a,b):
while b: a, b = (b, a % b)
return a
#------------------------------------------
# --- a ^ n % p ---
def modular_exponentation(a, n, p):
ans = 1
while n:
if n & 1 : ans = ans * a % p
a = a * a % p
n >>= 1
return ans
#------------------------------------------
# --- a ^ n ---
def fast_exponentation(a, n):
ans = 1
while n:
if n & 1 : ans = ans * a
a = a * a
n >>= 1
return ans
#------------------------------------------
# Determines whether n is a power a ^ b, O(lg n (lg lg n) ^ 2)
def is_power(n):
if (- n & n) == n: return True # 2 ^ k
lgn = 1 + ( len( bin ( abs ( n ) ) ) - 2) # ceil
for b in range(2,lgn):
# b lg a = lg n
lowa = 1L
higha = 1L << (lgn / b + 1)
while lowa < higha - 1:
mida = (lowa + higha) >> 1
ab = fast_exponentation(mida,b)
if ab > n: higha = mida
elif ab < n: lowa = mida
else: return True # mida ^ b
return False
#------------------------------------------
def order_r(n):
logn2 = (log(n) / log(2)) ** 2 # float
r = int(logn2 + 1)
while True:
if gcd(n,r) == 1:
ans = 1
for k in range(1,r):
ans = ans * n % r
if ans == 1:
if k > logn2:
return r
else:
break
r = r + 1
#------------------------------------------
def polyMult(a, b, r, p):
res = [0]*r;
for i, u in enumerate(a):
for j, v in enumerate(b):
idx = (i+j) % r
res[idx] = (res[idx] + u * v ) % p
return res
#------------------------------------------
# (a + x)^p == (a + x^p)
def test(a, p, r):
ans = [1]
poly = [a, 1]
n = p
while n != 0:
if n & 1 :
ans = polyMult(ans,poly,r,p)
n >>= 1
poly = polyMult(poly,poly,r,p)
check = [0] * r
check[0], check[p%r] = a, 1
#print ans
#print check
return ans == check
#------------------------------------------
# Determines if n is prime using the AKS algorithm
def is_prime(n):
if is_power(n): return False
r = order_r(n)
for a in range(2, min(n,r+1)):
if n % a == 0: return False
#if gcd(a, n) != 1: return False
if r >= n: return True
ed = int(floor(sqrt(r) * log(n) / log(2)))
print "n = %d r = %d [a] = %d" % (n,r,ed)
for a in range(1, ed + 1):
if not test(a, n, r):
return False
return True
#------------------------------------------
def main():
if len(sys.argv) > 1:
n = int(sys.argv[1]);
res = is_prime(n);
print [n,res]
#------------------------------------------
if __name__ == "__main__":
#print('-----------');
main();
#print('-----------');