Foursquare.py

import random
_mrpt_num_trials = 7
def is_probable_prime(n): #This one is copied from rosettacode
    if n == 1:
        return False
    if n == 2:
        return True
    if n % 2 == 0:
        return False
    s = 0
    d = n - 1
    while True:
        quotient, remainder = divmod(d,2)
        if remainder == 1:
            break
        s += 1
        d = quotient
    assert(2**s * d == n-1)
    def try_composite(a):
        if pow(a,d, n) == 1:
            return False
        for i in range(s):
            if pow(a, 2**i * d, n) == n-1:
                return False
        return True
    for i in range(_mrpt_num_trials):
        a = random.randrange(2,n)
        if try_composite(a):
            return False
    return True

def foursquare(n):
    def split(n):
        i = int(n**0.5)
        j = int((n - i**2)**0.5)
        while not i**2 + j**2 == n:
            i -= 1
            j = int((n - i**2)**0.5)
        return i,j
    def stepone(n):
        r = 0
        while not n % 4:
            n /= 4
            r += 1
        return r,n
    def stepthree(n):
        r = 0
        while not n % 2:
            n /= 2
            r += 1
        return r,n
    r, s = stepone(n)
    if (int(s**0.5))**2 == s:
        return int(n**0.5),0,0,0
    if (s % 8) < 7:
        arbsquare = 0
        residue = s
        while True:
            m, k = stepthree(residue)
            if k % 4 == 1:
                if is_probable_prime(k) or k == 1:
                    if not m % 2:
                        a, b = split(k)
                        a *= 2**(m/2)
                        b *= 2**(m/2)
                    else:
                        a, b = split(2*k)
                        a *= 2**((m-1)/2)
                        b *= 2**((m-1)/2)
                    return (2**r)*a, (2**r)*b, (2**r)*arbsquare,0
            arbsquare += 1
            residue = s - (arbsquare**2)
    else:
        c = 1
        d = 1
        residue = s - 2
        while True:
            m, k = stepthree(residue)
            if k % 4 == 1:
                if is_probable_prime(k) or k == 1:
                    if not m % 2:
                        a, b = split(k)
                        a *= 2**(m/2)
                        b *= 2**(m/2)
                    else:
                        a, b = split(2*k)
                        a *= 2**((m-1)/2)
                        b *= 2**((m-1)/2)
                    return (2**r)*a, (2**r)*b, (2**r)*c, (2**r)*d
            if c + 1 < n**0.5:
                c += 1
            else:
                c = 1
                d += 1
            residue = s - c**2 - d**2