content: -- |
  -- Module:      Math.NumberTheory.Primes.Testing.Certified
  -- Copyright:   (c) 2011 Daniel Fischer
  -- Licence:     MIT
  -- Maintainer:  Daniel Fischer <daniel.is.fischer@googlemail.com>
  --
  -- Deterministic primality testing.
  
  {-# LANGUAGE MagicHash           #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE UnboxedTuples       #-}
  {-# LANGUAGE ViewPatterns        #-}
  {-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-}
  
  module Math.NumberTheory.Primes.Testing.Certified
    ( isCertifiedPrime
    ) where
  
  import Prelude hiding (Foldable(..))
  import Data.Bits ((.&.))
  import Data.Foldable
  import Data.Mod
  import Data.Proxy
  import GHC.Num.Integer
  import GHC.TypeNats (SomeNat(..), someNatVal)
  
  import Math.NumberTheory.Roots (integerSquareRoot)
  import Math.NumberTheory.Primes (unPrime)
  import Math.NumberTheory.Primes.Factorisation.TrialDivision (trialDivisionPrimeTo, trialDivisionTo, trialDivisionWith)
  import Math.NumberTheory.Primes.Factorisation.Montgomery (montgomeryFactorisation, smallFactors, findParms)
  import Math.NumberTheory.Primes.Testing.Probabilistic (bailliePSW, isPrime, isStrongFermatPP, lucasTest)
  import Math.NumberTheory.Primes.Sieve.Eratosthenes (primeList, primeSieve)
  import Math.NumberTheory.Utils (splitOff)
  
  -- | @'isCertifiedPrime' n@ tests primality of @n@, first trial division
  --   by small primes is performed, then a Baillie PSW test and finally a
  --   prime certificate is constructed and verified, provided no step before
  --   found @n@ to be composite. Constructing prime certificates can take
  --   a /very/ long time, so use this with care.
  isCertifiedPrime :: Integer -> Bool
  isCertifiedPrime n
      | n < 0     = isCertifiedPrime (-n)
      | otherwise = isPrime n && ((n < bpbd) || checkPrimalityProof (certifyBPSW n))
        where
          bpbd = 100000000000000000
  -- Although it is known that there are no Baillie PSW pseudoprimes below 2^64,
  -- use the verified bound 10^17, I don't know whether Gilchrist's result has been
  -- verified yet.
  
  -- | A proof of primality of a positive number. The type is
  --   abstract to ensure the validity of proofs.
  data PrimalityProof
      = Pocklington { cprime :: !Integer          -- ^ The number whose primality is proved.
                    , _factorisedPart, _cofactor :: !Integer
                    , _knownFactors :: ![(Integer, Word, Integer, PrimalityProof)]
                    }
      | TrialDivision { cprime :: !Integer        -- ^ The number whose primality is proved.
                      , _tdLimit :: !Integer }
      | Trivial { cprime :: !Integer              -- ^ The number whose primality is proved.
                }
        deriving Show
  
  -- | Check the validity of a 'PrimalityProof'. Since it should be
  --   impossible to create invalid proofs by the public interface, this
  --   should never return 'False'.
  checkPrimalityProof :: PrimalityProof -> Bool
  checkPrimalityProof (Trivial n) = isTrivialPrime n
  checkPrimalityProof (TrialDivision p b) = p <= b*b && trialDivisionPrimeTo b p
  checkPrimalityProof (Pocklington p a b fcts) = b > 0 && a > b && a*b == pm1 && a == ppProd fcts && all verify fcts
    where
      pm1 = p-1
      ppProd pps = product [pf^e | (pf,e,_,_) <- pps]
      verify (pf,_,base,proof) = pf == cprime proof && crit pf base && checkPrimalityProof proof
      crit pf base = gcd p (toInteger x-1) == 1 && y == 1
        where
          (# x | #) = integerPowMod# base (pm1 `quot` pf) (fromInteger p)
          (# y | #) = integerPowMod# (toInteger x) pf (fromInteger p)
  
  -- | @'isTrivialPrime'@ checks whether its argument is a trivially
  --   known prime.
  isTrivialPrime :: Integer -> Bool
  isTrivialPrime n = n `elem` trivialPrimes
  
  -- | List of trivially known primes.
  trivialPrimes :: [Integer]
  trivialPrimes = [2,3,5,7,11,13,17,19,23,29]
  
  -- | Certify a small number. This is not exposed and should only
  --   be used where correct. It is always checked after use, though,
  --   so it shouldn't be able to lie.
  smallCert :: Integer -> PrimalityProof
  smallCert n
      | n < 30    = Trivial n
      | otherwise = TrialDivision n (integerSquareRoot n + 1)
  
  -- | @'certify' n@ constructs, for @n > 1@, a proof of either
  --   primality or compositeness of @n@. This may take a very long
  --   time if the number has no small(ish) prime divisors
  certify :: Integer -> Maybe PrimalityProof
  certify n
      | n < 2     = error "Only numbers larger than 1 can be certified"
      | n < 31    = case trialDivisionWith trivialPrimes n of
                      ((p,_):_) | p < n     -> Nothing
                                | otherwise -> Just (Trivial n)
                      _ -> error "Impossible"
      | n < billi = let r2 = integerSquareRoot n + 2 in
                    case trialDivisionTo r2 n of
                      ((p,_):_) | p < n       -> Nothing
                                | otherwise   -> Just (TrialDivision n r2)
                      _ -> error "Impossible"
      | otherwise = case smallFactors (fromInteger (abs n)) of
                      ([], Just _) | not (isStrongFermatPP n 2) -> Nothing
                                   | not (lucasTest n) -> Nothing
                                   | otherwise -> Just (certifyBPSW n)       -- if it isn't we error and ask for a report.
                      ((toInteger -> p,_):_, _)
                        | p == n -> Just (TrialDivision n (min 100000 n))
                        | otherwise -> Nothing
                      _ -> error ("***Error factorising " ++ show n ++ "! Please report this to maintainer of arithmoi.")
        where
          billi = 1000000000000
  
  -- | Certify a number known to be not too small, having no small prime divisors and having
  --   passed the Baillie PSW test. So we assume it's prime, erroring if not.
  --   Since it's presumably a large number, we don't bother with trial division and
  --   construct a Pocklington certificate.
  certifyBPSW :: Integer -> PrimalityProof
  certifyBPSW n = Pocklington n a b kfcts
    where
      nm1 = n-1
      h = nm1 `quot` 2
      m3 = fromInteger n .&. (3 :: Int) == 3
      (a,pp,b) = findDecomposition nm1
      kfcts0 = map check pp
      kfcts = foldl' force [] kfcts0
      force xs t@(_,_,_,prf) = prf `seq` (t:xs)
      check (p,e,byTD) = go 2
        where
          go bs
              | bs > h    = error (bpswMessage n)
              | x == 1    = if m3 && (p == 2) then (p,e,n-bs,Trivial 2) else go (bs+1)
              | g /= 1    = error (bpswMessage n ++ found g)
              | y /= 1    = error (bpswMessage n ++ fermat bs)
              | byTD      = (p,e,bs, smallCert p)
              | otherwise = case certify p of
                              Nothing -> error ("***Error in factorisation code: " ++ show p
                                                          ++ " was supposed to be prime but isn't.\n"
                                                          ++ "Please report this to the maintainer.\n\n")
                              Just ppr ->(p,e,bs,ppr)
                where
                  q = nm1 `quot` p
                  (# x | #) = integerPowMod# bs q (fromInteger n)
                  (# y | #) = integerPowMod# (toInteger x) p (fromInteger n)
                  g = gcd n (toInteger x-1)
  
  -- | Find a decomposition of p-1 for the pocklington certificate.
  --   Usually bloody slow if p-1 has two (or more) /large/ prime divisors.
  findDecomposition :: Integer -> (Integer, [(Integer, Word, Bool)], Integer)
  findDecomposition n = go 1 n [] prms
    where
      sr = integerSquareRoot n
      pbd = min 1000000 (sr+20)
      prms = map unPrime $ primeList (primeSieve pbd)
      go a b afs (p:ps)
          | a > b     = (a,afs,b)
          | otherwise = case splitOff p b of
                          (0,_) -> go a b afs ps
                          (e,q) -> go (a*p^e) q ((p,e,True):afs) ps
      go a b afs []
          | a > b     = (a,afs,b)
          | bailliePSW b  = (b,[(b,1,False)],a)   -- Until a Baillie PSW pseudoprime is found, I'm going with this
          | e == 0    = error ("Error in factorisation, " ++ show p ++ " was found as a factor of " ++ show b ++ " but isn't.")
          | otherwise = go (a*p^e) q ((p,e,False):afs) []
            where
              p = findFactor b 8 6
              (e,q) = splitOff p b
  
  -- | Find a factor of a known composite with approximately digits digits,
  --   starting with curve s. Actually, this may loop infinitely, but the
  --   loop should not be entered before the heat death of the universe.
  findFactor :: Integer -> Int -> Integer -> Integer
  findFactor n digits s = case findLoop n lo hi count s of
                            Left t  -> findFactor n (digits+5) t
                            Right f -> f
    where
      (lo,hi,count) = findParms digits
  
  -- | Find a factor or say with which curve to continue.
  findLoop :: Integer -> Word -> Word -> Word -> Integer -> Either Integer Integer
  findLoop _ _  _  0  s = Left s
  findLoop n lo hi ct s
      | n <= s+2  = Left 6
      | otherwise = case someNatVal (fromInteger n) of
                      SomeNat (_ :: Proxy t) -> case montgomeryFactorisation lo hi (fromInteger s :: Mod t) of
                        Nothing -> findLoop n lo hi (ct-1) (s+1)
                        Just fct
                          | bailliePSW fct -> Right fct
                          | otherwise -> Right (findFactor fct 8 (s+1))
  
  -- | Message in the unlikely case a Baillie PSW pseudoprime is found.
  bpswMessage :: Integer -> String
  bpswMessage n = unlines
                      [ "\n***Congratulations! You found a Baillie PSW pseudoprime!"
                      , "Please report this finding to the maintainers:"
                      , "<daniel.is.fischer@googlemail.com>,"
                      , "<andrew.lelechenko@gmail.com>"
                      , "The number in question is:\n"
                      , show n
                      , "\nOther parties like wikipedia might also be interested."
                      , "\nSorry for aborting your programm, but this is a major discovery."
                      ]
  
  -- | Found a factor
  found :: Integer -> String
  found g = "\nA nontrivial divisor is:\n" ++ show g
  
  -- | Fermat failure
  fermat :: Integer -> String
  fermat b = "\nThe Fermat test fails for base\n" ++ show b
  