{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE KindSignatures        #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE FlexibleInstances     #-}

module Generics.MultiRec.Rewriting.Rewriting (
    Rewrite
  , apply, prop_LeftRight
  , rewrite, rewriteM
  ) where

import Control.Monad.State
import qualified Data.Map as M
import Prelude hiding (zip)

import Generics.MultiRec.Base
import Generics.MultiRec.HFunctor
import Generics.MultiRec.Rewriting.HZip
import Generics.MultiRec.HFix

import Generics.MultiRec.Rewriting.Rules

data Dyn s   = forall ix. Dyn (s ix) ix
type Subst s = M.Map MVar (Dyn s)


-----------------------------------------------------------------------------
-- Type class synonym summarizing generic functions.
-----------------------------------------------------------------------------
-- | The @Rewrite@ is a type class synonym, hiding some of the implementation
-- details.
--
-- To be able to use the rewriting functions, the user is required to provide
-- an instance of this type class.
class (HFunctor s (PF s), HZip s (PF s), LR s (PF s), EqS s)
        => Rewrite s


-----------------------------------------------------------------------------
-- Rewriting a term.
-----------------------------------------------------------------------------
-- | Rewrite a term according to a rule. Result admits failure of matching.
rewriteM :: forall s ix m. (Rewrite s, MonadPlus m, Fam s)
         => s ix -> Rule s ix -> ix -> m ix
rewriteM p (lhs :~> rhs) t = match p lhs t >>= return . (\s -> inst s p rhs)

-- | Rewrite a term according to a rule. Returns the unchanged term in case of 
-- matching failure.
rewrite :: forall s ix. (Rewrite s, Fam s)
        => s ix -> Rule s ix -> ix -> ix
rewrite p f t = maybe t id (rewriteM p f t)

-----------------------------------------------------------------------------
-- Matching and application
-----------------------------------------------------------------------------
-- | Apply a substition to a scheme
apply :: forall s ix . (EqS s, HFunctor s (PF s), Fam s) 
      => Subst s -> s ix -> Scheme s ix -> Scheme s ix
apply s p sc
  = case viewScheme p sc of
      MVar x ->
        case M.lookup x s of
          Nothing          -> mvar p x
          Just (Dyn p' e) ->
            case eqS p p' of
              Just Refl -> toScheme p e
              Nothing   -> error "Coerce error in apply"
      PF r -> HIn (R (hmap (apply s) p r))

-- | Matching a term
match :: (Rewrite s, MonadPlus m, Fam s)
      => s ix -> Scheme s ix -> ix -> m (Subst s)
match p pat term = execStateT (matchM p pat (I0 term)) M.empty

matchM :: (Rewrite s, MonadPlus m, Fam s)
      => s ix -> Scheme s ix -> I0 ix -> StateT (Subst s) m ()
matchM p scheme (I0 e)
  = case viewScheme p scheme of
      MVar var ->
        do subst <- get
           case M.lookup var subst of
             Nothing     -> put (M.insert var (Dyn p e) subst)
             Just exTerm -> checkEqual p e exTerm
      PF r -> combine matchM p r (from p e)

checkEqual :: (HZip s (PF s), MonadPlus m, EqS s, Fam s)
           => s ix -> ix -> Dyn s -> m ()
checkEqual p e (Dyn p' e')
  = case eqS p p' of
      Just Refl -> eq' p (I0 e) (I0 e')
      Nothing   -> mzero

-- | Apply a substitution assuming that the result contains
-- no meta-variables
inst :: forall s ix . (EqS s, HFunctor s (PF s), Fam s) 
     => Subst s -> s ix -> Scheme s ix -> ix
inst s p sc
  = case viewScheme p sc of
      MVar x ->
        case M.lookup x s of
          Just (Dyn p' e)
               -> case eqS p p' of
                    Just Refl -> e
                    Nothing   -> error "inst: Coerce error"
          Nothing
               -> error "inst: Metavar not in substitution"
      PF r   ->
        (to p . hmap (\p' -> I0 . inst s p') p) r


-- | Checks left/right property (see van Noort et al.)
prop_LeftRight :: forall a s . (Show a, Rewrite s, EqS s, Fam s) 
               => s a -> IO ()
prop_LeftRight p
  = case hzipM fake (from p lvalue) (from p rvalue) of
    Nothing -> putStrLn "Left/Right property satisfied" >> showVals
    Just _  -> putStrLn "Left/Right property not satisfied: values are not different enough" >> showVals
  where
    fake :: forall ix r r' r''. s ix -> r ix -> r' ix -> Maybe (r'' ix)
    fake = undefined
    lvalue = left p
    rvalue = right p
    showVals = do
                 putStrLn (concat ["Left value  :    ",take 100 (show lvalue)])
                 putStrLn (concat ["Right value :    ",take 100 (show rvalue)])