FediPrimes/lucas.py

# Finding lucas pseudoprimes
from math import gcd, floor, log
from psimp import trd


def isqrt(n):
    """ Find the integer square root of n via newton's method, code via
    stackoverflow:

    (https://stackoverflow.com/questions/15390807/integer-square-root-in-python)
    """
    x = n
    y = (x + 1) // 2
    while y < x:
        x = y
        y = (x + n // x) // 2
    return x


def hasIntSQRT(n):
    """Detect whether the square root of n is an integer,
    i.e. whether the isqrt(n) is the true square root.
    """
    isq = isqrt(n)
    return isq * isq == n


def Dsequence():
    """Generate sequence 5, -7, 9, -11, 13, -15...
    """
    val = 5
    while True:
        if val % 4 == 1:
            yield val
        else:
            yield -val
        val = val + 2


def Legendre(a, p):
    """Function for calculating Legendre symbol.

    Note that this is only supposed to be defined if p is an odd prime, but we
    don't know in advance if that is true - will therefore only throw and error
    if p is even.

    Uses original power definition from
    <https://en.wikipedia.org/wiki/Legendre_symbol>
    """
    if (p % 2 == 0):
        raise ValueError("p must be odd, is {}".format(p))
    lv = pow(a, (p-1)//2, p)
    if lv == p - 1:
        lv = -1
    return lv


def Jacobi(a, n):
    """Function for calculating Jacobi symbol.

    Note that this is only defined for positive odd integers n.

    Uses algorithm from
    <https://en.wikipedia.org/wiki/Jacobi_symbol#Calculating_the_Jacobi_symbol>
    """
    if n < 1:
        raise ValueError("n must be positive")
    if (n % 2 == 0):
        raise ValueError("n must be odd")
    if a % n != a:
        return Jacobi(a % n, n)
    if a != 0 and a % 2 == 0:
        nm8 = n % 8
        if (nm8 == 3 or nm8 == 5):
            return -1 * Jacobi(a//2, n)
        else:
            return Jacobi(a//2, n)
    if a == 1:
        return 1
    if gcd(a, n) != 1:
        return 0
    if a == 0 and n == 1:
        return 1
    if n % 4 == 3 and a % 4 == 3:
        return -1 * Jacobi(n, a)
    return Jacobi(n, a)


def FirstD(n):
    """Return first D in the sequence 5, -7, 9, -11, 13, -15... for which the
    Jacobi symbol (D/n) is -1. Returns 0 if the number is a perfect square, or
    otherwise the value can't be found
    """
    if n < 1 or n % 2 == 0:
        raise ValueError("n must be a positive odd number")
    # Flag to see if it's been checked whether or not n is
    # square
    haschecked = False
    for D in Dsequence():
        Jv = Jacobi(D, n)
        # What we're looking for
        if Jv == -1:
            return D
        # Have found a factor
        if Jv == 0 and abs(D) != n:
            return 0
        # Check for square at 13
        if D > 10 and not haschecked:
            if hasIntSQRT(n):
                return 0
    # Shouldn't fire
    return 0


def Lucas(n, p, q):
    """Function for generating values of a lucas sequence with parameters p, q.
    Via formula at <https://en.wikipedia.org/wiki/Lucas_sequence>
    """
    if n < 0:
        raise ValueError("n must be a non-negative integer")
    Uc = 0
    Vc = 2
    for i in range(n):
        Un = (p * Uc + Vc) // 2
        Vn = ((p * p - 4 * q) * Uc + p * Vc) // 2
        Uc = Un
        Vc = Vn
    return Uc, Vc


def LucasUn(n, p, q):
    """Return only the U numbers from a lucas sequence Un(P, Q). For example if
    P = 1, and Q = -1 this will return the Fibonacci numbers; if P = 3 and Q =
    2 then it returns the Mersenne numbers etc.
    """
    return Lucas(n, p, q)[0]


def LucasVn(n, p, q):
    """Return only V numbers from a local sequence Vn(P, Q).
    """
    return Lucas(n, p, q)[1]


def LucasFast(k, p, q, n):
    """Faster Lucas algorithm, modulo a number
    """
    if k < 0:
        raise ValueError("k must be a non-negative integer")
    if n < 1:
        raise ValueError("n must be a postive integer")
    if n % 2 == 0:
        raise ValueError("n must be odd")
    uc, vc = Lucas(0, p, q)
    dk = floor(log(k, 2))
    tp = 2**dk
    krunn = k
    ks = 0
    for i in range(dk + 1):
        un = (uc * vc)
        vn = (pow(vc, 2, n) - 2 * pow(q, ks, n))
        uc = un % n
        vc = vn % n
        ks = ks * 2
        if krunn >= tp:
            krunn = krunn - tp
            ks = ks + 1
            un = p * uc + vc
            if un % 2 != 0:
                un = un + n
            vn = (p * p - 4 * q) * uc + p * vc
            if vn % 2 != 0:
                vn = vn + n
            uc = (un // 2) % n
            vc = (vn // 2) % n
        tp = tp // 2
    return(uc, vc)


def LucasPrime(n):
    """Lucas probable prime test
    (note: not the "strong" test.)
    """
    if n < 2:
        return False
    if n == 2:
        return True
    if n % 2 == 0:
        return False
    d = FirstD(n)
    if d == 0:
        return False
    q = (1 - d) // 4
    lnum = LucasFast(n + 1, 1, q, n)
    return lnum[0] == 0


def StrongLucasPrime(n):
    """Strong version of the Lucas probable prime test
    """
    if n < 2:
        return False
    if n == 2:
        return True
    if n % 2 == 0:
        return False
    d = FirstD(n)
    if d == 0:
        return False
    s, od = trd(n + 1)
    p = 1
    q = (1 - d) // 4
    uc, vc = Lucas(0, p, q)
    dk = floor(log(od, 2))
    tp = 2**dk
    krunn = od
    ks = 0
    for i in range(dk + 1):
        un = (uc * vc)
        vn = (pow(vc, 2, n) - 2 * pow(q, ks, n))
        uc = un % n
        vc = vn % n
        ks = ks * 2
        if krunn >= tp:
            krunn = krunn - tp
            ks = ks + 1
            un = p * uc + vc
            if un % 2 != 0:
                un = un + n
            vn = (p * p - 4 * q) * uc + p * vc
            if vn % 2 != 0:
                vn = vn + n
            uc = (un // 2) % n
            vc = (vn // 2) % n
        tp = tp // 2
    if uc == 0:
        return True
    if vc == 0:
        return True
    for r in range(1, s):
        un = (uc * vc)
        vn = (pow(vc, 2, n) - 2 * pow(q, ks, n))
        uc = un % n
        vc = vn % n
        if vc == 0:
            return True
        ks = ks * 2
    return False