#!/usr/bin/env nix-script
#!>haskell
#! haskell | mtl random 

import Control.Monad        (replicateM)
import Control.Monad.State  (State, state, evalState)
import System.Random        (Random, StdGen, newStdGen, randomR)
import Data.List            (genericLength)


-- Modular integral exponentiation
(^%) :: Integral a => a -> a -> a -> a
(^%) m b k = ex b k 1 where
  ex a k s
    | k == 0          = s
    | k `mod` 2 == 0  = ex (a * a `mod` m) (k `div` 2) s
    | otherwise       = ex (a * a `mod` m) (k `div` 2) (s * a `mod` m)


-- Deconstruct an integral n into two
-- numbers d, r such as d*2^r = n
decons :: Integral a => a -> (a, a)
decons n = (exp quots, rest quots)
  where
    rest  = head . dropWhile even
    exp   = genericLength . takeWhile even 
    quots = iterate (`quot` 2) n


-- Test a number given a witness
test :: Integer -> Integer -> Bool
test n a
  | n < 2   || even n       = False
  | b0 == 1 || b0 == pred n = True
  | otherwise = iter (tail b)
  where
    (r, d) = decons (pred n)
    b0 = (n ^% a)  d
    b = take (fromIntegral r) (iterate (n ^% 2) b0)
    iter [] = False
    iter (x:xs)
      | x == 1      = False
      | x == pred n = True
      | otherwise   = iter xs         


-- Check whether n is probably a prime number 
-- with un uncertainty of 2^-k
isPrime :: Integer -> Int -> State StdGen Bool
isPrime n k = all (test n) <$> replicateM k a
  where a = state (randomR (2, pred n))


main :: IO ()
main = do
  gen <- newStdGen
  let test n = evalState (isPrime n 64) gen
  print (test 9917340299)