-- efficient lazy evaluation for lambda calculus with constructors and delta functions
-- eliminators are used instead of case expressions
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RankNTypes #-}
import Control.Monad.ST
import Control.Monad.Fix
import System.Environment
import Data.STRef

--------------------------------------------------------------------------------

data Exp s
    = Ref !(STRef s (Exp s))
    | Lam !(Exp s -> Exp s)
    | App (Exp s) (Exp s)
    | Action !(ST s (Exp s))
    | Con !ConName [Exp s]
    | Int !Int
    | LetRec (Exp s -> Exp s) (Exp s -> Exp s)

infixl 1 `App`

type VarName = Int
data ConName = T2 | CTrue | CFalse | Nil | Cons
    deriving (Eq, Show)

instance Show (Exp s) where
    show = \case
        Int i -> show i

--------------------------------------------------------------------------------

pureEval :: (forall s . Exp s) -> String
pureEval e = runST (show <$> evalN e)

-- eval to some normal form (not needed in the current test case)
evalN :: Exp s -> ST s (Exp s)
evalN x = eval x >>= \case
    Con f es -> Con f <$> mapM evalN es
    c -> return c

-- eval to weak head normal form
eval :: Exp s -> ST s (Exp s)
eval = \case
    App (Lam f) x -> eval . f =<< addHeap x     -- optimization
    App l x -> eval l >>= \(Lam f) -> eval . f =<< addHeap x
    Action m -> m
    Ref i -> do
        x <- readSTRef i
        if reduced x then return x else do          -- optimization
            -- writeSTRef i $ error "cycle in spine"    -- optional test
            z <- eval x
            writeSTRef i z
            return z
    LetRec f g -> eval . g =<< mfix (fmap Ref . newSTRef . f)
    x -> return x

reduced = \case
    App{} -> False
    Action{} -> False
    _ -> True

addHeap :: Exp s -> ST s (Exp s)
addHeap e = if reduced e then return e else Ref <$> newSTRef e

iSqrt :: Int -> Int
iSqrt = round . sqrt . fromIntegral

-------------------------------------------------------------------------------- example codes

infixl 1 @@, @@.

(@@), (@@.) :: Exp s -> Exp s -> Exp s
(@@) = App
(@@.) = (@@)

f0 s = Action s
f1 s = Lam $ \v0 -> Action $ s v0
f2 s = Lam $ \v0 -> Lam $ \v1 -> Action $ s v0 v1
f3 s = Lam $ \v0 -> Lam $ \v1 -> Lam $ \v2 -> Action $ s v0 v1 v2
c0 s = Con s []
c1 s = Lam $ \v0 -> Con s [v0]
c2 s = Lam $ \v0 -> Lam $ \v1 -> Con s [v0, v1]
c3 s = Lam $ \v0 -> Lam $ \v1 -> Lam $ \v2 -> Con s [v0, v1, v2]

intOp g = f2 $ \a b -> do
    Int x <- eval a
    Int y <- eval b
    return $ Int $ g x y
boolOp g = f2 $ \a b -> do
    Int x <- eval a
    Int y <- eval b
    return $ if g x y then true else false

undef = f0 $ return $ error "undef"

-- tuples
t2 = c2 T2
elimT2 = f2 $ \t f -> eval t >>= eval . \(Con _ (x: y:_)) -> f `App` x `App` y


-- booleans
false = c0 CFalse
true = c0 CTrue
elimBool = f3 $ \xs f g -> eval xs >>= eval . \case
        Con CFalse _ -> f
        Con CTrue _ -> g
and' = Lam $ \v0 -> Lam $ \v1 -> elimBool @@ v0 @@ false @@ v1

-- lists
nil = c0 Nil
cons = c2 Cons
elimList = f3 $ \xs f g -> eval xs >>= eval . \case
        Con Nil _ -> f
        Con Cons (x: y:_) -> g `App` x `App` y

-- integers
zero = Int 0
one = Int 1
two = Int 2
add = intOp (+)
sub = intOp (-)
mod' = intOp mod
eq = boolOp (==)
neq = boolOp (/=)
leq = boolOp (<=)
geq = boolOp (>=)
iSqrt' = f1 $ \x -> eval x >>= \(Int n) -> return $ Int (iSqrt n)

nthPrime n =
    LetRec (\r -> Lam $ \v0 -> cons @@. v0 @@. (r @@ (add @@. one @@. v0))) $ \from ->
    LetRec (\r -> Lam $ \v0 -> Lam $ \v1 -> Lam $ \v2 ->
        elimList @@. v2 @@. v1 @@. (Lam $ \v3 -> Lam $ \v4 -> v0 @@ v3 @@ (r @@ v0 @@ v1 @@ v4))) $ \foldr' ->
    LetRec (\r -> Lam $ \v0 -> Lam $ \v1 ->
        elimList @@. v1 @@. nil @@. (Lam $ \v2 -> Lam $ \v3 -> elimBool @@. (v0 @@ v2) @@. nil @@. (cons @@. v2 @@. (r @@ v0 @@ v3)))) $ \takeWhile' ->
    LetRec (\r -> Lam $ \v0 -> Lam $ \v1 ->
        elimList @@. v0 @@. undef @@. (Lam $ \v2 -> Lam $ \v3 -> elimBool @@. (eq @@. v1 @@. zero) @@. (r @@ v3 @@ (sub @@. v1 @@. one)) @@. v2)) $ \nth ->
    let
        map_ = Lam $ \v0 -> foldr' @@ (Lam $ \v1 -> Lam $ \v2 -> cons @@. (v0 @@ v1) @@. v2) @@ nil
        and'' = foldr' @@ and' @@ true
        filter' = Lam $ \v0 -> foldr' @@ (Lam $ \v1 -> Lam $ \v2 -> elimBool @@. (v0 @@ v1) @@. v2 @@. (cons @@. v1 @@. v2)) @@ nil

    in LetRec (\r -> cons @@. Int 2 @@. (cons @@. Int 3 @@. (filter' @@ (Lam $ \v0 -> and'' @@ (map_ @@. (Lam $ \v1 -> neq @@. zero @@. (mod' @@. v0 @@. v1)) @@ (takeWhile' @@ (geq @@ (iSqrt' @@. v0)) @@ r))) @@ (from @@ Int 5)))) $ \primes ->
        nth @@ primes @@ Int n

main = getArgs >>= \case
    [n] -> putStrLn $ pureEval $ nthPrime $ read n
    [_, n] -> print $ primes !! read n

primes = 2:3: filter (\n -> and $ map (\p -> n `mod` p /= 0) (takeWhile (<= iSqrt n) primes)) [5..]