{-# LANGUAGE ExistentialQuantification #-}
{-| Zobrist keys compactly represent the state of perfect information games. -}
module Data.ZCache (
    ZSet(..), ZMap(..), zSet, zMap,
    flipPos, zArray) where
import Data.Array.Unboxed
import Data.Word
import Data.Bits
import System.Random.Mersenne.Pure64
import Data.Set (Set)
import qualified Data.Set as Set
import Data.Map (Map)
import Data.List
import qualified Data.Map as Map

type ZArray a = UArray a Word64

makeRandom :: (PureMT,Int) -> Maybe (Word64,(PureMT,Int))
makeRandom (_,0) = Nothing
makeRandom (p,i) = Just (w,(p',i-1)) where
    (w,p') = randomWord64 p

-- |Create the internal array used for the zobrist key creation
zArray :: Ix a => Word64 -> (a,a) -> Int -> ZArray a
zArray seed dim l = listArray dim (unfoldr makeRandom (pureMT seed,l))

-- |Sets of board states, indexed by zobrist keys
data ZSet a = ZSet Word64 (ZArray a) (Set Word64) deriving (Eq, Show)

-- |Construct a ZSet
zSet :: Ix a =>
    Word64 -- ^ Seed to use for random number generation
    -> (a,a) -- ^ Bounds for piece position values
    -> Int -- ^ Total number of piece positions
    -> ZSet a
zSet seed dim l = ZSet (0::Word64) (zArray seed dim l) Set.empty

flipPos :: Ix a =>
    a -- ^ Most recent move in the game
    -> ZSet a -- ^ The set to insert into
    -> Maybe (ZSet a) -- ^ Final set or Nothing if the board state exists already
flipPos x (ZSet c a s) = if Set.member c' s then Nothing else Just$ ZSet c' a (Set.insert c' s) where
    c' = xor c (a!x)

-- |Maps from board states, indexed by zobrist keys
data ZMap k v = ZMap Word64 (ZArray k) (Map Word64 v) deriving (Eq, Show)

-- |Construct a ZMap
zMap :: forall a b . Ix a =>
    Word64 -- ^ Seed to use for random number generation
    -> (a,a) -- ^ Bounds for piece position values
    -> Int -- ^ Total number of piece positions
    -> ZMap a b
zMap seed dim l = ZMap (0::Word64) (zArray seed dim l) Map.empty

-- |Store a value into a ZMap at a board state
storePos :: Ix a =>
    a -- ^ Most recent move in the game
    -> b -- ^ Value to be associated with the board state
    -> (ZMap a b) -- ^ ZMap to use for storage
    -> Maybe (ZMap a b) -- ^ Final map or Nothing if the board state exists already

storePos k v (ZMap c a m) = if Map.member c' m then Nothing else Just$ ZMap c' a (Map.insert c' v m) where
    c' = xor c (a!k)