{- |
Module      :  $Header$
Description :  Maps of sets
Copyright   :  (c) DFKI GmbH 2011
License     :  GPLv2 or higher, see LICENSE.txt

Maintainer  :  Christian.Maeder@dfki.de
Stability   :  provisional
Portability :  portable

supply total mappings from keys to sets of values based on Data.Map.
Undefined keys are mapped to the empty set. An abstract data type is needed
to avoid differences due to empty set values versus undefined map keys.

-}

module Common.Lib.MapSet
  ( rmNullSets
  , setLookup
  , setElems
  , setMember
  , setInsert
  , setAll
  , setDifference
  , setToMap
  , restrict
  , imageList
  , imageSet
  , MapSet
  , toMap
  , fromDistinctMap
  , fromMap
  , empty
  , null
  , fromList
  , toList
  , toPairList
  , keysSet
  , elems
  , insert
  , update
  , lookup
  , member
  , delete
  , union
  , difference
  , intersection
  , map
  , mapMonotonic
  , mapSet
  , foldWithKey
  , filter
  , filterWithKey
  , all
  , isSubmapOf
  , preImage
  , transpose
  , keysSet
  ) where

import Prelude hiding (all, filter, map, null, lookup)
import qualified Data.Map as Map
import qualified Data.Set as Set
import qualified Data.List as List

-- * functions directly working over unwrapped maps of sets

-- | remove empty elements from the map
rmNullSets :: Ord a => Map.Map a (Set.Set b) -> Map.Map a (Set.Set b)
rmNullSets = Map.filter (not . Set.null)

-- | get elements for a key
setLookup :: Ord a => a -> Map.Map a (Set.Set b) -> Set.Set b
setLookup = Map.findWithDefault Set.empty

-- | all elementes united
setElems :: Ord b => Map.Map a (Set.Set b) -> Set.Set b
setElems = Set.unions . Map.elems

-- | test for an element under a key
setMember :: (Ord a, Ord b) => a -> b -> Map.Map a (Set.Set b) -> Bool
setMember k v = Set.member v . setLookup k

-- | insert into a set of values
setInsert :: (Ord k, Ord a) => k -> a -> Map.Map k (Set.Set a)
  -> Map.Map k (Set.Set a)
setInsert k v m = Map.insert k (Set.insert v $ setLookup k m) m

-- | update an element set under the given key
setUpdate :: (Ord k, Ord a) => (Set.Set a -> Set.Set a) -> k
  -> Map.Map k (Set.Set a) -> Map.Map k (Set.Set a)
setUpdate f k m = let s = f $ setLookup k m in
  if Set.null s then Map.delete k m else Map.insert k s m

-- | test all elements of a set
setAll :: (a -> Bool) -> Set.Set a -> Bool
setAll p = List.all p . Set.toList

-- | difference function for differenceWith, returns Nothing for empty sets
setDifference :: Ord a => Set.Set a -> Set.Set a -> Maybe (Set.Set a)
setDifference s t = let d = Set.difference s t in
    if Set.null d then Nothing else Just d

-- | convert a set into an identity map
setToMap :: Ord a => Set.Set a -> Map.Map a a
setToMap = Map.fromDistinctAscList . List.map (\ a -> (a, a)) . Set.toList

-- | restrict a map by a keys set
restrict :: Ord k => Map.Map k a -> Set.Set k -> Map.Map k a
restrict m = Map.intersection m . setToMap

-- | the image of a map
imageList :: Ord k => Map.Map k a -> Set.Set k -> [a]
imageList m = Map.elems . restrict m

-- | the image of a map
imageSet :: (Ord k, Ord a) => Map.Map k a -> Set.Set k -> Set.Set a
imageSet m = Set.fromList . imageList m

-- * protected maps of set as a newtype

-- | a map to non-empty sets
newtype MapSet a b = MapSet { toMap :: Map.Map a (Set.Set b) }
  deriving (Eq, Ord)

instance (Show a, Show b) => Show (MapSet a b) where
    show = show . toMap

instance (Ord a, Read a, Ord b, Read b) => Read (MapSet a b) where
    readsPrec d = List.map (\ (m, r) -> (fromMap m , r)) . readsPrec d

-- | unsafe variant of fromMap (without removal of empty sets)
fromDistinctMap :: Map.Map a (Set.Set b) -> MapSet a b
fromDistinctMap = MapSet

-- | remove empty values from the map before applying wrapping constructor
fromMap :: Ord a => Map.Map a (Set.Set b) -> MapSet a b
fromMap = MapSet . rmNullSets

-- | the empty relation
empty :: MapSet a b
empty = MapSet Map.empty

-- | test for the empty mapping
null :: MapSet a b -> Bool
null (MapSet m) = Map.null m

-- | convert from a list
fromList :: (Ord a, Ord b) => [(a, [b])] -> MapSet a b
fromList = fromMap
  . Map.fromListWith Set.union
  . List.map (\ (a, bs) -> (a, Set.fromList bs))

-- | convert to a list
toList :: MapSet a b -> [(a, [b])]
toList = List.map (\ (a, bs) -> (a, Set.toList bs)) . Map.toList . toMap

toPairList :: MapSet a b -> [(a, b)]
toPairList = concatMap (\ (c, ts) -> List.map (\ t -> (c, t)) ts) . toList

-- | keys for non-empty elements
keysSet :: MapSet a b -> Set.Set a
keysSet = Map.keysSet . toMap

-- | all elementes united
elems :: Ord b => MapSet a b -> Set.Set b
elems = setElems . toMap

-- | get elements for a key
lookup :: Ord a => a -> MapSet a b -> Set.Set b
lookup k = setLookup k . toMap

-- | insert an element under the given key
insert :: (Ord a, Ord b) => a -> b -> MapSet a b -> MapSet a b
insert k v = MapSet . setInsert k v . toMap

-- | update an element set under the given key
update :: (Ord a, Ord b) => (Set.Set b -> Set.Set b) -> a -> MapSet a b
  -> MapSet a b
update f k = MapSet . setUpdate f k . toMap

-- | test for an element under a key
member :: (Ord a, Ord b) => a -> b -> MapSet a b -> Bool
member k v = setMember k v . toMap

-- | delete an element under the given key
delete :: (Ord a, Ord b) => a -> b -> MapSet a b -> MapSet a b
delete k v m@(MapSet r) = MapSet
  $ let s = Set.delete v $ lookup k m in
    if Set.null s then Map.delete k r else Map.insert k s r

-- | union of two maps
union :: (Ord a, Ord b) => MapSet a b -> MapSet a b -> MapSet a b
union (MapSet m) = MapSet . Map.unionWith Set.union m . toMap

-- | difference of two maps
difference :: (Ord a, Ord b) => MapSet a b -> MapSet a b -> MapSet a b
difference (MapSet m) = MapSet . Map.differenceWith setDifference m . toMap

-- | intersection of two maps
intersection :: (Ord a, Ord b) => MapSet a b -> MapSet a b -> MapSet a b
intersection (MapSet m) = fromMap
  . Map.intersectionWith Set.intersection m . toMap

-- | map a function to all elements
map :: (Ord b, Ord c) => (b -> c) -> MapSet a b -> MapSet a c
map f = MapSet . Map.map (Set.map f) . toMap

-- | unsafe map a function to all elements
mapMonotonic :: (b -> c) -> MapSet a b -> MapSet a c
mapMonotonic f = MapSet . Map.map (Set.mapMonotonic f) . toMap

-- | apply a function to all element sets
mapSet :: Ord a => (Set.Set b -> Set.Set c) -> MapSet a b -> MapSet a c
mapSet f = fromMap . Map.map f . toMap

-- | fold over all elements
foldWithKey :: (a -> b -> c -> c) -> c -> MapSet a b -> c
foldWithKey f e = Map.foldWithKey (\ a bs c -> Set.fold (f a) c bs) e . toMap

-- | filter elements
filter :: (Ord a, Ord b) => (b -> Bool) -> MapSet a b -> MapSet a b
filter p = fromMap . Map.map (Set.filter p) . toMap

-- | filter complete entries
filterWithKey :: Ord a => (a -> Set.Set b -> Bool) -> MapSet a b -> MapSet a b
filterWithKey p = MapSet . Map.filterWithKey p . toMap

-- | test all elements
all :: Ord b => (b -> Bool) -> MapSet a b -> Bool
all p = setAll p . elems

-- | submap test
isSubmapOf :: (Ord a, Ord b) => MapSet a b -> MapSet a b -> Bool
isSubmapOf (MapSet m) = Map.isSubmapOfBy Set.isSubsetOf m . toMap

-- | pre-image of a map
preImage :: (Ord a, Ord b) => Map.Map a b -> MapSet b a
preImage = Map.foldWithKey (flip insert) empty

-- | transpose a map set
transpose :: (Ord a, Ord b) => MapSet a b -> MapSet b a
transpose = foldWithKey (flip insert) empty
