{-# LANGUAGE TypeFamilies               #-}
{-# LANGUAGE GADTs                      #-}
{-# LANGUAGE RankNTypes                 #-}
{-# LANGUAGE TypeOperators              #-}
{-# LANGUAGE MultiParamTypeClasses      #-}
{-# LANGUAGE FlexibleInstances          #-}
{-# LANGUAGE FlexibleContexts           #-}
{-# LANGUAGE OverlappingInstances       #-}
{-# LANGUAGE UndecidableInstances       #-}
{-# LANGUAGE ScopedTypeVariables        #-}
{-# LANGUAGE DataKinds                  #-}
{-# LANGUAGE PolyKinds                  #-}

module Generics.MultiRec.Transformations.MemoTable where

import Generics.MultiRec hiding ( show, foldM )
import Control.Monad.State hiding ( foldM, mapM )
import qualified Data.Map as Map
import Data.Map ( Map )

import Generics.MultiRec.Transformations.Path
import Generics.MultiRec.Transformations.Children

--------------------------------------------------------------------------------
-- One memotable per type in the family
--------------------------------------------------------------------------------
type family Ixs (phi :: * -> *) :: [*] -- one instance per family; tells us
                                       -- which types are in the family

data HList (l :: [*]) where -- heterogeneous collections
  HNil  :: HList '[]
  HCons :: h -> HList t -> HList (h ': t)

type family HMap (h :: * -> * -> *) (f :: * -> *) (g :: * -> *) (l :: [*]) :: [*]
type instance HMap h f g '[]      = '[]
type instance HMap h f g (t ': ts) = h (f t) (g t) ': HMap h f g ts

-- Memotables are now heterogenous collections mapping keys to values, indexed
-- over each family type
-- type MemoTable' phi top ixs = HList (HMap Map (MemKey' top) (MemVal' top) ixs)
type MemoTable phi top ixs = HList (MemoTable' One phi top ixs)

data Type = One | Two

type family MemoTable' (x :: Type) (phi :: * -> *) top (ixs :: [*]) :: [*]
type instance MemoTable' x phi top '[]       = '[]
type instance MemoTable' x phi top (t ': ts) = MemCell x phi top t
                                            ': MemoTable' x phi top ts

type family MemCell (x :: Type) (phi :: * -> *) top (t :: *) :: *
-- type instance MemCell One phi top t = Map (MemKey t) (MemVal phi top  t)
type instance MemCell One phi top t = Map (MemKey t) (MemVal phi top  t)

type MemKey         t = (Bool, t, t)
type MemVal phi top t = [Insert phi top t]

-- Lookup on the right table
data Proxy (t :: k) = Proxy
class {- (ixs ~ Ixs phi) => -} Lookup phi (ixs :: [*]) ix where
  lookupMT :: Proxy '(phi, ixs)
           -- I don't like this Proxy too much, but I'm afraid it's necessary
           -> MemoTable phi top ixs -> MemKey ix -> Maybe (MemVal phi top ix)
  insertMT :: Proxy '(phi, ixs)
              -> MemKey ix -> MemVal phi top ix
              -> MemoTable phi top ixs -> MemoTable phi top ixs

-- this probably shouldn't happen...
instance Lookup phi '[] ix where
  lookupMT _ HNil _ = Nothing
  insertMT _ _ _ HNil = HNil

-- this is the right memotable for this type
instance (Ord t) => Lookup phi (t ': ts) t where
  lookupMT _ (HCons mt _) k = Map.lookup k mt
  insertMT _ k v (HCons mt tl) = HCons (Map.insert k v mt) tl

-- this is not the right memotable for this type, keep searching
instance (Lookup phi ts ix) => Lookup phi (t ': ts) ix where
  lookupMT _ (HCons _ mts) = lookupMT (Proxy :: Proxy '(phi, ts)) mts
  insertMT _ k v (HCons hd mts) = HCons hd (insertMT (Proxy :: Proxy '(phi, ts)) k v mts)

type Memo phi top a = State (MemoTable phi top (Ixs phi)) a

class EmptyMemo (phi :: * -> *) top (ixs :: [*]) where
  emptyMemo :: Proxy '(phi,top,ixs) -> MemoTable phi top ixs

instance EmptyMemo phi top '[] where
  emptyMemo _ = HNil

instance (EmptyMemo phi top t) => EmptyMemo phi top (h ': t) where
  emptyMemo _ = HCons Map.empty (emptyMemo (Proxy :: Proxy '(phi,top,t)))

runMemo :: forall phi top a. (EmptyMemo phi top (Ixs phi))
        => Proxy '(phi,top) -> Memo phi top a -> a
runMemo _ = flip evalState (emptyMemo (Proxy :: Proxy '(phi,top,Ixs phi)))

recMemo :: forall phi top ix.
           (Fam phi, Children phi (PF phi) ix, Lookup phi (Ixs phi) ix, Eq ix, GetChildrenTable phi (Ixs phi) ix) =>
           (forall ix. (Children phi (PF phi) ix, Lookup phi (Ixs phi) ix, Eq ix, GetChildrenTable phi (Ixs phi) ix) => Bool -> phi ix -> ix -> ix -> Memo phi top [ Insert phi top ix ])
           -> Bool -> phi ix -> ix -> ix -> Memo phi top [ Insert phi top ix ]
recMemo f a p b c = do
  mp <- get
  let k = (a,b,c)
  case lookupMT (Proxy :: Proxy '(phi, Ixs phi)) mp k of
    Just r -> return r
    Nothing -> do
      r <- f a p b c
      modify $ insertMT (Proxy :: Proxy '(phi, Ixs phi)) k r
      return r

--------------------------------------------------------------------------------
-- Memo table of all children
--------------------------------------------------------------------------------
-- type family ChildTable (phi :: * -> *) top (ixs :: [*]) :: [*]
-- type instance ChildTable phi top '[]       = '[]
-- type instance ChildTable phi top (t ': ts) = [(Path phi top t, t)] ': ChildTable phi top ts

type ChildTable phi top ixs = MemoTable' Two phi top ixs
type instance MemCell Two phi top t = [(Path phi t top, t)]

{-
This is a first step to merge ChildTable and MemoTable. However, it would be
really cool if we can also merge getChTable and lookupMT, for example. But for
that we have to change MemCell to either be a Map or a list of pairs.

Changing ChildTable to use a Map as its cell isn't nice, because we start
requiring |Ord (r ix)| constraints in AllChildren.

The other way around, changing MemoTable to use a list of pairs instead, is
probably possible, but I'm not sure if it's a good idea. I'd like Jeroen's
opinion on this.
-}

class ChildrenTable (phi :: * -> *) top (ixs :: [*]) where
  childrenTable :: Proxy '(phi,top,ixs) -> top -> HList (ChildTable phi top ixs)

instance ChildrenTable phi top '[] where
  childrenTable _ _ = HNil

-- For the toplevel type, also include full tree
instance (Fam phi, El phi top, Children phi (PF phi) top, ChildrenTable phi top t)
         => ChildrenTable phi top (top ': t) where
  childrenTable _ top = HCons ((Empty,top) : allChildren proof proof top)
                        (childrenTable (Proxy :: Proxy '(phi,top,t)) top)

-- For all other types
instance (Fam phi, El phi h, El phi top, Children phi (PF phi) h, ChildrenTable phi top t)
         => ChildrenTable phi top (h ': t) where
  childrenTable _ top = HCons (allChildren proof proof top)
                        (childrenTable (Proxy :: Proxy '(phi,top,t)) top)

class GetChildrenTable phi (ixs :: [*]) ix where
  getChTable :: Proxy '(phi, top, ixs)
           -- I don't like this Proxy too much, but I'm afraid it's necessary
           -> HList (ChildTable phi top ixs) -> MemCell Two phi top ix

instance GetChildrenTable phi '[] ix where
  getChTable _ HNil = error "This shouldn't happen"

instance (Ord t) => GetChildrenTable phi (t ': ts) t where
  getChTable _ (HCons mt _) = mt

-- this is not the right memotable for this type, keep searching
instance (GetChildrenTable phi ts ix) => GetChildrenTable phi (t ': ts) ix where
  getChTable _ (HCons _ mts) = getChTable (Proxy :: Proxy '(phi, top, ts)) mts