-- |
-- 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 PostfixOperators #-}
{-# 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)
import Data.List.Infinite (Infinite(..), (...))
import qualified Data.List.Infinite as Inf
import Data.List.NonEmpty (NonEmpty(..))
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
$c== :: GaussianInteger -> GaussianInteger -> Bool
== :: GaussianInteger -> GaussianInteger -> Bool
$c/= :: GaussianInteger -> GaussianInteger -> Bool
/= :: 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
$ccompare :: GaussianInteger -> GaussianInteger -> Ordering
compare :: GaussianInteger -> GaussianInteger -> Ordering
$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
>= :: GaussianInteger -> GaussianInteger -> Bool
$cmax :: GaussianInteger -> GaussianInteger -> GaussianInteger
max :: GaussianInteger -> GaussianInteger -> GaussianInteger
$cmin :: GaussianInteger -> GaussianInteger -> GaussianInteger
min :: GaussianInteger -> GaussianInteger -> 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
$cfrom :: forall x. GaussianInteger -> Rep GaussianInteger x
from :: forall x. GaussianInteger -> Rep GaussianInteger x
$cto :: forall x. Rep GaussianInteger x -> GaussianInteger
to :: forall x. Rep GaussianInteger x -> GaussianInteger
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 :: Infinite (U.Prime GaussianInteger)
primes :: Infinite (Prime GaussianInteger)
primes = Infinite GaussianInteger -> Infinite (Prime GaussianInteger)
forall a b. Coercible a b => a -> b
coerce (Infinite GaussianInteger -> Infinite (Prime GaussianInteger))
-> Infinite GaussianInteger -> Infinite (Prime GaussianInteger)
forall a b. (a -> b) -> a -> b
$ (Integer
1 Integer -> Integer -> GaussianInteger
:+ Integer
1) GaussianInteger
-> Infinite GaussianInteger -> Infinite GaussianInteger
forall a. a -> Infinite a -> Infinite a
:< (GaussianInteger -> GaussianInteger -> Ordering)
-> Infinite GaussianInteger
-> Infinite GaussianInteger
-> Infinite GaussianInteger
forall a.
(a -> a -> Ordering) -> Infinite a -> Infinite a -> Infinite a
mergeBy ((GaussianInteger -> Integer)
-> GaussianInteger -> GaussianInteger -> Ordering
forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing GaussianInteger -> Integer
norm) Infinite GaussianInteger
l Infinite GaussianInteger
r
  where
    leftPrimes, rightPrimes :: Infinite (Prime Integer)
    (Infinite (Prime Integer)
leftPrimes, Infinite (Prime Integer)
rightPrimes) = (Prime Integer -> Bool)
-> Infinite (Prime Integer)
-> (Infinite (Prime Integer), Infinite (Prime Integer))
forall a. (a -> Bool) -> Infinite a -> (Infinite a, Infinite a)
Inf.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 :: Infinite (GaussianInteger)
    l :: Infinite GaussianInteger
l = (Prime Integer -> GaussianInteger)
-> Infinite (Prime Integer) -> Infinite GaussianInteger
forall a b. (a -> b) -> Infinite a -> Infinite b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\Prime Integer
p -> Prime Integer -> Integer
forall a. Prime a -> a
unPrime Prime Integer
p Integer -> Integer -> GaussianInteger
:+ Integer
0) Infinite (Prime Integer)
leftPrimes

    r :: Infinite (GaussianInteger)
    r :: Infinite GaussianInteger
r = (Prime Integer -> NonEmpty GaussianInteger)
-> Infinite (Prime Integer) -> Infinite GaussianInteger
forall a b. (a -> NonEmpty b) -> Infinite a -> Infinite b
Inf.concatMap
        (\Prime Integer
p -> let Integer
x :+ Integer
y = Prime GaussianInteger -> GaussianInteger
forall a. Prime a -> a
unPrime (Prime Integer -> Prime GaussianInteger
findPrime Prime Integer
p) in (Integer
x Integer -> Integer -> GaussianInteger
:+ Integer
y) GaussianInteger -> [GaussianInteger] -> NonEmpty GaussianInteger
forall a. a -> [a] -> NonEmpty a
:| [Integer
y Integer -> Integer -> GaussianInteger
:+ Integer
x])
        Infinite (Prime Integer)
rightPrimes

-- |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 :: * -> *) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
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. HasCallStack => [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)]
forall a b. Coercible a b => a -> b
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