module PrimesIsInP
    ( isPrime
    ) where

import List (find, unfoldr)
import Maybe (fromMaybe)
import Array

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

isPrime :: Integer -> 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 :: Integer -> 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 n

findR :: Integer -> Maybe Integer
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 (r-1) / q')) `mod` r /= 1)
                    q  = largestPrimeFactor (r-1)
                    n' = fromInteger n
                    r' = fromInteger r
                    q' = fromInteger q

testIdentity :: Integer -> Integer -> Bool
testIdentity n r = all test [1 .. last]
    where test a = x == y
              where x = polyExp n r base n
                        where base = zero // [(1,1),(0,-a)]
                              zero = array (0,r-1) [(i,0) | i <- [0..(r-1)]]
                    y = accumArray (\a b -> (a + b) `mod` n) 0 (0,r-1)
                                   [(n `mod` r, 1), (0,-a)]
          last = floor (2 * sqrt (fromInteger r) * logBase 2 (fromInteger n))

-----------------------------------------------------------------------------
-- FIXME

primes :: Integer -> [Integer]
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 :: Integer -> Integer
largestPrimeFactor n =
    fromMaybe 1 (find (\x -> n `mod` x == 0) (reverse (primes n)))

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

type Poly = Array Integer Integer

polyExp :: Integer -> Integer -> Poly -> Integer -> Poly
polyExp n r base exp = func exp
    where func 1 = base
          func x = if m==0
                   then polyMult n r (func d) (func d)
                   else polyMult n r base (func (x-1))
              where (d,m) = x `divMod` 2

polyMult :: Integer -> Integer -> Poly -> Poly -> Poly
polyMult n r x y =
    accumArray (\a b -> (a+b) `mod` n) 0 (0,r-1)
               [((i+j) `mod` r, a*b) |
                (i,a) <- assocs x, (j,b) <- assocs y]

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