%html <h1><center>Алгоритм Агравала—Каяла—Саксены (AKS)
%html <h3> Test AKS
%html Если существует $r\in\mathbb{Z}$ такое, что $o_r(n)>\log^2 n$ и для любого $a$ от 1 до $\left\lfloor\sqrt{\varphi(r)}\log_2(n)\right\rfloor$ выполняется
%html $$ (x+a)^n\equiv(x^n+a)\pmod{x^r-1,\;n}, $$
%html то $n$ — либо простое число, либо степень простого числа.
%html $o_r(n)$ — показатель числа $n$ по модулю $r$ - такое $k$, что $n^k \equiv 1 \pmod n$
%html $\log_2$ — двоичный логарифм  1111
%html $\varphi(\cdot)$ — функция Эйлера
%html Сравнение по двум модулям вида $a(x)\equiv b(x)\pmod{h(x),\;n}$ для многочленов $a(x),\;b(x)\in\mathbb{Z}[x]$ означает, что существует $g(x)\in\mathbb{Z}[x]$ такой, что все коэффициенты многочлена $a(x)-b(x)-g(x)h(x)$ кратны $n$, где $\mathbb{Z}[x]$ — кольцо многочленов от $x$ над целыми числами.
%html Данное утверждение похоже на малую теорему Ферма, однако малую теорему Ферма можно проверить за $O(n)$, в то время как AKS можно проверить быстрее, выбрав правильное $r$.

import time

# Kostul, I hate this thing -- tell me how normal human beings should do it
def make_global(func):
    globals()[func.__name__] = func
make_global(make_global)

__global_time__ = 0
akslog_show = 0

def akslog_reset_time():
    global __global_time__
    __global_time__ = time.time()
make_global(akslog_reset_time)

def akslog(formatstr, *args):
    if 'akslog_show' not in globals() or akslog_show == 0:
        return

    global __global_time__
    print "%.2f" % (time.time() - __global_time__), formatstr.format(*args)
make_global(akslog)

akslog_reset_time()

%html Некоторые дополнительные функции для работы алгоритма.
%html Проверка того, что предоставленное число - полная степень числа

def check_comp_perfect_power(n):
    for b in range(2, log(n, 2) + 1):
        a = n ^ (1 / b)
        print a
        if a.is_integer():
            return True

    return False
make_global(check_comp_perfect_power)

print "31 полная степень числа -", check_comp_perfect_power(31)
print "25 полная степень числа -", check_comp_perfect_power(25)

%html Поиск $r$, такого что $o_r(n)>\log^2 n$ по определению
%html $o_r(n)$ — показатель числа $n$ по модулю $r$ - такое $k$, что $n^k \equiv 1 \pmod n$
%html При оценке сложности используется, что $r \leqslant max(3,\lceil\log^5 n\rceil)$, что может быть улучшено при помощи алгоритма Ленстры-Померанса до $\log^2 n$
def find_ord(n):
    r = 2
    log2n = int(log(n, 2) ^ 2)

    # ord is such k that foreach n^k = 1 mod r
    # we need k > log2(n)
    while (True):
        r += 1
        A = Integers(r)
        nr = A(n)

        is_ord = True
        for k in range(1, log2n + 1):
            # 1 means ord is done, 0 is not coprime
            if nr ** k == 1 or nr ** k == 0:
                is_ord = False
                break

        if is_ord:
            break

    return r
make_global(find_ord)

print find_ord(3571)

print "Подходящее r для числа 31 -", find_ord(1111111111111111111)
print "Подходящее r для числа 3571 -", find_ord(3571)
print "Подходящее r для числа 3469 -", find_ord(3469)
print "Подходящее r для числа 2 -", find_ord(2)
print "Подходящее r для числа -1 -", find_ord(-1)

%html Проверка наличия делителей у числа до заданного
def check_comp_gcd_foreach(n, r):
    for a in range(3, min(r, n-1) + 1, 2):
        if gcd(a, n) != 1:
            return True

    return False
make_global(check_comp_gcd_foreach)

print "Делители 31 до 29 -", check_comp_gcd_foreach(31, 29)
print "Делители 15 до 4 -", check_comp_gcd_foreach(15, 4)
print "Делители 15 до 2 -", check_comp_gcd_foreach(15, 2)

%html Проверка самого AKS теста до заданного значения
%html Расчет осуществляется при помощи SageMath кольца полиномов над $\mathbb{Z}_n$ с делителем $x^r-1$

def get_aks_poly(n, r, a):
    # Setup polynomial ring in Integers(n)
    R = PolynomialRing(Integers(n), 'x')
    x, = R.gens()

    divisor = x^r - 1

    # Define function that will calc remainder
    def rem(val):
        print(val)
        print(divisor)
        _, rem =  val.quo_rem(divisor)
        print(rem, "\n")
        return rem

    # Optimal power using method of divide and conquer - важно понять почему стандартное тормозит.
    def optpow(val, power):
        if power <= 8:
            return rem(val ^ power)
        else:
            if power % 2 == 0:
                powval = optpow(val, power/2)
                return rem(powval * powval)
            else:
                power -= 1
                powval = optpow(val, power/2)
                return rem(powval * powval * val)

    # Do it
    lhs = optpow(x + a, n)
    rhs = rem(x ^ n + a)

    return lhs - rhs

def check_comp_aks_test(n, r, maxval):
    for a in range(1, maxval + 1):
        poly = get_aks_poly(n, r, a)
        if poly:
            return True

    return False
make_global(check_comp_aks_test)

print "Проверка числа 31, r=29 до a=26, число составное -", check_comp_aks_test(31, 29, 26)

%html <h3>Алгоритм AKS</h3>
%html Ввод: целое число $n>1$.
%html Если $n=a^b$ для целых чисел $a>0$ и $b>1$, вернуть "составное".
%html Найдем наименьшее $r$, такое что $o_r(n)>(\log_2(n))^2$.
%html Если $1< gcd(a,n) < n $ для некоторого $a\leqslant r$, вернуть "составное".
%html Если $n\leqslant r$, вернуть "простое".
%html Если для всех $a$ от 1 до $\left\lfloor\sqrt{\varphi(r)}\log(n)\right\rfloor$ верно, что $(x+a)^n\equiv x^n+a\pmod{x^r-1,\;n}$, вернуть "простое".
%html Иначе вернуть "составное".
%html Сложность $\widetilde{O}(\log n^{10.5})$, Ленстра-Померанс $\widetilde{O}(\log n^6)$

def check_prime_aks(n):
    reset_time()

    if check_comp_perfect_power(n):
        akslog("AKS: Is Perfect Power")
        return False

    r = find_ord(n)
    akslog("AKS: Using r {0}", r)

    if check_comp_gcd_foreach(n, r):
        akslog("AKS: Divisor detected smaller than r {0}", r)
        return False

    if n <= r:
        akslog("AKS: {0} <= {1}", n, r)
        return True

    maxval = int(sqrt(euler_phi(r)) * log(n, 2))
    akslog("AKS: Maxval {0}", maxval)

    if check_comp_aks_test(n, r, maxval):
        akslog("AKS: AKS Test Failed")
        return False

    akslog("AKS: Prime")
    return True
make_global(check_prime_aks)

print "Число 11939 простое -", check_prime_aks(11939)

%html Тестирование для алгоритма AKS для первых 1000 чисел и сравнение с эталоном
count = 100

is_failed = False
for n in range(3, count, 2):
    answer = is_prime(n)
    aksAnswer = check_prime_aks(n)
    if answer != aksAnswer:
        is_failed = True
        print "Wrong answer for ", n

if is_failed:
    %html AKS выдал неверные ответы
else:
    %html Ошибок не найдет

data = []
P = Primes()
n = P.unrank(2000)

def nextprime(n, P):
    for i in range(20):
        n = P.next(n + 1)
    return n

for i in range(100):
    n = nextprime(n, P)
    timestart = time.time()
    result = check_prime_aks(n)
    timeend = time.time()

    t = timeend - timestart

    x = float(log(n))
    y = float(log(t))

    print result, x, y

    data.append((x, y))

# Fit data to linear model
_ = var('k, b, x')
model(x) = k * x + b

sol = find_fit(data, model)
f(x) = model(b=sol[0].rhs(), k=sol[1].rhs())

a = plot([])
a += plot(f(x), x, data[0][0], data[-1][0])
a += points(data, color='darkgreen', pointsize=50, aspect_ratio=1)
show(a)