module TraversableT where

import Data.Traversable
import Control.Monad (join,MonadPlus(..))
import Control.Monad.Trans.Class (MonadTrans(..))
import Control.Applicative
import Data.Foldable (Foldable(foldMap), toList)
import Data.Maybe (maybeToList)

-- | 
--
-- /Note:/ this does not yield a monad unless the argument monad is commutative.
newtype TraversableT t m a = TraversableT { runTraversableT :: m (t a) }

-- | Map between 'TraversableT' computations.
--
-- * @'runTraversableT' ('mapTraversableT' f m) = f ('runTraversableT' m)@
mapTraversableT :: (m (t a) -> n (t b)) -> TraversableT t m a -> TraversableT t n b
mapTraversableT f m = TraversableT $ f (runTraversableT m)


instance (Functor m, Functor t) => Functor (TraversableT t m) where
    fmap f = mapTraversableT $ fmap $ fmap f

instance (Foldable f, Foldable t) => Foldable (TraversableT t f) where
    foldMap f (TraversableT a) = foldMap (foldMap f) a

instance (Traversable f, Traversable t) => Traversable (TraversableT t f) where
    traverse f (TraversableT a) = TraversableT <$> traverse (traverse f) a

instance (Applicative m,Applicative t) => Applicative (TraversableT t m) where
    pure a  = TraversableT $ pure (pure a)
    f <*> v = TraversableT $ (<*>) <$> runTraversableT f <*> runTraversableT v

instance (Applicative m, Alternative t) => Alternative (TraversableT t m) where
    empty   = TraversableT $ pure empty
    m <|> n = TraversableT $ (<|>) <$> runTraversableT m <*> runTraversableT n

instance (Monad m, Traversable t, Monad t) => Monad (TraversableT t m) where
    return = TraversableT . return . return
    m >>= k  = TraversableT $ do
        a <- runTraversableT m
        b <- forM a $ runTraversableT . k
        return (join b)
    fail s = TraversableT $ return (fail s)

instance (Monad m, Traversable t, MonadPlus t) => MonadPlus (TraversableT t m) where
    mzero       = TraversableT $ return mzero
    m `mplus` n = TraversableT $ do
        a <- runTraversableT m
        b <- runTraversableT n
        return (a `mplus` b)

instance Applicative t => MonadTrans (TraversableT t) where
    lift m = TraversableT $ do
        a <- m
        return (pure a)

liftT :: Monad m => t a -> TraversableT t m a
liftT = TraversableT . return

liftMT :: m (t a) -> TraversableT t m a
liftMT = TraversableT

-- | Lift a @callCC@ operation to the new monad.
liftCallCC :: Applicative t =>
    ((( (t a) -> m (t b)) -> m (t a)) -> m (t a)) ->
    ((a -> TraversableT t m b) -> TraversableT t m a) -> TraversableT t m a
liftCallCC callCC f = TraversableT $
    callCC $ \c ->
    runTraversableT (f (\a -> TraversableT $ c (pure a)))

-- | Lift a @catchError@ operation to the new monad.
liftCatch :: (m (t a) -> (e -> m (t a)) -> m (t a)) ->
    TraversableT t m a -> (e -> TraversableT t m a) -> TraversableT t m a
liftCatch catchError m h = TraversableT $ runTraversableT m
    `catchError` \e -> runTraversableT (h e)

liftMaybe :: Monad m => Maybe a -> TraversableT [] m a
liftMaybe = liftT . maybeToList

liftIOMaybe :: IO (Maybe a) -> TraversableT [] IO a
liftIOMaybe = liftMT . fmap maybeToList

handleT :: ( Monad m
           , Traversable t ) =>
        (t a -> TraversableT t m a)
        -> TraversableT t m a
        -> TraversableT t m a
handleT catcher body = TraversableT $ do
    tx <- runTraversableT body
    if null (toList tx)
        then runTraversableT $ catcher tx
        else return tx