Primality test using Chebyshev polynomials of the first kind

chebyshev.py

# Authors: Peter Taylor and Pedja Terzic

import sys

print("                        ***** CHEBYSHEV *****\n\n\n")
while True:
	
    n=int(input("Enter a number : "))


    def is_prime_naive(n):
        if n <= 2 or n % 2 == 0:
            return n == 2
        return all(n % i != 0 for i in range(3, int(n ** 0.5 + 1), 2))


    def is_pseudoprime_chebyshev(n):
        r = smallest_r(n)
        if r == 0:
            return n == 2

        # We work with polynomials represented in little-endian format
        # (poly[i] is the coefficient of x^i)
        xp, _ = tpair_mod_n_xpowrm1(n, n, r)
        return normalise_poly(xp) == [0] * (n % r) + [1]


    def smallest_r(n):
        if n <= 1 or n % 2 == 0:
            return 0

        for r in range(5, sys.maxsize**10, 2):
            if not is_prime_naive(r):
                continue

            u = n % r
            if u == 0 and r < u:
                return 0
            if u > 1 and u < r - 1:
                return r


    def tpair_mod_n_xpowrm1(m, n, r):
        """Returns the tuple (T_m, T_{m+1}) modulo (n, x^r - 1)
           where T_i is the ith Chebyshev polynomial"""
        if m == 0:
            return ([1], [0, 1])

        T_k, T_kinc = tpair_mod_n_xpowrm1(m >> 1, n, r)

        # T_2kinc = 2 * T_k * T_kinc - x
        T_2kinc = [2 * coeff % n
                   for coeff in convolve_mod_n_xpowrm1(T_k, T_kinc, n, r)]
        T_2kinc = poly_add_mod_n(T_2kinc, [0, n - 1], n)

        if m & 1:
            # m = 2k+1, m+1 = 2k+2
            T_m = T_2kinc
            # T_minc = 2 * T_kinc * T_kinc - 1
            # TODO Optimised polynomial squaring
            T_minc = [2 * coeff % n
                      for coeff in convolve_mod_n_xpowrm1(T_kinc, T_kinc, n, r)]
            T_minc = poly_add_mod_n(T_minc, [n - 1], n)
        else:
            # m = 2k, m+1 = 2k+1
            # T_m = 2 * T_k * T_k - 1
            T_m = [2 * coeff % n
                   for coeff in convolve_mod_n_xpowrm1(T_k, T_k, n, r)]
            T_m = poly_add_mod_n(T_m, [n - 1], n)
            T_minc = T_2kinc
        return (T_m, T_minc)


    def convolve_mod_n_xpowrm1(p, q, n, r):
        """ Calculates the convolution p * q modulo (n, x^r - 1) """
        if p == [] or q == []:
            return []

        conv = [0] * min(len(p) + len(q) - 1, r)
        for i, p_i in enumerate(p):
            for j, q_j in enumerate(q):
                idx = (i + j) % r
                conv[idx] = (conv[idx] + p_i * q_j) % n

        return conv


    def poly_add_mod_n(p, q, n):
        r = [0] * max(len(p), len(q))

        for i, p_i in enumerate(p):
            r[i] = p_i

        for j, q_j in enumerate(q):
            r[j] = (r[j] + q_j) % n

        return r


    def normalise_poly(p):
        while len(p) > 0 and p[-1] == 0:
            p = p[:-1]

        return p

    if n < 2:
        print("Number must be greater than one")
    elif is_pseudoprime_chebyshev(n):
        print("prime")
    else:
	print("composite")
		
    try_again = ""
    # Loop until users opts to go again or quit
    while not(try_again == "1") and not(try_again == "0"):
        try_again = input("Press 1 to try again, 0 to exit. ")
        if try_again in ["1", "0"]:
            continue  # a valid entry found
        else:
            print("Invalid input- Press 1 to try again, 0 to exit.") 
    # at this point, try_again must be "0" or "1"
    if try_again == "0":
        break