{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE OverlappingInstances #-}


module Control.Search.Memo where

import Control.Monatron.Monatron hiding (Abort, L, state, cont)
import Control.Monatron.Zipper hiding (i,r)
import Control.Monatron.IdT
import Control.Monatron.MonadInfo

import Data.List (sort, nub, sortBy)
import Data.Maybe (fromJust)
import Data.Map (Map)
import qualified Data.Map as Map

import Control.Search.Language
import Control.Search.GeneratorInfo
import Control.Search.SStateT

data MemoKey  = MemoKey { memoFn :: String, memoInfo :: Maybe Info, memoStack :: Maybe String, memoExtra :: Maybe (Map Int String), memoStatement :: Maybe Statement, memoParams :: [String] }
  deriving (Eq, Ord)

data MemoValue = MemoValue { memoId :: Int, memoCode :: Statement, memoUsed :: Int, memoFields :: [(String,String)] }

data MemoInfo = MemoInfo { memoMap :: Map MemoKey MemoValue 
                         , memoCount :: Int
                         , memoRead :: Map Int String
                         }

initMemoInfo = MemoInfo { memoMap = Map.empty
                        , memoCount = 0
                        , memoRead = Map.empty
                        }

newtype MemoT m a = MemoT { unMemoT :: SStateT MemoInfo m a }
  deriving (MonadT,StateM MemoInfo,FMonadT)

instance MonadInfoT MemoT where
  tminfo x = miInc "MemoT" (minfo $ runMemoT x)

-- runMemoT :: Monad m => MemoT m a -> m (a,[(String,Statement,[(String,String)])])
runMemoT m = do (Tup2 a s) <- runSStateT initMemoInfo (unMemoT m)
                return (a, {- map (\(key,val) -> ( memoFn key ++ show (memoId val)
                                              , comment (" fn=" ++ memoFn key ++ " stack='" ++ show (memoStack key) ++ "' extra='" ++ show (memoExtra key) ++ "' used: " ++ show (memoUsed val)) >>> memoCode val
                                              , memoFields key
                                              )
                                 ) $ -} sortBy (\(ka,va) (kb,vb) -> compare (memoId va) (memoId vb)) $ Map.toList (memoMap s)
                       )

-- runReaderMemoT :: (ReaderM r m, ReaderMemoM r (MemoT m)) => MemoT m a -> m (a,[(String,Statement,Info)])
-- runReaderMemoT m = do val <- ask
--                      runMemoT (memoLocal (const val) m)

class Monad m => MemoM m where
  getMemo :: m MemoInfo 
  setMemo :: MemoInfo -> m ()

instance Monad m => MemoM (MemoT m) where
  getMemo  = MemoT $ get 
  setMemo  = MemoT . put

instance (MemoM m, FMonadT t) => MemoM (t m) where
  getMemo = lift $ getMemo
  setMemo = lift . setMemo