module Utils.Set
( setUnionAll
, setIntersectNotOne
, setIntersections
, setMapMaybe
, setSymmDiff
, setProduct
, setPartitionEithers
, setFromFunc
, mapIntersectNotOne
, getAllUserIdsSetList
) where 

import qualified Data.Set as Set
import qualified Data.Map.Strict()
import qualified Data.Map as Map
import ClassyPrelude
import Data.Universe
import Control.Lens.Prism
import Control.Lens

-- | cardinal number of a Set
setUnionAll :: Ord a => [Set.Set a] -> Int
setUnionAll [] = 0
setUnionAll [x] = Set.size x
setUnionAll (x:y:z) = setUnionAll (xy:z) where xy = Set.union x y

-- | cardinal number of an intersection of a set and a list of sets
setIntersectNotOne :: Ord a => Set.Set a -> [Set.Set a] -> Int 
setIntersectNotOne _ [] = 0
setIntersectNotOne k r = Set.size $ Set.intersection k others where others = setUnionOthers r

-- | Union of a list of sets
setUnionOthers :: Ord a => [Set.Set a] -> Set.Set a
setUnionOthers [] = Set.empty 
setUnionOthers [x] = x
setUnionOthers (x:y:z) = setUnionOthers (xy:z) where xy = Set.union x y

----------------------------------
-- Functions fro Particiants.hs --
----------------------------------

-- | list of values of a Map with Sets as values
getAllUserIdsSetList :: Map.Map a (Set b) -> [Set b]
getAllUserIdsSetList m = loop (Map.toList m) [] where
    loop [] uids = uids
    loop ((_,v):xs) uids = loop xs $ uids ++ [v] 

-- | value of a Map with given key
getUserIdsOfOneCourse :: (Ord a, Ord b) => Map.Map a (Set b) -> a -> Set b
getUserIdsOfOneCourse m cid
    |m == Map.empty = Set.empty 
    |otherwise = m Map.! cid

-- | extracts from a map a list of values (sets) without one specific entry (a)
getAllUserIdsWithoutOne :: (Ord a, Ord b) => Map.Map a (Set b) -> a -> [Set b]
getAllUserIdsWithoutOne m cid 
    |m == Map.empty = []
    |otherwise = getAllUserIdsSetList $ Map.delete cid m 

-- | transforms values (sets) of a map to integers. The number gives information about how many entreis are not only in this one 
mapIntersectNotOne :: (Ord a, Ord b) => Map.Map a (Set b) -> Map.Map a Int 
mapIntersectNotOne m = loop (Map.toList m) Map.empty where 
    loop [] ino = ino 
    loop ((k,_):xs) ino = loop xs $ Map.insert k (setIntersectNotOne (getUserIdsOfOneCourse m k) (getAllUserIdsWithoutOne m k)) ino

-----------------------------
-- Functions from Utils.hs --
-----------------------------

-- | Intersection of multiple sets. Returns empty set for empty input list
setIntersections :: Ord a => [Set a] -> Set a
setIntersections []    = Set.empty
setIntersections (h:t) = foldl' Set.intersection h t

setMapMaybe :: Ord b => (a -> Maybe b) -> Set a -> Set b
setMapMaybe f = Set.fromList . mapMaybe f . Set.toList

-- | Symmetric difference of two sets.
setSymmDiff :: Ord a => Set a -> Set a -> Set a
setSymmDiff x y = (x `Set.difference` y) `Set.union` (y `Set.difference` x)

setProduct :: Set a -> Set b -> Set (a, b)
-- ^ Depends on the valid internal structure of the given sets
setProduct (Set.toAscList -> as) (Set.toAscList -> bs) = Set.fromDistinctAscList $ (,) <$> as <*> bs

setPartitionEithers :: (Ord a, Ord b) => Set (Either a b) -> (Set a, Set b)
setPartitionEithers = (,) <$> setMapMaybe (preview _Left) <*> setMapMaybe (preview _Right)

setFromFunc :: (Finite k, Ord k) => (k -> Bool) -> Set k
setFromFunc = Set.fromList . flip filter universeF