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

millerRabinWitness :: Integer -> Integer -> Integer -> Integer -> Bool
millerRabinWitness n a d s =
  let x0 = modExp a d n
  in  x0 == 1 || x0 == n - 1 || loop x0 (s - 1)
  where
    loop _ 0 = False
    loop x r =
      let x' = (x * x) `mod` n
      in  x' == n - 1 || loop x' (r - 1)

millerRabin :: Integer -> Integer -> IO Bool
millerRabin n k
  | n < 2     = return False
  | n == 2    = return True
  | even n    = return False
  | otherwise = do
      let (d, s) = factorOutTwos (n - 1)
      go k d s
  where
    go 0 _ _ = return True
    go i d s = do
      a <- randomRIO (2, n - 2)
      if millerRabinWitness n a d s
        then go (i - 1) d s
        else return False