content: -----------------------------------------------------------------------------
  -- |
  -- Copyright   :  (C) 2013-15 Edward Kmett
  -- License     :  BSD-style (see the file LICENSE)
  -- Maintainer  :  Edward Kmett <ekmett@gmail.com>
  -- Stability   :  experimental
  -- Portability :  non-portable
  --
  -- Utility functions for the RRR class and offset representation
  -----------------------------------------------------------------------------
  module Succinct.Internal.Binomial
    ( binomial
    , logBinomial
    , bitmap
    , offset
    ) where
  
  import Control.Monad.ST
  import Data.Bits
  import Data.Foldable as F
  import Data.Function
  import Data.List as List (sortBy, scanl)
  import Data.Vector as V
  import Data.Vector.Primitive as P
  import Data.Vector.Primitive.Mutable as PM
  import Data.Word
  
  -- maximum binomial term
  _N :: Int
  _N = 15
  
  bits :: Word16 -> Word8
  bits = go 0 where
    go i 0 = i
    go i n = go (i + 1) (unsafeShiftR n 1)
  {-# INLINE bits #-}
  
  -- build a fixed sized table of binomial coeffients
  -- bins :: (V.Vector (P.Vector Word16), V.Vector (P.Vector Word8))
  
  bin :: V.Vector (P.Vector Word16)
  lbin :: V.Vector (P.Vector Word8)
  (bin, lbin) = runST $ do
    mbin  <- V.replicateM (_N+2) $ PM.replicate (_N+2) 0
    mlbin <- V.replicateM (_N+2) $ PM.replicate (_N+2) 0
    F.forM_ [0.._N] $ \i -> do
      mbini <- V.unsafeIndexM mbin i
      PM.unsafeWrite mbini 0 1
      PM.unsafeWrite mbini 1 1
      PM.unsafeWrite mbini i 1
      mlbini <- V.unsafeIndexM mlbin i
      PM.unsafeWrite mlbini 0 0
      PM.unsafeWrite mlbini 1 0
      PM.unsafeWrite mlbini i 0
    F.forM_ [1.._N+1] $ \j ->
      F.forM_ [j+1.._N+1] $ \i -> do
        mbinim1 <- V.unsafeIndexM mbin (i - 1)
        mbini   <- V.unsafeIndexM mbin i
        bx <- PM.unsafeRead mbinim1 (j-1)
        by <- PM.unsafeRead mbinim1 j
        let bz = bx + by
        PM.unsafeWrite mbini j bz
        mlbini <- V.unsafeIndexM mlbin i
        PM.unsafeWrite mlbini j (bits bz)
    xs <- V.forM mbin P.unsafeFreeze
    ys <- V.forM mlbin P.unsafeFreeze
    return (xs,ys)
  
  -- | The sortBy below could be replaced with judicious use of <http://alexbowe.com/popcount-permutations/>
  -- to reduce startup times.
  bitmaps :: P.Vector Word16
  bitmaps = P.fromListN 32768 $ sortBy (compare `on` popCount) [0..0x7fff]
  
  classOffsets :: P.Vector Word16
  classOffsets = P.fromListN (_N+1) $ List.scanl (\a k -> a + fromIntegral (binomial _N k)) 0 [0.._N]
  
  offsets :: P.Vector Word16
  offsets
    = P.fromListN (P.length bitmaps)
    $ fmap (\(w,o) -> o - (classOffsets P.! popCount w))
    $ sortBy (on compare fst)
    $ Prelude.zip (P.toList bitmaps) [0..]
  
  -- @'binomial' n k@ returns @n `choose` k@ for @n, k <= 15@
  binomial :: Int -> Int -> Int
  binomial n k
    | 0 <= n, n <= _N, 0 <= k, k <= _N = fromIntegral $ P.unsafeIndex (V.unsafeIndex bin n) k
    | otherwise = error "binomial: out of range"
  {-# INLINE binomial #-}
  
  logBinomial :: Int -> Int -> Int
  logBinomial n k
    | 0 <= n, n <= _N, 0 <= k, k <= _N = fromIntegral $ P.unsafeIndex (V.unsafeIndex lbin n) k
    | otherwise = error "logBinomial: out of range"
  {-# INLINE logBinomial #-}
  
  -- | bitmap by class and offset
  --
  -- There are 16 classes @k@ (based on popCount) each with @binomial 15 k@ possible offsets in a 'Word16' with the msb unset
  bitmap :: Int -> Word16 -> Word16
  bitmap k o = bitmaps P.! (fromIntegral (classOffsets P.! fromIntegral k) + fromIntegral o)
  {-# INLINE bitmap #-}
  
  -- | Calculate the offset of a Word16 into its class.
  --
  -- You can use @popCount@ to calculate the class, given @w <= 0x7fff@
  --
  -- @
  -- 'bitmap' ('popCount' w) ('offset' w) = w
  -- @
  offset :: Word16 -> Word16
  offset w = fromIntegral (offsets P.! fromIntegral w)
  {-# INLINE offset #-}
  