{-# LANGUAGE CPP                        #-}
{-# LANGUAGE ConstrainedClassMethods    #-}
{-# LANGUAGE DefaultSignatures          #-}
{-# LANGUAGE DeriveFoldable             #-}
{-# LANGUAGE DeriveFunctor              #-}
{-# LANGUAGE DeriveGeneric              #-}
{-# LANGUAGE DerivingVia                #-}
{-# LANGUAGE FlexibleContexts           #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE KindSignatures             #-}
{-# LANGUAGE LambdaCase                 #-}
{-# LANGUAGE ScopedTypeVariables        #-}
{-# LANGUAGE StandaloneDeriving         #-}
{-# LANGUAGE TypeOperators              #-}


module Control.Functor.Compactable
  (
  -- * Compact
    Compactable (..)
  , separate
  -- * Handly flips
  , fforMaybe
  , fforThese
  -- * More general lefts and rights
  , mfold'
  , mlefts
  , mrights
  -- * Monad Transformer utils
  , mapMaybeM
  , mapTheseM
  , fforMaybeM
  , fforTheseM
  , applyMaybeM
  , bindMaybeM
  , traverseMaybeM
  -- * Alternative Defaults
  , altDefaultCompact
  , altDefaultSeparate

  ) where


import           Control.Applicative         (Alternative (empty, (<|>)),
                                              Const (Const), WrappedMonad (..),
                                              ZipList (ZipList))
import           Control.Monad               (join, (<=<))
import           Control.Monad.Trans.Except  (ExceptT, runExceptT)
import           Control.Monad.Trans.Maybe   (MaybeT (runMaybeT))
import           Data.Bifunctor              (bimap)
import           Data.Foldable               as F (foldl', foldr')
import           Data.Functor.Compose        (Compose (Compose))
import           Data.Functor.Contravariant  (Contravariant (contramap))
import qualified Data.Functor.Product        as FP
import           Data.Kind                   (Type)
import qualified Data.List                   as List
import qualified Data.Map                    as Map
import qualified Data.Maybe                  as May
import           Data.Monoid                 (Alt (Alt))
import           Data.Proxy                  (Proxy (..))
import qualified Data.Sequence               as Seq
import qualified Data.Set                    as Set
import qualified Data.Vector                 as V
import           GHC.Conc                    (STM)
import           GHC.Generics                (M1 (M1), Rec1 (Rec1), U1 (U1),
                                              type (:*:) ((:*:)),
                                              type (:.:) (Comp1))

import           Data.These                  (These (..), these)

import           Control.Functor.Dichotomous (Dichotomous (dichotomy), hushLeft,
                                              hushRight, mfold', mlefts,
                                              mrights)
import qualified Data.IntMap                 as IntMap
#if __GLASGOW_HASKELL__ < 900
import           Data.Semigroup              (Option (Option))
#endif


separateMap :: (Dichotomous g, Functor f, Compactable f) => (a -> g l r) -> f a -> (f l, f r)
separateMap f = separate . mapMaybe (dichotomy . f)
{-# INLINABLE separateMap #-}


separate :: (Dichotomous g, Functor f, Compactable f) => f (g l r) -> (f l, f r)
separate = separateThese . mapMaybe dichotomy
{-# INLINABLE separate #-}


-- | A generalization of catMaybes
--
-- prop> compact . map Just = id
-- prop> compact . mapMaybe id
-- prop> compact (pure Just <*> a) = a
-- prop> applyMaybe (pure Just) = id
-- prop> applyMaybe (pure id) = compact
-- prop> bindMaybe (return . Just) = id
-- prop> bindMaybe return = compact
-- prop> compact (return . Just =<< a) = a
-- prop> mapMaybe (l <=< r) = mapMaybe l . mapMaybe r
-- prop> compact (Nothing <$ a) = empty
-- prop> compact (Nothing <$ a) = mempty
-- prop> compact empty = empty
-- prop> compact mempty = mempty
-- prop> traverseMaybe (Just . Just) = Just
-- prop> traverseMaybe (map Just . f) = traverse f
class Compactable (f :: Type -> Type) where
  {-# MINIMAL compact | separateThese #-}

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

  separateThese :: f (These l r) -> (f l, f r)
  default separateThese :: Functor f => f (These l r) -> (f l, f r)
  separateThese xs = (compact $ hushRight <$> xs, compact $ hushLeft <$> xs)
  {-# INLINABLE separateThese #-}

  filter :: (a -> Bool) -> f a -> f a
  default filter :: Functor f => (a -> Bool) -> f a -> f a
  filter f = mapMaybe $ \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 = separateMap $ \a -> if f a then Right a else Left a
  {-# INLINEABLE partition #-}

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

  contramapMaybe :: Contravariant f => (Maybe b -> a) -> f a -> f b
  contramapMaybe f = compact . contramap f
  {-# INLINABLE contramapMaybe #-}

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

  contramapThese :: Contravariant f => (These l r -> a) -> f a -> (f l, f r)
  contramapThese f = separateThese . contramap f
  {-# INLINEABLE contramapThese #-}

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

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

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

  bindThese :: Monad f => (a -> f (These l r)) -> f a -> (f l, f r)
  bindThese f x = separate $ x >>= f
  {-# INLINABLE bindThese #-}

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

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


instance Compactable Maybe where
    compact = join
    {-# INLINABLE compact #-}
    mapMaybe = (=<<)
    {-# INLINABLE mapMaybe #-}
    separateThese = \case
      Just x -> case x of
        This  l   -> (Just l, Nothing)
        That    r -> (Nothing, Just r)
        These l r -> (Just l, Just r)
      _ -> (Nothing, Nothing)
    {-# INLINABLE separateThese #-}


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

instance Monoid m => Compactable (These m) where
  compact = \case
    This x           -> This x
    That (Just x)    -> That x
    That Nothing     -> This mempty
    These x (Just y) -> These x y
    These x Nothing  -> This x
  {-# INLINABLE compact #-}

instance Compactable [] where
    compact = May.catMaybes
    {-# INLINABLE compact #-}
    mapMaybe _ []    = []
    mapMaybe f (h:t) = maybe (mapMaybe f t) (: mapMaybe f t) (f h)
    {-# INLINABLE mapMaybe #-}
    filter = Prelude.filter
    {-# INLINABLE filter #-}
    partition = List.partition
    {-# INLINABLE partition #-}
    separateThese = foldr (these l_ r_ lr_) ([],[])
      where
      l_  a   ~(l, r) = (a:l,   r)
      r_    b ~(l, r) = (  l, b:r)
      lr_ a b ~(l, r) = (a:l, b:r)
    {-# INLINABLE separateThese #-}
    mapThese f = foldr' deal ([],[])
      where deal a ~(bs, cs) = case f a of
             This  b   -> (b:bs,   cs)
             That    c -> (  bs, c:cs)
             These b c -> (b:bs, c:cs)
    {-# INLINABLE mapThese #-}
    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 x = x >>= maybe (error "compact called on (x :: IO (Maybe _)) where x = return Nothing") return
    {-# NOINLINE compact #-}

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

instance Compactable Proxy where
    compact _ = Proxy
    {-# INLINABLE compact #-}
    separateThese _ = (Proxy, Proxy)
    {-# INLINABLE separateThese #-}
    filter _ _ = Proxy
    {-# INLINABLE filter #-}
    partition _ _ = (Proxy, Proxy)
    {-# INLINABLE partition #-}
    mapMaybe _ _ = Proxy
    {-# INLINABLE mapMaybe #-}
    applyMaybe _ _ = Proxy
    {-# INLINABLE applyMaybe #-}
    bindMaybe _ _ = Proxy
    {-# INLINABLE bindMaybe #-}
    mapThese _ _ = (Proxy, Proxy)
    {-# INLINABLE mapThese #-}
    applyThese _ _ = (Proxy, Proxy)
    {-# INLINABLE applyThese #-}
    bindThese _ _ = (Proxy, Proxy)
    {-# INLINABLE bindThese #-}


instance Compactable U1
  where compact U1 = U1


#if __GLASGOW_HASKELL__ < 900
instance Compactable Option where
    compact (Option x) = Option (join x)
    {-# INLINABLE compact #-}
    mapMaybe f (Option (Just x)) = Option (f x)
    mapMaybe _ _                 = Option Nothing
    {-# INLINABLE mapMaybe #-}
    separateThese = altDefaultSeparate
    {-# INLINABLE separateThese #-}
#endif


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 f, Compactable g)
         => Compactable (Compose f g) where
    compact (Compose fg) = Compose $ compact <$> fg
    {-# INLINABLE compact #-}


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

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

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

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

instance Compactable (Const r) where
    compact (Const r) = Const r
    {-# INLINABLE compact #-}
    mapMaybe _ (Const r) = Const r
    {-# INLINABLE mapMaybe #-}
    applyMaybe _ (Const r) = Const r
    {-# INLINABLE applyMaybe #-}
    bindMaybe _ (Const r) = Const r
    {-# INLINABLE bindMaybe #-}
    mapThese _ (Const r) = (Const r, Const r)
    {-# INLINABLE mapThese #-}
    applyThese _ (Const r) = (Const r, Const r)
    {-# INLINABLE applyThese #-}
    bindThese _ (Const r) = (Const r, Const r)
    {-# INLINABLE bindThese #-}
    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 #-}
    separateThese = bimap Set.fromDistinctAscList Set.fromDistinctAscList . separate . Set.toAscList
    {-# INLINABLE separateThese #-}
    filter = Set.filter
    {-# INLINABLE filter #-}
    partition = Set.partition
    {-# INLINABLE partition #-}


instance (Compactable a, Monad a) => Compactable (WrappedMonad a)
  where compact (WrapMonad x) = WrapMonad $ compact x

instance (Compactable a, Functor a) => Compactable (Rec1 a)
  where compact (Rec1 x) = Rec1 $ compact x

instance (Compactable a, Functor a) => Compactable (Alt a)
  where compact (Alt a) = Alt $ compact a

instance (Compactable a, Functor a, Compactable b, Functor b) => Compactable (a :*: b)
  where compact (a :*: b) = compact a :*: compact b

instance (Compactable f, Functor f) => Compactable (M1 i c f)
  where compact (M1 x) = M1 $ compact x

instance (Functor f, Compactable g, Functor g) => Compactable (f :.: g)
  where compact (Comp1 x) = Comp1 $ compact <$> x


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


fforThese :: (Compactable f, Functor f) => f a -> (a -> These l r) -> (f l, f r)
fforThese = flip mapThese


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


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


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


fforTheseM :: (Compactable f, Monad f) => f a -> (a -> ExceptT l f r) -> (f l, f r)
fforTheseM = flip mapTheseM


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


bindMaybeM :: (Compactable f, Monad f) => f a -> (a -> f (MaybeT f b)) -> f b
bindMaybeM x = compact . 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 :: (Dichotomous d, Alternative f, Foldable f) => f (d l r) -> (f l, f r)
altDefaultSeparate = foldl' (\(l', r') d -> case dichotomy d of
  Nothing          -> (l', r')
  Just (This l)    -> (l' <|> pure l ,r')
  Just (That r)    -> (l', r' <|> pure r)
  Just (These l r) -> (l' <|> pure l, r' <|> pure r)) (empty, empty)
{-# INLINABLE altDefaultSeparate #-}