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 = (Boolean -> Boolean) -> m Boolean -> m Boolean
forall a b. (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Boolean -> Boolean
not (m Boolean -> m Boolean) -> m Boolean -> m Boolean
forall a b. (a -> b) -> a -> b
$ Int -> [Boolean] -> m Boolean
forall (m :: * -> *). MonadSAT m => Int -> [Boolean] -> m Boolean
atmost (Int
kInt -> Int -> Int
forall 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 <- Bool -> m Boolean
forall (m :: * -> *). MonadSAT m => Bool -> m Boolean
constant (Bool -> m Boolean) -> Bool -> m Boolean
forall a b. (a -> b) -> a -> b
$ Bool
True
    [Boolean] -> m [Boolean]
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ([Boolean] -> m [Boolean]) -> [Boolean] -> m [Boolean]
forall a b. (a -> b) -> a -> b
$ Int -> Boolean -> [Boolean]
forall a. Int -> a -> [a]
replicate (Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Boolean
t
atmost_block Int
k (Boolean
x:[Boolean]
xs) = do
    [Boolean]
cs <- Int -> [Boolean] -> m [Boolean]
forall (m :: * -> *). MonadSAT m => Int -> [Boolean] -> m [Boolean]
atmost_block Int
k [Boolean]
xs
    Boolean
f <- Bool -> m Boolean
forall (m :: * -> *). MonadSAT m => Bool -> m Boolean
constant Bool
False
    [m Boolean] -> m [Boolean]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
forall (m :: * -> *) a. Monad m => [m a] -> m [a]
sequence ([m Boolean] -> m [Boolean]) -> [m Boolean] -> m [Boolean]
forall a b. (a -> b) -> a -> b
$ do
        (Boolean
p,Boolean
q) <- [Boolean] -> [Boolean] -> [(Boolean, Boolean)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Boolean]
cs ( Boolean
f Boolean -> [Boolean] -> [Boolean]
forall a. a -> [a] -> [a]
: [Boolean]
cs )
        m Boolean -> [m Boolean]
forall a. a -> [a]
forall (m :: * -> *) a. Monad m => a -> m a
return (m Boolean -> [m Boolean]) -> m Boolean -> [m Boolean]
forall a b. (a -> b) -> a -> b
$ do
            (Bool -> Bool -> Bool -> Bool)
-> Boolean -> Boolean -> Boolean -> m Boolean
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 <- Int -> [Boolean] -> m [Boolean]
forall (m :: * -> *). MonadSAT m => Int -> [Boolean] -> m [Boolean]
atmost_block Int
k [Boolean]
xs
    Boolean -> m Boolean
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Boolean -> m Boolean) -> Boolean -> m Boolean
forall a b. (a -> b) -> a -> b
$ [Boolean]
cs [Boolean] -> Int -> Boolean
forall a. HasCallStack => [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 <- Bool -> m Boolean
forall (m :: * -> *). MonadSAT m => Bool -> m Boolean
constant Bool
True
    Boolean
f <- Bool -> m Boolean
forall (m :: * -> *). MonadSAT m => Bool -> m Boolean
constant Bool
False
    [Boolean] -> m [Boolean]
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ([Boolean] -> m [Boolean]) -> [Boolean] -> m [Boolean]
forall a b. (a -> b) -> a -> b
$ Boolean
t Boolean -> [Boolean] -> [Boolean]
forall a. a -> [a] -> [a]
: Int -> Boolean -> [Boolean]
forall a. Int -> a -> [a]
replicate Int
k Boolean
f
exactly_block Int
k (Boolean
x:[Boolean]
xs) = do
    Boolean
f <- Bool -> m Boolean
forall (m :: * -> *). MonadSAT m => Bool -> m Boolean
constant Bool
False
    [Boolean]
cs <- Int -> [Boolean] -> m [Boolean]
forall (m :: * -> *). MonadSAT m => Int -> [Boolean] -> m [Boolean]
exactly_block Int
k [Boolean]
xs
    [m Boolean] -> m [Boolean]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
forall (m :: * -> *) a. Monad m => [m a] -> m [a]
sequence ([m Boolean] -> m [Boolean]) -> [m Boolean] -> m [Boolean]
forall a b. (a -> b) -> a -> b
$ do
        (Boolean
p,Boolean
q) <- [Boolean] -> [Boolean] -> [(Boolean, Boolean)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Boolean]
cs ( Boolean
f Boolean -> [Boolean] -> [Boolean]
forall a. a -> [a] -> [a]
: [Boolean]
cs )
        m Boolean -> [m Boolean]
forall a. a -> [a]
forall (m :: * -> *) a. Monad m => a -> m a
return (m Boolean -> [m Boolean]) -> m Boolean -> [m Boolean]
forall a b. (a -> b) -> a -> b
$ do
            (Bool -> Bool -> Bool -> Bool)
-> Boolean -> Boolean -> Boolean -> m Boolean
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 <- Int -> [Boolean] -> m [Boolean]
forall (m :: * -> *). MonadSAT m => Int -> [Boolean] -> m [Boolean]
exactly_block Int
k [Boolean]
xs
    Boolean -> m Boolean
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Boolean -> m Boolean) -> Boolean -> m Boolean
forall a b. (a -> b) -> a -> b
$ [Boolean]
cs [Boolean] -> Int -> Boolean
forall a. HasCallStack => [a] -> Int -> a
!! Int
k