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('-----------');