import Bound
import Control.Applicative
import Control.Monad.Reader

-- Two is just "data Two e a = Two (e a) (e a)"
-- but we'll crib from ReaderT Bool for the instances

type Two = ReaderT Bool

two :: m a -> m a -> Two m a
two a b = ReaderT (\x -> if x then a else b)

unTwo :: Two m a -> (m a, m a)
unTwo t = let r = runReaderT t in (r True, r False)

-- Library

data Mu t a = V a | Roll (t (Mu t) a)

instance Bound t => Functor (Mu t) where
  fmap f (V a)    = V (f a)
  fmap f (Roll g) = Roll (fmap2 f g)

fmap2 :: (Monad f, Bound t) => (a -> c) -> t f a -> t f c
fmap2 f x = x >>>= return . f

instance Bound t => Applicative (Mu t) where pure = V; (<*>) = ap

instance Bound t => Monad (Mu t) where
  return = V
  V a    >>= f = f a
  Roll g >>= f = Roll (g >>>= f)

-- Client: This is all you need to specify for your language!
-- Bindings for free!

data T e a = Ap  (Two e a)
           | Lam (Scope () e a)

instance Bound T where
  Ap r  >>>= f = Ap  (r >>>= f)
  Lam l >>>= f = Lam (l >>>= f)

type Exp = Mu T

lam :: Eq a => a -> Exp a -> Exp a
lam v b = Roll (Lam (abstract1 v b))

app :: Exp a -> Exp a -> Exp a
app e1 e2 = Roll (Ap (two e1 e2))

-- Using the client

whnf :: Exp a -> Exp a
whnf (Roll (Ap a)) =
  let (f, x) = unTwo a
  in case whnf f of
    Roll (Lam b) -> whnf (instantiate1 x b)
    f'           -> Roll (Ap (two f' x))
whnf e = e

term :: Exp String
term = lam "f" (lam "x" (V "f" `app` V "x"))

pretty :: Int -> Exp String -> String
pretty i e = case e of
  V a -> a
  Roll (Ap r) -> let (f, x) = unTwo r
                 in "(" ++ pretty i f ++ ") (" ++ pretty i x ++ ")"
  Roll (Lam l) -> "(\\"
                  ++ show i
                  ++ "."
                  ++ pretty (i + 1) (instantiate1 (return (show i)) l)
                  ++ ")"

main :: IO ()
main = putStrLn (pretty 0 term)

-- > main
-- (\0.(\1.(0) (1)))