{-# LANGUAGE CPP #-}

module Data.Set.Extra
    ( module Set
    , mapM
    , mapM_
    , filterM
    , catMaybes
    , mapMaybe
    , flatten
    , concatMap
    , concatMapM
    , any
    , all
    , or
    , and
    , ss
    , toSS
    , fromSS
    , ssMapM
    , distrib
    , cartesianProduct
    , groupBy
    , powerset
    , partitionM
    , unzip
    , gFind
    ) where

import qualified Control.Monad as List (mapM, mapM_, filterM, foldM)
import Control.Monad.State ()
import qualified Data.Foldable as Foldable (all, any, and, or)
import Data.Map as Map (Map, insertWith, empty)
import Data.Set as Set
import Data.Set.ExtraG
import qualified Data.List as List
--import qualified Data.Maybe
import Prelude hiding (mapM, mapM_, unzip, all, any, map, filter, null, concatMap, and, or)
--import qualified Prelude

mapM :: (Monad m, Ord b) => (a -> m b) -> Set a -> m (Set b)
mapM :: forall (m :: * -> *) b a.
(Monad m, Ord b) =>
(a -> m b) -> Set a -> m (Set b)
mapM a -> m b
f Set a
s = (a -> m b) -> [a] -> m [b]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
List.mapM a -> m b
f (Set a -> [a]
forall a. Set a -> [a]
toList Set a
s) m [b] -> ([b] -> m (Set b)) -> m (Set b)
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Set b -> m (Set b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Set b -> m (Set b)) -> ([b] -> Set b) -> [b] -> m (Set b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [b] -> Set b
forall a. Ord a => [a] -> Set a
fromList

mapM_ :: (Monad m, Ord b) => (a -> m b) -> Set a -> m ()
mapM_ :: forall (m :: * -> *) b a.
(Monad m, Ord b) =>
(a -> m b) -> Set a -> m ()
mapM_ a -> m b
f Set a
s = (a -> m b) -> [a] -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
List.mapM_ a -> m b
f (Set a -> [a]
forall a. Set a -> [a]
toList Set a
s)

filterM :: (Ord a, Monad m) => (a -> m Bool) -> Set a -> m (Set a)
filterM :: forall a (m :: * -> *).
(Ord a, Monad m) =>
(a -> m Bool) -> Set a -> m (Set a)
filterM a -> m Bool
p Set a
s = (a -> m Bool) -> [a] -> m [a]
forall (m :: * -> *) a.
Applicative m =>
(a -> m Bool) -> [a] -> m [a]
List.filterM a -> m Bool
p (Set a -> [a]
forall a. Set a -> [a]
toList Set a
s) m [a] -> ([a] -> m (Set a)) -> m (Set a)
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Set a -> m (Set a)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Set a -> m (Set a)) -> ([a] -> Set a) -> [a] -> m (Set a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> Set a
forall a. Ord a => [a] -> Set a
fromList

catMaybes :: Ord a => Set (Maybe a) -> Set a
catMaybes :: forall a. Ord a => Set (Maybe a) -> Set a
catMaybes Set (Maybe a)
sm = (Maybe a -> Set a -> Set a) -> Set a -> Set (Maybe a) -> Set a
forall a b. (a -> b -> b) -> b -> Set a -> b
fold (\ Maybe a
mx Set a
s -> Set a -> (a -> Set a) -> Maybe a -> Set a
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Set a
s (a -> Set a -> Set a
forall a. Ord a => a -> Set a -> Set a
`insert` Set a
s) Maybe a
mx) Set a
forall a. Set a
Set.empty Set (Maybe a)
sm

mapMaybe :: (Ord a, Ord b) => (a -> Maybe b) -> Set a -> Set b
mapMaybe :: forall a b. (Ord a, Ord b) => (a -> Maybe b) -> Set a -> Set b
mapMaybe a -> Maybe b
f Set a
s = Set (Maybe b) -> Set b
forall a. Ord a => Set (Maybe a) -> Set a
catMaybes ((a -> Maybe b) -> Set a -> Set (Maybe b)
forall b a. Ord b => (a -> b) -> Set a -> Set b
Set.map a -> Maybe b
f Set a
s)

flatten :: Ord a => Set (Set a) -> Set a
flatten :: forall a. Ord a => Set (Set a) -> Set a
flatten Set (Set a)
ss' = (Set a -> Set a -> Set a) -> Set a -> Set (Set a) -> Set a
forall a b. (a -> b -> b) -> b -> Set a -> b
fold Set a -> Set a -> Set a
forall a. Ord a => Set a -> Set a -> Set a
union Set a
forall a. Set a
Set.empty Set (Set a)
ss'
--flatten = unions . toList

concatMap :: (Ord a, Ord b) => (a -> Set b) -> Set a -> Set b
concatMap :: forall a b. (Ord a, Ord b) => (a -> Set b) -> Set a -> Set b
concatMap a -> Set b
f Set a
s = Set (Set b) -> Set b
forall a. Ord a => Set (Set a) -> Set a
flatten ((a -> Set b) -> Set a -> Set (Set b)
forall b a. Ord b => (a -> b) -> Set a -> Set b
Set.map a -> Set b
f Set a
s)

concatMapM :: (Monad m, Ord a, Ord b) => (a -> m (Set b)) -> Set a -> m (Set b)
concatMapM :: forall (m :: * -> *) a b.
(Monad m, Ord a, Ord b) =>
(a -> m (Set b)) -> Set a -> m (Set b)
concatMapM a -> m (Set b)
f Set a
s = (a -> m (Set b)) -> Set a -> m (Set (Set b))
forall (m :: * -> *) b a.
(Monad m, Ord b) =>
(a -> m b) -> Set a -> m (Set b)
mapM a -> m (Set b)
f Set a
s m (Set (Set b)) -> (Set (Set b) -> m (Set b)) -> m (Set b)
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Set b -> m (Set b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Set b -> m (Set b))
-> (Set (Set b) -> Set b) -> Set (Set b) -> m (Set b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Set (Set b) -> Set b
forall a. Ord a => Set (Set a) -> Set a
flatten

any :: Ord a => (a -> Bool) -> Set a -> Bool
any :: forall a. Ord a => (a -> Bool) -> Set a -> Bool
any = (a -> Bool) -> Set a -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
Foldable.any

{-
anyM :: Monad m => (a -> m (Maybe Bool)) -> Set a -> m (Maybe Bool)
anyM p s =
    List.mapM p (toList s) >>= return . Data.Maybe.catMaybes >>= return . chk
    where chk [] = Nothing
          chk ys = Just (Prelude.or ys)
-}

all :: Ord a => (a -> Bool) -> Set a -> Bool
all :: forall a. Ord a => (a -> Bool) -> Set a -> Bool
all = (a -> Bool) -> Set a -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
Foldable.all

or :: Set Bool -> Bool
or :: Set Bool -> Bool
or = Set Bool -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
Foldable.or

and :: Set Bool -> Bool
and :: Set Bool -> Bool
and = Set Bool -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
Foldable.and

-- |Create a singleton set containing a singleton set of a.
ss :: Ord a => a -> Set (Set a)
ss :: forall a. Ord a => a -> Set (Set a)
ss = Set a -> Set (Set a)
forall a. a -> Set a
singleton (Set a -> Set (Set a)) -> (a -> Set a) -> a -> Set (Set a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Set a
forall a. a -> Set a
singleton

-- |Turn a list of lists into a set of sets.
toSS :: Ord a => [[a]] -> Set (Set a)
toSS :: forall a. Ord a => [[a]] -> Set (Set a)
toSS = [Set a] -> Set (Set a)
forall a. Ord a => [a] -> Set a
fromList ([Set a] -> Set (Set a))
-> ([[a]] -> [Set a]) -> [[a]] -> Set (Set a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([a] -> Set a) -> [[a]] -> [Set a]
forall a b. (a -> b) -> [a] -> [b]
List.map [a] -> Set a
forall a. Ord a => [a] -> Set a
fromList

fromSS :: Ord a => Set (Set a) -> [[a]]
fromSS :: forall a. Ord a => Set (Set a) -> [[a]]
fromSS = (Set a -> [a]) -> [Set a] -> [[a]]
forall a b. (a -> b) -> [a] -> [b]
List.map Set a -> [a]
forall a. Set a -> [a]
toList ([Set a] -> [[a]])
-> (Set (Set a) -> [Set a]) -> Set (Set a) -> [[a]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Set (Set a) -> [Set a]
forall a. Set a -> [a]
toList

ssMapM :: (Monad m, Ord a, Ord b) => (a -> m b) -> Set (Set a) -> m (Set (Set b))
ssMapM :: forall (m :: * -> *) a b.
(Monad m, Ord a, Ord b) =>
(a -> m b) -> Set (Set a) -> m (Set (Set b))
ssMapM a -> m b
f Set (Set a)
s = ([a] -> m [b]) -> [[a]] -> m [[b]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
List.mapM ((a -> m b) -> [a] -> m [b]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
List.mapM a -> m b
f) (Set (Set a) -> [[a]]
forall a. Ord a => Set (Set a) -> [[a]]
fromSS Set (Set a)
s) m [[b]] -> ([[b]] -> m (Set (Set b))) -> m (Set (Set b))
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Set (Set b) -> m (Set (Set b))
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Set (Set b) -> m (Set (Set b)))
-> ([[b]] -> Set (Set b)) -> [[b]] -> m (Set (Set b))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [[b]] -> Set (Set b)
forall a. Ord a => [[a]] -> Set (Set a)
toSS

-- | distrib {a, b, c} {d, e, f} -> {a+d, a+e, a+f, b+d, b+e, b+f, c+d, c+e, c+f}
distrib :: Ord a => Set (Set a) -> Set (Set a) -> Set (Set a)
distrib :: forall a. Ord a => Set (Set a) -> Set (Set a) -> Set (Set a)
distrib Set (Set a)
lss Set (Set a)
rss = Set (Set (Set a)) -> Set (Set a)
forall a. Ord a => Set (Set a) -> Set a
flatten (Set (Set (Set a)) -> Set (Set a))
-> Set (Set (Set a)) -> Set (Set a)
forall a b. (a -> b) -> a -> b
$ (Set a -> Set (Set a)) -> Set (Set a) -> Set (Set (Set a))
forall b a. Ord b => (a -> b) -> Set a -> Set b
map (\ Set a
rs -> ((Set a -> Set a) -> Set (Set a) -> Set (Set a)
forall b a. Ord b => (a -> b) -> Set a -> Set b
map (\ Set a
ls -> Set a -> Set a -> Set a
forall a. Ord a => Set a -> Set a -> Set a
union Set a
rs Set a
ls) Set (Set a)
lss)) Set (Set a)
rss

#if !MIN_VERSION_containers(0,5,11)
cartesianProduct :: (Ord a, Ord b) => Set a -> Set b -> Set (a, b)
cartesianProduct xs ys = flatten $ Set.map (\ x -> Set.map (\ y -> (x, y)) ys) xs
#endif

groupBy :: (Ord a, Ord b) => (a -> b) -> Set a -> Map.Map b (Set a)
groupBy :: forall a b. (Ord a, Ord b) => (a -> b) -> Set a -> Map b (Set a)
groupBy a -> b
f Set a
xs = (a -> Map b (Set a) -> Map b (Set a))
-> Map b (Set a) -> Set a -> Map b (Set a)
forall a b. (a -> b -> b) -> b -> Set a -> b
fold (\ a
x Map b (Set a)
m -> (Set a -> Set a -> Set a)
-> b -> Set a -> Map b (Set a) -> Map b (Set a)
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
Map.insertWith Set a -> Set a -> Set a
forall a. Ord a => Set a -> Set a -> Set a
union (a -> b
f a
x) (a -> Set a
forall a. a -> Set a
singleton a
x) Map b (Set a)
m) Map b (Set a)
forall k a. Map k a
Map.empty Set a
xs

powerset :: Ord a => Set a -> Set (Set a)
powerset :: forall a. Ord a => Set a -> Set (Set a)
powerset Set a
s
    | Set a
s Set a -> Set a -> Bool
forall a. Eq a => a -> a -> Bool
== Set a
forall a. Set a
Set.empty = Set a -> Set (Set a)
forall a. a -> Set a
singleton Set a
forall a. Set a
Set.empty
    | Bool
otherwise = (Set a -> Set a) -> Set (Set a) -> Set (Set a)
forall b a. Ord b => (a -> b) -> Set a -> Set b
Set.map (a -> Set a -> Set a
forall a. Ord a => a -> Set a -> Set a
insert a
x) Set (Set a)
pxs Set (Set a) -> Set (Set a) -> Set (Set a)
forall a. Ord a => Set a -> Set a -> Set a
`union` Set (Set a)
pxs
        where (a
x, Set a
xs) = Set a -> (a, Set a)
forall a. Set a -> (a, Set a)
deleteFindMin Set a
s
              pxs :: Set (Set a)
pxs = Set a -> Set (Set a)
forall a. Ord a => Set a -> Set (Set a)
powerset Set a
xs

partitionM :: (Monad m, Ord a) => (a -> m Bool) -> Set a -> m (Set a, Set a)
partitionM :: forall (m :: * -> *) a.
(Monad m, Ord a) =>
(a -> m Bool) -> Set a -> m (Set a, Set a)
partitionM a -> m Bool
p Set a
xs =
    ((Set a, Set a) -> a -> m (Set a, Set a))
-> (Set a, Set a) -> [a] -> m (Set a, Set a)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
List.foldM (Set a, Set a) -> a -> m (Set a, Set a)
f (Set a
forall a. Set a
Set.empty, Set a
forall a. Set a
Set.empty) (Set a -> [a]
forall a. Set a -> [a]
toList Set a
xs)
    where f :: (Set a, Set a) -> a -> m (Set a, Set a)
f (Set a
ts, Set a
fs) a
x = a -> m Bool
p a
x m Bool -> (Bool -> m (Set a, Set a)) -> m (Set a, Set a)
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \ Bool
flag -> (Set a, Set a) -> m (Set a, Set a)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ((Set a, Set a) -> m (Set a, Set a))
-> (Set a, Set a) -> m (Set a, Set a)
forall a b. (a -> b) -> a -> b
$ if Bool
flag then (a -> Set a -> Set a
forall a. Ord a => a -> Set a -> Set a
insert a
x Set a
ts, Set a
fs) else (Set a
ts, a -> Set a -> Set a
forall a. Ord a => a -> Set a -> Set a
insert a
x Set a
fs)

unzip :: (Ord a, Ord b) => Set (a, b) -> (Set a, Set b)
unzip :: forall a b. (Ord a, Ord b) => Set (a, b) -> (Set a, Set b)
unzip Set (a, b)
s = ((a, b) -> (Set a, Set b) -> (Set a, Set b))
-> (Set a, Set b) -> Set (a, b) -> (Set a, Set b)
forall a b. (a -> b -> b) -> b -> Set a -> b
Set.fold (\ (a
l, b
r) (Set a
ls, Set b
rs) -> (a -> Set a -> Set a
forall a. Ord a => a -> Set a -> Set a
Set.insert a
l Set a
ls, b -> Set b -> Set b
forall a. Ord a => a -> Set a -> Set a
Set.insert b
r Set b
rs)) (Set a
forall a. Set a
Set.empty, Set b
forall a. Set a
Set.empty) Set (a, b)
s