-- |
-- 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 = (:+) { real :: !Integer, imag :: !Integer }
    deriving (Eq, Ord, Generic)

instance NFData GaussianInteger

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

instance Show GaussianInteger where
    show (a :+ b)
        | b == 0     = show a
        | a == 0     = s ++ b'
        | otherwise  = show a ++ op ++ b'
        where
            b' = if abs b == 1 then "ι" else show (abs b) ++ "*ι"
            op = if b > 0 then "+" else "-"
            s  = if b > 0 then "" else "-"

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

instance S.Semiring GaussianInteger where
    plus          = (+)
    times         = (*)
    zero          = 0 :+ 0
    one           = 1 :+ 0
    fromNatural n = naturalToInteger n :+ 0

instance S.Ring GaussianInteger where
    negate = negate

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

instance GcdDomain GaussianInteger

instance Euclidean GaussianInteger where
  degree = fromInteger . norm
  quotRem x (d :+ 0) = quotRemInt x d
  quotRem x y = (q, x - q * y)
    where
      (q, _) = quotRemInt (x * conjugate y) (norm y)

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

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

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

-- |Compute whether a given Gaussian integer is prime.
isPrime :: GaussianInteger -> Bool
isPrime g@(x :+ y)
    | x == 0 && y /= 0 = abs y `mod` 4 == 3 && isJust (U.isPrime y)
    | y == 0 && x /= 0 = abs x `mod` 4 == 3 && isJust (U.isPrime x)
    | otherwise        = isJust $ U.isPrime $ norm 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 = coerce $ (1 :+ 1) :< mergeBy (comparing norm) l r
  where
    leftPrimes, rightPrimes :: Infinite (Prime Integer)
    (leftPrimes, rightPrimes) = Inf.partition (\p -> unPrime p `mod` 4 == 3) (U.nextPrime 3 ...)

    l :: Infinite (GaussianInteger)
    l = fmap (\p -> unPrime p :+ 0) leftPrimes

    r :: Infinite (GaussianInteger)
    r = Inf.concatMap
        (\p -> let x :+ y = unPrime (findPrime p) in (x :+ y) :| [y :+ x])
        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 p = case sqrtsModPrime (-1) p of
    []    -> error "findPrime: an argument must be prime p = 4k + 1"
    z : _ -> Prime $ go (unPrime p) z -- Effectively we calculate gcdG' (p :+ 0) (z :+ 1)
    where
        sqrtp :: Integer
        sqrtp = integerSquareRoot (unPrime p)

        go :: Integer -> Integer -> GaussianInteger
        go g h
            | g <= sqrtp = g :+ h
            | otherwise  = go h (g `mod` 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 g = concat $ snd $ mapAccumL go g (U.factorise $ norm g)
    where
        go :: GaussianInteger -> (Prime Integer, Word) -> (GaussianInteger, [(Prime GaussianInteger, Word)])
        go z (Prime 2, e) = (divideByTwo z, [(Prime (1 :+ 1), e)])
        go z (p, e)
            | unPrime p `mod` 4 == 3
            = let e' = e `quot` 2 in (z `quotI` (unPrime p ^ e'), [(Prime (unPrime p :+ 0), e')])
            | otherwise
            = (z', filter ((> 0) . snd) [(gp, k), (gp', k')])
                where
                    gp = findPrime p
                    (k, k', z') = divideByPrime gp (unPrime p) e z
                    gp' = Prime (abs (conjugate (unPrime gp)))

-- | Remove all (1:+1) factors from the argument,
-- avoiding complex division.
divideByTwo :: GaussianInteger -> GaussianInteger
divideByTwo z@(x :+ y)
    | even x, even y
    = divideByTwo $ z `quotI` 2
    | odd x, odd y
    = (x - y) `quot` 2 :+ (x + y) `quot` 2
    | otherwise
    = 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 p np k = go k 0
    where
        go :: Word -> Word -> GaussianInteger -> (Word, Word, GaussianInteger)
        go 0 d z = (d, d, z)
        go c d z
            | c >= 2
            , Just z' <- z `quotEvenI` np
            = go (c - 2) (d + 1) z'
        go c d z = (d + d1, d + d2, z'')
            where
                (d1, z') = go1 c 0 z
                d2 = c - d1
                z'' = iterate (\g -> fromMaybe err $ (g * unPrime p) `quotEvenI` np) z' !! wordToInt d2

        go1 :: Word -> Word -> GaussianInteger -> (Word, GaussianInteger)
        go1 0 d z = (d, z)
        go1 c d z
            | Just z' <- (z * conjugate (unPrime p)) `quotEvenI` np
            = go1 (c - 1) (d + 1) z'
            | otherwise
            = (d, z)

        err = error $ "divideByPrime: malformed arguments" ++ show (p, np, k)

quotI :: GaussianInteger -> Integer -> GaussianInteger
quotI (x :+ y) n = x `quot` n :+ y `quot` n

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

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

instance U.UniqueFactorisation GaussianInteger where
  factorise 0 = []
  factorise g = coerce $ factorise g

  isPrime g = if isPrime g then Just (Prime g) else Nothing