APR-CL素数判定法 - wacchoz’s note

APR-CL素数判定法(Adleman–Pomerance–Rumely-Cohen-Lenstra primality test, APR-CL primality test)について概説＆実装してみる。
この素数判定法はAdleman–Pomerance–Rumelyらによって開発された方法で、のちにCohenとLenstraによって改良されています。
数100桁程度の数に対しても現実的な時間で素数判定ができるすごいやつです。
ただし使う数学の道具も高度で、場合分けも多く、実装は重めです。確率的に「ほぼほぼ素数」と判定すればいいだけなら、Miller–Rabin法などを使う方が実装は楽です。

参考文献
Dirichlet指標とGauss和
Gauss和を用いた素数判定の原理
Jacobi和への書き換え
素数判定アルゴリズム
- 事前計算
- 素数判定
Python3による実装

参考文献

[APR] Adleman, Leonard M.; Pomerance, Carl; Rumely, Robert S. (1983). "On distinguishing prime numbers from composite numbers". Annals of Mathematics. 117 (1): 173–206.
[Coh-Len1] Cohen, Henri; Lenstra, Hendrik W., Jr. (1984). "Primality testing and Jacobi sums". Mathematics of Computation. 42 (165): 297–330.
[Coh-Len2] H. Cohen and A. K. Lenstra, Implementation of a new primality test. Math. Comp. 48 (1987), 103-121.
[和田] 和田秀男, コンピュータと素因子分解, 遊星社, 1987
[Cohen] H. Cohen, "A course in computational algebraic number theory". Springer-Verlag, Berlin, 1993
[Chmie] A Chmielowiec, "Primality proving with Gauss and Jacobi sums", 2004.

[APR]はおそらく読まなくても大丈夫。
[Coh-Len1]が[APR]からの改良です。
[Coh-Len2]も実装の上で有用なことが書いています。
和書でAPR法の解説をしているのは[和田]しか見たことがありません。ただ実装できるレベルまでの解説は無いので、実装するなら[Cohen]の本が良いかと思います。
[Chmie]は[Cohen]を簡潔にまとめた解説論文です。

Dirichlet指標とGauss和

まずDirichlet指標 $\chi$ とGauss和の導入をします。
乗法群 $(\mathbb{Z}/q\mathbb{Z})^\times$ から、（0でない）複素数全体の乗法群 $\mathbb{C}^\times$ への群準同型
$\displaystyle \chi : \mathbb{Z}/q\mathbb{Z}^\times\to \mathbb{C}^\times$
を素数 $q$ を法とするDirichlet指標と呼びます。
つまりmod $q$ での乗算構造をそのままに、複素数の乗法に置き換えたもの。
ディリクレ指標 - Wikipedia

具体的に指標の1つを構成します。
奇素数 $q$ の原始根の1つを $g$ とします。 $(a, q)\neq 1$ とすると
$\displaystyle a\equiv g^x (\text{mod } q)$
なる $x$ が $\text{mod } (q-1)$ で唯一に定まる。この $x$ を $\text{ind } a$ と表す。
イメージ的には $\log_g a$ のような対数。

$p$ を $q-1$ の素因子の1つとし、 $p^k \mid q-1$ であり $p^{k+1} \nmid q-1$ とします。
1の原始 $p^k$ 乗根を
$\displaystyle \zeta_{p^k}=e^{2\pi\sqrt{-1}/p^k}$
とすると、
$\displaystyle \chi(a) = \zeta_{p^k}{}^x =\zeta_{p^k}{}^{\text{ind }a}$
とすればDirichlet指標となっています。

あるいは以下のように書いた方が直感的でわかりやすいかも。
$\displaystyle \chi(g^x \text{ mod } q) = \zeta_{p^k}{}^x$

定義から
$\displaystyle \chi(ab)=\chi(a)\chi(b)$
が成立するのは明らかでしょう。

さてここで、Gauss和 $\tau(\chi)$ を
$\displaystyle \tau(\chi) = \sum_{x\in (\mathbb{Z}/q\mathbb{Z})^*} \chi(x) \zeta_q{}^x$
と定義します。

$\displaystyle \zeta_q^{q-1} + \zeta_q^{q-2} + \cdots + \zeta_q + 1 = 0$
$\displaystyle \zeta_{p^k}{}^{(p-1)p^{k-1}} + \zeta_{p^k}{}^{(p-2)p^{k-1}} + \cdots + \zeta_{p^k}{}^{p^{k-1}} + 1 = 0$
が成り立つため、Gauss和は
$\displaystyle \tau(\chi)=\sum_{i=0}^{(p-1)p^{k-1}-1} \sum_{j=0}^{q-2} a_{ij}\zeta_{p^k}{}^i\zeta_q{}^j$
という形に表せ、この形に表したときに $a_{ij}$ は一意に決まることが証明できる。
よって
$\displaystyle \sum_{i=0}^{(p-1)p^{k-1}-1} \sum_{j=0}^{q-2} a_{ij}\zeta_p{}^i\zeta_q{}^j \equiv \sum_{i=0}^{(p-1)p^{k-1}-1} \sum_{j=0}^{q-2} b_{ij}\zeta_p{}^i\zeta_q{}^j (\text{ mod } N)$
の意味を
$\text{すべての}i, j \text{に対して} a_{ij}\equiv b_{ij} (\text{mod }N)$
と定めてもいいです。

Gauss和を用いた素数判定の原理

具体的に話を進めるために
$t=5040=2^4 \cdot 3^2 \cdot 5 \cdot 7$
とする。 $q-1$ が $t$ の約数となる素数 $q$ はたくさんある。例えば
$11-1 = 2\times 5, \quad 43-1=2\times 3 \times 7$
であるので、 $11, 43$ などが $q$ の候補となる。
これらの $q$ についての積 $e(t)$ を
$\displaystyle e(t) = 2 \prod_{q \text{ is prime}, (q-1)\mid t} q^{v_q(t)+1}$
と定義する。ここで $v_q(t)$ は $t$ の素因子分解の中に $q$ が何回現れるかという回数である。

$t=5040$ のとき、
$\begin{eqnarray} e(t) &=& 2^6 \cdot 3^3 \cdot 5^2 \cdot 7^2 \cdot 11 \cdot 13\cdot 17\cdot 19\cdot 29\cdot 31\cdot 37\cdot 41\cdot 43\cdot 61\cdot 71\cdot 73\cdot 113\cdot 127\cdot 181\cdot 211\cdot 241\cdot 281\cdot 337\cdot 421\cdot 631\cdot 1009\cdot 2521\\ &=& 15321986788854443284662612735663611380010431225771200 \\ &\fallingdotseq& 1.53\times 10^{52} \end{eqnarray}$
となる。素数判定したい数を $N$ とすると、 $N\lt e(t)^2$ までの数が素数判定できる。
実装上は $e(t)$ はテーブルにしておき、条件を満たすtを探すことになる。

$N$ はあらかじめ試し割り算で小さな素因数は無いことがわかっているものとする。つまり、 $(N, t\cdot e(t))=1$ としてよい。

$q \mid e(t)$ なる $q$ を選び、さらに $p$ を $q-1$ の素因子の1つとし、 $p^k \mid q-1$ であり $p^{k+1} \nmid q-1$ とします。（最初の方のGauss和の説明と同じ記号にしています。）

このとき、 $N$ が素数であれば
$\displaystyle \begin{eqnarray} \tau(\chi)^N &=& \left( \sum \chi(x) \zeta_q{}^x \right)^N \\ & \equiv & \sum_x \chi(x)^N \zeta_q{}^{Nx} \\ & \equiv & \chi(N)^{-N} \sum_x \chi(Nx)^N \zeta_q{}^{Nx} \quad (\text{mod } N) \end{eqnarray}$
となる。 $x$ が $1$ から $q-1$ まで動けば、 $Nx$ も $1$ から $q-1$ まで動くので
$\displaystyle \begin{eqnarray} \tau(\chi)^N &\equiv& \chi(N)^{-N} \sum_x \chi(x)^N \zeta_q{}^{x} \quad (\text{mod } N) \\ &=& \chi(N)^{-N} \tau(\chi^N) \end{eqnarray}$
が成り立つ。

最後の
$\displaystyle \tau(\chi)^N \equiv \chi(N)^{-N} \tau(\chi^N) \quad (\text{mod } N)$
がポイントで、もしこの合同式が成立しなければ、 $N$ は合成数となる。
素数 $p, q$ の組合せは多数あるが、 $N$ が素数であれば、すべての組合せで成立しなければならず、一度でも成立しなければ合成数とわかる。
Fermatの小定理と考え方は同じ。

この合同式がすべての $p, q$ について成立しても、残念ながら素数とはいえず、次のステップに進むことになりますが、ひとまずそれは置いておきます。

この計算量を見積もってみます。 $N$ は100桁ぐらいを想定しておきます。 $p, q$ の組合せはいろいろありますが、一番大きいケースで $q=2521, p^k=3^2$ です。
Gauss和を
$\displaystyle \tau(\chi)=\sum_{i=0}^{(p-1)p^{k-1}-1} \sum_{j=0}^{q-2} a_{ij}\zeta_{p^k}{}^i\zeta_q{}^j$
と表したとき、メモリ量は $10^4$ のオーダーとなる。
Gauss和同士の積は、愚直に計算すると、その2乗となり $10^8$ 程度。
$\tau(\chi)^N$ の計算は繰り返し2乗法を使えるので、 $\log_2 N$ に比例する計算量であり、300回程度の繰り返しとなる。全体として $10^{10}$ となる。さらに $a_{ij}$ は $\text{mod } N$ で計算する必要があり、多倍長計算となる。
これを多数の $p, q$ の組合せについて計算するので、不可能ではないレベルだが、実用上は不満のある計算量となる。

これをJacobi和を用いた形に書き換えていく。

Jacobi和への書き換え

Gauss和を用いると計算量が増えるのは、Gauss和は $\mathbb{Z}\left[\zeta_q, \zeta_{p^k}\right]$ という大きな環に属していたことが原因といえる。
唐突だが、 $\chi_1$ と $\chi_2$ を $\text{mod }q$ でのDirichlet指標とするとき、 Jacobi和を
$\displaystyle j(\chi_1, \chi_2) = \sum_{x\in (\mathbb{Z}/q\mathbb{Z})^\times} \chi_1(x) \chi_2(1-x)$
と定義する。
このJacobi和は $\mathbb{Z}\left[\zeta_{p^k}\right]$ という小さな環に属しているので計算量が少ないです。先の例でいえば、 $p^k=2^4$ が最大である。

$\chi_0$ を $(x, q)=1$ のとき $\chi_0(x)=1$ 、それ以外のとき $\chi_0(x)=0$ となる自明な指標とする。
Jacobi和の性質として、以下が成り立つ。
(1) $\chi \neq \chi_0$ を $\text{mod }q$ のDirichlet指標とすると、
・ $\displaystyle \tau(\chi)\tau(\bar{\chi}) = \chi(-1)q$
・ $\displaystyle \left|\tau(\chi)\right| = \sqrt{q}$
(2) $\chi_1$ と $\chi_2$ を $\chi_1\chi_2\neq\chi_0$ なる $\text{mod }q$ のDirichlet指標とすると
・ $\displaystyle j(\chi_1, \chi_2) = \frac{\tau(\chi_1)\tau(\chi_2)}{\tau(\chi_1\chi_2)}$

例えば $p=5$ のとき
$\displaystyle \begin{eqnarray} \tau(\chi)^5 &=& (\tau(\chi)\cdot\tau(\chi))^2\cdot \tau(\chi) \\ &=& (j(\chi, \chi)\cdot\tau(\chi^2))^2\cdot \tau(\chi) \\ &=& j(\chi, \chi)^2\cdot j(\chi^2, \chi^2) \cdot \tau(\chi^4)\cdot \tau(\chi) \\ &=& j(\chi, \chi)^2\cdot j(\chi^2, \chi^2) \cdot \chi(-1)\cdot q \end{eqnarray}$
のようにJacobi和の積にできたりします。

新たな記号を導入します。
$(a,n)=1$ として、 $\sigma_a$ を $\sigma_a(\zeta_n)=\zeta_n{}^a$ とします。（ $\sigma_a \in \text{Gal}(\mathbb{Q}(\zeta_n))/\mathbb{Q}$ ）
$G=\{ \sigma_a: (a,n)=1, \sigma_a(\zeta_n)=\zeta_n{}^a \}$ とし、 $\sigma_a$ の整数係数の和からなる群環 $\mathbb{Z}[G]$ を定義します。
$f=\sum_{\sigma\in G} f_\sigma \sigma \quad (f_\sigma \in \mathbb{Z})$ と書くとき、
$f\in \mathbb{Z}[G]$ と $x\in \mathbb{Q}(\zeta_n)$ とするとき、 $f$ の $x$ への作用を、
$\displaystyle x^f=\prod_{\sigma\in G} \sigma(x)^{f_\sigma}$
で定義します。

少しわかりにくいですが、ただの表記だけの問題で、
$\displaystyle x^{2\sigma_a+3\sigma_b} = \sigma_a(x)^2\cdot \sigma_b(x)^3$
のように計算します。

素数判定式を以下のように言い換えます。
「 $\beta\in \mathbb{Z}[G]$ とするとき、 $N$ が素数であれば、1の $p^k$ 乗根 $\eta(\chi)$ が存在し、
$\displaystyle \tau(\chi)^{\beta(N-\sigma_N)} \equiv \eta(\chi)^{-\beta N}\quad (\text{mod }N)$
となる。」

ここからかなり技巧的に感じる式変形をします。
$a, b$ を $p \nmid ab(a+b)$ なる整数とし、 $E$ を $1\le x \lt p^k$ かつ $p\nmid x$ なる整数の集合とするとき、
$\displaystyle \alpha = \sum_{x\in E} \left\lfloor \frac{Nx}{p^k} \right\rfloor \sigma_x{}^{-1}$
と
$\displaystyle \beta = - \sum_{x\in E} \left( \left\lfloor \frac{xa}{p^k}\right\rfloor + \left\lfloor \frac{xb}{p^k}\right\rfloor - \left\lfloor \frac{x(a+b)}{p^k}\right\rfloor \right) \sigma_x{}^{-1}$
とすると、
$\displaystyle \tau(\chi)^{\beta(N-\sigma_N)} = j(\chi^a, \chi^b)^\alpha$
と変形でき、Gauss和がJacobi和に書き換えできました。

判定する際には、 $a=b=1$ をとります。ただし $p \nmid ab(a+b)$ という制約のため、 $p=2$ のときには場合分けが必要です。
ここから先の詳細は[Cohen]か[Coh-Len1]を参照ということで。

素数判定アルゴリズム

[Cohen]に従い、実際に実装できるレベルでアルゴリズムをまとめておきます。

事前計算

Jacobi和などを事前計算しておきます。
$B$ を素数判定したい数の上限とします。
このステップは $B$ のみにより、素数判定したい数 $N$ にはよりません。

$e(t)^2>B$ となる $t$ を見つける。
$e(t)$ を割りきる素数 $q\ge 3$ に対し以下の(a)から(c)を行う。

(a) $\text{mod }q$ での原始根 $g_q$ を見つける。
$1\le x\le q-2$ に対して $1-g_q{}^x=g_q{}^{f(x)}$ となるような $1\le f(x) \le q-2$ を計算する。
(b) $q-1$ を割りきる素数 $p$ に対し、 $k$ を $p^k \mid q-1$ かつ $p^{k+1} \nmid q-1$ なる整数とする。
Dirichlet指標 $\chi_{p,q}$ を、 $\chi_{p,q}(g_q{}^x) = \zeta_{p^k}{}^x$ で定義する。
(c) ・ $p \ge 3$ または $p=2$ かつ $k=2$ のとき、以下を計算する。
$\displaystyle J(p,q) = j(\chi_{p,q}, \chi_{p,q})= \sum_{1\le x \le q-2} \zeta_{p^k}{}^{x+f(x)}$

・ $p=2$ かつ $k\ge 3$ のとき、 $J(2,q)$ を上式で計算し、以下を計算する。
$\displaystyle J_3(q) = j(\chi_{2,q}, \chi_{2,q}) j(\chi_{2,q}, \chi_{2,q}{}^2)= J(2,q) \left( \sum_{1\le x \le q-2} \zeta_{2^k}{}^{2x+f(x)}\right)$

$\displaystyle J_2(q) = j(\chi_{2,q}{}^{2^{k-3}}, \chi_{2,q}{}^{3\cdot 2^{k-3}})^2 = \left( \sum_{1\le x \le q-2} \zeta_{8}{}^{3x+f(x)}\right)^2$

素数判定

素数判定したい数を $N$ ( $N\le B$ )とする。

Step1. $(t\cdot e(t), N)>1$ であれば、 $N$ は合成数である

Step2. $p\mid t$ に対して、
　　・ $p\ge 3$ かつ $N^{p-1}\not\equiv 1\quad \text{mod }p^2$ であれば、 $l_p\gets 1$
　　・それ以外であれば、 $l_p \gets 0$

Step3. $p^k \parallel (q-1) \mid t$ となる全ての素数のペア $(p,q)$ に対し
　　・ $p\ge 3$ のとき、Step 4aへ
　　・ $p=2$ かつ $k\ge 3$ のとき、Step 4bへ
　　・ $p=2$ かつ $k=2$ のとき、Step 4cへ
　　・ $p=2$ かつ $k=1$ のとき、Step 5へ

Step 4a. $E$ を $p^k$ 未満で $p$ と互いに素な自然数の集合とする。
　　・ $\Theta \gets \sum_{x\in E} x\sigma_x^{-1}$
　　・ $r\gets N \text{ mod }p^k$
　　・ $\alpha \gets \sum_{x\in E} \lfloor \frac{rx}{p^k} \rfloor \sigma_x^{-1}$
　　・ $s_1 \gets J(p,q)^\Theta \text{ mod } N$
　　・ $s_2 \gets s_1{}^{\lfloor N/p^k\rfloor} \text{ mod } N$
　　・ $S(p,q) \gets s_2 J(p,q)^\alpha \text{ mod } N$

　　このとき、 $S(p,q)\equiv \eta \text{ mod } N$ となるような $1$ の $p^k$ 乗根 $\eta$ が存在しないとき、 $N$ は合成数である。

　　もし $\eta$ が存在し、原始冪根であるとき、 $l_p\gets 1$ とする。

Step 4b. $E$ を $2^k$ 未満で $\text{mod }8$ で $1$ または $3$ に合同な自然数の集合とする。
　　・ $\Theta \gets \sum_{x\in E} x\sigma_x^{-1}$
　　・ $r\gets N \text{ mod }2^k$
　　・ $\alpha \gets \sum_{x\in E} \lfloor \frac{rx}{2^k} \rfloor \sigma_x^{-1}$
　　・ $s_1 \gets J_3(q)^\Theta \text{ mod } N$
　　・ $s_2 \gets s_1{}^{\lfloor N/2^k\rfloor} \text{ mod } N$
　　・ $S(2,q) \gets s_2 J_3(q)^\alpha J_2(q)^{\delta_N} \text{ mod } N$
　　　　　ここで $\delta_N$ は $N$ が $1$ か $3$ に合同であるときに $0$ 、それ以外は $1$ とする。

　　このとき、 $S(2,q)\equiv \eta \text{ mod } N$ となるような $1$ の $2^k$ 乗根 $\eta$ が存在しないとき、 $N$ は合成数である。

　　もし $\eta$ が存在し、原始冪根であり、かつ、 $q^{(N-1)/2}\equiv -1 \text{ mod }N$ のとき、 $l_2\gets 1$ とする。

Step 4c.
　　・ $s_1 \gets J(2,q)^2\cdot q \text{ mod } N$
　　・ $s_2 \gets s_1{}^{\lfloor N/4\rfloor} \text{ mod } N$
　　・ $N\equiv 1 \quad (\text{mod } 4)$ のとき、 $S(2,q) \gets s_2$
　　　 $N\equiv 3 \quad (\text{mod } 4)$ のとき、 $S(2,q)\gets s_2 J(2,q)^2$

　　このとき、 $S(2,q)\equiv \eta \text{ mod } N$ となるような $1$ の $4$ 乗根 $\eta$ が存在しないとき、 $N$ は合成数である。

　　もし $\eta$ が存在し、原始冪根であり、かつ、 $q^{(N-1)/2}\equiv -1 \text{ mod }N$ のとき、 $l_2\gets 1$ とする。

Step 4d.
　　・ $S(2,q)\gets (-q)^{(N-1)/2} \quad\text{mod }N$

　　このとき、 $S(2,q)\not\equiv \pm 1 \text{ mod } N$ であれば、 $N$ は合成数である。

　　もし $S(2,q)\equiv -1 \text{ mod } N$ かつ $N\equiv 1 \text{ mod } 4$ のとき、 $l_2\gets 1$ とする。

Step 5. $l_p=0$ なる $p\mid t$ に対して以下を行う。
　　・ $q\nmid e(t), p\mid(q-1)$ かつ $(q, N)=1$ となる $q$ をランダムに選ぶ。
　　・Step 4a-4dを $(p,q)$ に対して行う。
　　・何度か繰り返しても $0$ となる $l_p$ があれば「判定不能」を返す。（現実的にはこれはまず起こらない）

Step 6.
　　・ $i=1, \dots , t-1$ に対して、 $r_i\gets N^i\quad \text{mod } e(t)$ を計算する。
　　・もしある $i$ で $r_i$ が $N$ を割りきれば合成数である。そうでなければ素数と判定される。

以上がAPR-CL素数判定アルゴリズムです。
かなり場合分けが煩雑ですね。

上記アルゴリズムだけでは実装には悩む箇所があるので、実装上は[Coh-Len2]を参照するといいです。
$\sigma_x^{-1}$ を作用させる箇所や、1のべき根であることの判定法などが具体的に記載されています。

Python3による実装

[Cohen]とはできるだけ記号を合わせているので、比較しながら読めばわかるはず。
上のアルゴリズムで事前計算とされている部分は、事前計算にせずにその場で計算しています。
原始根なども保存せずに毎回計算したりしているので、少し工夫すればもう少しは早くなるかも。

手元の環境での実行時間は、100桁だと2秒、2^521-1（157桁）が18秒、2^1279-1（386桁）が5分55秒といったところで、まずまずの速度でしょうか。

GitHub - wacchoz/APR_CL: Adleman–Pomerance–Rumely-Cohen-Lenstra primality test (APR-CL)

# -*- coding: utf-8 -*-

import copy
import time
from math import gcd  # version >= 3.5



# primality test by trial division
def isprime_slow(n):
    if n<2:
        return False
    elif n==2 or n==3:
        return True
    elif n%2==0:
        return False
    else:
        i = 3
        while i*i <= n:
            if n%i == 0:
                return False
            i+=2
    return True        


# v_q(t): how many time is t divided by q
def v(q, t):
    ans = 0
    while(t % q == 0):
        ans +=1
        t//= q
    return ans


def prime_factorize(n):
    ret = []
    p=2
    while p*p <= n:
        if n%p==0:
            num = 0
            while n%p==0:
                num+=1
                n//= p
            ret.append((p,num))
        p+= 1

    if n!=1:
        ret.append((n,1))

    return ret


# calculate e(t)
def e(t):
    s = 1
    q_list = []
    for q in range(2, t+2):
        if t % (q-1) == 0 and isprime_slow(q):
            s *= q ** (1+v(q,t))
            q_list.append(q)
    return 2*s, q_list

# Jacobi sum
class JacobiSum(object):
    def __init__(self, p, k, q):
        self.p = p
        self.k = k
        self.q = q
        self.m = (p-1)*p**(k-1)
        self.pk = p**k
        self.coef = [0]*self.m

    # 1
    def one(self):
        self.coef[0] = 1
        for i in range(1,self.m):
            self.coef[i] = 0
        return self


    # product of JacobiSum
    # jac : JacobiSum
    def mul(self, jac):
        m = self.m
        pk = self.pk
        j_ret=JacobiSum(self.p, self.k, self.q)
        for i in range(m):
            for j in range(m):
                if (i+j)% pk < m:
                    j_ret.coef[(i+j)% pk] += self.coef[i] * jac.coef[j]
                else:
                    r = (i+j) % pk - self.p ** (self.k-1)                    
                    while r>=0:
                        j_ret.coef[r] -= self.coef[i] * jac.coef[j]
                        r-= self.p ** (self.k-1)

        return j_ret


    def __mul__(self, right):
        if type(right) is int:
            # product with integer
            j_ret=JacobiSum(self.p, self.k, self.q)
            for i in range(self.m):
                j_ret.coef[i] = self.coef[i] * right
            return j_ret
        else:
            # product with JacobiSum
            return self.mul(right)
    

    # power of JacobiSum（x-th power mod n）
    def modpow(self, x, n):
        j_ret=JacobiSum(self.p, self.k, self.q)
        j_ret.coef[0]=1
        j_a = copy.deepcopy(self)
        while x>0:
            if x%2==1:
                j_ret = (j_ret * j_a).mod(n)
            j_a = j_a*j_a
            j_a.mod(n)
            x //= 2
        return j_ret
    

    # applying "mod n" to coefficient of self
    def mod(self, n):
        for i in range(self.m):
            self.coef[i] %= n
        return self
    

    # operate sigma_x
    # verification for sigma_inv
    def sigma(self, x):
        m = self.m
        pk = self.pk
        j_ret=JacobiSum(self.p, self.k, self.q)
        for i in range(m):
            if (i*x) % pk < m:
                j_ret.coef[(i*x) % pk] += self.coef[i]
            else:
                r = (i*x) % pk - self.p ** (self.k-1)                    
                while r>=0:
                    j_ret.coef[r] -= self.coef[i]
                    r-= self.p ** (self.k-1)
        return j_ret
    
                
    # operate sigma_x^(-1)
    def sigma_inv(self, x):
        m = self.m
        pk = self.pk
        j_ret=JacobiSum(self.p, self.k, self.q)
        for i in range(pk):
            if i<m:
                if (i*x)%pk < m:
                    j_ret.coef[i] += self.coef[(i*x)%pk]
            else:
                r = i - self.p ** (self.k-1)
                while r>=0:
                    if (i*x)%pk < m:
                        j_ret.coef[r] -= self.coef[(i*x)%pk]
                    r-= self.p ** (self.k-1)

        return j_ret
    

    # Is self p^k-th root of unity (mod N)
    # if so, return h where self is zeta^h
    def is_root_of_unity(self, N):
        m = self.m
        p = self.p
        k = self.k

        # case of zeta^h (h<m)
        one = 0
        for i in range(m):
            if self.coef[i]==1:
                one += 1
                h = i
            elif self.coef[i] == 0:
                continue
            elif (self.coef[i] - (-1)) %N != 0:
                return False, None
        if one == 1:
            return True, h

        # case of zeta^h (h>=m)
        for i in range(m):
            if self.coef[i]!=0:
                break
        r = i % (p**(k-1))
        for i in range(m):
            if i % (p**(k-1)) == r:
                if (self.coef[i] - (-1))%N != 0:
                    return False, None
            else:
                if self.coef[i] !=0:
                    return False, None

        return True, (p-1)*p**(k-1)+ r
            

# find primitive root
def smallest_primitive_root(q):
    for r in range(2, q):
        s = set({})
        m = 1
        for i in range(1, q):
            m = (m*r) % q
            s.add(m)
        if len(s) == q-1:
            return r
    return None   # error


# calculate f_q(x)
def calc_f(q):
    g = smallest_primitive_root(q)
    m = {}
    for x in range(1,q-1):
        m[pow(g,x,q)] = x
    f = {}
    for x in range(1,q-1):
        f[x] = m[ (1-pow(g,x,q))%q ]

    return f


# sum zeta^(a*x+b*f(x))
def calc_J_ab(p, k, q, a, b):
    j_ret = JacobiSum(p,k,q)
    f = calc_f(q)
    for x in range(1,q-1):
        pk = p**k
        if (a*x+b*f[x]) % pk < j_ret.m:
            j_ret.coef[(a*x+b*f[x]) % pk] += 1
        else:
            r = (a*x+b*f[x]) % pk - p**(k-1)
            while r>=0:
                j_ret.coef[r] -= 1
                r-= p**(k-1)
    return j_ret


# calculate J(p,q)（p>=3 or p,q=2,2）
def calc_J(p, k, q):
    return calc_J_ab(p, k, q, 1, 1)

            
# calculate J_3(q)（p=2 and k>=3）
def calc_J3(p, k, q):
    j2q = calc_J(p, k, q)
    j21 = calc_J_ab(p, k, q, 2, 1)
    j_ret = j2q * j21
    return j_ret


# calculate J_2(q)（p=2 and k>=3）
def calc_J2(p, k, q):
    j31 = calc_J_ab(2, 3, q, 3, 1)
    j_conv = JacobiSum(p, k, q)
    for i in range(j31.m):
        j_conv.coef[i*(p**k)//8] = j31.coef[i]
    j_ret = j_conv * j_conv
    return j_ret


# in case of p>=3
def APRtest_step4a(p, k, q, N):

    print("Step 4a. (p^k, q = {0}^{1}, {2})".format(p,k,q))
    
    J = calc_J(p, k, q)
    # initialize s1=1
    s1 = JacobiSum(p,k,q).one()
    # J^Theta
    for x in range(p**k):
        if x % p == 0:
            continue
        t = J.sigma_inv(x)
        t = t.modpow(x, N)
        s1 = s1 * t
        s1.mod(N)

    # r = N mod p^k
    r = N % (p**k)

    # s2 = s1 ^ (N/p^k)
    s2 = s1.modpow(N//(p**k), N)

    # J^alpha
    J_alpha = JacobiSum(p,k,q).one()
    for x in range(p**k):
        if x % p == 0:
            continue
        t = J.sigma_inv(x)
        t = t.modpow((r*x)//(p**k), N)
        J_alpha = J_alpha * t
        J_alpha.mod(N)

    # S = s2 * J_alpha
    S = (s2 * J_alpha).mod(N)

    # Is S root of unity
    exist, h = S.is_root_of_unity(N)

    if not exist:
        # composite!
        return False, None
    else:
        # possible prime
        if h%p!=0:
            l_p = 1
        else:
            l_p = 0
        return True, l_p


# in case of p=2 and k>=3
def APRtest_step4b(p, k, q, N):

    print("Step 4b. (p^k, q = {0}^{1}, {2})".format(p,k,q))

    J = calc_J3(p, k, q)
    # initialize s1=1
    s1 = JacobiSum(p,k,q).one()
    # J3^Theta
    for x in range(p**k):
        if x % 8 not in [1,3]:
            continue
        t = J.sigma_inv(x)
        t = t.modpow(x, N)
        s1 = s1 * t
        s1.mod(N)

    # r = N mod p^k
    r = N % (p**k)

    # s2 = s1 ^ (N/p^k)
    s2 = s1.modpow(N//(p**k), N)

    # J3^alpha
    J_alpha = JacobiSum(p,k,q).one()
    for x in range(p**k):
        if x % 8 not in [1,3]:
            continue
        t = J.sigma_inv(x)
        t = t.modpow((r*x)//(p**k), N)
        J_alpha = J_alpha * t
        J_alpha.mod(N)

    # S = s2 * J_alpha * J2^delta
    if N%8 in [1,3]:
        S = (s2 * J_alpha ).mod(N)
    else:
        J2_delta = calc_J2(p,k,q)
        S = (s2 * J_alpha * J2_delta).mod(N)

    # Is S root of unity
    exist, h = S.is_root_of_unity(N)

    if not exist:
        # composite 
        return False, None
    else:
        # possible prime
        if h%p!=0 and (pow(q,(N-1)//2,N) + 1)%N==0:
            l_p = 1
        else:
            l_p = 0
        return True, l_p


# in case of p=2 and k=2
def APRtest_step4c(p, k, q, N):

    print("Step 4c. (p^k, q = {0}^{1}, {2})".format(p,k,q))

    J2q = calc_J(p, k, q)

    # s1 = J(2,q)^2 * q (mod N)
    s1 = (J2q * J2q * q).mod(N)

    # s2 = s1 ^ (N/4)
    s2 = s1.modpow(N//4, N)

    if N%4 == 1:
        S = s2
    elif N%4 == 3:
        S = (s2 * J2q * J2q).mod(N)
    else:
        print("Error")

    # Is S root of unity
    exist, h = S.is_root_of_unity(N)

    if not exist:
        # composite
        return False, None
    else:
        # possible prime
        if h%p!=0 and (pow(q,(N-1)//2,N) + 1)%N==0:
            l_p = 1
        else:
            l_p = 0
        return True, l_p


# in case of p=2 and k=1
def APRtest_step4d(p, k, q, N):

    print("Step 4d. (p^k, q = {0}^{1}, {2})".format(p,k,q))

    S2q = pow(-q, (N-1)//2, N)
    if (S2q-1)%N != 0 and (S2q+1)%N != 0:
        # composite
        return False, None
    else:
        # possible prime
        if (S2q + 1)%N == 0 and (N-1)%4==0:
            l_p=1
        else:
            l_p=0
        return True, l_p


# Step 4
def APRtest_step4(p, k, q, N):

    if p>=3:
        result, l_p = APRtest_step4a(p, k, q, N)
    elif p==2 and k>=3:
        result, l_p = APRtest_step4b(p, k, q, N)
    elif p==2 and k==2:
        result, l_p = APRtest_step4c(p, k, q, N)
    elif p==2 and k==1:
        result, l_p = APRtest_step4d(p, k, q, N)
    else:
        print("error")

    if not result:
        print("Composite")

    return result, l_p


def APRtest(N):
    t_list = [
        2,
        12,
        60,
        180,
        840,
        1260,
        1680,
        2520,
        5040,
        15120,
        55440,
        110880,
        720720,
        1441440,
        4324320,
        24504480,
        73513440
        ]

    print("N=", N)

    if N<=3:
        print("input should be greater than 3")
        return False
 
    # Select t
    for t in t_list:
        et, q_list = e(t)
        if N < et*et:
            break
    else:
        print("t not found")
        return False

    print("t=", t)
    print("e(t)=", et, q_list)

    # Step 1
    print("=== Step 1 ===")
    g = gcd(t*et, N)
    if g > 1:
        print("Composite")
        return False

    # Step 2
    print("=== Step 2 ===")
    l = {}
    fac_t = prime_factorize(t)
    for p, k in fac_t:
        if p>=3 and pow(N, p-1, p*p)!=1:
            l[p] = 1
        else:
            l[p] = 0
    print("l_p=", l)

    # Step 3 & Step 4
    print("=== Step 3&4 ===")
    for q in q_list:
        if q == 2:
            continue
        fac = prime_factorize(q-1)
        for p,k in fac:

            # Step 4
            result, l_p = APRtest_step4(p, k, q, N)

            if not result:
                # composite
                print("Composite")
                return False
            elif l_p==1:
                l[p] = 1

    # Step 5          
    print("=== Step 5 ===")
    print("l_p=", l)
    for p, value in l.items():
        if value==0:
            # try other pair of (p,q)
            print("Try other (p,q). p={}".format(p))
            count = 0
            i = 1
            found = False
            # try maximum 30 times
            while count < 30:
                q = p*i+1
                if N%q != 0 and isprime_slow(q) and (q not in q_list):
                    count += 1

                    k = v(p, q-1)
                    # Step 4
                    result, l_p = APRtest_step4(p, k, q, N)

                    if not result:
                        # composite
                        print("Composite")
                        return False
                    elif l_p == 1:
                        found = True
                        break
                i += 1
            if not found:
                print("error in Step 5")
                return False

    # Step 6
    print("=== Step 6 ===")
    r = 1
    for t in range(t-1):
        r = (r*N) % et
        if r!=1 and r!= N and N % r == 0:
            print("Composite", r)
            return False
    # prime!!
    print("Prime!")
    return True


if __name__ == '__main__':


    start_time = time.time()

    APRtest(2**521-1)   # 157 digit, 18 sec
#    APRtest(2**1279-1)  # 386 digit, 355 sec

    end_time = time.time()
    print(end_time - start_time, "sec")