{-# OPTIONS -fglasgow-exts #-}

module PrimesIsInP_2
    ( isPrime
    ) where

import List (find, unfoldr)
import Maybe (fromMaybe)
import Monad (when)
import Data.Array.ST
import GHC.ST

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

isPrime :: Int -> Bool
isPrime n
    | n <= 1    = False
    | isExp n   = False
    | otherwise =
        case findR n of
        Just r  -> testIdentity n r
        Nothing -> False

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

-- check if n == a^b for some natural number a and b > 1.
isExp :: Int -> Bool
isExp n = any f [2 .. floor (logBase 2 n')]
    where f b = a == fromInteger (floor a)
              where a  = exp (log n' / b')
                    b' = fromInteger b
          n' = fromInteger (toInteger n)

findR :: Int -> Maybe Int
findR n = f 2
    where f r | r >= n       = Just r
              | gcd n r /= 1 = Nothing
              | test         = Just r
              | otherwise    = f (r+1)
              where test  = (r `elem` primes r) &&
                            (q' >= 4 * sqrt r' * logBase 2 n') &&
                            ((n ^ floor (fromInteger (toInteger (r-1)) / q')) `mod` r /= 1)
                    q  = largestPrimeFactor (r-1)
                    n' = fromInteger (toInteger n)
                    r' = fromInteger (toInteger r)
                    q' = fromInteger (toInteger q)

allM :: (Monad m) => (a -> m Bool) -> [a] -> m Bool
allM _ []     = return True
allM f (a:ax) = do b <- f a
                   if b
                       then allM f ax
                       else return False

testIdentity :: Int -> Int -> Bool
testIdentity n r = runST st
    where st :: forall s. ST s Bool
          st = do base <- newArray (0,r-1) 0
                  tmp  <- newArray (0,r-1) 0
                  lhs  <- newArray (0,r-1) 0
                  rhs  <- newArray (0,r-1) 0
                  allM (test n r base tmp lhs rhs) [1..last]
          last = floor (2 * sqrt (fromInteger (toInteger r)) * logBase 2 (fromInteger (toInteger n)))

test :: Int -> Int ->
        STUArray s Int Int ->
        STUArray s Int Int ->
        STUArray s Int Int ->
        STUArray s Int Int ->
        Int ->
        ST s Bool
test n r base tmp lhs rhs a =
    do mapM_ (\i -> writeArray base i 0) [0..(r-1)]
       writeArray base 0 (-a `mod` n)
       writeArray base 1 1
       polyExp n r base n lhs tmp

       mapM_ (\i -> writeArray rhs i 0) [0..(r-1)]
       writeArray rhs (n `mod` r) 1
       rhs0 <- readArray rhs 0
       writeArray rhs 0 ((rhs0-a) `mod` n)

       allM (\i -> do x <- readArray lhs i
                      y <- readArray rhs i
                      return (x==y))
            [0..(r-1)]

polyExp :: (MArray a Int m) =>
           Int -> Int ->
           a Int Int ->
           Int ->
           a Int Int ->
           a Int Int ->
           m ()
polyExp n r base exp dest tmp = func exp tmp dest
    where func 1 _ dest =
              mapM_ (\i -> do e <- readArray base i; writeArray dest i e)
                    [0..(r-1)]
          func x tmp dest =
              case x `divMod` 2 of
              (d,0) ->
                  do func d dest tmp
                     polyMult n r tmp tmp dest
              _     ->
                  do func (x-1) dest tmp
                     polyMult n r base tmp dest

polyMult :: (MArray a Int m) =>
            Int -> Int ->
            a Int Int ->
            a Int Int ->
            a Int Int ->
            m ()
polyMult n r a b dest = seq n $ seq r $ seq a $ seq b $ seq dest $
    do mapM_ (\i -> writeArray dest i 0) [0..(r-1)]
       mapM_ f [0..(r-1)]
    where f i =
              do x <- readArray a i
                 when (x /= 0) (mapM_ (g x) [0..(r-1)])
              where g x j =
                        do y <- readArray b j
                           when (y /= 0) $
                                do k <- return ((i+j) `mod` r)
                                   z <- readArray dest k
                                   writeArray dest k ((z+(x*y)) `mod` n)

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

primes :: Int -> [Int]
primes n = unfoldr f [2..n]
    where f (a:ax) = Just (a, sieve a ax)
          f []     = Nothing
          sieve a ax = filter (\x -> x `mod` a /= 0) ax

largestPrimeFactor :: Int -> Int
largestPrimeFactor n =
    fromMaybe 1 (find (\x -> n `mod` x == 0) (reverse (primes n)))

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