module Composite.Ix.NonEmpty (groupBy, groupBy', groupBy'') where

import           Composite.Record
import           Control.Lens
import           Data.List.NonEmpty (NonEmpty ((:|)))
import qualified Data.List.NonEmpty as NE
import qualified Data.Map.NonEmpty  as MNE
import           Data.Proxy
import           Data.Vinyl         hiding (rlens, rlens')

-- | Group a NonEmpty list of records by a subtype of the record.
--
-- @since 0.0.1.0
groupBy :: forall ys xs f.
           (ys ⊆ xs)
        => Ord (Rec f ys)
        => NonEmpty (Rec f xs)
        -> MNE.NEMap (Rec f ys) (NonEmpty (Rec f xs))
groupBy xs = MNE.fromListWith (<>) $ fmap (\x -> (rcast x, x :| [])) xs

-- | Group a NonEmpty list of records by a single field.
--
-- @since 0.0.1.0
groupBy' :: forall k s y xs.
            k ~ (s :-> y)
         => k ∈ xs
         => Ord y
         => NonEmpty (Record xs)
         -> MNE.NEMap y (NonEmpty (Record xs))
groupBy' xs = MNE.fromListWith (<>) $ fmap (\x -> (view (rlens (Proxy @(s :-> y))) x, x :| [])) xs

-- | Group a NonEmpty list of records by a single field (HKD).
--
-- @since 0.0.1.0
groupBy'' :: forall k s y xs f.
             k ~ (s :-> y)
          => k ∈ xs
          => Functor f
          => Ord (f y)
          => NonEmpty (Rec f xs)
          -> MNE.NEMap (f y) (NonEmpty (Rec f xs))
groupBy'' xs = MNE.fromListWith (<>) $ fmap (\x -> (view (rlens' (Proxy @(s :-> y))) x, x :| [])) xs