Decomposes n into a,b such that a^2 + b^2 = n

ssd.py

from math import floor as fl
from math import sqrt
from random import randint

def ctuple(z):
	return z.real, z.imag

def floor(x):
	if x < 0:
		return -fl(-x)
	else:
		return fl(x)

def zmod(z,w):
	a,b = ctuple(z/w)
	a,b = a-floor(a), b-floor(b)
	q = w*(a + b*1j)
	j,k = ctuple(q)
	j,k = round(j,0),round(k,0)
	q = j + k*1j

	return q

def gcd(a,b):
	while b:
		a,b = b,a%b
	
	return a

def zgcd(z,w):
	while w != 0:
		z, w = w, zmod(z,w)
	
	return z

def rho(n):
	if n == 1:
		return 1
	
	if n % 2 == 0:
		return 2
	
	x = randint(2, n)
	y = x

	c = randint(1, n)
	d = 1

	def f(x):
		return (pow(x, 2, n) + c)%n

	while d == 1:
		x = f(x)
		y = f(f(y))

		d = gcd(abs(x-y), n)

	return d

def is_prime(n, k=7):
	q = n

	for x in range(k):
		q = rho(q)
	
	return q == n

def pfactor(n):
	f = []
	m = 0
	no = n

	while n != 1:
		r = rho(n)
		
		if n == r and is_prime(r):
			k = r
		elif n == r:
			k = 1
		else:
			k = r

		m = n/k

		if m != n:
			f.append(k)
		
		n = m
	
	p = True
	for x in f:
		p = is_prime(x) and p
	
	if not p:
		return pfactor(no)
	else:
		return f

def conj_fac(p):
	if p%4 == 1:
		for x in range(2,p):
			y = pow(x,(p-1)/2,p)

			if y == p-1:
				break
		
		q = pow(x, (p-1)/4) + 1j

		a,b = ctuple(zgcd(q, p))

		return complex(abs(a), abs(b))
	elif p%4 == 3:
		return sqrt(p)+0j
	elif p == 2:
		return 1+1j
	

def swap(z):
	return complex(z.imag, z.real)

def swap_combo(l):
	ll = []

	for i in range(1<<len(l)):
		b = bin(i)[2:].zfill(len(l))
		
		ll.append([swap(l[y]) if b[y] == "1" else l[y] for y in range(len(l))])
	
	return ll

def ssd(n):
	p = pfactor(n)
	c = map(conj_fac, p)
	s = swap_combo(c)
	pt = set()

	for l in s:
		a,b = ctuple(reduce(lambda z,w: z*w, l))
		pt.add((abs(a),abs(b)))
	
	pt = [tuple(x) for x in pt]
	pt = filter(lambda j: 0 not in j, pt)

	return pt

n = int(raw_input("n="))

pf = pfactor(n)

for x in pf:
	if x%4 == 3 and pf.count(x) % 2 != 0:
		print("No Sum of Squares Decomposition can be made")
		exit(0)

sp = ssd(n)

if len(sp) == 0:
	print("No Sum of Squares Decomposition can be made")
	exit(0)


for a,b in sp:
	print(str(int(round(a,0))) + "^2 + " + str(int(round(b,0))) + "^2 = " + str(n))