{-# Language GADTs, FlexibleInstances, FlexibleContexts, MultiParamTypeClasses, FunctionalDependencies, ScopedTypeVariables, TypeFamilies, NoMonomorphismRestriction, OverlappingInstances, IncoherentInstances #-}
module L where

import qualified Data.SBV as SBV
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 -}
type Arr repr a b = repr a -> repr b

class (Num (LInt l)) => Lambda l where
    type LInt l :: *
    labs    :: (l a -> l b) -> l (Arr l a b)
    app     :: l (Arr l a b) -> l a -> l b
    lit     :: (LInt l) -> l (LInt l)
    op      :: l (LInt l) -> OpName -> l (LInt l) -> l (LInt l)

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

instance (Lambda l, LInt l ~ a) => Num (l a) 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 -> Int -> ShowS }

instance Lambda LPrint where
  type LInt LPrint = Int
  lit x      = LPrint $ \_ -> 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 v -> let
    var = LPrint $ \_ -> return $ showVar v
    l = unPrint (e var) 0 $ succ v
    in showParen (p>0) $ showChar 'λ' . showVar v . showChar '.' . l

instance SizeExp LPrint where
  type List LPrint = []
  aabs f   = LPrint $ \p v -> let 
    v2 = succ v
    var1 = LPrint $ \_ _ -> showVar v
    var2 = LPrint $ \_ _ -> showVar v2
    l = unPrint (f var1 var2) 0 (v+2)
    in showParen (p>0) $ showChar 'Λ' . showVar v . 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'

view :: LPrint a -> LPrint a
view = id
instance Show (LPrint a) where
    showsPrec _ e = unPrint e 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
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))

v = (labs $ \s -> s+1) `app` (lit 1+2)
w1 = (labs $ \s -> s+ lit 1)
w2 = (labs $ \s -> s+ lit 2)
q1 :: (Lambda l, a ~ LInt l, Num (l a)) => l (Arr l b (Arr l a a))
q1 = (labs $ \q -> labs $ \s -> s + lit 1)
q2 :: (Lambda l, a ~ LInt l, Num (l a)) => l (Arr l a (Arr l a a))
q2 = (labs $ \q -> labs $ \s -> s + q)
q3 = (labs $ \q -> labs $ \s -> q + s)

{- Evaluating -}
eval = unR
newtype R a = R { unR :: a }
instance Lambda R where
  type LInt R = Int
  labs f      = R $ f
  app e1 e2   = (unR e1) 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 . (unR f) . R )
  aabs e = R $ ( \(RList (s,f)) -> (e (R s) (R (R . f . unR))) ) . unR
  shift a s c = R $ \i -> if (unR i)<(unR s) then unR a i else unR c i

newtype E a = E { unE :: a }
instance Lambda E where
  type LInt E = SBV.SInteger
  labs f = E $ f
  app e1 e2   = (unE e1) e2
  lit         = E 
  op e1 op e2 = let opm = case op of
                                  PLUS  -> (+)
                                  MINUS -> (-)
                                  MUL   -> (*)
                in E $ opm (unE e1) (unE e2)

class Sizeable' a where
  getl' :: E (Arr E a b) -> SBV.Symbolic b

instance Sizeable' (SBV.SInteger) where
  getl' f = do
    x <- SBV.sInteger "x"
    return $ unE $ (unE f) $ E x

{-
getEq f = SBV.proveWith (SBV.z3 {SBV.verbose=True}) $ do
  x <- getl f
  return $ x SBV..== x
-}
data ST  = ST

class Instantiateable a r | a -> r where
  instantiate :: a -> SBV.Symbolic r

instance Instantiateable SBV.SInteger SBV.SInteger where
  instantiate _ = do
    liftIO $ putStrLn " INT VARIABLE x"
    SBV.sInteger "x"

instance Instantiateable a ST where
  instantiate _ = do
    liftIO $ putStrLn " ? VARIABLE x"
    return $ undefined -- SBV.uninterpret "f"

class Sizeable a where
  type Ret a :: *
  getl :: a -> a -> SBV.Symbolic (Ret a, Ret a)

instance (Instantiateable a r, Sizeable b) => Sizeable (Arr E a b) where
  type Ret (Arr E a b) = Ret b
  getl f g = do
    liftIO $ putStrLn " -> x"
    x <- instantiate (undefined :: a)
    let f' = unE $ f $ E x
    let g' = unE $ g $ E x
    getl f' g'

instance Sizeable (SBV.SInteger) where
  type Ret SBV.SInteger = SBV.SInteger
  getl f g = return (f,g)

getEq f g = SBV.proveWith (SBV.z3 {SBV.verbose=True}) $ do
  (x,y) <- getl (unE f) (unE g)
  return $ x SBV..== y