{-# LANGUAGE Safe          #-}
{-# LANGUAGE TypeFamilies  #-}
{-# LANGUAGE TypeOperators #-}

{- |
Module                  : Relude.Extra.Group
Copyright               : (c) 2018-2023 Kowainik
SPDX-License-Identifier : MIT
Maintainer              : Kowainik <xrom.xkov@gmail.com>
Stability               : Stable
Portability             : Portable

Polymorphic grouping functions.

@since 0.1.0
-}

module Relude.Extra.Group
    ( groupBy
    , groupOneBy
    ) where

import Relude
import Relude.Extra.Map

import Data.List.NonEmpty ((<|))


{- | Groups elements using results of the given function as keys.

>>> groupBy even [1..6] :: HashMap Bool (NonEmpty Int)
fromList [(False,5 :| [3,1]),(True,6 :| [4,2])]

@since 0.1.0
-}
groupBy :: forall f t a . (Foldable f, DynamicMap t, Val t ~ NonEmpty a, Monoid t)
        => (a -> Key t) -> f a -> t
groupBy f = flipfoldl' hmGroup mempty
  where
    hmGroup :: a -> t -> t
    hmGroup x =
        let val :: Maybe (NonEmpty a) -> NonEmpty a
            val Nothing   = one x
            val (Just xs) = x <| xs
        in alter (Just . val) (f x)

{- | Similar to 'groupBy' but keeps only one element as value.

>>> groupOneBy even [1 .. 6] :: HashMap Bool Int
fromList [(False,1),(True,2)]

@since 0.1.0
-}
groupOneBy :: forall f t a . (Foldable f, DynamicMap t, Val t ~ a, Monoid t)
           => (a -> Key t) -> f a -> t
groupOneBy f = flipfoldl' hmGroup mempty
  where
    hmGroup :: a -> t -> t
    hmGroup val m = let key = f val in
        case lookup key m of
            Nothing -> insert key val m
            Just _  -> m
