-- | A collection of a number of data types and type classes shared by all
-- bitset variants.

module Data.PrimitiveArray.Index.BitSetClasses where

import           Control.DeepSeq (NFData(..))
import           Data.Aeson (FromJSON,ToJSON,FromJSONKey,ToJSONKey)
import           Data.Binary (Binary)
import           Data.Hashable (Hashable)
import           Data.Serialize (Serialize)
import           Data.Vector.Unboxed.Deriving
import           GHC.Generics (Generic)
import qualified Data.Vector.Fusion.Stream.Monadic as SM
import qualified Data.Vector.Unboxed as VU

import           Data.Bits.Ordered
import           Data.PrimitiveArray.Index.Class
import           Data.PrimitiveArray.Index.IOC



-- * Boundaries, the interface(s) for bitsets.

-- | Certain sets have an interface, a particular element with special
-- meaning. In this module, certain ``meanings'' are already provided.
-- These include a @First@ element and a @Last@ element. We phantom-type
-- these to reduce programming overhead.

newtype Boundary boundaryType ioc = Boundary { getBoundary ∷ Int }
  deriving stock (Eq,Ord,Generic)
  deriving newtype (Num)

-- | Whenever we can not set the boundary we should have for a set, we use this
-- pattern. All legal boundaries are @>=0@. We also need to set the undefined
-- boundary to @0@, since the @linearIndex@ will use this value to add, which
-- for empty sets would reduce to @0 - UndefBoundary === 0@.

pattern UndefBoundary ∷ Boundary boundaryType ioc
pattern UndefBoundary = Boundary 0

instance Show (Boundary i t) where
  show (Boundary i) = "(I:" ++ show i ++ ")"

derivingUnbox "Boundary"
  [t| forall i t . Boundary i t → Int |]
  [| \(Boundary i) → i                |]
  [| Boundary                         |]

instance Binary    (Boundary i t)
instance Serialize (Boundary i t)
instance ToJSON    (Boundary i t)
instance FromJSON  (Boundary i t)
instance Hashable  (Boundary i t)

instance NFData (Boundary i t) where
  rnf (Boundary i) = rnf i
  {-# Inline rnf #-}

instance Index (Boundary i t) where
  newtype LimitType (Boundary i t) = LtBoundary Int
  linearIndex _ (Boundary z) = z
  {-# INLINE linearIndex #-}
  size (LtBoundary h) = h + 1
  {-# INLINE size #-}
  inBounds (LtBoundary h) z = 0 <= z && getBoundary z <= h
  {-# INLINE inBounds #-}
  zeroBound = Boundary 0
  {-# Inline zeroBound #-}
  zeroBound' = LtBoundary 0
  {-# Inline zeroBound' #-}
  totalSize (LtBoundary n) = [fromIntegral n]
  {-# Inline totalSize #-}
  fromLinearIndex _ = Boundary
  {-# Inline fromLinearIndex #-}
  showBound (LtBoundary b) = ["LtBoundary " ++ show b]
  showIndex (Boundary b) = ["Boundary " ++ show b]

instance IndexStream z ⇒ IndexStream (z:.Boundary k I) where
  streamUp   (ls:..LtBoundary l) (hs:..LtBoundary h) = SM.flatten (streamUpBndMk   l h) (streamUpBndStep   l h) $ streamUp   ls hs
  streamDown (ls:..LtBoundary l) (hs:..LtBoundary h) = SM.flatten (streamDownBndMk l h) (streamDownBndStep l h) $ streamDown ls hs
  {-# Inline streamUp   #-}
  {-# Inline streamDown #-}

instance IndexStream (Z:.Boundary k I) ⇒ IndexStream (Boundary k I) where
  streamUp l h = SM.map (\(Z:.i) -> i) $ streamUp (ZZ:..l) (ZZ:..h)
  {-# Inline streamUp #-}
  streamDown l h = SM.map (\(Z:.i) -> i) $ streamDown (ZZ:..l) (ZZ:..h)
  {-# Inline streamDown #-}

streamUpBndMk l h z = return (z, l)
{-# Inline [0] streamUpBndMk #-}

streamUpBndStep l h (z , k)
  | k > h     = return $ SM.Done
  | otherwise = return $ SM.Yield (z:.Boundary k) (z, k+1)
{-# Inline [0] streamUpBndStep #-}

streamDownBndMk l h z = return (z, h)
{-# Inline [0] streamDownBndMk #-}

streamDownBndStep l h (z , k)
  | k < l     = return $ SM.Done
  | otherwise = return $ SM.Yield (z:.Boundary k) (z,k-1)
{-# Inline [0] streamDownBndStep #-}

-- | Declare the interface to be the start of a path.

data First

-- | Declare the interface to be the end of a path.

data Last

-- | Declare the interface to match anything.
--
-- TODO needed? want to use later in ADPfusion

data Any



-- * Moving indices within sets.

-- | Successor and Predecessor for sets. Designed as a class to accomodate
-- sets with interfaces and without interfaces with one function.
--
-- The functions are not written recursively, as we currently only have
-- three cases, and we do not want to "reset" while generating successors
-- and predecessors.
--
-- Note that sets have a partial order. Within the group of element with
-- the same @popCount@, we use @popPermutation@ which has the same stepping
-- order for both, @setSucc@ and @setPred@.

class SetPredSucc s where
  -- | Set successor. The first argument is the lower set limit, the second
  -- the upper set limit, the third the current set.
  setSucc ∷ Int → Int → s → Maybe s
  -- | Set predecessor. The first argument is the lower set limit, the
  -- second the upper set limit, the third the current set.
  setPred ∷ Int → Int → s → Maybe s

-- | Masks are used quite often for different types of bitsets. We liberate
-- them as a type family.

type family Mask s ∷ *

-- | @Fixed@ allows us to fix some or all bits of a bitset, thereby
-- providing @succ/pred@ operations which are only partially free.
--
-- @f = getFixedMask .&. getFixed@ are the fixed bits.
-- @n = getFixed .&. complement getFixedMask@ are the free bits.
-- @to = complement getFixed@ is the to move mask
-- @n' = popShiftR to n@ yields the population after the move
-- @p = popPermutation undefined n'@ yields the new population permutation
-- @p' = popShiftL to p@ yields the population moved back
-- @final = p' .|. f@

data FixedMask t = FixedMask { getMask ∷ (Mask t) , getFixed ∷ !t }

-- | Assuming a bitset on bits @[0 .. highbit]@, we can apply a mask that
-- stretches out those bits over @[0 .. higherBit]@ with @highbit <=
-- higherBit@. Any active interfaces are correctly set as well.

class ApplyMask s where
  applyMask :: Mask s → s → s



-- | for 'Test.QuickCheck.Arbitrary'

arbitraryBitSetMax ∷ Int
arbitraryBitSetMax = 6