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

module Rewriting where

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

import Base
import HFunctor
import HZip
import LeftRight
import Repr
import TEq
import Rules

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

rewriteM :: forall a s . (Rewrite s, Ix s a)
         => Rule s a -> a -> Maybe a
rewriteM (lhs :~> rhs) term = match lhs term >>= return . (\s -> inst s index rhs)

rewrite :: forall a s . (Rewrite s, Ix s a)
        => Rule s a -> a -> a
rewrite f term = maybe term id (rewriteM f term)

-- | Matching a term

match :: (Rewrite s,Monad m,Ix s ix)
      => Scheme s ix -> ix -> m (Subst s)
match pat term = execStateT (matchM index pat (I0 term)) M.empty

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

checkEqual :: (HZip (PF s),Monad m,Ix s ix,Repr s)
           => s ix -> ix -> Dyn s -> m ()
checkEqual ix e (Dyn ix' e')
  = do
    Refl <- ixeq ix ix'
    geq' ix (I0 e) (I0 e')


class (HFunctor (PF s),HZip (PF s),LR (PF s) s,Repr s) => Rewrite s


-- | Apply a substition to a scheme

apply :: forall s ix . (Repr s,HFunctor (PF s)) =>
         Subst s -> s ix -> Scheme s ix -> Scheme s ix
apply s ix p
  = case viewScheme p of
      MVar x ->
        case M.lookup x s of
          Nothing -> mvar x
          Just (Dyn ix' e) ->
            case ixeq ix ix' of
              Just Refl -> toScheme e
              Nothing -> error "Coerce error in apply"
      PF r -> HIn (R (hmap (apply s) r))

-- | Apply a substitution assuming that the result contains
-- no meta-variables

inst :: forall s ix . (Repr s,HFunctor (PF s),Ix s ix) =>
        Subst s -> s ix -> Scheme s ix -> ix
inst s ix p
  = case viewScheme p of
      MVar x ->
        case M.lookup x s of
          Just (Dyn ix' e)
            -> case ixeq ix ix' of
                 Just Refl -> e
                 Nothing -> error "Coerce error in inst"
      PF r ->
        (to . hmap (\ix -> I0 . inst s ix)) r


-- | Checks left/right property (see van Noort et al.)

prop_LeftRight :: forall a s . (Show a,Rewrite s,Repr s,Ix s a) =>
                  s a -> IO ()
prop_LeftRight witness
  = case hzipM fake (from'' lvalue) (from'' 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''. Ix s ix => s ix -> r ix -> r' ix -> Maybe (r'' ix)
    fake = undefined
    lvalue = left witness
    rvalue = right witness
    showVals = do
                 putStrLn (concat ["Left value  :    ",take 100 (show lvalue)])
                 putStrLn (concat ["Right value :    ",take 100 (show rvalue)])


-- The horrible stuff below should be erased when the more general "rule" is fixed
-- Actually "rule" is already fixed, but it only works in GHC HEAD, see ASTExamples.
--     -- Alexey

liftTerm2 :: forall s ix' ix'' ix''' .
             (HZip (PF s),LR (PF s) s,Ix s ix',Ix s ix'',Ix s ix''')
          => MVar -> MVar -> (ix' -> ix'' -> ix''') -> Scheme s ix'''
liftTerm2 name1 name2 f =
   let pat1 = insertMVar name1 ix''' 
                (I0 $ f (left ix') (left ix'')) (I0 $ f (right ix') (left ix''))
       pat2 = insertMVar name2 ix''' 
                (I0 $ f (left ix') (left ix'')) (I0 $ f (left ix') (right ix''))
   in  pat1 `mergeSchemes` pat2
  where
    ix'   = index :: s ix'
    ix''  = index :: s ix''
    ix''' = index :: s ix'''

liftRule2 :: (Ix s t1, Ix s ix'', Ix s ix', LR (PF s) s, HZip (PF s)) =>
             s t1 -> (ix' -> ix'' -> RuleSpec t1) -> Rule s t1
liftRule2 _ f = liftTerm2 0 1 (\x y -> lhsR (f x y)) :~>
                liftTerm2 0 1 (\x y -> rhsR (f x y))

rewrite2 :: (Ix s t1,
             Ix s ix'',
             Ix s ix',
             LR (PF s) s,
             Rewrite s) =>
            s t1 -> (ix' -> ix'' -> RuleSpec t1) -> t1 -> Maybe t1
rewrite2 proxy f = rewriteM (liftRule2 proxy f)