{-# LANGUAGE FlexibleContexts     #-}
{-# LANGUAGE TypeFamilies         #-}
{-# LANGUAGE ScopedTypeVariables  #-}
{-# LANGUAGE TypeOperators        #-}

module ASTExamples where

import Data.Maybe (fromJust)
import Control.Arrow ((>>>))
import Control.Monad ((>=>))

-- Replace ASTUse with ASTTHUse below if you want
-- to test TH code generation.
import AST
import ASTUse
-- import ASTTHUse

import Generics.MultiRec.Elems.None
import Generics.MultiRec.Compos
import qualified Generics.MultiRec.Fold as F
import Generics.MultiRec.Fold (con, tag)
import Generics.MultiRec.FoldAlg as FA
import Generics.MultiRec.Eq
import Generics.MultiRec.Show as GS

-- | Example expression

example = Let (Seq ("x" := Mul (Const 6) (Const 9)) None)
              (Add (EVar "x") (EVar "y"))

-- | Renaming variables using 'compos'

renameVar :: Expr -> Expr
renameVar = renameVar' Expr
  where
    renameVar' :: AST NoElems a ix -> a -> a
    renameVar' Var x = x ++ "_"
    renameVar' p   x = compos renameVar' p x

-- | Test for 'renameVar'

testRename :: Expr
testRename = renameVar example

-- | Result of evaluating an expression

data family Value aT :: *
data instance Value Zero               =  EV  (Env -> Int)
data instance Value (Suc Zero)         =  DV  (Env -> Env)
data instance Value (Suc (Suc Zero))   =  VV  Var

type Env = [(Var, Int)]

infixr 5 &.

(&.) = (F.&)

type R n = Case NoElems Value (Rec n)

evalAlgebra1 :: F.Algebra AST NoElems Value
evalAlgebra1 _ = 
      tag  (   con (\ (K x)                   -> EV (const x))
           &.  con (\ (I (CR (EV x) :: R Zero)
                   :*: I (CR (EV y) :: R Zero))
                   -> EV (\ env -> x env + y env))
           &.  con (\ (I (CR (EV x) :: R Zero)
                   :*: I (CR (EV y) :: R Zero))
                   -> EV (\ env -> x env * y env))
           &.  con (\ (I (CR (VV x) :: R (Suc (Suc Zero))))
                   -> EV (fromJust . lookup x))
           &.  con (\ (I (CR (DV e) :: R (Suc Zero))
                   :*: I (CR (EV x) :: R Zero))
                   -> EV (\ env -> x (e env)))
           )
  &.  tag  (   con ((\ (I (CR (VV x) :: R (Suc (Suc Zero)))
                   :*:  I (CR (EV v) :: R Zero))
                   -> DV (\ env -> (x, v env) : env)) 
                   :: ((I (Rec (Suc (Suc Zero))) :*: I (Rec (Zero))) (Case NoElems Value) (Rec (Suc Zero))) -> Value (Suc Zero))
                   {- Why do I have to give the type signature only in this specific case? -}
           &.  con (\ (I (CR (DV f) :: R (Suc Zero))
                   :*: I (CR (DV g) :: R (Suc Zero)))
                   -> DV (g . f))
           &.  con (\ U -> DV id)
           )
  &.  tag          (\ (K x) -> VV x)

-- | More convenient algebra for evaluating an expression

evalAlgebra2 :: FA.Algebra AST NoElems Value
evalAlgebra2 _ =

     (  (\ x             -> EV (const x))
     &  (\ (EV x) (EV y) -> EV (\ env -> x env  +  y env))
     &  (\ (EV x) (EV y) -> EV (\ env -> x env  *  y env))
     &  (\ (VV x)        -> EV (fromJust . lookup x))
     &  (\ (DV e) (EV x) -> EV (\ env -> x (e env)))
     )
  &  (  (\ (VV x) (EV v) -> DV (\ env -> (x, v env) : env ))
     &  (\ (DV f) (DV g) -> DV (g . f))
     &  (                   DV id)
     )
  &     (\ x             -> VV x)

-- | Evaluator

eval1 :: Expr -> Env -> Int
eval1 x = let (EV f) = F.fold evalAlgebra1 Expr x in f

-- | Evaluator

eval2 :: Expr -> Env -> Int
eval2 x = let (EV f) = FA.fold evalAlgebra2 Expr x in f

-- | Test for 'eval1'

testEval1 :: Int
testEval1 = eval1 example [("y", -12)] 

-- | Test for 'eval2'

testEval2 :: Int
testEval2 = eval2 example [("y", -12)] 

-- | Equality instance for 'Expr'

instance Eq Expr where
  (==) = eq Expr

-- | Test for equality

testEq :: (Bool, Bool)
testEq = (example == example, example == testRename)

-- | Test for generic show

testShow :: IO ()
testShow = putStrLn $ GS.show Expr example