-- | functions in this module have no extra variables but exponential cost.

module Satchmo.Counting.Direct 

( atleast
, atmost
, exactly
, assert_implies_atmost
, assert_implies_exactly
)

where

import Satchmo.Boolean ( Boolean, MonadSAT )  
import qualified Satchmo.Boolean as B

import Control.Monad ( forM, forM_ )

select :: Int -> [a] -> [[a]]
select :: forall a. Int -> [a] -> [[a]]
select Int
0 [a]
xs = [[]]
select Int
k [] = []
select Int
k (a
x:[a]
xs) =
  forall a. Int -> [a] -> [[a]]
select Int
k [a]
xs forall a. [a] -> [a] -> [a]
++ (forall a b. (a -> b) -> [a] -> [b]
map (a
xforall a. a -> [a] -> [a]
:) forall a b. (a -> b) -> a -> b
$ forall a. Int -> [a] -> [[a]]
select (Int
kforall a. Num a => a -> a -> a
-Int
1) [a]
xs)

atleast :: MonadSAT m => Int -> [ Boolean ] -> m Boolean
atleast :: forall (m :: * -> *). MonadSAT m => Int -> [Boolean] -> m Boolean
atleast Int
k [Boolean]
xs = forall (m :: * -> *). MonadSAT m => [Boolean] -> m Boolean
B.or forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM (forall a. Int -> [a] -> [[a]]
select Int
k [Boolean]
xs) forall (m :: * -> *). MonadSAT m => [Boolean] -> m Boolean
B.and

atmost :: MonadSAT m => Int -> [ Boolean ] -> m Boolean
atmost :: forall (m :: * -> *). MonadSAT m => Int -> [Boolean] -> m Boolean
atmost Int
k [Boolean]
xs = forall (m :: * -> *). MonadSAT m => Int -> [Boolean] -> m Boolean
atleast (forall (t :: * -> *) a. Foldable t => t a -> Int
length [Boolean]
xs forall a. Num a => a -> a -> a
- Int
k) forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map Boolean -> Boolean
B.not [Boolean]
xs

exactly :: MonadSAT m => Int -> [ Boolean ] -> m Boolean
exactly :: forall (m :: * -> *). MonadSAT m => Int -> [Boolean] -> m Boolean
exactly Int
k [Boolean]
xs = do
  Boolean
this <- forall (m :: * -> *). MonadSAT m => Int -> [Boolean] -> m Boolean
atleast Int
k [Boolean]
xs
  Boolean
that <- forall (m :: * -> *). MonadSAT m => Int -> [Boolean] -> m Boolean
atmost Int
k [Boolean]
xs
  Boolean
this forall {m :: * -> *}. MonadSAT m => Boolean -> Boolean -> m Boolean
B.&& Boolean
that

-- | (and ys) implies (atmost k xs)
assert_implies_atmost :: [Boolean] -> Int -> [Boolean] -> m ()
assert_implies_atmost [Boolean]
ys Int
k [Boolean]
xs | Int
k forall a. Ord a => a -> a -> Bool
>= Int
0 = 
  forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ (forall a. Int -> [a] -> [[a]]
select (Int
kforall a. Num a => a -> a -> a
+Int
1) [Boolean]
xs) forall a b. (a -> b) -> a -> b
$ \ [Boolean]
sub -> do
    forall (m :: * -> *). MonadSAT m => [Boolean] -> m ()
B.assert forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map Boolean -> Boolean
B.not [Boolean]
ys forall a. [a] -> [a] -> [a]
++ forall a b. (a -> b) -> [a] -> [b]
map Boolean -> Boolean
B.not [Boolean]
sub
assert_implies_atmost [Boolean]
ys Int
k [Boolean]
_ =
  forall (m :: * -> *). MonadSAT m => [Boolean] -> m ()
B.assert forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map Boolean -> Boolean
B.not [Boolean]
ys

assert_implies_atleast :: [Boolean] -> Int -> [Boolean] -> m ()
assert_implies_atleast [Boolean]
ys Int
k [Boolean]
xs =
  forall {m :: * -> *}.
MonadSAT m =>
[Boolean] -> Int -> [Boolean] -> m ()
assert_implies_atmost [Boolean]
ys (forall (t :: * -> *) a. Foldable t => t a -> Int
length [Boolean]
xs forall a. Num a => a -> a -> a
- Int
k) (forall a b. (a -> b) -> [a] -> [b]
map Boolean -> Boolean
B.not [Boolean]
xs)

-- | asserting that  (and ys)  implies  (exactly k xs)
assert_implies_exactly :: [Boolean] -> Int -> [Boolean] -> m ()
assert_implies_exactly [Boolean]
ys Int
k [Boolean]
xs = do
  forall {m :: * -> *}.
MonadSAT m =>
[Boolean] -> Int -> [Boolean] -> m ()
assert_implies_atmost [Boolean]
ys Int
k [Boolean]
xs
  forall {m :: * -> *}.
MonadSAT m =>
[Boolean] -> Int -> [Boolean] -> m ()
assert_implies_atleast [Boolean]
ys Int
k [Boolean]
xs

-- | (atmost k xs) implies (or ys)
assert_atmost_implies :: [Boolean] -> Int -> [Boolean] -> m ()
assert_atmost_implies [Boolean]
xs Int
k [Boolean]
ys =
  forall {m :: * -> *}.
MonadSAT m =>
[Boolean] -> Int -> [Boolean] -> m ()
assert_implies_atleast (forall a b. (a -> b) -> [a] -> [b]
map Boolean -> Boolean
B.not [Boolean]
ys) (Int
kforall a. Num a => a -> a -> a
+Int
1) [Boolean]
xs

assert_atleast_implies :: [Boolean] -> Int -> [Boolean] -> m ()
assert_atleast_implies [Boolean]
xs Int
k [Boolean]
ys =
  forall {m :: * -> *}.
MonadSAT m =>
[Boolean] -> Int -> [Boolean] -> m ()
assert_implies_atmost (forall a b. (a -> b) -> [a] -> [b]
map Boolean -> Boolean
B.not [Boolean]
ys) (Int
kforall a. Num a => a -> a -> a
+Int
1) [Boolean]
xs