-- Copyright (c) 2006-2011, David Amos. All rights reserved.

{-# LANGUAGE NoMonomorphismRestriction #-}

-- |A module providing functions to test for primality, and find next and previous primes.
module Math.NumberTheory.Prime (primes, isTrialDivisionPrime, isMillerRabinPrime,
                                isPrime, notPrime, prevPrime, nextPrime) where

import System.Random
import System.IO.Unsafe


isTrialDivisionPrime n
    | n > 1 = not $ any (\p -> n `rem` p == 0) (takeWhile (\p -> p*p <= n) primes)
    | otherwise = False

-- |A (lazy) list of the primes
primes :: [Integer]
primes = 2 : filter isPrime [3,5..] where
    isPrime n = not $ any (\p -> n `rem` p == 0) (takeWhile (\p -> p*p <= n) primes)

{-
-- This is just marginally faster, but less elegant
primes2 :: [Integer]
primes2 = 2 : 3 : 5 : 7 : filter isPrime
    (concat [ [m30+11,m30+13,m30+17,m30+19,m30+23,m30+29,m30+31,m30+37] | m30 <- [0,30..] ])
    where isPrime n = not $ any (\p -> n `rem` p == 0) (takeWhile (\p -> p*p <= n) primes2')
          primes2' = drop 3 primes2
-}

{-
-- initial version. This isn't going to be very good if n has any "large" prime factors (eg > 10000)
pfactors1 n | n > 0 = pfactors' n primes
            | n < 0 = -1 : pfactors' (-n) primes
    where pfactors' n (d:ds) | n == 1 = []
                             | n < d*d = [n]
                             | r == 0 = d : pfactors' q (d:ds)
                             | otherwise = pfactors' n ds
                             where (q,r) = quotRem n d
-}

-- MILLER-RABIN TEST
-- Cohen, A Course in Computational Algebraic Number Theory, p422
-- Koblitz, A Course in Number Theory and Cryptography


-- Let n-1 = 2^s * q, q odd
-- Then n is a strong pseudoprime to base b if
-- either b^q == 1 (mod n)
-- or b^(2^r * q) == -1 (mod n) for some 0 <= r < s
-- (For we know that if n is prime, then b^(n-1) == 1 (mod n)

isStrongPseudoPrime n b =
    let (s,q) = split2s 0 (n-1)  -- n-1 == 2^s * q, with q odd
    in isStrongPseudoPrime' n (s,q) b

isStrongPseudoPrime' n (s,q) b
    | b' == 1 = True
    | otherwise = n-1 `elem` squarings
    where b' = power_mod b q n     -- b' = b^q `mod` n
          squarings = take s $ iterate (\x -> x*x `mod` n) b' -- b^(2^r *q) for 0 <= r < s

-- split2s 0 m returns (s,t) such that 2^s * t == m, t odd
split2s s t = let (q,r) = t `quotRem` 2
              in if r == 0 then split2s (s+1) q else (s,t)

-- power_mod b t n == b^t mod n
power_mod b t n = powerMod' b 1 t
    where powerMod' x y 0 = y
          powerMod' x y t = let (q,r) = t `quotRem` 2
                            in powerMod' (x*x `rem` n) (if r == 0 then y else x*y `rem` n) q

isMillerRabinPrime' n
    | n >= 4 =
        let (s,q) = split2s 0 (n-1) -- n-1 == 2^s * q, with q odd
        in do g <- getStdGen
              let rs = randomRs (2,n-1) g
              return $ all (isStrongPseudoPrime' n (s,q)) (take 25 rs)
    | n >= 2 = return True
    | otherwise = return False
-- Cohen states that if we restrict our rs to single word numbers, we can use a more efficient powering algorithm

-- isMillerRabinPrime :: Integer -> Bool
isMillerRabinPrime n = unsafePerformIO (isMillerRabinPrime' n)


-- |Is this number prime? The algorithm consists of using trial division to test for very small factors,
-- followed if necessary by the Miller-Rabin probabilistic test.
isPrime :: Integer -> Bool
isPrime n | n > 1 = isPrime' $ takeWhile (< 100) primes
          | otherwise = False
    where isPrime' (d:ds) | n < d*d = True
                          | otherwise = let (q,r) = quotRem n d
                                        in if r == 0 then False else isPrime' ds
          isPrime' [] = isMillerRabinPrime n
-- the < 100 is found heuristically to be about the point at which trial division stops being worthwhile

notPrime :: Integer -> Bool
notPrime = not . isPrime

-- |Given n, @prevPrime n@ returns the greatest p, p < n, such that p is prime
prevPrime :: Integer -> Integer
prevPrime n | n > 5 = head $ filter isPrime $ candidates
            | n < 3 = error "prevPrime: no previous primes"
            | n == 3 = 2
            | otherwise = 3
            where n6 = (n `div` 6) * 6
                  candidates = dropWhile (>= n) $ concat [ [m6+5,m6+1] | m6 <- [n6, n6-6..] ]

-- |Given n, @nextPrime n@ returns the least p, p > n, such that p is prime
nextPrime :: Integer -> Integer
nextPrime n | n < 2 = 2
            | n < 3 = 3
            | otherwise = head $ filter isPrime $ candidates
            where n6 = (n `div` 6) * 6
                  candidates = dropWhile (<= n) $ concat [ [m6+1,m6+5] | m6 <- [n6, n6+6..] ]

{-
-- slightly better version. This is okay so long as n has at most one "large" prime factor (> 10000)
-- if it has more, it does at least tell you, via an error message, that it has run into difficulties
pfactors2 n | n > 0 = pfactors' n $ takeWhile (< 10000) primes
            | n < 0 = -1 : pfactors' (-n) (takeWhile (< 10000) primes)
    where pfactors' n (d:ds) | n == 1 = []
                             | n < d*d = [n]
                             | r == 0 = d : pfactors' q (d:ds)
                             | otherwise = pfactors' n ds
                             where (q,r) = quotRem n d
          pfactors' n [] = if isMillerRabinPrime n then [n] else error ("pfactors2: can't factor " ++ show n)
-}