import System.Random (randomRIO)
import ModularArithmeticUtils (modExp)

fermatPrimeTest :: Integer -> Integer -> IO Bool
fermatPrimeTest n k
  | n <= 3    = return (n == 2 || n == 3)
  | even n    = return False
  | otherwise = go k
  where
    go 0 = return True
    go i = do
      a <- randomRIO (2, n - 2)
      if modExp a (n - 1) n /= 1
        then return False
        else go (i - 1)