{-# OPTIONS_GHC -Wunused-imports #-}

-- | Create clusters of non-overlapping things.

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

import Control.Monad

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

-- NB: We keep this module independent of Agda.Utils.List1
import Data.List.NonEmpty     ( NonEmpty(..), nonEmpty, toList )
import Data.Maybe             ( fromMaybe )

import qualified Data.Map.Strict as MapS

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 = (NonEmpty a -> NonEmpty (NonEmpty a)) -> [a] -> [NonEmpty a]
forall a b. (NonEmpty a -> NonEmpty b) -> [a] -> [b]
liftList1 ((NonEmpty a -> NonEmpty (NonEmpty a)) -> [a] -> [NonEmpty a])
-> ((a -> NonEmpty c) -> NonEmpty a -> NonEmpty (NonEmpty a))
-> (a -> NonEmpty c)
-> [a]
-> [NonEmpty a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> NonEmpty c) -> NonEmpty a -> NonEmpty (NonEmpty a)
forall c a.
Ord c =>
(a -> NonEmpty c) -> NonEmpty a -> NonEmpty (NonEmpty a)
cluster1

-- | 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' = (NonEmpty (a, NonEmpty c) -> NonEmpty (NonEmpty a))
-> [(a, NonEmpty c)] -> [NonEmpty a]
forall a b. (NonEmpty a -> NonEmpty b) -> [a] -> [b]
liftList1 NonEmpty (a, NonEmpty c) -> NonEmpty (NonEmpty a)
forall c a.
Ord c =>
NonEmpty (a, NonEmpty c) -> NonEmpty (NonEmpty a)
cluster1'

-- | Lift a function on non-empty lists to a function on lists.
--
-- Duplicate of 'Agda.Utils.List1.liftList1'.
liftList1 :: (NonEmpty a -> NonEmpty b) -> [a] -> [b]
liftList1 :: forall a b. (NonEmpty a -> NonEmpty b) -> [a] -> [b]
liftList1 NonEmpty a -> NonEmpty b
f = \case
  []     -> []
  a
a : [a]
as -> NonEmpty b -> [b]
forall a. NonEmpty a -> [a]
toList (NonEmpty b -> [b]) -> NonEmpty b -> [b]
forall a b. (a -> b) -> a -> b
$ NonEmpty a -> NonEmpty b
f (NonEmpty a -> NonEmpty b) -> NonEmpty a -> NonEmpty b
forall a b. (a -> b) -> a -> b
$ a
a a -> [a] -> NonEmpty a
forall a. a -> [a] -> NonEmpty a
:| [a]
as

-- | Given a function @f :: a -> NonEmpty c@ which returns a non-empty list of
--   characteristics of @a@, partition a non-empty 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.
cluster1 :: Ord c => (a -> NonEmpty c) -> NonEmpty a -> NonEmpty (NonEmpty a)
cluster1 :: forall c a.
Ord c =>
(a -> NonEmpty c) -> NonEmpty a -> NonEmpty (NonEmpty a)
cluster1 a -> NonEmpty c
f NonEmpty a
as = NonEmpty (a, NonEmpty c) -> NonEmpty (NonEmpty a)
forall c a.
Ord c =>
NonEmpty (a, NonEmpty c) -> NonEmpty (NonEmpty a)
cluster1' (NonEmpty (a, NonEmpty c) -> NonEmpty (NonEmpty a))
-> NonEmpty (a, NonEmpty c) -> NonEmpty (NonEmpty a)
forall a b. (a -> b) -> a -> b
$ (a -> (a, NonEmpty c)) -> NonEmpty a -> NonEmpty (a, NonEmpty c)
forall a b. (a -> b) -> NonEmpty a -> NonEmpty b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\ a
a -> (a
a, a -> NonEmpty c
f a
a)) NonEmpty a
as

-- | Partition a non-empty 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.
cluster1' :: Ord c => NonEmpty (a, NonEmpty c) -> NonEmpty (NonEmpty a)
cluster1' :: forall c a.
Ord c =>
NonEmpty (a, NonEmpty c) -> NonEmpty (NonEmpty a)
cluster1' NonEmpty (a, NonEmpty c)
acs = Fail (NonEmpty (NonEmpty a)) -> NonEmpty (NonEmpty a)
forall a. Fail a -> a
runFail_ (Fail (NonEmpty (NonEmpty a)) -> NonEmpty (NonEmpty a))
-> Fail (NonEmpty (NonEmpty a)) -> NonEmpty (NonEmpty a)
forall a b. (a -> b) -> a -> b
$ (c -> c)
-> (c -> c -> c)
-> (forall {s}. EquivT s c c Fail (NonEmpty (NonEmpty a)))
-> Fail (NonEmpty (NonEmpty a))
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 c -> c
forall a. a -> a
id c -> c -> c
forall a b. a -> b -> a
const ((forall {s}. EquivT s c c Fail (NonEmpty (NonEmpty a)))
 -> Fail (NonEmpty (NonEmpty a)))
-> (forall {s}. EquivT s c c Fail (NonEmpty (NonEmpty a)))
-> Fail (NonEmpty (NonEmpty a))
forall a b. (a -> b) -> a -> b
$ do
  -- Construct the equivalence classes of characteristics.
  NonEmpty (a, NonEmpty c)
-> ((a, NonEmpty c) -> EquivT s c c Fail ())
-> EquivT s c c Fail ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ NonEmpty (a, NonEmpty c)
acs (((a, NonEmpty c) -> EquivT s c c Fail ()) -> EquivT s c c Fail ())
-> ((a, NonEmpty c) -> EquivT s c c Fail ())
-> EquivT s c c Fail ()
forall a b. (a -> b) -> a -> b
$ \ (a
_, c
c :| [c]
cs) -> [c] -> EquivT s c c Fail ()
forall c v d (m :: * -> *). MonadEquiv c v d m => [v] -> m ()
equateAll ([c] -> EquivT s c c Fail ()) -> [c] -> EquivT s c c Fail ()
forall a b. (a -> b) -> a -> b
$ c
cc -> [c] -> [c]
forall a. a -> [a] -> [a]
:[c]
cs
  -- Pair each element with its class.
  NonEmpty (Map c (NonEmpty a))
cas <- NonEmpty (a, NonEmpty c)
-> ((a, NonEmpty c) -> EquivT s c c Fail (Map c (NonEmpty a)))
-> EquivT s c c Fail (NonEmpty (Map c (NonEmpty a)))
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM NonEmpty (a, NonEmpty c)
acs (((a, NonEmpty c) -> EquivT s c c Fail (Map c (NonEmpty a)))
 -> EquivT s c c Fail (NonEmpty (Map c (NonEmpty a))))
-> ((a, NonEmpty c) -> EquivT s c c Fail (Map c (NonEmpty a)))
-> EquivT s c c Fail (NonEmpty (Map c (NonEmpty a)))
forall a b. (a -> b) -> a -> b
$ \ (a
a, c
c :| [c]
_) -> c -> EquivT s c c Fail c
forall c v d (m :: * -> *). MonadEquiv c v d m => v -> m d
classDesc c
c EquivT s c c Fail c
-> (c -> Map c (NonEmpty a))
-> EquivT s c c Fail (Map c (NonEmpty a))
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \ c
k -> c -> NonEmpty a -> Map c (NonEmpty a)
forall k a. k -> a -> Map k a
MapS.singleton c
k (a -> NonEmpty a
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 = (NonEmpty a -> NonEmpty a -> NonEmpty a)
-> NonEmpty (Map c (NonEmpty a)) -> Map c (NonEmpty a)
forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
(a -> a -> a) -> f (Map k a) -> Map k a
MapS.unionsWith NonEmpty a -> NonEmpty a -> NonEmpty a
forall a. Semigroup a => a -> a -> a
(<>) NonEmpty (Map c (NonEmpty a))
cas
  -- Return the values of the map
  NonEmpty (NonEmpty a) -> EquivT s c c Fail (NonEmpty (NonEmpty a))
forall a. a -> EquivT s c c Fail a
forall (m :: * -> *) a. Monad m => a -> m a
return (NonEmpty (NonEmpty a)
 -> EquivT s c c Fail (NonEmpty (NonEmpty a)))
-> NonEmpty (NonEmpty a)
-> EquivT s c c Fail (NonEmpty (NonEmpty a))
forall a b. (a -> b) -> a -> b
$ NonEmpty (NonEmpty a)
-> Maybe (NonEmpty (NonEmpty a)) -> NonEmpty (NonEmpty a)
forall a. a -> Maybe a -> a
fromMaybe ([Char] -> NonEmpty (NonEmpty a)
forall a. HasCallStack => [Char] -> a
error [Char]
"impossibility at Agda.Utils.Cluster.cluster'") (Maybe (NonEmpty (NonEmpty a)) -> NonEmpty (NonEmpty a))
-> Maybe (NonEmpty (NonEmpty a)) -> NonEmpty (NonEmpty a)
forall a b. (a -> b) -> a -> b
$ [NonEmpty a] -> Maybe (NonEmpty (NonEmpty a))
forall a. [a] -> Maybe (NonEmpty a)
nonEmpty ([NonEmpty a] -> Maybe (NonEmpty (NonEmpty a)))
-> [NonEmpty a] -> Maybe (NonEmpty (NonEmpty a))
forall a b. (a -> b) -> a -> b
$
    Map c (NonEmpty a) -> [NonEmpty a]
forall k a. Map k a -> [a]
MapS.elems Map c (NonEmpty a)
m