import Control.DeepSeq
import Numeric.Natural
import Data.List

data Val
    = PIN !Val
    | LAW !Natural !Natural !Val
    | APP !Val Val
    | NAT !Natural
  deriving (Show)

instance NFData Val where
    rnf (APP f x) = rnf f `seq` rnf x
    rnf _         = ()

f % x | arity f == 1 = apply f [x]
f % x | otherwise    = APP f x

arity (APP f _)     = arity f - 1
arity (PIN (NAT n)) = case n of 1->3; 3->6; _->1
arity (PIN l@LAW{}) = arity l
arity (PIN i)       = 1
arity (LAW _ a _)   = fromIntegral a :: Integer
arity (NAT n)       = 0

cas p _ _ _ _ (PIN x)     = (p % x)
cas _ l _ _ _ (LAW n a b) = (l % NAT n % NAT a % b)
cas _ _ a _ _ (APP f x)   = (a % f % x)
cas _ _ _ z _ (NAT 0)     = z
cas _ _ _ _ m (NAT n)     = (m % NAT (n-1))

kal n e (NAT b) | b<=n          = e !! fromIntegral (n-b)
kal n e (NAT 0 `APP` f `APP` x) = (kal n e f % kal n e x)
kal _ _ (NAT 0 `APP` x)         = x
kal _ _ x                       = x

run :: Natural -> [Val] -> Val -> Val
run arity ie body = res
  where (n, e, res::Val) = go arity ie body
        go i acc (NAT 1 `APP` v `APP` k') = go (i+1) (kal n e v : acc) k'
        go i acc x                        = (i, acc, kal n e x)

nat (NAT n) = n
nat _       = 0

apply (APP f x)           xs                  = apply f (x:xs)
apply f@(LAW n a b)       xs                  = run a (reverse (f:xs)) b
apply f@(PIN (LAW n a b)) xs                  = run a (reverse (f:xs)) b
apply (PIN (NAT 0))       [x]                 = PIN (force x)
apply (PIN (NAT 1))       [n,a,b] | nat a > 0 = LAW (nat n) (nat a) (force b)
apply (PIN (NAT 2))       [x]                 = NAT (nat x + 1)
apply (PIN (NAT 3))       [p,l,a,z,m,o]       = cas p l a z m o
apply f                   e                   = error (show (f,e))