module Data.List.MultiValuedGrouping
  ( GroupingStrategy
  , GroupResult
  , isGroup
  , memberGroupingStrategy
  ) where

import Control.Applicative ((<$>), (<*>))
import Control.Arrow (first)
import Data.Either (partitionEithers)
import qualified Data.Map as M
import Data.Maybe (catMaybes)
import Data.Set (Set, singleton, toList)

type GroupResult k v = (Set k, v, [v])

isGroup :: GroupResult k v -> Bool
isGroup (_, _, xs) = not $ null xs

type GroupingStrategy k v = [(Set k, v)] -> [GroupResult k v]

partitionMaybe :: (a -> Maybe b) -> [a] -> ([a], [(a, b)])
partitionMaybe f = partitionEithers . ((flip (maybe Left ((Right .) . flip (,))) <*> f) <$>)

set :: Int -> a -> [a] -> [a]
set i x = uncurry (++) . ((\(_:t) -> x:t) <$>) . splitAt i

-- | Works like in SQL, except that unkeyed items are left in their original positions
memberGroupingStrategy :: Ord k => GroupingStrategy k v
memberGroupingStrategy = catMaybes . fst . foldr combiner ([], M.empty)
    where
      combiner (keys, value) acc@(res, seen)
          = flip (foldr (\key s -> M.insert key (length res, [value]) s)) ungrouped
            <$> ((case grouped of
                    [] -> first ((return (keys, value, [])) :)
                    _  -> id
                 ) $ foldr group acc grouped)
          where
            group (k, (i, oldValues)) (r, s)
                = ( return (singleton k, value, oldValues) : set (length r - (i + 1)) Nothing r
                  , M.insert k (length r, value : oldValues) s
                  )
            (ungrouped, grouped) = partitionMaybe (`M.lookup` seen) $ toList keys