-- |
-- Module:      Math.NumberTheory.GaussianIntegers
-- Copyright:   (c) 2016 Chris Fredrickson, Google Inc.
-- Licence:     MIT
-- Maintainer:  Chris Fredrickson <chris.p.fredrickson@gmail.com>
--
-- This module exports functions for manipulating Gaussian integers, including
-- computing their prime factorisations.
--

{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE TypeFamilies  #-}

module Math.NumberTheory.Quadratic.GaussianIntegers (
    GaussianInteger(..),
    ι,
    conjugate,
    norm,
    primes,
    findPrime,
) where

import Prelude hiding (quot, quotRem)
import Control.DeepSeq (NFData)
import Data.Coerce
import Data.Euclidean
import Data.List (mapAccumL, partition)
import Data.Maybe
import Data.Ord (comparing)
import qualified Data.Semiring as S
import GHC.Generics

import Math.NumberTheory.Moduli.Sqrt
import Math.NumberTheory.Roots (integerSquareRoot)
import Math.NumberTheory.Primes.Types
import qualified Math.NumberTheory.Primes as U
import Math.NumberTheory.Utils              (mergeBy)
import Math.NumberTheory.Utils.FromIntegral

infix 6 :+
-- |A Gaussian integer is a+bi, where a and b are both integers.
data GaussianInteger = (:+) { GaussianInteger -> Integer
real :: !Integer, GaussianInteger -> Integer
imag :: !Integer }
    deriving (GaussianInteger -> GaussianInteger -> Bool
(GaussianInteger -> GaussianInteger -> Bool)
-> (GaussianInteger -> GaussianInteger -> Bool)
-> Eq GaussianInteger
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: GaussianInteger -> GaussianInteger -> Bool
$c/= :: GaussianInteger -> GaussianInteger -> Bool
== :: GaussianInteger -> GaussianInteger -> Bool
$c== :: GaussianInteger -> GaussianInteger -> Bool
Eq, Eq GaussianInteger
Eq GaussianInteger
-> (GaussianInteger -> GaussianInteger -> Ordering)
-> (GaussianInteger -> GaussianInteger -> Bool)
-> (GaussianInteger -> GaussianInteger -> Bool)
-> (GaussianInteger -> GaussianInteger -> Bool)
-> (GaussianInteger -> GaussianInteger -> Bool)
-> (GaussianInteger -> GaussianInteger -> GaussianInteger)
-> (GaussianInteger -> GaussianInteger -> GaussianInteger)
-> Ord GaussianInteger
GaussianInteger -> GaussianInteger -> Bool
GaussianInteger -> GaussianInteger -> Ordering
GaussianInteger -> GaussianInteger -> GaussianInteger
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: GaussianInteger -> GaussianInteger -> GaussianInteger
$cmin :: GaussianInteger -> GaussianInteger -> GaussianInteger
max :: GaussianInteger -> GaussianInteger -> GaussianInteger
$cmax :: GaussianInteger -> GaussianInteger -> GaussianInteger
>= :: GaussianInteger -> GaussianInteger -> Bool
$c>= :: GaussianInteger -> GaussianInteger -> Bool
> :: GaussianInteger -> GaussianInteger -> Bool
$c> :: GaussianInteger -> GaussianInteger -> Bool
<= :: GaussianInteger -> GaussianInteger -> Bool
$c<= :: GaussianInteger -> GaussianInteger -> Bool
< :: GaussianInteger -> GaussianInteger -> Bool
$c< :: GaussianInteger -> GaussianInteger -> Bool
compare :: GaussianInteger -> GaussianInteger -> Ordering
$ccompare :: GaussianInteger -> GaussianInteger -> Ordering
$cp1Ord :: Eq GaussianInteger
Ord, (forall x. GaussianInteger -> Rep GaussianInteger x)
-> (forall x. Rep GaussianInteger x -> GaussianInteger)
-> Generic GaussianInteger
forall x. Rep GaussianInteger x -> GaussianInteger
forall x. GaussianInteger -> Rep GaussianInteger x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep GaussianInteger x -> GaussianInteger
$cfrom :: forall x. GaussianInteger -> Rep GaussianInteger x
Generic)

instance NFData GaussianInteger

-- |The imaginary unit, where
--
-- > ι .^ 2 == -1
ι :: GaussianInteger
ι :: GaussianInteger
ι = Integer
0 Integer -> Integer -> GaussianInteger
:+ Integer
1

instance Show GaussianInteger where
    show :: GaussianInteger -> String
show (Integer
a :+ Integer
b)
        | Integer
b Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
0     = Integer -> String
forall a. Show a => a -> String
show Integer
a
        | Integer
a Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
0     = String
s String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
b'
        | Bool
otherwise  = Integer -> String
forall a. Show a => a -> String
show Integer
a String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
op String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
b'
        where
            b' :: String
b' = if Integer -> Integer
forall a. Num a => a -> a
abs Integer
b Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
1 then String
"ι" else Integer -> String
forall a. Show a => a -> String
show (Integer -> Integer
forall a. Num a => a -> a
abs Integer
b) String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"*ι"
            op :: String
op = if Integer
b Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Integer
0 then String
"+" else String
"-"
            s :: String
s  = if Integer
b Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Integer
0 then String
"" else String
"-"

instance Num GaussianInteger where
    + :: GaussianInteger -> GaussianInteger -> GaussianInteger
(+) (Integer
a :+ Integer
b) (Integer
c :+ Integer
d) = (Integer
a Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
c) Integer -> Integer -> GaussianInteger
:+ (Integer
b Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
d)
    * :: GaussianInteger -> GaussianInteger -> GaussianInteger
(*) (Integer
a :+ Integer
b) (Integer
c :+ Integer
d) = (Integer
a Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
c Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
b Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
d) Integer -> Integer -> GaussianInteger
:+ (Integer
a Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
d Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
b Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
c)
    abs :: GaussianInteger -> GaussianInteger
abs = (GaussianInteger, GaussianInteger) -> GaussianInteger
forall a b. (a, b) -> a
fst ((GaussianInteger, GaussianInteger) -> GaussianInteger)
-> (GaussianInteger -> (GaussianInteger, GaussianInteger))
-> GaussianInteger
-> GaussianInteger
forall b c a. (b -> c) -> (a -> b) -> a -> c
. GaussianInteger -> (GaussianInteger, GaussianInteger)
absSignum
    negate :: GaussianInteger -> GaussianInteger
negate (Integer
a :+ Integer
b) = (-Integer
a) Integer -> Integer -> GaussianInteger
:+ (-Integer
b)
    fromInteger :: Integer -> GaussianInteger
fromInteger Integer
n = Integer
n Integer -> Integer -> GaussianInteger
:+ Integer
0
    signum :: GaussianInteger -> GaussianInteger
signum = (GaussianInteger, GaussianInteger) -> GaussianInteger
forall a b. (a, b) -> b
snd ((GaussianInteger, GaussianInteger) -> GaussianInteger)
-> (GaussianInteger -> (GaussianInteger, GaussianInteger))
-> GaussianInteger
-> GaussianInteger
forall b c a. (b -> c) -> (a -> b) -> a -> c
. GaussianInteger -> (GaussianInteger, GaussianInteger)
absSignum

instance S.Semiring GaussianInteger where
    plus :: GaussianInteger -> GaussianInteger -> GaussianInteger
plus          = GaussianInteger -> GaussianInteger -> GaussianInteger
forall a. Num a => a -> a -> a
(+)
    times :: GaussianInteger -> GaussianInteger -> GaussianInteger
times         = GaussianInteger -> GaussianInteger -> GaussianInteger
forall a. Num a => a -> a -> a
(*)
    zero :: GaussianInteger
zero          = Integer
0 Integer -> Integer -> GaussianInteger
:+ Integer
0
    one :: GaussianInteger
one           = Integer
1 Integer -> Integer -> GaussianInteger
:+ Integer
0
    fromNatural :: Natural -> GaussianInteger
fromNatural Natural
n = Natural -> Integer
naturalToInteger Natural
n Integer -> Integer -> GaussianInteger
:+ Integer
0

instance S.Ring GaussianInteger where
    negate :: GaussianInteger -> GaussianInteger
negate = GaussianInteger -> GaussianInteger
forall a. Num a => a -> a
negate

absSignum :: GaussianInteger -> (GaussianInteger, GaussianInteger)
absSignum :: GaussianInteger -> (GaussianInteger, GaussianInteger)
absSignum GaussianInteger
0 = (GaussianInteger
0, GaussianInteger
0)
absSignum z :: GaussianInteger
z@(Integer
a :+ Integer
b)
  -- first quadrant: (0, inf) x [0, inf)i
  | Integer
a Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
>  Integer
0 Bool -> Bool -> Bool
&& Integer
b Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
>= Integer
0 = (GaussianInteger
z, GaussianInteger
1)
  -- second quadrant: (-inf, 0] x (0, inf)i
  | Integer
a Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
0 Bool -> Bool -> Bool
&& Integer
b Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
>  Integer
0 = (Integer
b Integer -> Integer -> GaussianInteger
:+ (-Integer
a), GaussianInteger
ι)
  -- third quadrant: (-inf, 0) x (-inf, 0]i
  | Integer
a Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<  Integer
0 Bool -> Bool -> Bool
&& Integer
b Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
0 = (-GaussianInteger
z, -GaussianInteger
1)
  -- fourth quadrant: [0, inf) x (-inf, 0)i
  | Bool
otherwise        = ((-Integer
b) Integer -> Integer -> GaussianInteger
:+ Integer
a, -GaussianInteger
ι)

instance GcdDomain GaussianInteger

instance Euclidean GaussianInteger where
  degree :: GaussianInteger -> Natural
degree = Integer -> Natural
forall a. Num a => Integer -> a
fromInteger (Integer -> Natural)
-> (GaussianInteger -> Integer) -> GaussianInteger -> Natural
forall b c a. (b -> c) -> (a -> b) -> a -> c
. GaussianInteger -> Integer
norm
  quotRem :: GaussianInteger
-> GaussianInteger -> (GaussianInteger, GaussianInteger)
quotRem GaussianInteger
x (Integer
d :+ Integer
0) = GaussianInteger -> Integer -> (GaussianInteger, GaussianInteger)
quotRemInt GaussianInteger
x Integer
d
  quotRem GaussianInteger
x GaussianInteger
y = (GaussianInteger
q, GaussianInteger
x GaussianInteger -> GaussianInteger -> GaussianInteger
forall a. Num a => a -> a -> a
- GaussianInteger
q GaussianInteger -> GaussianInteger -> GaussianInteger
forall a. Num a => a -> a -> a
* GaussianInteger
y)
    where
      (GaussianInteger
q, GaussianInteger
_) = GaussianInteger -> Integer -> (GaussianInteger, GaussianInteger)
quotRemInt (GaussianInteger
x GaussianInteger -> GaussianInteger -> GaussianInteger
forall a. Num a => a -> a -> a
* GaussianInteger -> GaussianInteger
conjugate GaussianInteger
y) (GaussianInteger -> Integer
norm GaussianInteger
y)

quotRemInt :: GaussianInteger -> Integer -> (GaussianInteger, GaussianInteger)
quotRemInt :: GaussianInteger -> Integer -> (GaussianInteger, GaussianInteger)
quotRemInt GaussianInteger
z   Integer
1  = ( GaussianInteger
z, GaussianInteger
0)
quotRemInt GaussianInteger
z (-1) = (-GaussianInteger
z, GaussianInteger
0)
quotRemInt (Integer
a :+ Integer
b) Integer
c = (Integer
qa Integer -> Integer -> GaussianInteger
:+ Integer
qb, (Integer
ra Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
bumpA) Integer -> Integer -> GaussianInteger
:+ (Integer
rb Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
bumpB))
  where
    halfC :: Integer
halfC    = Integer -> Integer
forall a. Num a => a -> a
abs Integer
c Integer -> Integer -> Integer
forall a. Euclidean a => a -> a -> a
`quot` Integer
2
    bumpA :: Integer
bumpA    = Integer -> Integer
forall a. Num a => a -> a
signum Integer
a Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
halfC
    bumpB :: Integer
bumpB    = Integer -> Integer
forall a. Num a => a -> a
signum Integer
b Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
halfC
    (Integer
qa, Integer
ra) = (Integer
a Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
bumpA) Integer -> Integer -> (Integer, Integer)
forall a. Euclidean a => a -> a -> (a, a)
`quotRem` Integer
c
    (Integer
qb, Integer
rb) = (Integer
b Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
bumpB) Integer -> Integer -> (Integer, Integer)
forall a. Euclidean a => a -> a -> (a, a)
`quotRem` Integer
c

-- |Conjugate a Gaussian integer.
conjugate :: GaussianInteger -> GaussianInteger
conjugate :: GaussianInteger -> GaussianInteger
conjugate (Integer
r :+ Integer
i) = Integer
r Integer -> Integer -> GaussianInteger
:+ (-Integer
i)

-- |The square of the magnitude of a Gaussian integer.
norm :: GaussianInteger -> Integer
norm :: GaussianInteger -> Integer
norm (Integer
x :+ Integer
y) = Integer
x Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
x Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
y Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
y

-- |Compute whether a given Gaussian integer is prime.
isPrime :: GaussianInteger -> Bool
isPrime :: GaussianInteger -> Bool
isPrime g :: GaussianInteger
g@(Integer
x :+ Integer
y)
    | Integer
x Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
0 Bool -> Bool -> Bool
&& Integer
y Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
/= Integer
0 = Integer -> Integer
forall a. Num a => a -> a
abs Integer
y Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
4 Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
3 Bool -> Bool -> Bool
&& Maybe (Prime Integer) -> Bool
forall a. Maybe a -> Bool
isJust (Integer -> Maybe (Prime Integer)
forall a. UniqueFactorisation a => a -> Maybe (Prime a)
U.isPrime Integer
y)
    | Integer
y Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
0 Bool -> Bool -> Bool
&& Integer
x Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
/= Integer
0 = Integer -> Integer
forall a. Num a => a -> a
abs Integer
x Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
4 Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
3 Bool -> Bool -> Bool
&& Maybe (Prime Integer) -> Bool
forall a. Maybe a -> Bool
isJust (Integer -> Maybe (Prime Integer)
forall a. UniqueFactorisation a => a -> Maybe (Prime a)
U.isPrime Integer
x)
    | Bool
otherwise        = Maybe (Prime Integer) -> Bool
forall a. Maybe a -> Bool
isJust (Maybe (Prime Integer) -> Bool) -> Maybe (Prime Integer) -> Bool
forall a b. (a -> b) -> a -> b
$ Integer -> Maybe (Prime Integer)
forall a. UniqueFactorisation a => a -> Maybe (Prime a)
U.isPrime (Integer -> Maybe (Prime Integer))
-> Integer -> Maybe (Prime Integer)
forall a b. (a -> b) -> a -> b
$ GaussianInteger -> Integer
norm GaussianInteger
g

-- |An infinite list of the Gaussian primes. Uses primes in Z to exhaustively
-- generate all Gaussian primes (up to associates), in order of ascending
-- magnitude.
--
-- >>> take 10 primes
-- [Prime 1+ι,Prime 2+ι,Prime 1+2*ι,Prime 3,Prime 3+2*ι,Prime 2+3*ι,Prime 4+ι,Prime 1+4*ι,Prime 5+2*ι,Prime 2+5*ι]
primes :: [U.Prime GaussianInteger]
primes :: [Prime GaussianInteger]
primes = [GaussianInteger] -> [Prime GaussianInteger]
coerce ([GaussianInteger] -> [Prime GaussianInteger])
-> [GaussianInteger] -> [Prime GaussianInteger]
forall a b. (a -> b) -> a -> b
$ (Integer
1 Integer -> Integer -> GaussianInteger
:+ Integer
1) GaussianInteger -> [GaussianInteger] -> [GaussianInteger]
forall a. a -> [a] -> [a]
: (GaussianInteger -> GaussianInteger -> Ordering)
-> [GaussianInteger] -> [GaussianInteger] -> [GaussianInteger]
forall a. (a -> a -> Ordering) -> [a] -> [a] -> [a]
mergeBy ((GaussianInteger -> Integer)
-> GaussianInteger -> GaussianInteger -> Ordering
forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing GaussianInteger -> Integer
norm) [GaussianInteger]
l [GaussianInteger]
r
  where
    leftPrimes, rightPrimes :: [Prime Integer]
    ([Prime Integer]
leftPrimes, [Prime Integer]
rightPrimes) = (Prime Integer -> Bool)
-> [Prime Integer] -> ([Prime Integer], [Prime Integer])
forall a. (a -> Bool) -> [a] -> ([a], [a])
partition (\Prime Integer
p -> Prime Integer -> Integer
forall a. Prime a -> a
unPrime Prime Integer
p Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
4 Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
3) [Integer -> Prime Integer
forall a.
(Bits a, Integral a, UniqueFactorisation a) =>
a -> Prime a
U.nextPrime Integer
3 ..]
    l :: [GaussianInteger]
l = [Prime Integer -> Integer
forall a. Prime a -> a
unPrime Prime Integer
p Integer -> Integer -> GaussianInteger
:+ Integer
0 | Prime Integer
p <- [Prime Integer]
leftPrimes]
    r :: [GaussianInteger]
r = [GaussianInteger
g | Prime Integer
p <- [Prime Integer]
rightPrimes, let Prime (Integer
x :+ Integer
y) = Prime Integer -> Prime GaussianInteger
findPrime Prime Integer
p, GaussianInteger
g <- [Integer
x Integer -> Integer -> GaussianInteger
:+ Integer
y, Integer
y Integer -> Integer -> GaussianInteger
:+ Integer
x]]


-- |Find a Gaussian integer whose norm is the given prime number
-- of form 4k + 1 using
-- <http://www.ams.org/journals/mcom/1972-26-120/S0025-5718-1972-0314745-6/S0025-5718-1972-0314745-6.pdf Hermite-Serret algorithm>.
--
-- >>> import Math.NumberTheory.Primes (nextPrime)
-- >>> findPrime (nextPrime 5)
-- Prime 2+ι
findPrime :: Prime Integer -> U.Prime GaussianInteger
findPrime :: Prime Integer -> Prime GaussianInteger
findPrime Prime Integer
p = case Integer -> Prime Integer -> [Integer]
sqrtsModPrime (-Integer
1) Prime Integer
p of
    []    -> String -> Prime GaussianInteger
forall a. HasCallStack => String -> a
error String
"findPrime: an argument must be prime p = 4k + 1"
    Integer
z : [Integer]
_ -> GaussianInteger -> Prime GaussianInteger
forall a. a -> Prime a
Prime (GaussianInteger -> Prime GaussianInteger)
-> GaussianInteger -> Prime GaussianInteger
forall a b. (a -> b) -> a -> b
$ Integer -> Integer -> GaussianInteger
go (Prime Integer -> Integer
forall a. Prime a -> a
unPrime Prime Integer
p) Integer
z -- Effectively we calculate gcdG' (p :+ 0) (z :+ 1)
    where
        sqrtp :: Integer
        sqrtp :: Integer
sqrtp = Integer -> Integer
forall a. Integral a => a -> a
integerSquareRoot (Prime Integer -> Integer
forall a. Prime a -> a
unPrime Prime Integer
p)

        go :: Integer -> Integer -> GaussianInteger
        go :: Integer -> Integer -> GaussianInteger
go Integer
g Integer
h
            | Integer
g Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
sqrtp = Integer
g Integer -> Integer -> GaussianInteger
:+ Integer
h
            | Bool
otherwise  = Integer -> Integer -> GaussianInteger
go Integer
h (Integer
g Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
h)

-- | Compute the prime factorisation of a Gaussian integer. This is unique up to units (+/- 1, +/- i).
-- Unit factors are not included in the result.
factorise :: GaussianInteger -> [(Prime GaussianInteger, Word)]
factorise :: GaussianInteger -> [(Prime GaussianInteger, Word)]
factorise GaussianInteger
g = [[(Prime GaussianInteger, Word)]]
-> [(Prime GaussianInteger, Word)]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[(Prime GaussianInteger, Word)]]
 -> [(Prime GaussianInteger, Word)])
-> [[(Prime GaussianInteger, Word)]]
-> [(Prime GaussianInteger, Word)]
forall a b. (a -> b) -> a -> b
$ (GaussianInteger, [[(Prime GaussianInteger, Word)]])
-> [[(Prime GaussianInteger, Word)]]
forall a b. (a, b) -> b
snd ((GaussianInteger, [[(Prime GaussianInteger, Word)]])
 -> [[(Prime GaussianInteger, Word)]])
-> (GaussianInteger, [[(Prime GaussianInteger, Word)]])
-> [[(Prime GaussianInteger, Word)]]
forall a b. (a -> b) -> a -> b
$ (GaussianInteger
 -> (Prime Integer, Word)
 -> (GaussianInteger, [(Prime GaussianInteger, Word)]))
-> GaussianInteger
-> [(Prime Integer, Word)]
-> (GaussianInteger, [[(Prime GaussianInteger, Word)]])
forall (t :: * -> *) a b c.
Traversable t =>
(a -> b -> (a, c)) -> a -> t b -> (a, t c)
mapAccumL GaussianInteger
-> (Prime Integer, Word)
-> (GaussianInteger, [(Prime GaussianInteger, Word)])
go GaussianInteger
g (Integer -> [(Prime Integer, Word)]
forall a. UniqueFactorisation a => a -> [(Prime a, Word)]
U.factorise (Integer -> [(Prime Integer, Word)])
-> Integer -> [(Prime Integer, Word)]
forall a b. (a -> b) -> a -> b
$ GaussianInteger -> Integer
norm GaussianInteger
g)
    where
        go :: GaussianInteger -> (Prime Integer, Word) -> (GaussianInteger, [(Prime GaussianInteger, Word)])
        go :: GaussianInteger
-> (Prime Integer, Word)
-> (GaussianInteger, [(Prime GaussianInteger, Word)])
go GaussianInteger
z (Prime Integer
2, Word
e) = (GaussianInteger -> GaussianInteger
divideByTwo GaussianInteger
z, [(GaussianInteger -> Prime GaussianInteger
forall a. a -> Prime a
Prime (Integer
1 Integer -> Integer -> GaussianInteger
:+ Integer
1), Word
e)])
        go GaussianInteger
z (Prime Integer
p, Word
e)
            | Prime Integer -> Integer
forall a. Prime a -> a
unPrime Prime Integer
p Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
4 Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
3
            = let e' :: Word
e' = Word
e Word -> Word -> Word
forall a. Euclidean a => a -> a -> a
`quot` Word
2 in (GaussianInteger
z GaussianInteger -> Integer -> GaussianInteger
`quotI` (Prime Integer -> Integer
forall a. Prime a -> a
unPrime Prime Integer
p Integer -> Word -> Integer
forall a b. (Num a, Integral b) => a -> b -> a
^ Word
e'), [(GaussianInteger -> Prime GaussianInteger
forall a. a -> Prime a
Prime (Prime Integer -> Integer
forall a. Prime a -> a
unPrime Prime Integer
p Integer -> Integer -> GaussianInteger
:+ Integer
0), Word
e')])
            | Bool
otherwise
            = (GaussianInteger
z', ((Prime GaussianInteger, Word) -> Bool)
-> [(Prime GaussianInteger, Word)]
-> [(Prime GaussianInteger, Word)]
forall a. (a -> Bool) -> [a] -> [a]
filter ((Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
> Word
0) (Word -> Bool)
-> ((Prime GaussianInteger, Word) -> Word)
-> (Prime GaussianInteger, Word)
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Prime GaussianInteger, Word) -> Word
forall a b. (a, b) -> b
snd) [(Prime GaussianInteger
gp, Word
k), (Prime GaussianInteger
gp', Word
k')])
                where
                    gp :: Prime GaussianInteger
gp = Prime Integer -> Prime GaussianInteger
findPrime Prime Integer
p
                    (Word
k, Word
k', GaussianInteger
z') = Prime GaussianInteger
-> Integer
-> Word
-> GaussianInteger
-> (Word, Word, GaussianInteger)
divideByPrime Prime GaussianInteger
gp (Prime Integer -> Integer
forall a. Prime a -> a
unPrime Prime Integer
p) Word
e GaussianInteger
z
                    gp' :: Prime GaussianInteger
gp' = GaussianInteger -> Prime GaussianInteger
forall a. a -> Prime a
Prime (GaussianInteger -> GaussianInteger
forall a. Num a => a -> a
abs (GaussianInteger -> GaussianInteger
conjugate (Prime GaussianInteger -> GaussianInteger
forall a. Prime a -> a
unPrime Prime GaussianInteger
gp)))

-- | Remove all (1:+1) factors from the argument,
-- avoiding complex division.
divideByTwo :: GaussianInteger -> GaussianInteger
divideByTwo :: GaussianInteger -> GaussianInteger
divideByTwo z :: GaussianInteger
z@(Integer
x :+ Integer
y)
    | Integer -> Bool
forall a. Integral a => a -> Bool
even Integer
x, Integer -> Bool
forall a. Integral a => a -> Bool
even Integer
y
    = GaussianInteger -> GaussianInteger
divideByTwo (GaussianInteger -> GaussianInteger)
-> GaussianInteger -> GaussianInteger
forall a b. (a -> b) -> a -> b
$ GaussianInteger
z GaussianInteger -> Integer -> GaussianInteger
`quotI` Integer
2
    | Integer -> Bool
forall a. Integral a => a -> Bool
odd Integer
x, Integer -> Bool
forall a. Integral a => a -> Bool
odd Integer
y
    = (Integer
x Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
y) Integer -> Integer -> Integer
forall a. Euclidean a => a -> a -> a
`quot` Integer
2 Integer -> Integer -> GaussianInteger
:+ (Integer
x Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
y) Integer -> Integer -> Integer
forall a. Euclidean a => a -> a -> a
`quot` Integer
2
    | Bool
otherwise
    = GaussianInteger
z

-- | Remove p and conj p factors from the argument,
-- avoiding complex division.
divideByPrime
    :: Prime GaussianInteger -- ^ Gaussian prime p
    -> Integer               -- ^ Precomputed norm of p, of form 4k + 1
    -> Word                  -- ^ Expected number of factors (either p or conj p)
                             --   in Gaussian integer z
    -> GaussianInteger       -- ^ Gaussian integer z
    -> ( Word                -- Multiplicity of factor p in z
       , Word                -- Multiplicity of factor conj p in z
       , GaussianInteger     -- Remaining Gaussian integer
       )
divideByPrime :: Prime GaussianInteger
-> Integer
-> Word
-> GaussianInteger
-> (Word, Word, GaussianInteger)
divideByPrime Prime GaussianInteger
p Integer
np Word
k = Word -> Word -> GaussianInteger -> (Word, Word, GaussianInteger)
go Word
k Word
0
    where
        go :: Word -> Word -> GaussianInteger -> (Word, Word, GaussianInteger)
        go :: Word -> Word -> GaussianInteger -> (Word, Word, GaussianInteger)
go Word
0 Word
d GaussianInteger
z = (Word
d, Word
d, GaussianInteger
z)
        go Word
c Word
d GaussianInteger
z
            | Word
c Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
>= Word
2
            , Just GaussianInteger
z' <- GaussianInteger
z GaussianInteger -> Integer -> Maybe GaussianInteger
`quotEvenI` Integer
np
            = Word -> Word -> GaussianInteger -> (Word, Word, GaussianInteger)
go (Word
c Word -> Word -> Word
forall a. Num a => a -> a -> a
- Word
2) (Word
d Word -> Word -> Word
forall a. Num a => a -> a -> a
+ Word
1) GaussianInteger
z'
        go Word
c Word
d GaussianInteger
z = (Word
d Word -> Word -> Word
forall a. Num a => a -> a -> a
+ Word
d1, Word
d Word -> Word -> Word
forall a. Num a => a -> a -> a
+ Word
d2, GaussianInteger
z'')
            where
                (Word
d1, GaussianInteger
z') = Word -> Word -> GaussianInteger -> (Word, GaussianInteger)
go1 Word
c Word
0 GaussianInteger
z
                d2 :: Word
d2 = Word
c Word -> Word -> Word
forall a. Num a => a -> a -> a
- Word
d1
                z'' :: GaussianInteger
z'' = (GaussianInteger -> GaussianInteger)
-> GaussianInteger -> [GaussianInteger]
forall a. (a -> a) -> a -> [a]
iterate (\GaussianInteger
g -> GaussianInteger -> Maybe GaussianInteger -> GaussianInteger
forall a. a -> Maybe a -> a
fromMaybe GaussianInteger
err (Maybe GaussianInteger -> GaussianInteger)
-> Maybe GaussianInteger -> GaussianInteger
forall a b. (a -> b) -> a -> b
$ (GaussianInteger
g GaussianInteger -> GaussianInteger -> GaussianInteger
forall a. Num a => a -> a -> a
* Prime GaussianInteger -> GaussianInteger
forall a. Prime a -> a
unPrime Prime GaussianInteger
p) GaussianInteger -> Integer -> Maybe GaussianInteger
`quotEvenI` Integer
np) GaussianInteger
z' [GaussianInteger] -> Int -> GaussianInteger
forall a. [a] -> Int -> a
!! Word -> Int
wordToInt Word
d2

        go1 :: Word -> Word -> GaussianInteger -> (Word, GaussianInteger)
        go1 :: Word -> Word -> GaussianInteger -> (Word, GaussianInteger)
go1 Word
0 Word
d GaussianInteger
z = (Word
d, GaussianInteger
z)
        go1 Word
c Word
d GaussianInteger
z
            | Just GaussianInteger
z' <- (GaussianInteger
z GaussianInteger -> GaussianInteger -> GaussianInteger
forall a. Num a => a -> a -> a
* GaussianInteger -> GaussianInteger
conjugate (Prime GaussianInteger -> GaussianInteger
forall a. Prime a -> a
unPrime Prime GaussianInteger
p)) GaussianInteger -> Integer -> Maybe GaussianInteger
`quotEvenI` Integer
np
            = Word -> Word -> GaussianInteger -> (Word, GaussianInteger)
go1 (Word
c Word -> Word -> Word
forall a. Num a => a -> a -> a
- Word
1) (Word
d Word -> Word -> Word
forall a. Num a => a -> a -> a
+ Word
1) GaussianInteger
z'
            | Bool
otherwise
            = (Word
d, GaussianInteger
z)

        err :: GaussianInteger
err = String -> GaussianInteger
forall a. HasCallStack => String -> a
error (String -> GaussianInteger) -> String -> GaussianInteger
forall a b. (a -> b) -> a -> b
$ String
"divideByPrime: malformed arguments" String -> ShowS
forall a. [a] -> [a] -> [a]
++ (Prime GaussianInteger, Integer, Word) -> String
forall a. Show a => a -> String
show (Prime GaussianInteger
p, Integer
np, Word
k)

quotI :: GaussianInteger -> Integer -> GaussianInteger
quotI :: GaussianInteger -> Integer -> GaussianInteger
quotI (Integer
x :+ Integer
y) Integer
n = Integer
x Integer -> Integer -> Integer
forall a. Euclidean a => a -> a -> a
`quot` Integer
n Integer -> Integer -> GaussianInteger
:+ Integer
y Integer -> Integer -> Integer
forall a. Euclidean a => a -> a -> a
`quot` Integer
n

quotEvenI :: GaussianInteger -> Integer -> Maybe GaussianInteger
quotEvenI :: GaussianInteger -> Integer -> Maybe GaussianInteger
quotEvenI (Integer
x :+ Integer
y) Integer
n
    | Integer
xr Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
0
    , Integer
yr Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
0
    = GaussianInteger -> Maybe GaussianInteger
forall a. a -> Maybe a
Just (Integer
xq Integer -> Integer -> GaussianInteger
:+ Integer
yq)
    | Bool
otherwise
    = Maybe GaussianInteger
forall a. Maybe a
Nothing
    where
        (Integer
xq, Integer
xr) = Integer
x Integer -> Integer -> (Integer, Integer)
forall a. Euclidean a => a -> a -> (a, a)
`quotRem` Integer
n
        (Integer
yq, Integer
yr) = Integer
y Integer -> Integer -> (Integer, Integer)
forall a. Euclidean a => a -> a -> (a, a)
`quotRem` Integer
n

-------------------------------------------------------------------------------

instance U.UniqueFactorisation GaussianInteger where
  factorise :: GaussianInteger -> [(Prime GaussianInteger, Word)]
factorise GaussianInteger
0 = []
  factorise GaussianInteger
g = [(Prime GaussianInteger, Word)] -> [(Prime GaussianInteger, Word)]
coerce ([(Prime GaussianInteger, Word)]
 -> [(Prime GaussianInteger, Word)])
-> [(Prime GaussianInteger, Word)]
-> [(Prime GaussianInteger, Word)]
forall a b. (a -> b) -> a -> b
$ GaussianInteger -> [(Prime GaussianInteger, Word)]
factorise GaussianInteger
g

  isPrime :: GaussianInteger -> Maybe (Prime GaussianInteger)
isPrime GaussianInteger
g = if GaussianInteger -> Bool
isPrime GaussianInteger
g then Prime GaussianInteger -> Maybe (Prime GaussianInteger)
forall a. a -> Maybe a
Just (GaussianInteger -> Prime GaussianInteger
forall a. a -> Prime a
Prime GaussianInteger
g) else Maybe (Prime GaussianInteger)
forall a. Maybe a
Nothing