module Data.Homeomorphic.Memo where

import Data.Homeomorphic.Internal
import Data.Maybe
import Data.IORef
import System.IO.Unsafe
import qualified Data.Map as Map
import Control.Monad.State


data Homeomorphic k v = Homeomorphic (IORef (Cache k)) [(Shell k, v)]

type Cache k = Map.Map (Shell k) (Map.Map (Shell k) Bool)


empty :: Homeomorphic k v
empty = Homeomorphic (unsafePerformIO $ newIORef Map.empty) []

insert :: Ord k => Shell k -> v -> Homeomorphic k v -> Homeomorphic k v
insert k v (Homeomorphic a b) = Homeomorphic a ((k,v):b)

findOne :: Ord k => Shell k -> Homeomorphic k v -> Maybe v
findOne k = listToMaybe . find k

find :: Ord k => Shell k -> Homeomorphic k v -> [v]
find y (Homeomorphic a b) = concatMap f b
    where
        f (x,v) = unsafePerformIO $ do
            s <- readIORef a
            let (r,s2) = runState (test x y) s
            writeIORef a s2
            return [v | r]


getCache :: Ord k => Shell k -> Shell k -> State (Cache k) (Maybe Bool)
getCache x y = do
    s <- get
    return $ Map.lookup y s >>= Map.lookup x


addCache :: Ord k => Shell k -> Shell k -> Bool -> State (Cache k) ()
addCache x y b = modify $ Map.insertWith add y (Map.singleton x b)
    where add _ old = Map.insert x b old


-- given a cache, test if it exists, then work the way back
test :: Ord k => Shell k -> Shell k -> State (Cache k) Bool
test x y = do
    v <- getCache x y
    case v of
        Just b -> return b
        Nothing -> do
            b <- diveM x y `orM` coupleM x y
            addCache x y b
            return b

diveM :: Ord k => Shell k -> Shell k -> State (Cache k) Bool
diveM x (Shell _ _ ys) = anyM (x `test`) ys

coupleM :: Ord k => Shell k -> Shell k -> State (Cache k) Bool
coupleM (Shell x1 x2 x3) (Shell y1 y2 y3)
    | x1 == y1 && x2 == y2 = andsM (zipWith test x3 y3)
    | otherwise = return False



orM a b = do a <- a; if a then return True else b
orsM x = foldr orM (return False) x
anyM f = orsM . map f

andM a b = do a <- a; if a then b else return False
andsM x = foldr andM (return True) x