-- This Haskell program calculates A_p and B_p of theorem 2 of Robert Langlands's "A little bit of number theory."

-- author: Anthony V. Pulido. Last modified: 03/16/2017 18:49:39.

-- This file was originally named lbnt.hs.

module Main where
	
import Math.NumberTheory.Primes.Sieve

-- This gives us the list of primes called "primes" used in the last line.

n = 100

-- n is the number of primes the program considers. Change this to have the program do fewer or more (odd) primes.

pOpt p = floor . sqrt $ (fromIntegral p :: Float)

-- When decomposing a prime p into sums of squares, we check only integers of absolute value less than the square root of p.

ap p
	| p `mod` 4 == 3 = 0
	| p `mod` 4 == 1 = apmod1 p
	| otherwise = 0


apmod1 p = head [ 2 * (x^3 - 3 * x * y^2) |
	x <- [-pOpt p..pOpt p], y <- [-pOpt p..pOpt p], p == x^2 + y^2, 
	([1, 0] == map (`mod` 4) [x,y] || 
	 [3, 2] == map (`mod` 4) [x,y])]
	 
-- These two functions make up the definition of A_p.
	 
-- 2 == -2 mod 4, and so the list comprehension in the function apmod1 will generate pairs (p,y) and (p,-y). The second entries, y and -y, are always squared, and so the result is the same. We can thus use "head."
	 
-- Haskell prefers 3 == a mod 4 to -1 == a mod 4. The number -1 is used in "A little bit of number theory."

bp p
	| p `mod` 4 == 1 = sAmod1 p - sBmod1 p
	| p `mod` 4 == 3 = sAmod3 p - sBmod3 p
	| otherwise = 0

-- Definition of B_p. sAmod1,3 and sBmod1,3 defined below correspond to A and B in the note "Examples."
	
sAmod1 p = sum [ a^2 + b^2 - g^2 - d^2 | a <- [-pOpt p..pOpt p], b <- [-pOpt p..pOpt p], g <- [-pOpt p..pOpt p], d <- [-pOpt p..pOpt p], p == a^2 + b^2 + g^2 + d^2, 1 == a `mod` 4, 
	([0,0,0] == map (`mod` 4) [b,g,d] ||
	 [2,0,0] == map (`mod` 4) [b,g,d] || 
	 [0,2,2] == map (`mod` 4) [b,g,d] || 
	 [2,2,2] == map (`mod` 4) [b,g,d])]
	 
sBmod1 p = sum [ a^2 + b^2 - g^2 - d^2 | a <- [-pOpt p..pOpt p], b <- [-pOpt p..pOpt p], g <- [-pOpt p..pOpt p], d <- [-pOpt p..pOpt p], p == a^2 + b^2 + g^2 + d^2, 1 == a `mod` 4, 
	([2,2,0] == map (`mod` 4) [b,g,d] ||
	 [0,2,0] == map (`mod` 4) [b,g,d] || 
	 [2,0,2] == map (`mod` 4) [b,g,d] || 
	 [0,0,2] == map (`mod` 4) [b,g,d])]


sAmod3 p = sum [ a^2 + b^2 - g^2 - d^2 | a <- [-pOpt p..pOpt p], b <- [-pOpt p..pOpt p], g <- [-pOpt p..pOpt p], d <- [-pOpt p..pOpt p], p == a^2 + b^2 + g^2 + d^2, 1 == a `mod` 4, 
	([0,1,1] == map (`mod` 4) [b,g,d] ||
	 [2,3,3] == map (`mod` 4) [b,g,d] || 
	 [0,3,3] == map (`mod` 4) [b,g,d] || 
	 [2,1,1] == map (`mod` 4) [b,g,d])]
	 
sBmod3 p = sum [ a^2 + b^2 - g^2 - d^2 | a <- [-pOpt p..pOpt p], b <- [-pOpt p..pOpt p], g <- [-pOpt p..pOpt p], d <- [-pOpt p..pOpt p], p == a^2 + b^2 + g^2 + d^2, 1 == a `mod` 4, 
	([0,3,1] == map (`mod` 4) [b,g,d] ||
	 [2,1,3] == map (`mod` 4) [b,g,d] || 
	 [0,1,3] == map (`mod` 4) [b,g,d] || 
	 [2,3,1] == map (`mod` 4) [b,g,d])]

main = print $ [(p, p `mod` 4, ap p, bp p) | p <- take n (filter odd primes)]