-- | Venn diagrams. See <https://en.wikipedia.org/wiki/Venn_diagram>
--
-- TODO: write a more efficient implementation (for example an array of size @2^n@)
--

{-# LANGUAGE BangPatterns #-}
module Math.Combinat.Sets.VennDiagrams where

--------------------------------------------------------------------------------

import Data.List

import GHC.TypeLits
import Data.Proxy

import qualified Data.Map as Map
import Data.Map (Map)

import Math.Combinat.Compositions
import Math.Combinat.ASCII

--------------------------------------------------------------------------------

-- | Venn diagrams of @n@ sets. Each possible zone is annotated with a value
-- of type @a@. A typical use case is to annotate with the cardinality of the
-- given zone.
--
-- Internally this is representated by a map from @[Bool]@, where @True@ means element 
-- of the set, @False@ means not.
--
-- TODO: write a more efficient implementation (for example an array of size 2^n)
newtype VennDiagram a = VennDiagram { VennDiagram a -> Map [Bool] a
vennTable :: Map [Bool] a } deriving (VennDiagram a -> VennDiagram a -> Bool
(VennDiagram a -> VennDiagram a -> Bool)
-> (VennDiagram a -> VennDiagram a -> Bool) -> Eq (VennDiagram a)
forall a. Eq a => VennDiagram a -> VennDiagram a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: VennDiagram a -> VennDiagram a -> Bool
$c/= :: forall a. Eq a => VennDiagram a -> VennDiagram a -> Bool
== :: VennDiagram a -> VennDiagram a -> Bool
$c== :: forall a. Eq a => VennDiagram a -> VennDiagram a -> Bool
Eq,Eq (VennDiagram a)
Eq (VennDiagram a)
-> (VennDiagram a -> VennDiagram a -> Ordering)
-> (VennDiagram a -> VennDiagram a -> Bool)
-> (VennDiagram a -> VennDiagram a -> Bool)
-> (VennDiagram a -> VennDiagram a -> Bool)
-> (VennDiagram a -> VennDiagram a -> Bool)
-> (VennDiagram a -> VennDiagram a -> VennDiagram a)
-> (VennDiagram a -> VennDiagram a -> VennDiagram a)
-> Ord (VennDiagram a)
VennDiagram a -> VennDiagram a -> Bool
VennDiagram a -> VennDiagram a -> Ordering
VennDiagram a -> VennDiagram a -> VennDiagram a
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
forall a. Ord a => Eq (VennDiagram a)
forall a. Ord a => VennDiagram a -> VennDiagram a -> Bool
forall a. Ord a => VennDiagram a -> VennDiagram a -> Ordering
forall a. Ord a => VennDiagram a -> VennDiagram a -> VennDiagram a
min :: VennDiagram a -> VennDiagram a -> VennDiagram a
$cmin :: forall a. Ord a => VennDiagram a -> VennDiagram a -> VennDiagram a
max :: VennDiagram a -> VennDiagram a -> VennDiagram a
$cmax :: forall a. Ord a => VennDiagram a -> VennDiagram a -> VennDiagram a
>= :: VennDiagram a -> VennDiagram a -> Bool
$c>= :: forall a. Ord a => VennDiagram a -> VennDiagram a -> Bool
> :: VennDiagram a -> VennDiagram a -> Bool
$c> :: forall a. Ord a => VennDiagram a -> VennDiagram a -> Bool
<= :: VennDiagram a -> VennDiagram a -> Bool
$c<= :: forall a. Ord a => VennDiagram a -> VennDiagram a -> Bool
< :: VennDiagram a -> VennDiagram a -> Bool
$c< :: forall a. Ord a => VennDiagram a -> VennDiagram a -> Bool
compare :: VennDiagram a -> VennDiagram a -> Ordering
$ccompare :: forall a. Ord a => VennDiagram a -> VennDiagram a -> Ordering
$cp1Ord :: forall a. Ord a => Eq (VennDiagram a)
Ord,Int -> VennDiagram a -> ShowS
[VennDiagram a] -> ShowS
VennDiagram a -> String
(Int -> VennDiagram a -> ShowS)
-> (VennDiagram a -> String)
-> ([VennDiagram a] -> ShowS)
-> Show (VennDiagram a)
forall a. Show a => Int -> VennDiagram a -> ShowS
forall a. Show a => [VennDiagram a] -> ShowS
forall a. Show a => VennDiagram a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [VennDiagram a] -> ShowS
$cshowList :: forall a. Show a => [VennDiagram a] -> ShowS
show :: VennDiagram a -> String
$cshow :: forall a. Show a => VennDiagram a -> String
showsPrec :: Int -> VennDiagram a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> VennDiagram a -> ShowS
Show)

-- | How many sets are in the Venn diagram
vennDiagramNumberOfSets :: VennDiagram a -> Int
vennDiagramNumberOfSets :: VennDiagram a -> Int
vennDiagramNumberOfSets (VennDiagram Map [Bool] a
table) = [Bool] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([Bool] -> Int) -> [Bool] -> Int
forall a b. (a -> b) -> a -> b
$ ([Bool], a) -> [Bool]
forall a b. (a, b) -> a
fst (([Bool], a) -> [Bool]) -> ([Bool], a) -> [Bool]
forall a b. (a -> b) -> a -> b
$ Map [Bool] a -> ([Bool], a)
forall k a. Map k a -> (k, a)
Map.findMin Map [Bool] a
table

-- | How many zones are in the Venn diagram
--
-- > vennDiagramNumberOfZones v == 2 ^ (vennDiagramNumberOfSets v)
--
vennDiagramNumberOfZones :: VennDiagram a -> Int
vennDiagramNumberOfZones :: VennDiagram a -> Int
vennDiagramNumberOfZones VennDiagram a
venn = Int
2 Int -> Int -> Int
forall a b. (Num a, Integral b) => a -> b -> a
^ (VennDiagram a -> Int
forall a. VennDiagram a -> Int
vennDiagramNumberOfSets VennDiagram a
venn)

-- | How many /nonempty/ zones are in the Venn diagram
vennDiagramNumberOfNonemptyZones :: VennDiagram Int -> Int
vennDiagramNumberOfNonemptyZones :: VennDiagram Int -> Int
vennDiagramNumberOfNonemptyZones (VennDiagram Map [Bool] Int
table) = [Int] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([Int] -> Int) -> [Int] -> Int
forall a b. (a -> b) -> a -> b
$ (Int -> Bool) -> [Int] -> [Int]
forall a. (a -> Bool) -> [a] -> [a]
filter (Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/=Int
0) ([Int] -> [Int]) -> [Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ Map [Bool] Int -> [Int]
forall k a. Map k a -> [a]
Map.elems Map [Bool] Int
table

unsafeMakeVennDiagram :: [([Bool],a)] -> VennDiagram a
unsafeMakeVennDiagram :: [([Bool], a)] -> VennDiagram a
unsafeMakeVennDiagram = Map [Bool] a -> VennDiagram a
forall a. Map [Bool] a -> VennDiagram a
VennDiagram (Map [Bool] a -> VennDiagram a)
-> ([([Bool], a)] -> Map [Bool] a)
-> [([Bool], a)]
-> VennDiagram a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [([Bool], a)] -> Map [Bool] a
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList

-- | We call venn diagram trivial if all the intersection zones has zero cardinality
-- (that is, the original sets are all disjoint)
isTrivialVennDiagram :: VennDiagram Int -> Bool
isTrivialVennDiagram :: VennDiagram Int -> Bool
isTrivialVennDiagram (VennDiagram Map [Bool] Int
table) = [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and [ Int
c Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 | ([Bool]
bs,Int
c) <- Map [Bool] Int -> [([Bool], Int)]
forall k a. Map k a -> [(k, a)]
Map.toList Map [Bool] Int
table , [Bool] -> Bool
isIntersection [Bool]
bs ] where
  isIntersection :: [Bool] -> Bool
isIntersection [Bool]
bs = case (Bool -> Bool) -> [Bool] -> [Bool]
forall a. (a -> Bool) -> [a] -> [a]
filter Bool -> Bool
forall a. a -> a
id [Bool]
bs of
    []  -> Bool
False
    [Bool
_] -> Bool
False
    [Bool]
_   -> Bool
True

printVennDiagram :: Show a => VennDiagram a -> IO ()
printVennDiagram :: VennDiagram a -> IO ()
printVennDiagram = String -> IO ()
putStrLn (String -> IO ())
-> (VennDiagram a -> String) -> VennDiagram a -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VennDiagram a -> String
forall a. Show a => VennDiagram a -> String
prettyVennDiagram

prettyVennDiagram :: Show a => VennDiagram a -> String
prettyVennDiagram :: VennDiagram a -> String
prettyVennDiagram = [String] -> String
unlines ([String] -> String)
-> (VennDiagram a -> [String]) -> VennDiagram a -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ASCII -> [String]
asciiLines (ASCII -> [String])
-> (VennDiagram a -> ASCII) -> VennDiagram a -> [String]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VennDiagram a -> ASCII
forall a. Show a => VennDiagram a -> ASCII
asciiVennDiagram

asciiVennDiagram :: Show a => VennDiagram a -> ASCII
asciiVennDiagram :: VennDiagram a -> ASCII
asciiVennDiagram (VennDiagram Map [Bool] a
table) = [String] -> ASCII
asciiFromLines ([String] -> ASCII) -> [String] -> ASCII
forall a b. (a -> b) -> a -> b
$ (([Bool], a) -> String) -> [([Bool], a)] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map ([Bool], a) -> String
forall a. Show a => ([Bool], a) -> String
f (Map [Bool] a -> [([Bool], a)]
forall k a. Map k a -> [(k, a)]
Map.toList Map [Bool] a
table) where
  f :: ([Bool], a) -> String
f ([Bool]
bs,a
a) = String
"{" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> ShowS
extendTo ([Bool] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Bool]
bs) [ if Bool
b then Char
z else Char
' ' | (Bool
b,Char
z) <- [Bool] -> String -> [(Bool, Char)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Bool]
bs String
abc ] String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"} -> " String -> ShowS
forall a. [a] -> [a] -> [a]
++ a -> String
forall a. Show a => a -> String
show a
a
  extendTo :: Int -> ShowS
extendTo Int
k String
str = String
str String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> Char -> String
forall a. Int -> a -> [a]
replicate (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- String -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length String
str) Char
' '
  abc :: String
abc = [Char
'A'..Char
'Z']

instance Show a => DrawASCII (VennDiagram a) where
  ascii :: VennDiagram a -> ASCII
ascii = VennDiagram a -> ASCII
forall a. Show a => VennDiagram a -> ASCII
asciiVennDiagram

-- | Given a Venn diagram of cardinalities, we compute the cardinalities of the
-- original sets (note: this is slow!)
vennDiagramSetCardinalities :: VennDiagram Int -> [Int]
vennDiagramSetCardinalities :: VennDiagram Int -> [Int]
vennDiagramSetCardinalities (VennDiagram Map [Bool] Int
table) = Int -> [([Bool], Int)] -> [Int]
go Int
n [([Bool], Int)]
list where
  list :: [([Bool], Int)]
list = Map [Bool] Int -> [([Bool], Int)]
forall k a. Map k a -> [(k, a)]
Map.toList Map [Bool] Int
table
  n :: Int
n = [Bool] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([Bool] -> Int) -> [Bool] -> Int
forall a b. (a -> b) -> a -> b
$ ([Bool], Int) -> [Bool]
forall a b. (a, b) -> a
fst (([Bool], Int) -> [Bool]) -> ([Bool], Int) -> [Bool]
forall a b. (a -> b) -> a -> b
$ [([Bool], Int)] -> ([Bool], Int)
forall a. [a] -> a
head [([Bool], Int)]
list
  go :: Int -> [([Bool],Int)] -> [Int]
  go :: Int -> [([Bool], Int)] -> [Int]
go !Int
0 [([Bool], Int)]
_  = []
  go !Int
k [([Bool], Int)]
xs = Int
this Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: Int -> [([Bool], Int)] -> [Int]
go (Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) ((([Bool], Int) -> ([Bool], Int))
-> [([Bool], Int)] -> [([Bool], Int)]
forall a b. (a -> b) -> [a] -> [b]
map ([Bool], Int) -> ([Bool], Int)
forall a b. ([a], b) -> ([a], b)
xtail [([Bool], Int)]
xs) where
    this :: Int
this = (Int -> Int -> Int) -> Int -> [Int] -> Int
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) Int
0 [ Int
c | ((Bool
True:[Bool]
_) , Int
c) <- [([Bool], Int)]
xs ]
  xtail :: ([a], b) -> ([a], b)
xtail ([a]
bs,b
c) = ([a] -> [a]
forall a. [a] -> [a]
tail [a]
bs,b
c)

--------------------------------------------------------------------------------

-- | Given the cardinalities of some finite sets, we list all possible
-- Venn diagrams.
--
-- Note: we don't include the empty zone in the tables, because it's always empty.
--
-- Remark: if each sets is a singleton set, we get back set partitions:
--
-- > > [ length $ enumerateVennDiagrams $ replicate k 1 | k<-[1..8] ]
-- > [1,2,5,15,52,203,877,4140]
-- >
-- > > [ countSetPartitions k | k<-[1..8] ]
-- > [1,2,5,15,52,203,877,4140]
--
-- Maybe this could be called multiset-partitions?
--
-- Example:
--
-- > autoTabulate RowMajor (Right 6) $ map ascii $ enumerateVennDiagrams [2,3,3]
--
enumerateVennDiagrams :: [Int] -> [VennDiagram Int]
enumerateVennDiagrams :: [Int] -> [VennDiagram Int]
enumerateVennDiagrams [Int]
dims = 
  case [Int]
dims of
    []     -> []
    [Int
d]    -> Int -> [VennDiagram Int]
venns1 Int
d
    (Int
d:[Int]
ds) -> (VennDiagram Int -> [VennDiagram Int])
-> [VennDiagram Int] -> [VennDiagram Int]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (Int -> Int -> VennDiagram Int -> [VennDiagram Int]
worker ([Int] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
ds) Int
d) ([VennDiagram Int] -> [VennDiagram Int])
-> [VennDiagram Int] -> [VennDiagram Int]
forall a b. (a -> b) -> a -> b
$ [Int] -> [VennDiagram Int]
enumerateVennDiagrams [Int]
ds
  where

    worker :: Int -> Int -> VennDiagram Int -> [VennDiagram Int]
worker !Int
n !Int
d (VennDiagram Map [Bool] Int
table) = [VennDiagram Int]
result where

      list :: [([Bool], Int)]
list   = Map [Bool] Int -> [([Bool], Int)]
forall k a. Map k a -> [(k, a)]
Map.toList Map [Bool] Int
table
      falses :: [Bool]
falses = Int -> Bool -> [Bool]
forall a. Int -> a -> [a]
replicate Int
n Bool
False

      comps :: Int -> [[Int]]
comps Int
k = [Int] -> Int -> [[Int]]
compositions' ((([Bool], Int) -> Int) -> [([Bool], Int)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map ([Bool], Int) -> Int
forall a b. (a, b) -> b
snd [([Bool], Int)]
list) Int
k
      result :: [VennDiagram Int]
result = 
        [ [([Bool], Int)] -> VennDiagram Int
forall a. [([Bool], a)] -> VennDiagram a
unsafeMakeVennDiagram ([([Bool], Int)] -> VennDiagram Int)
-> [([Bool], Int)] -> VennDiagram Int
forall a b. (a -> b) -> a -> b
$ 
            [ (Bool
FalseBool -> [Bool] -> [Bool]
forall a. a -> [a] -> [a]
:[Bool]
tfs    , Int
mInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
c) | (([Bool]
tfs,Int
m),Int
c) <- [([Bool], Int)] -> [Int] -> [(([Bool], Int), Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [([Bool], Int)]
list [Int]
comp ] [([Bool], Int)] -> [([Bool], Int)] -> [([Bool], Int)]
forall a. [a] -> [a] -> [a]
++
            [ (Bool
True Bool -> [Bool] -> [Bool]
forall a. a -> [a] -> [a]
:[Bool]
tfs    ,   Int
c) | (([Bool]
tfs,Int
m),Int
c) <- [([Bool], Int)] -> [Int] -> [(([Bool], Int), Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [([Bool], Int)]
list [Int]
comp ] [([Bool], Int)] -> [([Bool], Int)] -> [([Bool], Int)]
forall a. [a] -> [a] -> [a]
++
            [ (Bool
True Bool -> [Bool] -> [Bool]
forall a. a -> [a] -> [a]
:[Bool]
falses , Int
dInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
k) ]
        | Int
k <- [Int
0..Int
d]
        , [Int]
comp <- Int -> [[Int]]
comps Int
k
        ]

    venns1 :: Int -> [VennDiagram Int]
    venns1 :: Int -> [VennDiagram Int]
venns1 Int
p = [ VennDiagram Int
theVenn ] where 
      theVenn :: VennDiagram Int
theVenn = [([Bool], Int)] -> VennDiagram Int
forall a. [([Bool], a)] -> VennDiagram a
unsafeMakeVennDiagram [ ([Bool
True],Int
p) ] 

--------------------------------------------------------------------------------

{-

-- | for testing only
venns2 :: Int -> Int -> [Venn Int]
venns2 p q = 
  [ mkVenn [ ([t,f],p-k) , ([f,t],q-k) , ([t,t],k) ]
  | k <- [0..min p q] 
  ]
  where
    t = True
    f = False
-}