{-# LANGUAGE CPP                        #-}
{-# LANGUAGE ConstrainedClassMethods    #-}
{-# LANGUAGE DefaultSignatures          #-}
{-# LANGUAGE FlexibleContexts           #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE KindSignatures             #-}
{-# LANGUAGE LambdaCase                 #-}
{-# LANGUAGE StandaloneDeriving         #-}
{-# LANGUAGE TypeOperators              #-}

module Control.Compactable
  (
  -- * Compact
    Compactable (..)
  -- * Compact Fold
  , CompactFold (..)
  -- * Handly flips
  , fforMaybe
  , fforFold
  , fforEither
  , fforBifold
  -- * More general lefts and rights
  , mfold'
  , mlefts
  , mrights
  -- * Monad Transformer utils
  , fmapMaybeM
  , fmapEitherM
  , fforMaybeM
  , fforEitherM
  , applyMaybeM
  , bindMaybeM
  , traverseMaybeM
  -- * Alternative Defaults
  , altDefaultCompact
  , altDefaultSeparate
  ) where

import           Control.Applicative
import           Control.Arrow
import           Control.Monad                   (MonadPlus, join)
import           Control.Monad.Trans.Except
import           Control.Monad.Trans.Maybe
import           Data.Bifoldable
import           Data.Bifunctor                  (bimap)
import           Data.Either                     (partitionEithers)
import           Data.Foldable                   as F (foldl', toList)
import           Data.Functor.Compose
import qualified Data.Functor.Product            as FP
import qualified Data.IntMap                     as IntMap
import qualified Data.List                       as List
import qualified Data.Map                        as Map
import           Data.Maybe
import           Data.Monoid
import           Data.Proxy
import           Data.Semigroup
import qualified Data.Sequence                   as Seq
import qualified Data.Set                        as Set
import qualified Data.Vector                     as V
import           GHC.Conc
import           GHC.Generics
import           Text.ParserCombinators.ReadP
import           Text.ParserCombinators.ReadPrec

{-|
Class 'Compactable' provides two methods which can be writen in terms of each other, compact and separate.

is generalization of catMaybes as a new function. Compact
has relations with Functor, Applicative, Monad, Alternative, and Traversable.
In that we can use these class to provide the ability to operate on a data type
by throwing away intermediate Nothings. This is useful for representing
stripping out values or failure.

To be compactable alone, no laws must be satisfied other than the type signature.

If the data type is also a Functor the following should hold:

[/Kleisli composition/]

    @fmapMaybe (l <=< r) = fmapMaybe l . fmapMaybe r@

[/Functor identity 1/]

    @compact . fmap Just = id@

[/Functor identity 2/]

    @fmapMaybe Just = id@

[/Functor relation/]

    @compact = fmapMaybe id@

According to Kmett, (Compactable f, Functor f) is a functor from the
kleisli category of Maybe to the category of haskell data types.
@Kleisli Maybe -> Hask@.

If the data type is also Applicative the following should hold:

[/Applicative left identity/]

    @compact . (pure Just \<*\>) = id@

[/Applicative right identity/]

    @applyMaybe (pure Just) = id@

[/Applicative relation/]

    @compact = applyMaybe (pure id)@

If the data type is also a Monad the following should hold:

[/Monad left identity/]

    @flip bindMaybe (return . Just) = id@

[/Monad right identity/]

    @compact . (return . Just =<<) = id@

[/Monad relation/]

    @compact = flip bindMaybe return@

If the data type is also Alternative the following should hold:

[/Alternative identity/]

    @compact empty = empty@

[/Alternative annihilation/]

    @compact (const Nothing \<$\> xs) = empty@

If the data type is also Traversable the following should hold:

[/Traversable Applicative relation/]

    @traverseMaybe (pure . Just) = pure@

[/Traversable composition/]

    @Compose . fmap (traverseMaybe f) . traverseMaybe g = traverseMaybe (Compose . fmap (traverseMaybe f) . g)@

[/Traversable Functor relation/]

    @traverse f = traverseMaybe (fmap Just . f)@

[/Traversable naturality/]

    @t . traverseMaybe f = traverseMaybe (t . f)@

= Separate and filter
have recently elevated roles in this typeclass, and is not as well explored as compact. Here are the laws known today:

[/Functor identity 3/]

    @fst . separate . fmap Right = id@

[/Functor identity 4/]

    @snd . separate . fmap Left = id@

[/Applicative left identity 2/]

    @snd . separate . (pure Right \<*\>) = id@

[/Applicative right identity 2/]

    @fst . separate . (pure Left \<*\>) = id@

[/Alternative annihilation left/]

    @snd . separate . fmap (const Left) = empty@

[/Alternative annihilation right/]

    @fst , separate . fmap (const Right) = empty@

Docs for relationships between these functions and, a cleanup of laws will happen at some point.

If you know of more useful laws, or have better names for the ones above
(especially those marked "name me"). Please let me know.
-}

class Compactable (f :: * -> *) where
    {-# MINIMAL compact | separate #-}

    compact :: f (Maybe a) -> f a
    default compact :: Functor f => f (Maybe a) -> f a
    compact = snd . separate . fmap (\case Just x -> Right x; _ -> Left ())
    {-# INLINABLE compact #-}

    separate :: f (Either l r) -> (f l, f r)
    default separate :: Functor f => f (Either l r) -> (f l, f r)
    separate xs = (compact $ hush . flipEither <$> xs, compact $ hush <$> xs)
    {-# INLINABLE separate #-}

    filter :: (a -> Bool) -> f a -> f a
    default filter :: Functor f => (a -> Bool) -> f a -> f a
    filter f = fmapMaybe $ \a -> if f a then Just a else Nothing
    {-# INLINABLE filter #-}

    partition :: (a -> Bool) -> f a -> (f a, f a)
    default partition :: Functor f => (a -> Bool) -> f a -> (f a, f a)
    partition f = fmapEither $ \a -> if f a then Right a else Left a
    {-# INLINEABLE partition #-}

    fmapMaybe :: Functor f => (a -> Maybe b) -> f a -> f b
    fmapMaybe f = compact . fmap f
    {-# INLINABLE fmapMaybe #-}

    fmapEither :: Functor f => (a -> Either l r) -> f a -> (f l, f r)
    fmapEither f = separate . fmap f
    {-# INLINABLE fmapEither #-}

    applyMaybe :: Applicative f => f (a -> Maybe b) -> f a -> f b
    applyMaybe fa = compact . (fa <*>)
    {-# INLINABLE applyMaybe #-}

    applyEither :: Applicative f => f (a -> Either l r) -> f a -> (f l, f r)
    applyEither fa = separate . (fa <*>)
    {-# INLINABLE applyEither #-}

    bindMaybe :: Monad f => f a -> (a -> f (Maybe b)) -> f b
    bindMaybe x = compact . (x >>=)
    {-# INLINABLE bindMaybe #-}

    bindEither :: Monad f => f a -> (a -> f (Either l r)) -> (f l, f r)
    bindEither x = separate . (x >>=)
    {-# INLINABLE bindEither #-}

    traverseMaybe :: (Applicative g, Traversable f)
                  => (a -> g (Maybe b)) -> f a -> g (f b)
    traverseMaybe f = fmap compact . traverse f
    {-# INLINABLE traverseMaybe #-}

    traverseEither :: (Applicative g, Traversable f)
                   => (a -> g (Either l r)) -> f a -> g (f l, f r)
    traverseEither f = fmap separate . traverse f
    {-# INLINABLE traverseEither #-}


instance Compactable Maybe where
    compact = join
    {-# INLINABLE compact #-}
    fmapMaybe = (=<<)
    {-# INLINABLE fmapMaybe #-}
    separate = \case
      Just x -> case x of
         Left l  -> (Just l, Nothing)
         Right r -> (Nothing, Just r)
      _ -> (Nothing, Nothing)
    {-# INLINABLE separate #-}

instance Monoid m => Compactable (Either m) where
    compact (Right (Just x)) = Right x
    compact (Right _)        = Left mempty
    compact (Left x)         = Left x
    {-# INLINABLE compact #-}
    fmapMaybe f (Right x) = maybe (Left mempty) Right (f x)
    fmapMaybe _ (Left x)  = Left x
    {-# INLINABLE fmapMaybe #-}
    separate = \case
      Right (Left l) -> (Right l, Left mempty)
      Right (Right r) -> (Left mempty, Right r)
      Left x -> (Left x, Left x)
    {-# INLINABLE separate #-}

instance Compactable [] where
    compact = catMaybes
    {-# INLINABLE compact #-}
    fmapMaybe _ []    = []
    fmapMaybe f (h:t) = maybe (fmapMaybe f t) (: fmapMaybe f t) (f h)
    {-# INLINABLE fmapMaybe #-}
    filter = Prelude.filter
    {-# INLINABLE filter #-}
    partition = List.partition
    {-# INLINABLE partition #-}
    separate = partitionEithers
    {-# INLINABLE separate #-}
    fmapEither f = foldl' (deal f) ([],[])
      where deal g ~(bs, cs) a = case g a of
             Left b  -> (b:bs, cs)
             Right c -> (bs, c:cs)
    {-# INLINABLE fmapEither #-}
    traverseMaybe f = go where
      go (x:xs) = maybe id (:) <$> f x <*> go xs
      go []     = pure []
    {-# INLINE traverseMaybe #-}

instance Compactable ZipList where
  compact (ZipList xs) = ZipList $ compact xs

instance Compactable IO where
    compact = altDefaultCompact
    {-# INLINABLE compact #-}

instance Compactable STM where
    compact = altDefaultCompact
    {-# INLINABLE compact #-}

instance Compactable Proxy where
    compact _ = Proxy
    {-# INLINABLE compact #-}
    separate _ = (Proxy, Proxy)
    {-# INLINABLE separate #-}
    filter _ _ = Proxy
    {-# INLINABLE filter #-}
    partition _ _ = (Proxy, Proxy)
    {-# INLINABLE partition #-}
    fmapMaybe _ _ = Proxy
    {-# INLINABLE fmapMaybe #-}
    applyMaybe _ _ = Proxy
    {-# INLINABLE applyMaybe #-}
    bindMaybe _ _ = Proxy
    {-# INLINABLE bindMaybe #-}
    fmapEither _ _ = (Proxy, Proxy)
    {-# INLINABLE fmapEither #-}
    applyEither _ _ = (Proxy, Proxy)
    {-# INLINABLE applyEither #-}
    bindEither _ _ = (Proxy, Proxy)
    {-# INLINABLE bindEither #-}

instance Compactable U1

instance Compactable Option where
    compact (Option x) = Option (join x)
    {-# INLINABLE compact #-}
    fmapMaybe f (Option (Just x)) = Option (f x)
    fmapMaybe _ _                 = Option Nothing
    {-# INLINABLE fmapMaybe #-}
    separate = altDefaultSeparate
    {-# INLINABLE separate #-}

instance Compactable ReadP

instance Compactable ReadPrec

instance ( Functor f, Functor g, Compactable f, Compactable g )
         => Compactable (FP.Product f g) where
    compact (FP.Pair x y) = FP.Pair (compact x) (compact y)
    {-# INLINABLE compact #-}

instance ( Functor f, Functor g, Compactable g )
         => Compactable (Compose f g) where
    compact = fmapMaybe id
    {-# INLINABLE compact #-}
    fmapMaybe f (Compose fg) = Compose $ fmapMaybe f <$> fg
    {-# INLINABLE fmapMaybe #-}

instance Compactable IntMap.IntMap where
    compact = IntMap.mapMaybe id
    {-# INLINABLE compact #-}
    fmapMaybe = IntMap.mapMaybe
    {-# INLINABLE fmapMaybe #-}
    filter = IntMap.filter
    {-# INLINABLE filter #-}
    partition = IntMap.partition
    {-# INLINABLE partition #-}
    separate = IntMap.mapEither id
    {-# INLINABLE separate #-}
    fmapEither = IntMap.mapEither
    {-# INLINABLE fmapEither #-}

instance Compactable (Map.Map k) where
    compact = Map.mapMaybe id
    {-# INLINABLE compact #-}
    fmapMaybe = Map.mapMaybe
    {-# INLINABLE fmapMaybe #-}
    filter = Map.filter
    {-# INLINABLE filter #-}
    partition = Map.partition
    {-# INLINABLE partition #-}
    separate = Map.mapEither id
    {-# INLINABLE separate #-}
    fmapEither = Map.mapEither
    {-# INLINABLE fmapEither #-}

instance Compactable Seq.Seq where
    compact = fmap fromJust . Seq.filter isJust
    {-# INLINABLE compact #-}
    separate = altDefaultSeparate
    {-# INLINABLE separate #-}
    filter = Seq.filter
    {-# INLINABLE filter #-}
    partition = Seq.partition
    {-# INLINABLE partition #-}


instance Compactable V.Vector where
    compact = altDefaultCompact
    {-# INLINABLE compact #-}
    separate = altDefaultSeparate
    {-# INLINABLE separate #-}
    filter = V.filter
    {-# INLINABLE filter #-}
    partition = V.partition
    {-# INLINABLE partition #-}

instance Compactable (Const r) where
    compact (Const r) = Const r
    {-# INLINABLE compact #-}
    fmapMaybe _ (Const r) = Const r
    {-# INLINABLE fmapMaybe #-}
    applyMaybe _ (Const r) = Const r
    {-# INLINABLE applyMaybe #-}
    bindMaybe (Const r) _ = Const r
    {-# INLINABLE bindMaybe #-}
    fmapEither _ (Const r) = (Const r, Const r)
    {-# INLINABLE fmapEither #-}
    applyEither _ (Const r) = (Const r, Const r)
    {-# INLINABLE applyEither #-}
    bindEither (Const r) _ = (Const r, Const r)
    {-# INLINABLE bindEither #-}
    filter _ (Const r) = Const r
    {-# INLINABLE filter #-}
    partition _ (Const r) = (Const r, Const r)
    {-# INLINABLE partition #-}

instance Compactable Set.Set where
    compact = Set.fromDistinctAscList . compact . Set.toAscList
    {-# INLINABLE compact #-}
    separate = bimap Set.fromDistinctAscList Set.fromDistinctAscList . separate . Set.toAscList
    {-# INLINABLE separate #-}
    filter = Set.filter
    {-# INLINABLE filter #-}
    partition = Set.partition
    {-# INLINABLE partition #-}

instance (ArrowPlus a, ArrowApply a) => Compactable (ArrowMonad a) where
instance Monad a => Compactable (WrappedMonad a) where
instance Functor a => Compactable (Rec1 a) where
instance Functor a => Compactable (Alt a) where
instance (Functor a, Functor b) => Compactable (a :*: b)
instance Functor f => Compactable (M1 i c f)
instance (Functor f, Functor g) => Compactable (f :.: g)

newtype AltSum f a = AltSum { unAltSum :: f a }
    deriving (Functor, Applicative, Alternative)
#if __GLASGOW_HASKELL__ > 840
instance Alternative f => Monoid (AltSum f a) where
    mempty = empty
    AltSum a `mappend` AltSum b = AltSum (a <|> b)
#else
instance Alternative f => Semigroup (AltSum f a) where
    AltSum a <> AltSum b = AltSum (a <|> b)
instance Alternative f => Monoid (AltSum f a) where
    mappend = (Data.Semigroup.<>)
    mempty = empty
#endif


{-|
class `CompactFold` provides the same methods as `Compactable` but generalized to work on any `Foldable`.

When a type has Alternative (or similar) properties, we can extract the Maybe and the Either, and generalize to Foldable and Bifoldable.

Compactable can always be described in terms of CompactFold, because

  @compact = compactFold@

and

  @separate = separateFold@

as it's just a specialization. More exploration is needed on the relationship here.
-}
class Compactable f => CompactFold (f :: * -> *) where
    compactFold :: Foldable g => f (g a) -> f a
    default compactFold :: (Monad f, Alternative f, Foldable g) => f (g a) -> f a
    compactFold = (>>= mfold')
    {-# INLINEABLE compactFold #-}

    separateFold :: Bifoldable g => f (g a b) -> (f a, f b)
    default separateFold :: (Monad f, Alternative f, Bifoldable g) => f (g a b) -> (f a, f b)
    separateFold xs = (xs >>= mlefts, xs >>= mrights)
    {-# INLINEABLE separateFold #-}

    fmapFold :: (Functor f, Foldable g) => (a -> g b) -> f a -> f b
    fmapFold f = compactFold . fmap f
    {-# INLINABLE fmapFold #-}

    fmapBifold :: (Functor f, Bifoldable g) => (a -> g l r) -> f a -> (f l, f r)
    fmapBifold f = separateFold . fmap f
    {-# INLINABLE fmapBifold #-}

    applyFold :: (Applicative f, Foldable g) => f (a -> g b) -> f a -> f b
    applyFold f = compactFold . (f <*>)
    {-# INLINABLE applyFold #-}

    applyBifold :: (Applicative f, Bifoldable g) => f (a -> g l r) -> f a -> (f l, f r)
    applyBifold fa = separateFold . (fa <*>)
    {-# INLINABLE applyBifold #-}

    bindFold :: (Monad f, Foldable g) => f a -> (a -> f (g b)) -> f b
    bindFold f = compactFold . (f >>=)
    {-# INLINABLE bindFold #-}

    bindBifold :: (Monad f, Bifoldable g) => f a -> (a -> f (g l r)) -> (f l, f r)
    bindBifold f = separateFold . (f >>=)
    {-# INLINABLE bindBifold  #-}

    traverseFold :: (Applicative h, Foldable g, Traversable f) => (a -> h (g b)) -> f a -> h (f b)
    traverseFold f = fmap compactFold . traverse f
    {-# INLINABLE traverseFold #-}

    traverseBifold :: (Applicative h, Bifoldable g, Traversable f) => (a -> h (g l r)) -> f a -> h (f l, f r)
    traverseBifold f = fmap separateFold . traverse f
    {-# INLINABLE traverseBifold #-}


mfold' :: (Foldable f, Alternative m) => f a -> m a
mfold' = unAltSum . foldMap (AltSum . pure)

mlefts :: (Bifoldable f, Alternative m) => f a b -> m a
mlefts = unAltSum . bifoldMap (AltSum . pure) (const mempty)

mrights :: (Bifoldable f, Alternative m) => f a b -> m b
mrights = unAltSum . bifoldMap (const mempty) (AltSum . pure)


instance CompactFold [] where
    compactFold = (>>= F.toList)
    {-# INLINEABLE compactFold #-}

instance CompactFold Maybe
instance CompactFold IO
instance CompactFold ReadP
instance CompactFold ReadPrec
instance CompactFold STM
instance CompactFold ZipList where
  compactFold (ZipList xs) = ZipList $ compactFold xs
  separateFold (ZipList xs) = bimap ZipList ZipList $ separateFold xs
instance CompactFold Option
instance CompactFold U1
instance CompactFold Proxy
instance (ArrowPlus a, ArrowApply a) => CompactFold (ArrowMonad a)
instance MonadPlus a => CompactFold (WrappedMonad a)
instance (Alternative a, Monad a) => CompactFold (Rec1 a)
instance (Alternative a, Monad a) => CompactFold (Alt a)
instance (Alternative f, Monad f, Alternative g, Monad g) => CompactFold (f :*: g)
instance (Compactable f, Alternative f, Monad f, Compactable g, Alternative g, Monad g) => CompactFold (FP.Product f g)
instance (Alternative f, Monad f) => CompactFold (M1 i c f)


fforMaybe :: (Compactable f, Functor f) => f a -> (a -> Maybe b) -> f b
fforMaybe = flip fmapMaybe


fforFold :: (CompactFold f, Functor f, Foldable g) => f a -> (a -> g b) -> f b
fforFold = flip fmapFold


fforEither :: (Compactable f, Functor f) => f a -> (a -> Either l r) -> (f l, f r)
fforEither = flip fmapEither


fforBifold :: (CompactFold f, Functor f, Bifoldable g) => f a -> (a -> g l r) -> (f l, f r)
fforBifold = flip fmapBifold


fmapMaybeM :: (Compactable f, Monad f) => (a -> MaybeT f b) -> f a -> f b
fmapMaybeM f = (>>= compact . runMaybeT . f)


fforMaybeM :: (Compactable f, Monad f) => f a -> (a -> MaybeT f b) -> f b
fforMaybeM = flip fmapMaybeM


fmapEitherM :: (Compactable f, Monad f) => (a -> ExceptT l f r) -> f a -> (f l, f r)
fmapEitherM f x = separate $ runExceptT . f =<< x


fforEitherM :: (Compactable f, Monad f) => f a -> (a -> ExceptT l f r) -> (f l, f r)
fforEitherM = flip fmapEitherM


applyMaybeM :: (Compactable f, Monad f) => f (a -> MaybeT f b) -> f a -> f b
applyMaybeM fa = compact . join . fmap runMaybeT . (fa <*>)


bindMaybeM :: (Compactable f, Monad f) => f a -> (a -> f (MaybeT f b)) -> f b
bindMaybeM x = compact . join . fmap runMaybeT . (x >>=)


traverseMaybeM :: (Monad m, Compactable t, Traversable t) => (a -> MaybeT m b) -> t a -> m (t b)
traverseMaybeM f = unwrapMonad . traverseMaybe (WrapMonad . runMaybeT . f)

-- | While more constrained, when available, this default is going to be faster than the one provided in the typeclass
altDefaultCompact :: (Alternative f, Monad f) => f (Maybe a) -> f a
altDefaultCompact = (>>= maybe empty return)
{-# INLINABLE altDefaultCompact #-}

-- | While more constrained, when available, this default is going to be faster than the one provided in the typeclass
altDefaultSeparate :: (Alternative f, Foldable f) => f (Either l r) -> (f l, f r)
altDefaultSeparate = foldl' (\(l', r') -> \case
  Left l  -> (l' <|> pure l ,r')
  Right r -> (l', r' <|> pure r)) (empty, empty)
{-# INLINABLE altDefaultSeparate #-}


hush :: Either l r -> Maybe r
hush = \case (Right x) -> Just x; _ -> Nothing


flipEither :: Either a b -> Either b a
flipEither = \case (Right x) -> Left x; (Left x) -> Right x