{-# LANGUAGE CPP #-}

-- | Create clusters of non-overlapping things.

module Agda.Utils.Cluster
  ( cluster
  , cluster'
  ) where

import Control.Monad

-- An imperative union-find library:
import Data.Equivalence.Monad (runEquivT, equateAll, classDesc)
import Data.List.NonEmpty (NonEmpty(..))

import qualified Data.Map.Strict as MapS
#if __GLASGOW_HASKELL__ < 804
import Data.Semigroup
#endif

import Agda.Utils.Functor
import Agda.Utils.Singleton
import Agda.Utils.Fail

-- | Given a function @f :: a -> NonEmpty c@ which returns a non-empty list of
--   characteristics of @a@, partition a list of @a@s into groups such
--   that each element in a group shares at least one characteristic
--   with at least one other element of the group.
cluster :: Ord c => (a -> NonEmpty c) -> [a] -> [NonEmpty a]
cluster :: forall c a. Ord c => (a -> NonEmpty c) -> [a] -> [NonEmpty a]
cluster a -> NonEmpty c
f [a]
as = forall c a. Ord c => [(a, NonEmpty c)] -> [NonEmpty a]
cluster' forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (\ a
a -> (a
a, a -> NonEmpty c
f a
a)) [a]
as

-- | Partition a list of @a@s paired with a non-empty list of
--   characteristics into groups such that each element in a group
--   shares at least one characteristic with at least one other
--   element of the group.
cluster' :: Ord c => [(a, NonEmpty c)] -> [NonEmpty a]
cluster' :: forall c a. Ord c => [(a, NonEmpty c)] -> [NonEmpty a]
cluster' [(a, NonEmpty c)]
acs = forall a. Fail a -> a
runFail_ forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) v c a.
(Monad m, Applicative m) =>
(v -> c) -> (c -> c -> c) -> (forall s. EquivT s c v m a) -> m a
runEquivT forall a. a -> a
id forall a b. a -> b -> a
const forall a b. (a -> b) -> a -> b
$ do
  -- Construct the equivalence classes of characteristics.
  forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [(a, NonEmpty c)]
acs forall a b. (a -> b) -> a -> b
$ \ (a
_, c
c :| [c]
cs) -> forall c v d (m :: * -> *). MonadEquiv c v d m => [v] -> m ()
equateAll forall a b. (a -> b) -> a -> b
$ c
cforall a. a -> [a] -> [a]
:[c]
cs
  -- Pair each element with its class.
  [Map c (NonEmpty a)]
cas <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [(a, NonEmpty c)]
acs forall a b. (a -> b) -> a -> b
$ \ (a
a, c
c :| [c]
_) -> forall c v d (m :: * -> *). MonadEquiv c v d m => v -> m d
classDesc c
c forall (m :: * -> *) a b. Functor m => m a -> (a -> b) -> m b
<&> \ c
k -> forall k a. k -> a -> Map k a
MapS.singleton c
k (forall el coll. Singleton el coll => el -> coll
singleton a
a)
  -- Create a map from class to elements.
  let m :: Map c (NonEmpty a)
m = forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
(a -> a -> a) -> f (Map k a) -> Map k a
MapS.unionsWith forall a. Semigroup a => a -> a -> a
(<>) [Map c (NonEmpty a)]
cas
  -- Return the values of the map
  forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall k a. Map k a -> [a]
MapS.elems Map c (NonEmpty a)
m