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

module Generics.MultiRec.Rewriting.Rules where

import Prelude hiding (zip)

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

import Control.Monad (join)

-- * Rule types

infix 5 :~>
data RuleSpec a = a :~> a

lhsR :: RuleSpec a -> a
lhsR (lhs :~> _) = lhs

rhsR :: RuleSpec a -> a
rhsR (_ :~> rhs) = rhs

type Rule   phi a = RuleSpec (Scheme phi a)

-- * Schemes

type MVar = Int
type Scheme phi = HFix (K MVar :+: PF phi)

mvar :: phi es a ix -> MVar -> Scheme phi ix
mvar _ = HIn . L . K

pf :: phi es a ix -> (PF phi) (Scheme phi) ix -> Scheme phi ix
pf _ = HIn . R

toScheme :: (HFunctor (FamRel phi es) (FamRel phi es) (PF phi), Fam phi es) 
         => phi es a ix -> ix -> Scheme phi ix
toScheme p x = pf p $ hmap (\p' -> toScheme p' . unI0) p $ from p x

-- This can probably be written in a cleaner way
fromScheme :: (HFunctor (FamRel phi es) (FamRel phi es) (PF phi), Fam phi es)
           => phi es a ix -> Scheme phi ix -> Maybe ix
fromScheme p = fmap (to p) . join
             . fmap (hmapM (\p' -> fmap I0 . fromScheme p') p) . unpf p where
  unpf :: phi ix -> Scheme phi ix -> Maybe ((PF phi) (Scheme phi) ix)
  unpf _ (HIn (R x)) = Just x
  unpf _ _           = Nothing
{-
data ViewScheme phi ix = MVar MVar | PF (PF phi (Scheme phi) ix)

viewScheme :: phi ix -> Scheme phi ix -> ViewScheme phi ix
viewScheme _ (HIn (L (K x))) = MVar x
viewScheme _ (HIn (R r))     = PF r

foldScheme :: (HFunctor phi (PF phi))
           => phi ix
           -> (forall ix'. MVar -> r ix')
           -> (forall ix'. () => PF phi r ix' -> r ix')
           -> Scheme phi ix -> r ix
foldScheme p f g scheme =
  case viewScheme p scheme of
    MVar x -> f x
    PF r   -> g (hmap (\p' -> foldScheme p' f g) p r)

-- * Making a rule

class Builder phi a where
  type Target a :: *
  base          :: phi (Target a) -> a -> RuleSpec (Target a)
  diag          :: phi (Target a) -> a -> [RuleSpec (Target a)]

instance Builder phi (RuleSpec a) where
  type Target (RuleSpec a) = a
  base _ x                 = x
  diag _ x                 = [x]

instance (Builder phi a, El phi b, Fam phi) => Builder phi (b -> a) where
  type Target (b -> a) = Target a
  base ix f            = base ix (f (gleft  (proof::phi b)))
  diag ix f            = base ix (f (gright (proof::phi b))) :
                         diag ix (f (gleft  (proof::phi b)))

rule :: (Builder phi r, HZip phi (PF phi), Fam phi) =>
        phi (Target r) -> r -> Rule phi (Target r)
rule p f = foldr1 mergeRules rules
  where
    mergeRules x y = 
      mergeSchemes p (lhsR x) (lhsR y) :~>
      mergeSchemes p (rhsR x) (rhsR y)
    rules          = zipWith (ins (base p f)) (diag p f) [0..]
    ins x y v      = 
      insertMVar v p (I0 (lhsR x)) (I0 (lhsR y)) :~>
      insertMVar v p (I0 (rhsR x)) (I0 (rhsR y))

-- | Merge two schemes by accepting a meta-variable case in one of the 
-- schemes, and otherwise zip the remaining parts.
mergeSchemes :: HZip phi (PF phi) 
             => phi ix -> Scheme phi ix -> Scheme phi ix -> Scheme phi ix
mergeSchemes w p@(HIn x) q@(HIn y)
  = case viewScheme w p of
      MVar _ -> p
      _      -> case viewScheme w q of
                  MVar _ -> q
                  _      -> HIn (hzip' mergeSchemes w x y)

-- | Extend a value with a meta-variable case. Two structures are traversed 
-- in parallel, whenever the structures mismatch, insert a meta-variable.
-- Otherwise, lift the matching parts.
insertMVar :: forall phi ix .(HZip phi (PF phi), Fam phi) 
           => MVar -> phi ix -> I0 ix -> I0 ix -> Scheme phi ix
insertMVar name w (I0 x) (I0 y)
  = case hzip (insertMVar name) w (from w x) (from w y) of
      Just struc -> pf w struc
      Nothing    -> mvar w name
-}