{-# LANGUAGE GADTs, FlexibleInstances, TypeFamilies, StandaloneDeriving, ScopedTypeVariables #-}

import Data.Maybe
import Data.Reify
import Data.Reify.Graph
import Data.Typeable
import Control.Applicative

data Ast e where
  IntLit :: Int -> Ast Int
  Add :: Ast Int -> Ast Int -> Ast Int
  BoolLit :: Bool -> Ast Bool
  IfThenElse :: Ast Bool -> Ast e -> Ast e -> Ast e
  
data AstNode s where
  IntLitN :: Int -> AstNode s
  AddN :: s -> s -> AstNode s
  BoolLitN :: Bool -> AstNode s
  IfThenElseN :: TypeRep -> s -> s -> s -> AstNode s
  
data Ast2 e where
  IntLit2 :: Int -> Ast2 Int
  Add2 :: Unique -> Unique -> Ast2 Int
  BoolLit2 :: Bool -> Ast2 Bool
  IfThenElse2 :: Unique -> Unique -> Unique -> Ast2 e

deriving instance Show (Ast e)    
deriving instance Show (AstNode Int)  
deriving instance Show (Ast2 e)

instance Typeable e => MuRef (Ast e) where
  type DeRef (Ast e) = AstNode
  mapDeRef f (IntLit a) = pure $ IntLitN a
  mapDeRef f (Add a b) = AddN <$> f a <*> f b
  mapDeRef f (BoolLit a) = pure $ BoolLitN a
  mapDeRef f (IfThenElse a b c :: Ast e) = IfThenElseN (typeOf (undefined::e)) <$> f a <*> f b <*> f c
  
data Graph2 = Graph2 [(Unique, Ast2 Int)] [(Unique, Ast2 Bool)] Unique deriving Show
  
recoverTypes :: Graph AstNode -> Graph2
recoverTypes (Graph xs x) = Graph2 (catMaybes $ map (f toAst2Int) xs) (catMaybes $ map (f toAst2Bool) xs) x where
  f g (u,an) = do a2 <- g an
                  return (u,a2)
                  
  toAst2Int (IntLitN a) = Just $ IntLit2 a
  toAst2Int (AddN a b) = Just $ Add2 a b
  toAst2Int (IfThenElseN t a b c) | t == typeOf (undefined :: Int) = Just $ IfThenElse2 a b c
  toAst2Int _ = Nothing
  
  toAst2Bool (BoolLitN a) = Just $ BoolLit2 a
  toAst2Bool (IfThenElseN t a b c) | t == typeOf (undefined :: Bool) = Just $ IfThenElse2 a b c
  toAst2Bool _ = Nothing

  
expr = Add (IntLit 42) expr  
  
test = do
  graph <- reifyGraph expr
  print graph
  print $ recoverTypes graph