-- | 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 { vennTable :: Map [Bool] a } deriving (Eq,Ord,Show)

-- | How many sets are in the Venn diagram
vennDiagramNumberOfSets :: VennDiagram a -> Int
vennDiagramNumberOfSets (VennDiagram table) = length $ fst $ Map.findMin table

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

-- | How many /nonempty/ zones are in the Venn diagram
vennDiagramNumberOfNonemptyZones :: VennDiagram Int -> Int
vennDiagramNumberOfNonemptyZones (VennDiagram table) = length $ filter (/=0) $ Map.elems table

unsafeMakeVennDiagram :: [([Bool],a)] -> VennDiagram a
unsafeMakeVennDiagram = VennDiagram . 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 table) = and [ c == 0 | (bs,c) <- Map.toList table , isIntersection bs ] where
  isIntersection bs = case filter id bs of
    []  -> False
    [_] -> False
    _   -> True

printVennDiagram :: Show a => VennDiagram a -> IO ()
printVennDiagram = putStrLn . prettyVennDiagram

prettyVennDiagram :: Show a => VennDiagram a -> String
prettyVennDiagram = unlines . asciiLines . asciiVennDiagram

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

instance Show a => DrawASCII (VennDiagram a) where
  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 table) = go n list where
  list = Map.toList table
  n = length $ fst $ head list
  go :: Int -> [([Bool],Int)] -> [Int]
  go !0 _  = []
  go !k xs = this : go (k-1) (map xtail xs) where
    this = foldl' (+) 0 [ c | ((True:_) , c) <- xs ]
  xtail (bs,c) = (tail bs,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 dims = 
  case dims of
    []     -> []
    [d]    -> venns1 d
    (d:ds) -> concatMap (worker (length ds) d) $ enumerateVennDiagrams ds
  where

    worker !n !d (VennDiagram table) = result where

      list   = Map.toList table
      falses = replicate n False

      comps k = compositions' (map snd list) k
      result = 
        [ unsafeMakeVennDiagram $ 
            [ (False:tfs    , m-c) | ((tfs,m),c) <- zip list comp ] ++
            [ (True :tfs    ,   c) | ((tfs,m),c) <- zip list comp ] ++
            [ (True :falses , d-k) ]
        | k <- [0..d]
        , comp <- comps k
        ]

    venns1 :: Int -> [VennDiagram Int]
    venns1 p = [ theVenn ] where 
      theVenn = unsafeMakeVennDiagram [ ([True],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
-}