{-# LANGUAGE MultiParamTypeClasses, ScopedTypeVariables, AllowAmbiguousTypes, DeriveDataTypeable #-}
module Language.Memo.Safe.AGMemo where 

import Data.Generics.Zipper
import Data.Generics.Aliases
import Language.ZipperAG
import Language.StrategicData
import Data.Data
import Data.Maybe (fromJust)
import Data.List (union)

mkAG :: Data x => x -> Zipper x
mkAG = toZipper

-- MemoAG

atChild eval t i = let (v,t',d) = eval (t.$i)
                   in  (v, parent t',[Child i]) 

atParent eval t = let n = arity t 
                      (v,t',d) = eval (parent t)
                  in  (v, t'.$n, [Parent])

atRight eval t  = let (v,t',d) = eval (t.$>1)
                  in  (v, (t'.$<1),[AtRight])

atLeft eval t   = let (v,t',d) = eval (t.$<1)
                  in  (v, (t'.$>1),[AtLeft])

class MemoTable m => Memo att m a where
  mlookup :: att -> m -> Maybe a
  massign :: att -> a -> m -> m

data Dependency = Parent | Child Int | AtRight | AtLeft
 deriving (Data, Eq, Show)

class Typeable m => MemoTable m where 
  isValidMemoTable :: m -> Bool
  invalidateMemoTable :: m -> m 
  validateMemoTable :: m -> m 
  getDependencies :: m -> [Dependency]
  addDependency :: Dependency -> m -> m
  -- consider deleting memo class for this here:
  -- mstore :: 
  -- mlookup :: 

addDependencies :: MemoTable m => [Dependency] -> m -> m 
addDependencies d m = validateMemoTable $ foldr addDependency m d

class (Typeable dtype, MemoTable m) => Memoizable dtype m where 
  getMemoTable :: dtype m -> m 
  updMemoTable :: (m -> m) -> dtype m -> dtype m
  invalidateDependencies :: Zipper (dtype m) -> Zipper (dtype m)
  invalidateDependencies z = 
    if isValidMemoTable (memoTable z)
      then let z'  = upd' invalidateMemoTable z
               dep = getDependencies (memoTable z)
               fromDep Parent r = parent r 
               fromDep (Child n) r = r.$n
               fromDep AtRight r = r.$>1
               fromDep AtLeft  r = r.$<1
               unFromDep Parent r = r.$(arity z') 
               unFromDep (Child n) r = parent r
               unFromDep AtRight r = r.$<1
               unFromDep AtLeft  r = r.$>1
           in foldr (\d x -> unFromDep d $ invalidateDependencies (fromDep d x)) z' dep
      else z


-- type AGTree_m dtype m a = Zipper (dtype m) -> (a, Zipper (dtype m)) 
type AGTree_m dtype m a =  (Zipper (dtype m) -> (a, Zipper (dtype m), [Dependency]))


memo :: (Memoizable dtype m, Memo attr m a) => 
        attr -> 
        AGTree_m dtype m a -> 
        AGTree_m dtype m a
memo attr eval = \z ->
   case (isValidMemoTable (memoTable z), mlookup attr (memoTable z)) of
    (True, Just v)   -> (v,z, getDependencies (memoTable z))
    _                -> let (v,z', d) = eval z
                        in  (v, upd' (addDependencies d . massign attr v) z', d) 



memoTable :: forall dtype m. (Memoizable dtype m) => Zipper (dtype m) -> m
memoTable zx = let a' = (fromJust $ getHole zx) :: dtype m
               in getMemoTable a' 


-- "forced typechecking"
upd' :: Memoizable dtype m => (m -> m) -> Zipper (dtype m) -> Zipper (dtype m)
upd' f z = setHole (aux f z) z
 where aux :: (Memoizable dtype m) => (m -> m) -> Zipper (dtype m) -> dtype m
       aux f = updMemoTable f . fromJust . getHole