module Satchmo.Counting.Unary

( atleast
, atmost
, exactly
)

where

import Prelude hiding ( and, or, not )

import Satchmo.Boolean

import Satchmo.SAT ( SAT) -- for specializations

{-# specialize inline atleast :: Int -> [ Boolean] -> SAT Boolean #-}
{-# specialize inline atmost  :: Int -> [ Boolean] -> SAT Boolean #-}
{-# specialize inline exactly :: Int -> [ Boolean] -> SAT Boolean #-}

atleast :: MonadSAT m => Int -> [ Boolean ] -> m Boolean
atleast :: forall (m :: * -> *). MonadSAT m => Int -> [Boolean] -> m Boolean
atleast Int
k [Boolean]
xs = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Boolean -> Boolean
not forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). MonadSAT m => Int -> [Boolean] -> m Boolean
atmost (Int
kforall a. Num a => a -> a -> a
-Int
1) [Boolean]
xs
        

atmost_block :: MonadSAT m => Int -> [ Boolean ] -> m [ Boolean ]
atmost_block :: forall (m :: * -> *). MonadSAT m => Int -> [Boolean] -> m [Boolean]
atmost_block Int
k [] = do
    Boolean
t <- forall (m :: * -> *). MonadSAT m => Bool -> m Boolean
constant forall a b. (a -> b) -> a -> b
$ Bool
True
    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. Int -> a -> [a]
replicate (Int
kforall a. Num a => a -> a -> a
+Int
1) Boolean
t
atmost_block Int
k (Boolean
x:[Boolean]
xs) = do
    [Boolean]
cs <- forall (m :: * -> *). MonadSAT m => Int -> [Boolean] -> m [Boolean]
atmost_block Int
k [Boolean]
xs
    Boolean
f <- forall (m :: * -> *). MonadSAT m => Bool -> m Boolean
constant Bool
False
    forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence forall a b. (a -> b) -> a -> b
$ do
        (Boolean
p,Boolean
q) <- forall a b. [a] -> [b] -> [(a, b)]
zip [Boolean]
cs ( Boolean
f forall a. a -> [a] -> [a]
: [Boolean]
cs )
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ do
            forall (m :: * -> *).
MonadSAT m =>
(Bool -> Bool -> Bool -> Bool)
-> Boolean -> Boolean -> Boolean -> m Boolean
fun3  ( \ Bool
x Bool
p Bool
q -> if Bool
x then Bool
q else Bool
p ) Boolean
x Boolean
p Boolean
q

atmost :: MonadSAT m => Int -> [ Boolean ] -> m Boolean
atmost :: forall (m :: * -> *). MonadSAT m => Int -> [Boolean] -> m Boolean
atmost Int
k [Boolean]
xs = do
    [Boolean]
cs <- forall (m :: * -> *). MonadSAT m => Int -> [Boolean] -> m [Boolean]
atmost_block Int
k [Boolean]
xs
    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ [Boolean]
cs forall a. [a] -> Int -> a
!! Int
k
        

exactly_block :: MonadSAT m => Int -> [ Boolean ] -> m [ Boolean ]
exactly_block :: forall (m :: * -> *). MonadSAT m => Int -> [Boolean] -> m [Boolean]
exactly_block Int
k [] = do
    Boolean
t <- forall (m :: * -> *). MonadSAT m => Bool -> m Boolean
constant Bool
True
    Boolean
f <- forall (m :: * -> *). MonadSAT m => Bool -> m Boolean
constant Bool
False
    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ Boolean
t forall a. a -> [a] -> [a]
: forall a. Int -> a -> [a]
replicate Int
k Boolean
f
exactly_block Int
k (Boolean
x:[Boolean]
xs) = do
    Boolean
f <- forall (m :: * -> *). MonadSAT m => Bool -> m Boolean
constant Bool
False
    [Boolean]
cs <- forall (m :: * -> *). MonadSAT m => Int -> [Boolean] -> m [Boolean]
exactly_block Int
k [Boolean]
xs
    forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence forall a b. (a -> b) -> a -> b
$ do
        (Boolean
p,Boolean
q) <- forall a b. [a] -> [b] -> [(a, b)]
zip [Boolean]
cs ( Boolean
f forall a. a -> [a] -> [a]
: [Boolean]
cs )
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ do
            forall (m :: * -> *).
MonadSAT m =>
(Bool -> Bool -> Bool -> Bool)
-> Boolean -> Boolean -> Boolean -> m Boolean
fun3 ( \ Bool
x Bool
p Bool
q -> if Bool
x then Bool
q else Bool
p ) Boolean
x Boolean
p Boolean
q

exactly :: MonadSAT m => Int -> [ Boolean ] -> m Boolean
exactly :: forall (m :: * -> *). MonadSAT m => Int -> [Boolean] -> m Boolean
exactly Int
k [Boolean]
xs = do
    [Boolean]
cs <- forall (m :: * -> *). MonadSAT m => Int -> [Boolean] -> m [Boolean]
exactly_block Int
k [Boolean]
xs
    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ [Boolean]
cs forall a. [a] -> Int -> a
!! Int
k