{-# LANGUAGE BangPatterns #-}
module Data.HBlock where
import Prelude hiding (foldr, foldl)
import Data.Foldable (Foldable(..))
import Control.Monad.ST
import Data.Vector (Vector)
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as VM
import Data.IntSet (IntSet)
import qualified Data.IntSet as IS
import Data.SafeCopy
import GHC.Exts (build)
import qualified Data.List as L
-- import Data.Indexation.IxData (IxData(IxData, ixDataPos, ixData))

data HBlock a = HBlock { set      :: !IntSet
                       , values   :: (Vector a)
                       , size     :: !Int
                       , numItems :: !Int
                       , blockLen :: !Int
                       , emptyVal :: !a
                       } 

instance Functor HBlock where
    fmap f (HBlock s v len num blen e) = 
        if num == 0 
            then emptyhb
            else IS.foldl transf emptyhb s
        where 
            emptyhb = empty len blen (f e)
            transf hb pos = insert (f $ V.unsafeIndex v pos) hb

instance Foldable HBlock where
    foldr f z (HBlock s v _ num _ _) =
        if num == 0 then z else IS.foldr' foldit z s
        where 
            foldit pos item = f (V.unsafeIndex v pos) item
    foldl f z (HBlock s v _ num _ _) = 
        if num == 0 then z else IS.foldl foldit z s
        where 
            foldit item pos = f item (V.unsafeIndex v pos) 

instance (SafeCopy a) => SafeCopy (HBlock a) where
    version = 0
    kind = base
    putCopy h@(HBlock _ _ _ n bl e) = contain $ do safePut n
                                                   safePut bl
                                                   safePut e
                                                   safePut (toList h)
                                                   return ()
    getCopy = contain $ do n <- safeGet 
                           bl <- safeGet
                           e <- safeGet 
                           lst <- safeGet 
                           return $ fromList n bl e lst

-- | Creates a new HBlock 
-- receives the initial size, the ammount of items that the vector gets 
-- incremented when it is almost full, and an empty item to place on the growth
-- vector
empty :: Int -> Int -> a -> HBlock a
empty initSize blockSize e = 
    HBlock { set = IS.empty
           , values = V.replicate initSize e
           , size = initSize
           , numItems = 0
           , blockLen = blockSize
           , emptyVal = e
           }
{-# INLINABLE empty #-}

length :: HBlock a -> Int
length b = numItems b
{-# INLINE length#-}

-- | similar to insert but returns the index where the item was placed
insertGetIx :: a -> HBlock a -> (Int, HBlock a)
insertGetIx item b = 
    (pos, b { set = IS.insert pos (set b)
          , values = newvec
          , size = if growthNeeded then blockLen b else (size b) - 1
          , numItems = numItems b + 1
          })
    where
        pos = if IS.null (set b) then 0 else IS.findMax (set b) + 1
        oldvec = values b
        growthNeeded = size b == 0
        growthSize = blockLen b
        newvec = runST $ 
            do v <- V.unsafeThaw oldvec
               -- grow the vector by the extendAmmount
               let extendIxs = [pos..(pos+growthSize)]
               v' <- if growthNeeded
                        then do newv <- VM.unsafeGrow v growthSize
                                (mapM_ (\i -> VM.write newv i (emptyVal b)) extendIxs)
                                return newv
                        else return v
               VM.write v' pos item
               V.unsafeFreeze v'
{-# INLINABLE insertGetIx #-}

insert :: a -> HBlock a -> HBlock a
insert i b = snd $ insertGetIx i b
{-# INLINE insert #-}

adjust :: (a -> a) -> Int -> HBlock a -> HBlock a
adjust f pos b = b { values = adjusted }
    where
        oldvec = values b
        adjusted = runST $ 
            do v <- V.unsafeThaw oldvec
               val <- VM.read v pos
               VM.write v pos (f val)
               V.unsafeFreeze v
{-# INLINABLE adjust #-}

update :: a -> Int -> HBlock a -> HBlock a
update new pos b = b { values = adjusted }
    where
        oldvec = values b
        adjusted = runST $ 
            do v <- V.unsafeThaw oldvec
               VM.write v pos new
               V.unsafeFreeze v
{-# INLINABLE update #-}

occupied :: Int -> HBlock a -> Bool
occupied pos b = IS.member pos (set b)
{-# INLINE occupied #-}

getPos :: HBlock a -> Int -> Maybe a
getPos hb pos = if occupied pos hb then Just $ V.unsafeIndex (values hb) pos
                                   else Nothing
{-# INLINABLE getPos #-}

getSet :: HBlock a -> IntSet -> HBlock a
getSet hb s = hb { set = newSet
                 , numItems = IS.size newSet
                 }
    where
        newSet = IS.intersection s (set hb)
        -- ^TODO map and copy the whole vector into the new hb ?
        -- test growing the vector and then using the returned HB from this function

delete :: Int -> HBlock a -> HBlock a
delete pos b = b { set = IS.delete pos (set b) 
                 , numItems = if occupied pos b then numItems b - 1 
                                                else numItems b
                 }
{-# INLINABLE delete #-}

-- | /O(n)/ Return a list of this hblock's elements.  The list is
-- produced lazily.
-- toList :: HashMap k v -> [(k, v)]
-- toList t = build (\ c z -> foldrWithKey (curry c) z t)
toList :: HBlock a -> [a]
toList hb = build (\c z -> foldr c z hb)
{-# INLINE toList #-}

-- | /O(n)/ Construct a hblock with the values from a list.
fromList :: Int -> Int -> a -> [a] -> HBlock a
fromList s bs e lst = L.foldl' (flip insert) (empty s bs e) lst
{-# INLINABLE fromList #-}