{-# OPTIONS 
  -fglasgow-exts
  #-}
-- Search for -fallow-undecidable-instances to see why this is needed

{-# LANGUAGE ScopedTypeVariables, UndecidableInstances, OverlappingInstances #-}

-----------------------------------------------------------------------------
-- Module      :  Control.Monad.WriterX.Lazy
-- Copyright   :  (c) Mark Snyder 2008.
-- License     :  BSD-style
-- Maintainer  :  Mark Snyder, marks@ittc.ku.edu
-- Stability   :  experimental
-- Portability :  non-portable (multi-param classes, functional dependencies)
-----------------------------------------------------------------------------

module Control.Monad.WriterX.Lazy (
    module Control.Monad.WriterX.Class,
    WriterX(..),
    runWriterX,
    execWriterX,
    mapWriterX,

    WriterTX(..),
    runWriterTX,
    execWriterTX,
    mapWriterTX,
    module Control.Monad,
    module Control.Monad.Fix,
    module Control.Monad.Trans,
    module Data.Monoid,
  ) where

import Control.Monad
import Control.Monad.Cont.Class
import Control.Monad.Error.Class
import Control.Monad.Fix
import Control.Monad.Reader.Class
import Control.Monad.State.Class
import Control.Monad.Trans
import Control.Monad.Writer.Class
import Data.Monoid

import Control.Monad.Index

import Control.Monad.ErrorX.Class
import Control.Monad.ReaderX.Class
import Control.Monad.StateX.Class
import Control.Monad.WriterX.Class

-- ---------------------------------------------------------------------------

data (Index ix) => WriterX ix w a = WriterX ix (a,w)
runWriterX :: (Index ix) => ix -> WriterX ix w a -> (a,w)
runWriterX (_::ix) (WriterX (_::ix) f) = f
--instead of... newtype Writer w a = Writer { runWriter :: (a, w) }

execWriterX :: (Index ix) => ix -> WriterX ix w a -> w
execWriterX (ixv::ix) m = snd (runWriterX ixv m)

mapWriterX :: (Index ix) => ix -> ((a, w) -> (b, w')) -> WriterX ix w a -> WriterX ix w' b
mapWriterX (ixv::ix) f m = WriterX ixv $ f (runWriterX ixv m)

instance (Index ix) => Functor (WriterX ix w) where
    fmap f m = WriterX (getVal::ix) $ let (a, w) = runWriterX (getVal::ix) m in (f a, w)

instance (Monoid w, Index ix) => Monad (WriterX ix w) where
    return a = WriterX (getVal::ix) (a, mempty)
    m >>= k  = WriterX (getVal::ix) $ let
        (a, w)  = runWriterX (getVal::ix) m
        (b, w') = runWriterX (getVal::ix) (k a)
        in (b, w `mappend` w')

instance (Monoid w, Index ix) => MonadFix (WriterX ix w) where
    mfix m = WriterX (getVal::ix) $ let (a, w) = runWriterX (getVal::ix) (m a) in (a, w)

instance (Monoid w, Index ix) => MonadWriterX ix w (WriterX ix w) where
    tellx (ixv::ix)   w = WriterX ixv ((), w)
    listenx (ixv::ix) m = WriterX ixv $ let (a, w)      = runWriterX ixv m in ((a, w), w)
    passx (ixv::ix)   m = WriterX ixv $ let ((a, f), w) = runWriterX ixv m in (a, f w)

-- ---------------------------------------------------------------------------

data (Index ix) => WriterTX ix w m a = WriterTX ix (m (a,w))
runWriterTX :: (Index ix) => ix -> WriterTX ix w m a -> (m (a,w))
runWriterTX (_::ix) (WriterTX (_::ix) f) = f
--instead of... newtype WriterT w m a = WriterT { runWriterT :: m (a, w) }

execWriterTX :: (Index ix, Monad m) => ix -> WriterTX ix w m a -> m w
execWriterTX (ixv::ix) m = do
    ~(_, w) <- runWriterTX ixv m
    return w

mapWriterTX :: (Index ix) => ix -> (m (a, w) -> n (b, w')) -> WriterTX ix w m a -> WriterTX ix w' n b
mapWriterTX (ixv::ix) f m = WriterTX ixv $ f (runWriterTX ixv m)

instance (Monad m, Index ix) => Functor (WriterTX ix w m) where
    fmap f m = WriterTX (getVal::ix) $ do
        ~(a, w) <- runWriterTX (getVal::ix) m
        return (f a, w)

instance (Monoid w, Monad m, Index ix) => Monad (WriterTX ix w m) where
    return a = WriterTX (getVal::ix) $ return (a, mempty)
    m >>= k  = WriterTX (getVal::ix) $ do
        ~(a, w)  <- runWriterTX (getVal::ix) m
        ~(b, w') <- runWriterTX (getVal::ix) (k a)
        return (b, w `mappend` w')
    fail msg = WriterTX (getVal::ix) $ fail msg

instance (Monoid w, MonadPlus m, Index ix) => MonadPlus (WriterTX ix w m) where
    mzero       = WriterTX (getVal::ix) mzero
    m `mplus` n = WriterTX (getVal::ix) $ runWriterTX (getVal::ix) m `mplus` runWriterTX (getVal::ix) n

instance (Monoid w, MonadFix m, Index ix) => MonadFix (WriterTX ix w m) where
    mfix m = WriterTX (getVal::ix) $ mfix $ \ ~(a, _) -> runWriterTX (getVal::ix) (m a)

instance (Monoid w, Monad m, Index ix) => MonadWriterX ix w (WriterTX ix w m) where
    tellx (ixv::ix)   w = WriterTX ixv $ return ((), w)
    listenx (ixv::ix) m = WriterTX ixv $ do
        ~(a, w) <- runWriterTX ixv m
        return ((a, w), w)
    passx (ixv::ix)   m = WriterTX ixv $ do
        ~((a, f), w) <- runWriterTX ixv m
        return (a, f w)

-- ---------------------------------------------------------------------------
instance (Monoid w, Index ix) => MonadTrans (WriterTX ix w) where
    lift m = WriterTX (getVal::ix) $ do
        a <- m
        return (a, mempty)

instance (Monoid w, MonadIO m, Index ix) => MonadIO (WriterTX ix w m) where
    liftIO = lift . liftIO

instance (Monoid w, MonadCont m, Index ix) => MonadCont (WriterTX ix w m) where
    callCC f = WriterTX (getVal::ix) $
        callCC $ \c ->
        runWriterTX (getVal::ix) (f (\a -> WriterTX (getVal::ix) $ c (a, mempty)))




-- Error
instance (Monoid w, MonadError e m, Index ix) => MonadError e (WriterTX ix w m) where
    throwError       = lift . throwError
    m `catchError` h = WriterTX (getVal::ix) $ runWriterTX (getVal::ix) m
        `catchError` \e -> runWriterTX (getVal::ix) (h e)




-- ErrorX
instance (Monoid w, Index ixe, Index ixw, MonadErrorX ixe e m) => MonadErrorX ixe e (WriterTX ixw w m) where
    throwErrorx (ixv::ixe)       = lift . throwErrorx ixv
    catchErrorx (ixv::ixe) m h = WriterTX (getVal::ixw) $ 
                    catchErrorx
                    (ixv::ixe)
                    (runWriterTX (getVal::ixw) m)
                    (\e -> runWriterTX (getVal::ixw) (h e))




--Reader
-- This instance needs -fallow-undecidable-instances, because -- it does not satisfy the coverage condition
instance (Monoid w, MonadReader r m, Index ixw) => MonadReader r (WriterTX ixw w m) where
    ask       = lift ask
    local f m = WriterTX (getVal::ixw) $ local f (runWriterTX (getVal::ixw) m)




--ReaderX
-- This instance needs -fallow-undecidable-instances, because -- it does not satisfy the coverage condition
instance (Monoid w, Index ixr, Index ixw, MonadReaderX ixr r m) => MonadReaderX ixr r (WriterTX ixw w m) where
    askx (ixv::ixr)       = lift $ askx ixv
    localx (ixv::ixr) f m = WriterTX (getVal::ixw) $ localx ixv f (runWriterTX (getVal::ixw) m)




--State
-- Needs -fallow-undecidable-instances
instance (Monoid w, MonadState s m, Index ixw) => MonadState s (WriterTX ixw w m) where
    get = lift get
    put = lift . put




--StateX
-- Needs -fallow-undecidable-instances
instance (Monoid w, Index ixw, MonadStateX ixs s m) => MonadStateX ixs s (WriterTX ixw w m) where
    getx (ixv::ixs) = lift $ getx ixv
    putx (ixv::ixs) = lift . putx ixv




-- Writer
instance (Index ixw2, Monoid w, Monoid s, MonadWriter w m) => MonadWriter w (WriterTX ixw2 s m) where
   tell     = lift . tell
   listen m = WriterTX (getVal::ixw2) $ do
       ~((a,s'),w) <- listen (runWriterTX (getVal::ixw2) m)
       return ((a,w),s')
   pass   m = WriterTX (getVal::ixw2) $ pass $ do
       ~((a,f),s') <- runWriterTX (getVal::ixw2) m
       return ((a,s'),f)



-- WriterX
instance (Index ixw1, Index ixw2, Monoid w1, Monoid w2, MonadWriterX ixw1 w1 m)
    => MonadWriterX ixw1 w1 (WriterTX ixw2 w2 m) where
   tellx (ixv::ixw1)     = lift . tellx ixv
   listenx (ixv::ixw1) m = WriterTX (getVal::ixw2) $ do
       ~((a,s'),w) <- listenx ixv (runWriterTX (getVal::ixw2) m)
       return ((a,w),s')
   passx (ixv::ixw1)   m = WriterTX (getVal::ixw2) $ passx ixv $ do
       ~((a,f),s') <- runWriterTX (getVal::ixw2) m
       return ((a,s'),f)