Create a gist now

Instantly share code, notes, and snippets.

What would you like to do?
import Data.List
import Test.QuickCheck
import qualified Data.Map as Map
nats = [1..]
divides a b = b `mod` a == 0
sieve (n:t) m = case Map.lookup n m of
Nothing -> n : sieve t (Map.insert (n*n) [n] m)
Just ps -> sieve t (foldl reinsert (Map.delete n m) ps)
where
reinsert m p = Map.insertWith (++) (n+p) [p] m
primes = 2 : sieve (map (\n -> n*2+1) nats) Map.empty
factor n =
let
factor' 1 _ = []
factor' n ps@(p:pt) = if p `divides` n then p : factor' (n `div` p) ps else factor' n pt
in factor' n primes
modexp a 0 p = 1
modexp a n p =
let a2 = modexp (a*a `mod` p) (n `div` 2) p
in if n `mod` 2 == 0 then a2 else a2 * a `mod` p
modexp1 a 0 p b = b
modexp1 a n p b = modexp1 a (n-1) p ((b * a) `mod` p)
test0 a n p = n >= 0 && p > 0 ==> modexp a n p == modexp1 a n p 1
combine (a, b, c, d) (a', b', c', d') = (w, x, y, z)
where
w = a*a' + b*b' + c*c' + d*d'
x = a*b' - b*a' - c*d' + d*c'
y = a*c' + b*d' - c*a' - d*b'
z = a*d' - b*c' + c*b' - d*a'
sumOfSquares (a, b, c, d) = a*a + b*b + c*c + d*d
test1 a b = (sumOfSquares a) * (sumOfSquares b) == sumOfSquares (combine a b)
a |> f = f a
sortWith f = sortBy (\a b -> compare (f a) (f b))
map4 f (a, b, c, d) = (f a, f b, f c, f d)
foursquare 0 = (0, 0, 0, 0)
foursquare 1 = (1, 0, 0, 0)
foursquare n = factor n |> map foursquare' |> foldl1 combine
where
foursquare' p
| p == 2 = (1, 1, 0, 0)
| p `mod` 4 == 1 =
let
findSqrtMinus1 (a:ps) =
let b = modexp a ((p-1) `div` 4) p
in if b*b `mod` p == p-1 then b else findSqrtMinus1 ps
a = findSqrtMinus1 primes
in (a, 1, 0, 0) |> reduce
| p `mod` 4 == 3 =
let
findNumberWithSqrt (a:ns) =
let
x = (-1 - a*a) `mod` p
b = modexp x ((p+1) `div` 4) p
in if b*b `mod` p == x then (a, b) else findNumberWithSqrt ns
(a, b) = findNumberWithSqrt [1..p]
in (a, b, 1, 0) |> reduce
where
reduce abcd | sumOfSquares abcd == p = abcd
reduce abcd@(a, b, c, d) | otherwise =
let k = (sumOfSquares abcd) `div` p
in case k `mod` 2 of
0 ->
let [a', b', c', d'] = sortWith (`mod` 2) [a, b, c, d]
in (a'+b', a'-b', c'+d', c'-d') |> map4 (`div` 2) |> reduce
1 ->
abcd |> map4 (absoluteLeastResidue k) |> combine abcd |> map4 (`div` k) |> reduce
absoluteLeastResidue k m = if n <= k `div` 2 then n else n - k
where n = m `mod` k
test2 n = n >= 0 ==> sumOfSquares (foursquare n) == n
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment