{-# LANGUAGE DeriveTraversable #-}
module Data.Bag (
  Bag,
  bagFromList,
  bagFilter,
  bagPartition,
) where

import Data.Bifunctor (bimap)
import Data.Foldable (toList)
import Data.List.NonEmpty (NonEmpty((:|)))
import Data.Maybe (mapMaybe)
import Data.Semigroup


data Bag a
  = BZero
  | BOne a
  | BTwo (Bag a) (Bag a)
  | BList [Bag a]  -- make mconcat efficient
  | BList' [a]  -- make bagFromList efficient
  deriving (Functor, Traversable)

instance Show a => Show (Bag a) where
  showsPrec d b = showParen (d > 10) $
    showString "Bag " . showList (toList b)

instance Semigroup (Bag a) where
  BZero <> b = b
  b <> BZero = b
  b1 <> b2 = BTwo b1 b2

  sconcat (b :| bs) = b <> mconcat bs
  stimes n b = mconcat (replicate (fromIntegral n) b)

instance Monoid (Bag a) where
  mempty = BZero
  mconcat = BList

instance Applicative Bag where
  pure = BOne

  BZero <*> _ = BZero
  _ <*> BZero = BZero
  BOne f <*> b = f <$> b
  BTwo b1 b2 <*> b = BTwo (b1 <*> b) (b2 <*> b)
  BList bs <*> b = BList (map (<*> b) bs)
  BList' xs <*> b = BList (map BOne xs) <*> b

instance Foldable Bag where
  foldMap _ BZero = mempty
  foldMap f (BOne x) = f x
  foldMap f (BTwo b1 b2) = foldMap f b1 <> foldMap f b2
  foldMap f (BList l) = foldMap (foldMap f) l
  foldMap f (BList' l) = foldMap f l

  toList (BList' xs) = xs
  toList b = foldr (:) [] b

  null BZero = True
  null BOne{} = False
  null (BTwo b1 b2) = null b1 && null b2
  null (BList l) = all null l
  null (BList' l) = null l

bagFromList :: [a] -> Bag a
bagFromList = BList'

bagFilter :: (a -> Maybe b) -> Bag a -> Bag b
bagFilter _ BZero = BZero
bagFilter f (BOne x)
  | Just y <- f x = BOne y
  | otherwise = BZero
bagFilter f (BTwo b1 b2) = bagFilter f b1 <> bagFilter f b2
bagFilter f (BList bs) = BList (map (bagFilter f) bs)
bagFilter f (BList' xs) = BList' (mapMaybe f xs)

bagPartition :: (a -> Maybe b) -> Bag a -> (Bag b, Bag a)
bagPartition _ BZero = (BZero, BZero)
bagPartition f (BOne x)
  | Just y <- f x = (BOne y, BZero)
  | otherwise = (BZero, BOne x)
bagPartition f (BTwo b1 b2) = bagPartition f b1 <> bagPartition f b2
bagPartition f (BList bs) = foldMap (bagPartition f) bs
bagPartition f (BList' xs) =
  bimap bagFromList bagFromList $
    foldr (\x (l1,l2) -> case f x of
                           Just y -> (y : l1, l2)
                           Nothing -> (l1, x : l2))
          ([], []) xs