{-# LANGUAGE ForeignFunctionInterface #-}

-- |
-- Copyright   : Anders Claesson 2013
-- Maintainer  : Anders Claesson <anders.claesson@gmail.com>
--

module Sym.Internal.SubSeq
    (
      module Sym.Internal.CLongArray
    , SubSeq
    , choose
    ) where

import Sym.Internal.CLongArray
import Foreign
import Foreign.C.Types
import System.IO.Unsafe

-- | A SubSeq is represented by an increasing array of non-negative
-- integers.
type SubSeq = CLongArray

-- Bitmasks
-- --------

-- A sub-class of 'Bits' used internally. Minimal complete definiton: 'next'.
class (Bits a, Integral a) => Bitmask a where
    -- | Lexicographically, the next bitmask with the same Hamming weight.
    next :: a -> a

    -- | @ones k m@ is the set / subsequence of indices whose bits are
    -- set in @m@. Default implementation:
    -- 
    -- > ones m = fromListN (popCount m) $ filter (testBit m) [0..]
    -- 
    ones :: a -> SubSeq
    ones a
m = [Int] -> SubSeq
fromList ([Int] -> SubSeq) -> ([Int] -> [Int]) -> [Int] -> SubSeq
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
take (a -> Int
forall a. Bits a => a -> Int
popCount a
m) ([Int] -> SubSeq) -> [Int] -> SubSeq
forall a b. (a -> b) -> a -> b
$ (Int -> Bool) -> [Int] -> [Int]
forall a. (a -> Bool) -> [a] -> [a]
filter (a -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit a
m) [Int
0..]

instance Bitmask CLong where
    next :: CLong -> CLong
next = CLong -> CLong
nextCLong
    ones :: CLong -> SubSeq
ones = CLong -> SubSeq
onesCLong

instance Bitmask Integer where
    next :: Integer -> Integer
next = Integer -> Integer
forall a. (Integral a, Bits a) => a -> a
nextIntegral

-- @bitmasks n k@ is the list of bitmasks with Hamming weight @k@ and
-- size less than @2^n@.
bitmasks :: Bitmask a => Int -> Int -> [a]
bitmasks :: forall a. Bitmask a => Int -> Int -> [a]
bitmasks Int
n Int
k = Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
take Int
binomial ((a -> a) -> a -> [a]
forall a. (a -> a) -> a -> [a]
iterate a -> a
forall a. Bitmask a => a -> a
next ((a
1 a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftL` Int
k) a -> a -> a
forall a. Num a => a -> a -> a
- a
1))
    where
      n' :: Integer
n' = Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
n
      k' :: Integer
k' = Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
k
      binomial :: Int
binomial = Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer -> Int) -> Integer -> Int
forall a b. (a -> b) -> a -> b
$ [Integer] -> Integer
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [Integer
n', Integer
n'Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
-Integer
1 .. Integer
n'Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
-Integer
k'Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+Integer
1] Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` [Integer] -> Integer
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [Integer
1..Integer
k']

-- | @n \`choose\` k@ is the list of subsequences of @[0..n-1]@ with @k@
-- elements.
choose :: Int -> Int -> [SubSeq]
choose :: Int -> Int -> [SubSeq]
choose Int
n Int
k
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
32   = (CLong -> SubSeq) -> [CLong] -> [SubSeq]
forall a b. (a -> b) -> [a] -> [b]
map CLong -> SubSeq
forall a. Bitmask a => a -> SubSeq
ones (Int -> Int -> [CLong]
forall a. Bitmask a => Int -> Int -> [a]
bitmasks Int
n Int
k :: [CLong])
    | Bool
otherwise = (Integer -> SubSeq) -> [Integer] -> [SubSeq]
forall a b. (a -> b) -> [a] -> [b]
map Integer -> SubSeq
forall a. Bitmask a => a -> SubSeq
ones (Int -> Int -> [Integer]
forall a. Bitmask a => Int -> Int -> [a]
bitmasks Int
n Int
k :: [Integer])

foreign import ccall unsafe "bit.h next" c_next :: CLong -> CLong

-- | Lexicographically, the next 'CLong' with the same Hamming weight.
nextCLong :: CLong -> CLong
nextCLong :: CLong -> CLong
nextCLong = CLong -> CLong
c_next

foreign import ccall unsafe "bit.h ones" c_ones :: Ptr CLong -> CLong -> IO ()

-- | @onesCLong m@ gives the indices whose bits are set in @m@.
onesCLong :: CLong -> CLongArray
onesCLong :: CLong -> SubSeq
onesCLong CLong
m = IO SubSeq -> SubSeq
forall a. IO a -> a
unsafePerformIO (IO SubSeq -> SubSeq)
-> ((Ptr CLong -> IO ()) -> IO SubSeq)
-> (Ptr CLong -> IO ())
-> SubSeq
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> (Ptr CLong -> IO ()) -> IO SubSeq
unsafeNew (CLong -> Int
forall a. Bits a => a -> Int
popCount CLong
m) ((Ptr CLong -> IO ()) -> SubSeq) -> (Ptr CLong -> IO ()) -> SubSeq
forall a b. (a -> b) -> a -> b
$ (Ptr CLong -> CLong -> IO ()) -> CLong -> Ptr CLong -> IO ()
forall a b c. (a -> b -> c) -> b -> a -> c
flip Ptr CLong -> CLong -> IO ()
c_ones CLong
m

-- | Lexicographically, the next integral number with the same Hamming weight.
nextIntegral :: (Integral a, Bits a) => a -> a
nextIntegral :: forall a. (Integral a, Bits a) => a -> a
nextIntegral a
a =
    let b :: a
b = (a
a a -> a -> a
forall a. Bits a => a -> a -> a
.|. (a
a a -> a -> a
forall a. Num a => a -> a -> a
- a
1)) a -> a -> a
forall a. Num a => a -> a -> a
+ a
1
    in  a
b a -> a -> a
forall a. Bits a => a -> a -> a
.|. ((((a
b a -> a -> a
forall a. Bits a => a -> a -> a
.&. (-a
b)) a -> a -> a
forall a. Integral a => a -> a -> a
`div` (a
a a -> a -> a
forall a. Bits a => a -> a -> a
.&. (-a
a))) a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftR` Int
1) a -> a -> a
forall a. Num a => a -> a -> a
- a
1)