-- SPDX-FileCopyrightText: 2022 Winnie Ros <winnie.ros@campus.lmu.de>
--
-- SPDX-License-Identifier: AGPL-3.0-or-later

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

import qualified Data.List.NonEmpty as NonEmpty
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 an intersection of a set and a list of sets
setIntersectNotOne :: Ord a => Set a -> [Set a] -> Int 
setIntersectNotOne _ [] = 0
setIntersectNotOne k r = Set.size $ Set.intersection k others where others = Set.unions r

----------------------------------------
-- Functions for Handler.Participants --
----------------------------------------

-- | extracts from a map a list of values (sets) without one specific entry (a)
getAllElemsWithoutOne :: (Ord a) => Map a (Set b) -> a -> [Set b]
getAllElemsWithoutOne m cid = Map.elems $ 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 :: forall a b. (Ord a, Ord b) => Map a (Set b) -> Map a Int 
mapIntersectNotOne m = Map.mapWithKey f m where 
    f :: a -> Set b -> Int
    f k _ = setIntersectNotOne (Map.findWithDefault Set.empty k m) (getAllElemsWithoutOne m k)
 
--------------------------
-- Functions from Utils --
--------------------------

-- | 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


-- | convert a Set to NonEmpty, inserting a default value if necessary
set2NonEmpty :: a -> Set a -> NonEmpty.NonEmpty a
set2NonEmpty _ (Set.toList -> h:t) = h NonEmpty.:| t
set2NonEmpty d _ = d NonEmpty.:| []