    {-# LANGUAGE
      RankNTypes,
      FlexibleInstances,
      FlexibleContexts,
      MultiParamTypeClasses,
      UndecidableInstances,
      FunctionalDependencies #-}

    -- Suppress warnings about ''Control.Monad.Error'' being deprecated.

    {-# OPTIONS_GHC -fno-warn-deprecations #-}

    --------------------------------------------------------------------------------
    -- |
    -- Module      : Data.Equivalence.Monad
    -- Copyright   : Patrick Bahr, 2010
    -- License     : All Rights Reserved
    --
    -- Maintainer  :  Patrick Bahr
    -- Stability   :  unknown
    -- Portability :  unknown
    --
    -- This is an alternative interface to the union-find implementation
    -- in ''Data.Equivalence.STT''. It is wrapped into the monad
    -- transformer 'EquivT'.
    --
    --------------------------------------------------------------------------------

    module Data.Equivalence.Monad
        (
         MonadEquiv(..),
         EquivT(..),
         EquivT',
         EquivM,
         EquivM',
         runEquivT,
         runEquivT',
         runEquivM,
         runEquivM'
         ) where

    import Data.Equivalence.STT hiding (equate, equateAll, equivalent, classDesc, removeClass,
                                        getClass , combine, combineAll, same , desc , remove )
    import qualified Data.Equivalence.STT  as S

    import Control.Monad.Writer
    import Control.Monad.Reader
    import Control.Monad.Error.Class
    import Control.Monad.State
    import Control.Monad.Trans
    import Control.Monad.Identity
    import Control.Monad.ST.Trans
    import Control.Monad.Trans.Error (ErrorT)
    import Control.Monad.Trans.Except (ExceptT)
    import Control.Applicative

    import qualified Control.Monad.Fail as Fail

    {-| This monad transformer encapsulates computations maintaining an
    equivalence relation. A monadic computation of type 'EquivT' @s c v m
    a@ maintains a state space indexed by type @s@, maintains an
    equivalence relation over elements of type @v@ with equivalence class
    descriptors of type @c@ and contains an internal monadic computation
    of type @m a@. -}

    newtype EquivT s c v m a = EquivT {unEquivT :: ReaderT (Equiv s c v) (STT s m) a}

    {-| This monad transformer is a special case of 'EquivT' that only
    maintains trivial equivalence class descriptors of type @()@. -}

    type EquivT' s = EquivT s ()

    {-| This monad encapsulates computations maintaining an equivalence
    relation. A monadic computation of type 'EquivM' @s c v a@ maintains a
    state space indexed by type @s@, maintains an equivalence relation
    over elements of type @v@ with equivalence class descriptors of type
    @c@ and returns a value of type @a@.  -}

    type EquivM s c v = EquivT s c v Identity

    {-| This monad is a special case of 'EquivM' that only maintains
    trivial equivalence class descriptors of type @()@. -}

    type EquivM' s v = EquivM s () v

    instance Functor m => Functor (EquivT s c v m) where
      fmap f (EquivT m) = EquivT $ fmap f m

    instance (Applicative m, Monad m) => Applicative (EquivT s c v m) where
      pure = EquivT . pure
      (EquivT f) <*> (EquivT a) = EquivT (f <*> a)

    instance (Monad m) => Monad (EquivT s c v m) where
        EquivT m >>= f = EquivT (m >>= (unEquivT . f))
        return = EquivT . return

    instance MonadTrans (EquivT s c v) where
        lift = EquivT . lift . lift

    instance Monad m => Fail.MonadFail (EquivT s c v m) where
        fail = error

    instance (MonadReader r m) => MonadReader r (EquivT s c v m) where
        ask = EquivT $ lift ask
        local f (EquivT (ReaderT m)) = EquivT $ ReaderT $ (\ r -> local f (m r))

    instance (Monoid w, MonadWriter w m) => MonadWriter w (EquivT s c v m) where
        tell w = EquivT $ tell w
        listen (EquivT m) = EquivT $ listen m
        pass (EquivT m) = EquivT $ pass m

    instance (MonadState st m) => MonadState st (EquivT s c v m) where
        get = EquivT get
        put s = EquivT $ put s

    instance (MonadError e m) => MonadError e (EquivT s c v m) where
        throwError e = lift $ throwError e
        catchError (EquivT m) f = EquivT $ catchError m (unEquivT . f)

    {-| This function runs a monadic computation that maintains an
    equivalence relation. The first tow arguments specify how to construct
    an equivalence class descriptor for a singleton class and how to
    combine two equivalence class descriptors. -}

    runEquivT :: (Monad m, Applicative m)
              => (v -> c) -- ^ used to construct an equivalence class descriptor for a singleton class
              -> (c -> c -> c) -- ^ used to combine the equivalence class descriptor of two classes
                               --   which are meant to be combined.
              -> (forall s. EquivT s c v m a)
              -> m a
    runEquivT mk com m = runST $ do
      p <- leastEquiv mk com
      (`runReaderT` p) $ unEquivT m

    {-| This function is a special case of 'runEquivT' that only maintains
    trivial equivalence class descriptors of type @()@. -}

    runEquivT' :: (Monad m, Applicative m) => (forall s. EquivT' s v m a) -> m a
    runEquivT' = runEquivT (const ()) (\_ _-> ())

    {-| This function runs a monadic computation that maintains an
    equivalence relation. The first tow arguments specify how to construct
    an equivalence class descriptor for a singleton class and how to
    combine two equivalence class descriptors. -}
    runEquivM :: (v -> c) -- ^ used to construct an equivalence class descriptor for a singleton class
              -> (c -> c -> c) -- ^ used to combine the equivalence class descriptor of two classes
                               --   which are meant to be combined.
              -> (forall s. EquivM s c v a)
              -> a
    runEquivM sing comb m = runIdentity $ runEquivT sing comb m

    {-| This function is a special case of 'runEquivM' that only maintains
    trivial equivalence class descriptors of type @()@. -}

    runEquivM' :: (forall s. EquivM' s v a) -> a
    runEquivM' = runEquivM (const ()) (\_ _ -> ())

    {-| This class specifies the interface for a monadic computation that
    maintains an equivalence relation.  -}

    class (Monad m, Applicative m, Ord v) => MonadEquiv c v d m | m -> v, m -> c, m -> d where
        {-| This function decides whether the two given elements are
            equivalent in the current equivalence relation -}

        equivalent :: v -> v -> m Bool
        {-| This function obtains the descriptor of the given element's
            equivalence class. -}

        classDesc :: v -> m d

        {-| This function equates the element in the given list. That is, it
          unions the equivalence classes of the elements and combines their
          descriptor. -}

        equateAll :: [v] -> m ()

        {-| This function equates the given two elements. That is it
            unions the equivalence classes of the two elements. -}

        equate :: v -> v -> m ()
        equate x y = equateAll [x,y]

        {-| This function removes the equivalence class of the given
          element. If there is no corresponding equivalence class, @False@ is
          returned; otherwise @True@. -}
        removeClass :: v -> m Bool

        {-| This function provides the equivalence class the given element
          is contained in. -}

        getClass :: v -> m c

        {-| This function combines all equivalence classes in the given
          list. Afterwards all elements in the argument list represent the same
          equivalence class! -}

        combineAll :: [c] -> m ()

        {-| This function combines the two given equivalence
          classes. Afterwards both arguments represent the same equivalence
          class! One of it is returned in order to represent the new combined
          equivalence class. -}

        combine :: c -> c -> m c
        combine x y = combineAll [x,y] >> return x

        {-| This function decides whether the two given equivalence classes
          are the same. -}

        (===) :: c -> c -> m Bool

        {-| This function returns the descriptor of the given
          equivalence class. -}

        desc :: c -> m d

        {-| This function removes the given equivalence class. If the
          equivalence class does not exists anymore @False@ is returned;
          otherwise @True@. -}

        remove :: c -> m Bool

    instance (Monad m, Applicative m, Ord v) => MonadEquiv (Class s d v) v d (EquivT s d v m) where
        equivalent x y = EquivT $ do
          part <- ask
          lift $ S.equivalent part x y

        classDesc x = EquivT $ do
          part <- ask
          lift $ S.classDesc part x

        equateAll x = EquivT $ do
          part <- ask
          lift $ S.equateAll part x

        equate x y = EquivT $ do
          part <- ask
          lift $ S.equate part x y

        removeClass x = EquivT $ do
          part <- ask
          lift $ S.removeClass part x

        getClass x = EquivT $ do
          part <- ask
          lift $ S.getClass part x

        combineAll x = EquivT $ do
          part <- ask
          lift $ S.combineAll part x

        combine x y = EquivT $ do
          part <- ask
          lift $ S.combine part x y

        x === y = EquivT $ do
          part <- ask
          lift $ S.same part x y

        desc x = EquivT $ do
          part <- ask
          lift $ S.desc part x

        remove x = EquivT $ do
          part <- ask
          lift $ S.remove part x

    instance (MonadEquiv c v d m, Monoid w) => MonadEquiv c v d (WriterT w m) where
        equivalent x y = lift $ equivalent x y
        classDesc = lift . classDesc
        equateAll x = lift $ equateAll x
        equate x y = lift $ equate x y
        removeClass x = lift $ removeClass x
        getClass x = lift $ getClass x
        combineAll x = lift $ combineAll x
        combine x y = lift $ combine x y
        x === y = lift $ (===) x y
        desc x = lift $ desc x
        remove x = lift $ remove x

    instance (MonadEquiv c v d m, Error e) => MonadEquiv c v d (ErrorT e m) where
        equivalent x y = lift $ equivalent x y
        classDesc = lift . classDesc
        equateAll x = lift $ equateAll x
        equate x y = lift $ equate x y
        removeClass x = lift $ removeClass x
        getClass x = lift $ getClass x
        combineAll x = lift $ combineAll x
        combine x y = lift $ combine x y
        x === y = lift $ (===) x y
        desc x = lift $ desc x
        remove x = lift $ remove x

    instance (MonadEquiv c v d m) => MonadEquiv c v d (ExceptT e m) where
        equivalent x y = lift $ equivalent x y
        classDesc = lift . classDesc
        equateAll x = lift $ equateAll x
        equate x y = lift $ equate x y
        removeClass x = lift $ removeClass x
        getClass x = lift $ getClass x
        combineAll x = lift $ combineAll x
        combine x y = lift $ combine x y
        x === y = lift $ (===) x y
        desc x = lift $ desc x
        remove x = lift $ remove x

    instance (MonadEquiv c v d m) => MonadEquiv c v d (StateT s m) where
        equivalent x y = lift $ equivalent x y
        classDesc = lift . classDesc
        equateAll x = lift $ equateAll x
        equate x y = lift $ equate x y
        removeClass x = lift $ removeClass x
        getClass x = lift $ getClass x
        combineAll x = lift $ combineAll x
        combine x y = lift $ combine x y
        x === y = lift $ (===) x y
        desc x = lift $ desc x
        remove x = lift $ remove x

    instance (MonadEquiv c v d m) => MonadEquiv c v d (ReaderT r m) where
        equivalent x y = lift $ equivalent x y
        classDesc = lift . classDesc
        equateAll x = lift $ equateAll x
        equate x y = lift $ equate x y
        removeClass x = lift $ removeClass x
        getClass x = lift $ getClass x
        combineAll x = lift $ combineAll x
        combine x y = lift $ combine x y
        x === y = lift $ (===) x y
        desc x = lift $ desc x
        remove x = lift $ remove x