{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE KindSignatures        #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE TypeSynonymInstances  #-}
{-# LANGUAGE EmptyDataDecls        #-}
{-# LANGUAGE TemplateHaskell       #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE PolyKinds             #-}
{-# LANGUAGE TypeFamilies          #-}

module MultiRec where

import Datatypes
import Generics.MultiRec.Any
import Generics.MultiRec.Transformations.RewriteRules as RR
import Generics.MultiRec.Transformations.ZipperState
import Generics.MultiRec.Transformations.Explicit as Ex
import Generics.MultiRec.Rewriting
import Generics.MultiRec.Zipper

import Generics.MultiRec hiding (show)
import Generics.MultiRec.TH

import Control.Monad ( (>=>) )

--------------------------------------------------------------------------------
-- Multirec representations for the example datatypes
--------------------------------------------------------------------------------
data TreeAST :: * -> * where
  Tree :: TreeAST Tree

$(deriveAll ''TreeAST)

data ListAST :: * -> * -> * where
  List :: ListAST a (List a)

$(deriveAll ''ListAST)

data XAST :: * -> * where
  X :: XAST X

$(deriveAll ''XAST)

data ZigZag :: * -> * where
  Zig :: ZigZag Zig
  Zag :: ZigZag Zag

$(deriveAll ''ZigZag)

data AST i where
  BExpr  :: AST BExpr
  AExpr  :: AST AExpr
  Stmt   :: AST Stmt

$(deriveAll ''AST)

--------------------------------------------------------------------------------
-- Rewrite rules solution
--------------------------------------------------------------------------------
instance RR.Transform AST

-- Now we can simply do the above transformation in a nice way!
rr = RR.apply [insert (down >=> right >=> right) change] Stmt prog1 == Just prog2
  where
    change = rule $ \e a b -> If e a b :~> If (Not e) b a

-- The same one in two steps, which illustrates that rules can be of different
-- types
rr2 = RR.apply [ insert (down >=> right >=> right) swap
               , insert down addNot] Stmt prog1          == Just prog2
  where
    swap   :: Rule AST Stmt
    swap   = rule $ \e a b -> If e a b :~> If e b a
    addNot :: Rule AST BExpr
    addNot = rule $ \e -> e :~> Not e

--------------------------------------------------------------------------------
-- Zipper with state
--------------------------------------------------------------------------------
zs = navigate Stmt prog1 $ do
  downMonad >> rightMonad >> rightMonad
  -- Swap
  l <- downMonad >> rightMonad
  r <- rightMonad
  updateMonad (\p _ -> matchAny p l)
  leftMonad
  updateMonad (\p _ -> matchAny p r)
  -- Add the not
  leftMonad
  updateMonad (\p e -> case p of
                  BExpr -> Just (Not e)
                  _     -> Nothing)

--------------------------------------------------------------------------------
-- Explicit
--------------------------------------------------------------------------------
instance Ex.Transform AST

-- Ordering index of AST as AExpr < BExpr < Stmt
instance OrdI AST where
  compareI AExpr AExpr = EQ
  compareI AExpr _     = LT
  compareI BExpr AExpr = GT
  compareI BExpr BExpr = EQ
  compareI BExpr _     = LT
  compareI Stmt  Stmt  = EQ
  compareI Stmt  _     = GT

-- Family with references
class HasRef phi where
  type RefRep phi ix
  toRef   :: phi ix -> HFix (WithRef phi) ix -> RefRep phi ix
  fromRef :: phi ix -> RefRep phi ix -> HFix (WithRef phi) ix

data NiceInsert phi where
  NiceInsert :: phi ix -> Path -> RefRep phi ix -> NiceInsert phi

type NiceTransformation phi = [ NiceInsert phi]

toNiceTransformation :: HasRef phi => Ex.Transformation phi -> NiceTransformation phi
toNiceTransformation = map f
  where f (Ex.AnyInsert p l x) = NiceInsert p l (toRef p x)

-- Instances for example
data AExprEH = VarEH String
             | ConstEH Integer
             | NegEH AExprEH
             | AddEH AExprEH AExprEH
             | AExprRef Path
           deriving (Show, Eq)

data BExprEH = BConstEH Bool
             | NotEH BExprEH
             | AndEH BExprEH BExprEH
             | GreaterEH AExprEH AExprEH
             | BExprRef Path
             deriving (Show, Eq)

data StmtEH = SeqEH [StmtEH]
            | AssignEH String AExprEH
            | IfEH BExpr StmtEH StmtEH
            | WhileEH BExprEH StmtEH
            | SkipEH
            | StmtRef Path
            deriving (Show, Eq)

instance HasRef AST where
  type RefRep AST AExpr = AExprEH
  type RefRep AST BExpr = BExprEH
  type RefRep AST Stmt  = StmtEH
  
  -- Not complete, but enough for example below
  toRef AExpr (HIn (Ref p)) = AExprRef p
  toRef BExpr (HIn (Ref p)) = BExprRef p
  toRef BExpr (HIn (InR (L (Tag (R (L (C (I x)))))))) = NotEH (toRef BExpr x)
  toRef Stmt (HIn (Ref p)) = StmtRef p

  fromRef AExpr (AExprRef p) = HIn (Ref p)

-- Show existentials
instance Show (NiceInsert AST) where
  show (NiceInsert AExpr x l) = "NiceInsert AExpr (" ++ show x ++ ") " ++ show l
  show (NiceInsert BExpr x l) = "NiceInsert BExpr (" ++ show x ++ ") " ++ show l
  show (NiceInsert Stmt x l)  = "NiceInsert Stmt (" ++ show x ++ ") " ++ show l

-- Actual example
{- This prints: (note the different reference types here)
  [ NiceInsert BExpr (NotEH (BExprRef [2,0])) [2,0]
  , NiceInsert Stmt (StmtRef [2,2]) [2,1]
  , NiceInsert Stmt (StmtRef [2,1]) [2,2] ]
-}
expl1 = print $ toNiceTransformation $ diff Stmt prog1 prog2
expl2 = print $ toNiceTransformation $ diff Stmt prog3 prog4
expl3 = print $ toNiceTransformation $ diff Stmt prog5 prog6