{-# OPTIONS_GHC -fglasgow-exts #-}

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 Indices
import Base

-- --------------------------------------
-- Schemes
-- --------------------------------------
type MVar   = Int
type Scheme l = HFix (Sum (K MVar) (PF l) l)

var :: MVar -> Scheme l ix
var = HIn . Inl . K

pf :: (PF l) l (Scheme l) ix -> Scheme l ix
pf = HIn . Inr

toScheme :: (GMap (PF l), Ix l ix) => ix -> Scheme l ix
toScheme = pf . gmap (toScheme . unI) . from

data ViewScheme l ix = MVar MVar | PF (PF l l (Scheme l) ix)

viewScheme :: Scheme l ix -> ViewScheme l ix
viewScheme (HIn (Inl (K x))) = MVar x
viewScheme (HIn (Inr r))     = PF r

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

-- --------------------------------------
-- Matching a term
-- --------------------------------------
--

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

matchM :: (Rewrite l,Monad m,Ix l ix)
      => l ix -> Scheme l ix -> I ix -> StateT (Subst l) m ()
matchM ix scheme (I 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 ix matchM r (from'' e)

checkEqual :: (Zip (PF l),Monad m,Ix l ix,Repr l)
           => l ix -> ix -> Dyn l -> m ()
checkEqual ix e (Dyn ix' e')
  = do
    TEq <- ixeq ix ix'
    geq' ix (I e) (I e')


class (GMap (PF l),Zip (PF l),LR (PF l) l,Repr l) => Rewrite l


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

-- | Apply a substition to a pattern
apply :: forall l ix . (Repr l,GMap (PF l))
      => Subst l -> l ix -> Scheme l ix -> Scheme l ix
apply s ix p
  = case viewScheme p of
      MVar x ->
        case M.lookup x s of
          Nothing -> var x
          Just (Dyn ix' e) ->
            case ixeq ix ix' of
              Just TEq -> toScheme e
              Nothing -> error "Coerce error in apply"
      PF r -> HIn (Inr (gmapW (apply s) r))

inst :: forall l ix . (Repr l,GMap (PF l),Ix l ix)
     => Subst l -> l ix -> Scheme l 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 TEq -> e
                 Nothing -> error "Coerce error in inst"
      PF r ->
        (to . gmapW (\ix -> I . inst s ix)) r

rewrite :: forall a l . (Rewrite l, Ix l a)
        => Rule (Scheme l a) -> a -> Maybe a
rewrite (lhs :~> rhs) term = match lhs term >>= return . (\s -> inst s ix rhs)

data Rule a = a :~> a
lhsR (lhs :~> rhs) = lhs
rhsR (lhs :~> rhs) = rhs


diff :: (Zip (PF l),Ix l ix) => MVar -> l ix -> I ix -> I ix -> Scheme l ix
diff name ix (I x) (I y)
  = case zipW (diff name) (from'' x) (from'' y) of
      Just struc -> pf struc
      Nothing    -> var name

merge :: Zip (PF l) => Scheme l ix -> Scheme l ix -> Scheme l ix
merge p@(HIn x) q@(HIn y)
  = case viewScheme p of
      MVar i -> p
      _     -> case viewScheme q of
                 MVar j -> q
                 _     -> HIn (zip merge x y)



liftTerm2 :: forall l ix' ix'' ix''' .
             (Zip (PF l),LR (PF l) l,Ix l ix',Ix l ix'',Ix l ix''')
          => MVar -> MVar -> (ix' -> ix'' -> ix''') -> Scheme l ix'''
liftTerm2 name1 name2 f = let pat1 = diff name1 ix''' (I $ f (left ix') (left ix'')) (I $ f (right ix') (left ix''))
                              pat2 = diff name2 ix''' (I $ f (left ix') (left ix'')) (I $ f (left ix') (right ix''))
                          in  pat1 `merge` pat2
  where
    ix' = ix :: l ix'
    ix'' = ix :: l ix''
    ix''' = ix :: l ix'''

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

rewrite2 :: (Ix l t1,
             Ix l ix'',
             Ix l ix',
             LR (PF l) l,
             Rewrite l) =>
            l t1 -> (ix' -> ix'' -> Rule t1) -> t1 -> Maybe t1
rewrite2 proxy f = rewrite (liftRule2 proxy f)