-- |
-- Module      : Foundation.Array.Bitmap
-- License     : BSD-style
-- Maintainer  : Vincent Hanquez <vincent@snarc.org>
-- Stability   : experimental
-- Portability : portable
--
-- A simple abstraction to a set of Bits (Bitmap)
--
-- Largely a placeholder for a more performant implementation,
-- most operation goes through the List representation (e.g. [Bool])
-- to conduct even the most trivial operation, leading to a lots of
-- unnecessary churn.
--
{-# LANGUAGE BangPatterns #-}
module Foundation.Array.Bitmap
    ( Bitmap
    , MutableBitmap
    , empty
    , append
    , concat
    , unsafeIndex
    , index
    , read
    , unsafeRead
    , write
    , unsafeWrite
    , snoc
    , cons
    ) where

import           Basement.UArray (UArray)
import qualified Basement.UArray as A
import           Basement.UArray.Mutable (MUArray)
import           Basement.Compat.Bifunctor (first, second, bimap)
import           Basement.Compat.Semigroup
import           Basement.Exception
import           Basement.Compat.Base
import           Basement.Types.OffsetSize
import           Basement.Monad
import qualified Foundation.Collection as C
import           Foundation.Numerical
import           Data.Bits
import           Foundation.Bits
import           GHC.ST
import qualified Data.List

data Bitmap = Bitmap (CountOf Bool) (UArray Word32)

data MutableBitmap st = MutableBitmap (CountOf Bool) (MUArray Word32 st)

bitsPerTy :: Int
bitsPerTy = 32

shiftPerTy :: Int
shiftPerTy = 5

maskPerTy :: Int
maskPerTy = 0x1f

instance Show Bitmap where
    show v = show (toList v)
instance Eq Bitmap where
    (==) = equal
instance Ord Bitmap where
    compare = vCompare
instance Semigroup Bitmap where
    (<>) = append
instance Monoid Bitmap where
    mempty  = empty
    mappend = append
    mconcat = concat

type instance C.Element Bitmap = Bool

instance IsList Bitmap where
    type Item Bitmap = Bool
    fromList = vFromList
    toList = vToList

instance C.InnerFunctor Bitmap where
    imap = map

instance C.Foldable Bitmap where
    foldr = foldr
    foldl' = foldl'
    foldr' = foldr'

instance C.Collection Bitmap where
    null = null
    length = length
    elem e = Data.List.elem e . toList
    maximum = any id . C.getNonEmpty
    minimum = all id . C.getNonEmpty
    all = all
    any = any

instance C.Sequential Bitmap where
    take = take
    drop = drop
    splitAt = splitAt
    revTake n = unoptimised (C.revTake n)
    revDrop n = unoptimised (C.revDrop n)
    splitOn = splitOn
    break = break
    breakEnd = breakEnd
    span = span
    filter = filter
    reverse = reverse
    snoc = snoc
    cons = cons
    unsnoc = unsnoc
    uncons = uncons
    intersperse = intersperse
    find = find
    sortBy = sortBy
    singleton = fromList . (:[])
    replicate n = fromList . C.replicate n

instance C.IndexedCollection Bitmap where
    (!) l n
        | isOutOfBound n (length l) = Nothing
        | otherwise                     = Just $ index l n
    findIndex predicate c = loop 0
      where
        !len = length c
        loop i
            | i .==# len                  = Nothing
            | predicate (unsafeIndex c i) = Just i
            | otherwise                   = Nothing

instance C.MutableCollection MutableBitmap where
    type MutableFreezed MutableBitmap = Bitmap
    type MutableKey MutableBitmap = Offset Bool
    type MutableValue MutableBitmap = Bool

    thaw = thaw
    freeze = freeze
    unsafeThaw = unsafeThaw
    unsafeFreeze = unsafeFreeze

    mutNew = new
    mutUnsafeWrite = unsafeWrite
    mutUnsafeRead = unsafeRead
    mutWrite = write
    mutRead = read

bitmapIndex :: Offset Bool -> (Offset Word32, Int)
bitmapIndex (Offset !i) = (Offset (i .>>. shiftPerTy), i .&. maskPerTy)
{-# INLINE bitmapIndex #-}

-- return the index in word32 quantity and mask to a bit in a bitmap
{-
bitmapAddr :: Int -> (# Int , Word #)
bitmapAddr !i = (# idx, mask #)
  where (!idx, !bitIdx) = bitmapIndex i
        !mask = case bitIdx of
                    0  -> 0x1
                    1  -> 0x2
                    2  -> 0x4
                    3  -> 0x8
                    4  -> 0x10
                    5  -> 0x20
                    6  -> 0x40
                    7  -> 0x80
                    8  -> 0x100
                    9  -> 0x200
                    10 -> 0x400
                    11 -> 0x800
                    12 -> 0x1000
                    13 -> 0x2000
                    14 -> 0x4000
                    15 -> 0x8000
                    16 -> 0x10000
                    17 -> 0x20000
                    18 -> 0x40000
                    19 -> 0x80000
                    20 -> 0x100000
                    21 -> 0x200000
                    22 -> 0x400000
                    23 -> 0x800000
                    24 -> 0x1000000
                    25 -> 0x2000000
                    26 -> 0x4000000
                    27 -> 0x8000000
                    28 -> 0x10000000
                    29 -> 0x20000000
                    30 -> 0x40000000
                    _  -> 0x80000000
-}

thaw :: PrimMonad prim => Bitmap -> prim (MutableBitmap (PrimState prim))
thaw (Bitmap len ba) = MutableBitmap len `fmap` C.thaw ba

freeze :: PrimMonad prim => MutableBitmap (PrimState prim) -> prim Bitmap
freeze (MutableBitmap len mba) = Bitmap len `fmap` C.freeze mba

unsafeThaw :: PrimMonad prim => Bitmap -> prim (MutableBitmap (PrimState prim))
unsafeThaw (Bitmap len ba) = MutableBitmap len `fmap` C.unsafeThaw ba

unsafeFreeze :: PrimMonad prim => MutableBitmap (PrimState prim) -> prim Bitmap
unsafeFreeze (MutableBitmap len mba) = Bitmap len `fmap` C.unsafeFreeze mba

unsafeWrite :: PrimMonad prim => MutableBitmap (PrimState prim) -> Offset Bool -> Bool -> prim ()
unsafeWrite (MutableBitmap _ ma) i v = do
    let (idx, bitIdx) = bitmapIndex i
    w <- A.unsafeRead ma idx
    let w' = if v then setBit w bitIdx else clearBit w bitIdx
    A.unsafeWrite ma idx w'
{-# INLINE unsafeWrite #-}

unsafeRead :: PrimMonad prim => MutableBitmap (PrimState prim) -> Offset Bool -> prim Bool
unsafeRead (MutableBitmap _ ma) i = do
    let (idx, bitIdx) = bitmapIndex i
    flip testBit bitIdx `fmap` A.unsafeRead ma idx
{-# INLINE unsafeRead #-}

write :: PrimMonad prim => MutableBitmap (PrimState prim) -> Offset Bool -> Bool -> prim ()
write mb n val
    | isOutOfBound n len = primOutOfBound OOB_Write n len
    | otherwise          = unsafeWrite mb n val
  where
    len = mutableLength mb
{-# INLINE write #-}

read :: PrimMonad prim => MutableBitmap (PrimState prim) -> Offset Bool -> prim Bool
read mb n
    | isOutOfBound n len = primOutOfBound OOB_Read n len
    | otherwise        = unsafeRead mb n
  where len = mutableLength mb
{-# INLINE read #-}

-- | Return the element at a specific index from a Bitmap.
--
-- If the index @n is out of bounds, an error is raised.
index :: Bitmap -> Offset Bool -> Bool
index bits n
    | isOutOfBound n len = outOfBound OOB_Index n len
    | otherwise          = unsafeIndex bits n
  where len = length bits
{-# INLINE index #-}

-- | Return the element at a specific index from an array without bounds checking.
--
-- Reading from invalid memory can return unpredictable and invalid values.
-- use 'index' if unsure.
unsafeIndex :: Bitmap -> Offset Bool -> Bool
unsafeIndex (Bitmap _ ba) n =
    let (idx, bitIdx) = bitmapIndex n
     in testBit (A.unsafeIndex ba idx) bitIdx

{-# INLINE unsafeIndex #-}

-----------------------------------------------------------------------
-- higher level collection implementation
-----------------------------------------------------------------------
length :: Bitmap -> CountOf Bool
length (Bitmap sz _) = sz

mutableLength :: MutableBitmap st -> CountOf Bool
mutableLength (MutableBitmap sz _) = sz

empty :: Bitmap
empty = Bitmap 0 mempty

new :: PrimMonad prim => CountOf Bool -> prim (MutableBitmap (PrimState prim))
new sz@(CountOf len) =
    MutableBitmap sz <$> A.new nbElements
  where
    nbElements :: CountOf Word32
    nbElements = CountOf ((len `alignRoundUp` bitsPerTy) .>>. shiftPerTy)

-- | make an array from a list of elements.
vFromList :: [Bool] -> Bitmap
vFromList allBools = runST $ do
    mbitmap <- new len
    loop mbitmap 0 allBools
  where
    loop mb _ []     = unsafeFreeze mb
    loop mb i (x:xs) = unsafeWrite mb i x >> loop mb (i+1) xs

{-
    runST $ do
    mba <- A.new nbElements
    ba  <- loop mba (0 :: Int) allBools
    pure (Bitmap len ba)
  where
    loop mba _ [] = A.unsafeFreeze mba
    loop mba i l  = do
        let (l1, l2) = C.splitAt bitsPerTy l
            w = toPacked l1
        A.unsafeWrite mba i w
        loop mba (i+1) l2

    toPacked :: [Bool] -> Word32
    toPacked l =
        C.foldl' (.|.) 0 $ Prelude.zipWith (\b w -> if b then (1 `shiftL` w) else 0) l (C.reverse [0..31])
-}
    len        = C.length allBools

-- | transform an array to a list.
vToList :: Bitmap -> [Bool]
vToList a = loop 0
  where len = length a
        loop i | i .==# len  = []
               | otherwise = unsafeIndex a i : loop (i+1)

-- | Check if two vectors are identical
equal :: Bitmap -> Bitmap -> Bool
equal a b
    | la /= lb  = False
    | otherwise = loop 0
  where
    !la = length a
    !lb = length b
    loop n | n .==# la = True
           | otherwise = (unsafeIndex a n == unsafeIndex b n) && loop (n+1)

-- | Compare 2 vectors
vCompare :: Bitmap -> Bitmap -> Ordering
vCompare a b = loop 0
  where
    !la = length a
    !lb = length b
    loop n
        | n .==# la = if la == lb then EQ else LT
        | n .==# lb = GT
        | otherwise =
            case unsafeIndex a n `compare` unsafeIndex b n of
                EQ -> loop (n+1)
                r  -> r

-- | Append 2 arrays together by creating a new bigger array
--
-- TODO completely non optimized
append :: Bitmap -> Bitmap -> Bitmap
append a b = fromList $ toList a `mappend` toList b

-- TODO completely non optimized
concat :: [Bitmap] -> Bitmap
concat l = fromList $ mconcat $ fmap toList l

null :: Bitmap -> Bool
null (Bitmap nbBits _) = nbBits == 0

take :: CountOf Bool -> Bitmap -> Bitmap
take nbElems bits@(Bitmap nbBits ba)
    | nbElems <= 0      = empty
    | nbElems >= nbBits = bits
    | otherwise         = Bitmap nbElems ba -- TODO : although it work right now, take on the underlaying ba too

drop :: CountOf Bool -> Bitmap -> Bitmap
drop nbElems bits@(Bitmap nbBits _)
    | nbElems <= 0      = bits
    | nbElems >= nbBits = empty
    | otherwise         = unoptimised (C.drop nbElems) bits
        -- TODO: decide if we have drop easy by having a bit offset in the data structure
        -- or if we need to shift stuff around making all the indexing slighlty more complicated

splitAt :: CountOf Bool -> Bitmap -> (Bitmap, Bitmap)
splitAt n v = (take n v, drop n v)

-- unoptimised
splitOn :: (Bool -> Bool) -> Bitmap -> [Bitmap]
splitOn f bits = fmap fromList $ C.splitOn f $ toList bits

-- unoptimised
break :: (Bool -> Bool) -> Bitmap -> (Bitmap, Bitmap)
break predicate v = findBreak 0
  where
    len = length v
    findBreak i
        | i .==# len = (v, empty)
        | otherwise  =
            if predicate (unsafeIndex v i)
                then splitAt (offsetAsSize i) v
                else findBreak (i+1)

breakEnd :: (Bool -> Bool) -> Bitmap -> (Bitmap, Bitmap)
breakEnd predicate = bimap fromList fromList . C.breakEnd predicate . toList

span :: (Bool -> Bool) -> Bitmap -> (Bitmap, Bitmap)
span p = break (not . p)

map :: (Bool -> Bool) -> Bitmap -> Bitmap
map f bits = unoptimised (fmap f) bits

--mapIndex :: (Int -> Bool -> Bool) -> Bitmap -> Bitmap
--mapIndex f Bitmap =

cons :: Bool -> Bitmap -> Bitmap
cons v l = unoptimised (C.cons v) l

snoc :: Bitmap -> Bool -> Bitmap
snoc l v = unoptimised (flip C.snoc v) l

-- unoptimised
uncons :: Bitmap -> Maybe (Bool, Bitmap)
uncons b = fmap (second fromList) $ C.uncons $ toList b

-- unoptimised
unsnoc :: Bitmap -> Maybe (Bitmap, Bool)
unsnoc b = fmap (first fromList) $ C.unsnoc $ toList b

intersperse :: Bool -> Bitmap -> Bitmap
intersperse b = unoptimised (C.intersperse b)

find :: (Bool -> Bool) -> Bitmap -> Maybe Bool
find predicate vec = loop 0
  where
    !len = length vec
    loop i
        | i .==# len = Nothing
        | otherwise  =
            let e = unsafeIndex vec i
             in if predicate e then Just e else loop (i+1)

sortBy :: (Bool -> Bool -> Ordering) -> Bitmap -> Bitmap
sortBy by bits = unoptimised (C.sortBy by) bits

filter :: (Bool -> Bool) -> Bitmap -> Bitmap
filter predicate vec = unoptimised (Data.List.filter predicate) vec

reverse :: Bitmap -> Bitmap
reverse bits = unoptimised C.reverse bits

foldr :: (Bool -> a -> a) -> a -> Bitmap -> a
foldr f initialAcc vec = loop 0
  where
    len = length vec
    loop i
        | i .==# len = initialAcc
        | otherwise  = unsafeIndex vec i `f` loop (i+1)

foldr' :: (Bool -> a -> a) -> a -> Bitmap -> a
foldr' = foldr

foldl' :: (a -> Bool -> a) -> a -> Bitmap -> a
foldl' f initialAcc vec = loop 0 initialAcc
  where
    len = length vec
    loop i !acc
        | i .==# len = acc
        | otherwise  = loop (i+1) (f acc (unsafeIndex vec i))

all :: (Bool -> Bool) -> Bitmap -> Bool
all p bm = loop 0
  where
    len = length bm
    loop !i
      | i .==# len = True
      | not $ p (unsafeIndex bm i) = False
      | otherwise = loop (i + 1)

any :: (Bool -> Bool) -> Bitmap -> Bool
any p bm = loop 0
  where
    len = length bm
    loop !i
      | i .==# len = False
      | p (unsafeIndex bm i) = True
      | otherwise = loop (i + 1)

unoptimised :: ([Bool] -> [Bool]) -> Bitmap -> Bitmap
unoptimised f = vFromList . f . vToList