{-# Language GADTs, FlexibleInstances, FlexibleContexts, MultiParamTypeClasses, FunctionalDependencies, ScopedTypeVariables, TypeFamilies, NoMonomorphismRestriction #-}

import Data.Char(ord, chr)
import Control.Monad.State
import Data.Supply

data OpName = PLUS | MINUS | MUL
data Bottom = Bottom
  deriving Show

{- Lambda calculus without free variables -}
class Lambda l where
    labs    :: (l a -> l b) -> l (a -> b)
    app     :: l (a -> b) -> l a -> l b
    lit     :: Int -> l Int
    op      :: l Int -> OpName -> l Int -> l Int

class SizeExp l where
    type List l :: * -> *
    list    :: l Int -> l (Int -> b) -> l (List l b)
    aabs    :: (l Int -> l (Int -> a) -> l b) -> l (List l a -> b)
    shift   :: l (Int -> a) -> l Int -> l (Int -> a) -> l (Int -> a)
    undef   :: l Bottom
    unsized :: l ()

instance Lambda l => Num (l Int) where
    fromInteger = lit . fromIntegral
    lhs + rhs   = op lhs PLUS rhs
    lhs - rhs   = op lhs MINUS rhs
    lhs * rhs   = op lhs MUL rhs
    abs = error "abs is not implemented"
    signum = error "signum is not implemented"

{- Printing -}
newtype LPrint a = LPrint { unPrint :: Int -> State Int ShowS }

instance Lambda LPrint where
  lit x      = LPrint $ \p -> return $ shows x
  op m opc n = LPrint $ \p -> do
    let (prec, lprec, rprec, c) = getPrec opc
    l1 <- (unPrint m lprec)
    l2 <- (unPrint n rprec)
    return $ showParen (p>prec) $  l1 . showChar c . l2
  app f v    = LPrint $ \p -> do
    l1 <- unPrint f 6
    l2 <- unPrint v 7
    return $ showParen (p>6) $ l1 . showChar ' ' . l2
  labs e      = LPrint $ \p -> do
    v <- getVar
    let var = LPrint $ \_ -> return $ showVar v
    l <- unPrint (e var) 0 
    return $ showParen (p>0) $ showChar 'λ' . showVar v . showChar '.' . l

instance SizeExp LPrint where
  type List LPrint = []
  aabs f   = LPrint $ \p -> do
    v1 <- getVar
    v2 <- getVar
    let var1 = LPrint $ \_ -> return $ showVar v1
        var2 = LPrint $ \_ -> return $ showVar v2
    l <- unPrint (f var1 var2) 0
    return $ showParen (p>0) $ showChar 'Λ' . showVar v1 . showChar ',' . showVar v2 . showChar '.' . l
  list s f = LPrint $ \p -> do
    l1 <- unPrint s 9
    l2 <- unPrint f 9
    return $ showParen (p>0) $ showString "List " . l1 . showChar ' ' . l2
  shift e1 s e2 = LPrint $ \p -> do
    l1 <- unPrint e1 2
    l2 <- unPrint s 2
    l3 <- unPrint e2 2
    return $ showParen (p>0) $ showString "Shift " . l1 .  showChar ' ' . l2 . showChar ' ' . l3
  undef = LPrint $ \_ -> return $ showChar '┴'
  unsized = LPrint $ \_ -> return $ showChar 'U'

mprint x = putStrLn $ (fst $ runState (unPrint x 0) 0) ""

showVar x = if x>28 
    then showVar (x `div` 29) . showChar (chr $ ord 'a' + (x `mod` 29))
    else showChar $ chr $ ord 'a' + x

getPrec :: OpName -> (Int,Int,Int,Char)
getPrec PLUS = (4,4,5,'+')
getPrec MINUS = (4,4,5,'-')
getPrec MUL = (5,5,6,'*')

getVar :: State Int Int
getVar = do { x <- get; put (x+1); return x }

x = aabs (\s _ -> s) `app` list (lit 0) (labs $ \_ -> undef)
y = list 0 $ labs $ \_ -> undef
z = (labs $ \s -> s+1) `app` (lit 1+2)
conss = labs $ \x -> aabs $ \ys yf -> list (ys + 1) $ shift yf ys (labs $ \_ -> x)
concats = aabs $ \xs xf -> (aabs $ \ys yf -> list (xs+ys) $ shift xf xs yf)
maps = labs $ \x -> aabs $ \ys yf -> list ys ( x `app` yf )
l0 = list (lit 0) (labs $ \_ -> unsized)
l1 = list (lit 1) (labs $ \_ -> unsized)
l2 = list (lit 2) (labs $ \_ -> unsized)
l00 = list (lit 0) (labs $ \_ -> list (lit 0) (labs $ \_ -> unsized))
l10 = list (lit 1) (labs $ \_ -> list (lit 0) (labs $ \_ -> unsized))
l01 = list (lit 0) (labs $ \_ -> list (lit 1) (labs $ \_ -> unsized))
l11 = list (lit 1) (labs $ \_ -> list (lit 1) (labs $ \_ -> unsized))
l12 = list (lit 1) (labs $ \_ -> list (lit 2) (labs $ \_ -> unsized))

{- Evaluating -}
newtype R a = R { unR :: a }
instance Lambda R where
  labs f      = R $ unR . f . R
  app e1 e2   = R $ (unR e1) (unR e2)
  lit i       = R i
  op e1 op e2 = let opm = case op of
                                  PLUS  -> (+)
                                  MINUS -> (-)
                                  MUL   -> (*)
                in R $ opm (unR e1) (unR e2)

newtype RList a = RList { unList :: (Int, Int -> a)  }
instance (Show a) => Show (RList a) where
  show (RList (s,f)) = show (map f [0..s-1])

instance SizeExp R where
  type List R = RList
  undef       = R $ Bottom
  unsized     = R $ ()
  list s f    = R $ RList (unR s, unR f)
  aabs e      = R $ \(RList (s,f)) -> unR (e (R s) (R f))
  shift a s c = R $ \i -> if i<(unR s) then unR a i else unR c i

{- Create a restricted deep embedding to work with the expression -}

{- sized [a] to prove -}
data Nil
data Cons a
data SLInt a where
  SLNil  :: SLInt Nil
  SLCons :: Int -> SLInt a -> SLInt (Cons a)

class MGS a where
  mgs :: Supply Int -> a -> a
  mgs = mgs' []
  mgs' :: [Int] -> Supply Int -> a -> a


instance MGS () where
  mgs' _ _ _ = ()
instance MGS a => MGS (RList a) where
  mgs' bound supp  _ = RList undefined
    where
    (s1, s2, s3) = split3 supp
    listsize = mgs' bound supp (undefined :: a)