{-# OPTIONS -fglasgow-exts
            -fno-warn-orphans
  #-}

{-# LANGUAGE ScopedTypeVariables, UndecidableInstances #-}

{- |
Module      :  Control.Monad.ReaderX
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.ReaderX (
    module Control.Monad.ReaderX.Class,
    ReaderX(..),
    runReaderX,
    mapReaderx,
    withReaderx,

    ReaderTX(..),
    runReaderTX,
    mapReaderTX,
    withReaderTX,
    module Control.Monad,
    module Control.Monad.Fix,
    module Control.Monad.Trans,    
    module Control.Monad.Index
    ) where

import Control.Monad
import Control.Monad.Cont.Class
import Control.Monad.Error
import Control.Monad.Fix
import Control.Monad.Instances ()
import Control.Monad.Trans

import Control.Monad.Reader
import Control.Monad.State
import Control.Monad.Writer

import Control.Monad.Index

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

-- ----------------------------------------------------------------------------
-- The partially applied function type is a simple reader monad

instance (Index ix) => MonadReaderX ix r ((->) r) where
    askx (_::ix)       = id
    localx (_::ix) f m = m . f

data (Index ix) =>  ReaderX ix r a = ReaderX ix (r -> a)
runReaderX :: (Index ix) => ix -> ReaderX ix r a -> (r -> a)
runReaderX (_::ix) (ReaderX (_::ix) f) r = f r

--instead of...
-- newtype Reader r a = Reader {runReader :: r -> a}

mapReaderx :: (Index ix) => ix -> (a -> b) -> ReaderX ix r a -> ReaderX ix r b
mapReaderx (ixv::ix) f m = ReaderX ixv $ f . runReaderX ixv m

-- | A more general version of 'local'.

withReaderx :: (Index ix) => ix -> (r' -> r) -> ReaderX ix r a -> ReaderX ix r' a
withReaderx (ixv::ix) f m = ReaderX ixv $ runReaderX ixv m . f

instance (Index ix) => Functor (ReaderX ix r) where
    fmap f m = ReaderX (getVal::ix) $ \r -> f (runReaderX (getVal::ix) m r)

instance (Index ix) => Monad (ReaderX ix r) where
    return a = ReaderX (getVal::ix) $ \_ -> a
    m >>= k  = ReaderX (getVal::ix) $ \r -> runReaderX (getVal::ix) (k (runReaderX (getVal::ix) m r)) r

instance (Index ix) => MonadFix (ReaderX ix r) where
    mfix f = ReaderX (getVal::ix) $ \r -> let a = runReaderX (getVal::ix) (f a) r in a

instance (Index ix) => MonadReaderX ix r (ReaderX ix r) where
    askx (ixv::ix)       = ReaderX ixv id
    localx (ixv::ix) f m = ReaderX ixv $ runReaderX ixv m . f

{- |
The reader monad transformer.
Can be used to add environment reading functionality to other monads.
-}

data (Index ix) => ReaderTX ix r m a = ReaderTX ix (r -> m a)
runReaderTX :: (Index ix) => ix -> ReaderTX ix r m a -> (r -> m a)
runReaderTX (_::ix) (ReaderTX (_::ix) comp) r = comp r
--instead of...  newtype ReaderT r m a = ReaderT { runReaderT :: r -> m a }

mapReaderTX :: (Index ix) => ix -> (m a -> n b) -> ReaderTX ix w m a -> ReaderTX ix w n b
mapReaderTX (ixv::ix) f m = ReaderTX ixv $ f . runReaderTX ixv m

withReaderTX :: (Index ix) => ix -> (r' -> r) -> ReaderTX ix r m a -> ReaderTX ix r' m a
withReaderTX (ixv::ix) f m = ReaderTX ixv $ runReaderTX ixv m . f

instance (Monad m, Index ix) => Functor (ReaderTX ix r m) where
    fmap f m = ReaderTX (getVal::ix) $ \r -> do
        a <- runReaderTX (getVal::ix) m r
        return (f a)

instance (Monad m, Index ix) => Monad (ReaderTX ix r m) where
    return a = ReaderTX (getVal::ix) $ \_ -> return a
    m >>= k  = ReaderTX (getVal::ix) $ \r -> do
        a <- runReaderTX (getVal::ix) m r
        runReaderTX (getVal::ix) (k a) r
    fail msg = ReaderTX (getVal::ix) $ \_ -> fail msg

instance (MonadPlus m, Index ix) => MonadPlus (ReaderTX ix r m) where
    mzero       = ReaderTX (getVal::ix) $ \_ -> mzero
    m `mplus` n = ReaderTX (getVal::ix) $ \r -> runReaderTX (getVal::ix) m r `mplus` runReaderTX (getVal::ix) n r

instance (MonadFix m, Index ix) => MonadFix (ReaderTX ix r m) where
    mfix f = ReaderTX (getVal::ix) $ \r -> mfix $ \a -> runReaderTX (getVal::ix) (f a) r

instance (Monad m, Index ix) => MonadReaderX ix r (ReaderTX ix r m) where
    askx (ixv::ix)       = ReaderTX ixv return
    localx (ixv::ix) f m = ReaderTX ixv $ \r -> runReaderTX ixv m (f r)

-- ---------------------------------------------------------------------------
-- Instances for other mtl transformers

instance (Index ix) => MonadTrans (ReaderTX ix r) where
    lift m = ReaderTX (getVal::ix) $ \_ -> m

instance (MonadIO m, Index ix) => MonadIO (ReaderTX ix r m) where
    liftIO = lift . liftIO

instance (MonadCont m, Index ix) => MonadCont (ReaderTX ix r m) where
    callCC f = ReaderTX (getVal::ix) $ \r ->
        callCC $ \c ->
        runReaderTX (getVal::ix) (f (\a -> ReaderTX (getVal::ix) $ \_ -> c a)) r




-- Error
instance (MonadError e m, Index ix) => MonadError e (ReaderTX ix r m) where
    throwError       = lift . throwError
    m `catchError` h = ReaderTX (getVal::ix) $ \r -> runReaderTX (getVal::ix) m r
        `catchError` \e -> runReaderTX (getVal::ix) (h e) r

instance (Index ix, Error e, MonadReaderX ix r m) => MonadReaderX ix r (ErrorT e m) where
    askx (ixv::ix) = lift $ askx ixv
    localx (ixv::ix) f m = ErrorT $ localx ixv f (runErrorT m)




--ErrorX
instance (MonadErrorX ixe e m, Index ixe, Index ixr) => MonadErrorX ixe e (ReaderTX ixr r m) where
    throwErrorx (ixv::ixe)       = lift . throwErrorx ixv
    catchErrorx (ixv::ixe) m h = ReaderTX (getVal::ixr) $ \r -> 
                          catchErrorx 
                          ixv
                          (runReaderTX (getVal::ixr) m r)
                          (\e -> runReaderTX (getVal::ixr) (h e) r)


{-
instance MonadReaderX ixr r (ErrorTX ixe e m) where
    askx (ixv::ixr) = lift $ askx ixv
    localx (ixv::ixr) f m = ErrorTX (getVal::ixe) $ localx ixv f (runErrorTX (getVal::ixe) m)
-}



--Reader
-- Needs -fallow-undecidable-instances
instance (MonadReader r m, Index ix) => MonadReader r (ReaderTX ix r2 m) where
    ask = ReaderTX (getVal::ix) $ \_ -> ask
    local f (ReaderTX (ixv::ix) comp) = ReaderTX ixv $ \e -> local f (comp e)

instance (Monad m, MonadReaderX ix r1 m, Index ix) => 
    MonadReaderX ix r1 (ReaderT r2 m) where
    askx (ixv::ix) = ReaderT $ \_ -> askx ixv
    localx (ixv::ix) f (ReaderT comp) =  ReaderT $ \e -> localx ixv f (comp e)




--ReaderX
--a transformer should support other-numbered readers.
instance (Index ix1, Index ix2, MonadReaderX ix1 r1 m) => 
    MonadReaderX ix1 r1 (ReaderTX ix2 r2 m) where
   askx (ixv::ix1)  = ReaderTX (getVal::ix2) $ \(_::r2) -> askx ixv
   localx (ixv::ix1) (f::r1->r1) (ReaderTX (_::ix2) comp) = 
                      ReaderTX (getVal::ix2) $ \x -> localx ixv f (comp x)




-- State
-- Needs -fallow-undecidable-instances
instance (Index ix,  MonadState s m) => MonadState s (ReaderTX ix r m) where
    get   = lift $ get
    put s = lift $ put s

instance (Index ix, MonadReaderX ix r m) => MonadReaderX ix r (StateT s m) where
    askx (ixv::ix) = lift $ askx ixv
    localx (ixv::ix) f (StateT g) = StateT  $ \s -> localx ixv f (g s)




-- StateX
-- Needs -fallow-undecidable-instances
instance (Index ixr, Index ixs, MonadStateX ixs s m) => MonadStateX ixs s (ReaderTX ixr r m) where
    getx (ixv::ixs)   = lift $ getx ixv
    putx (ixv::ixs) s = lift $ putx ixv s

{-
instance (Index ixr, Index ixs, MonadReaderX ixr r m) => MonadReaderX ixr r (StateTX ixs s m) where
    askx (ixv::ixr) = lift $ askx ixv
    localx (ixvr::ixr) f (StateTX (ixvs::ixs) g) = StateTX ixvs  $ \s -> localx ixvr f (g s)
-}



-- Writer
-- This instance needs -fallow-undecidable-instances, because it does not satisfy the coverage condition
instance (Index ix, MonadWriter w m) => MonadWriter w (ReaderTX ix r m) where
    tell     = lift  . tell
    listen m = ReaderTX (getVal::ix) $ \w -> listen (runReaderTX (getVal::ix) m w)
    pass   m = ReaderTX (getVal::ix) $ \w -> pass   (runReaderTX (getVal::ix) m w)

instance (Index ix, MonadReaderX ix r m, Monoid w) => MonadReaderX ix r (WriterT w m) where
    askx   (ixv::ix)     = lift $ askx ixv
    localx (ixv::ix) f m = WriterT $ localx ixv f (runWriterT m)




-- WriterX
-- This instance needs -fallow-undecidable-instances, because it does not satisfy the coverage condition
instance (Index ixr, MonadWriterX ixw w m) => MonadWriterX ixw w (ReaderTX ixr r m) where
    tellx   (ixv::ixw)   = lift  . tellx ixv
    listenx (ixv::ixw) m = ReaderTX (getVal::ixr) $ \w -> listenx ixv (runReaderTX (getVal::ixr) m w)
    passx   (ixv::ixw) m = ReaderTX (getVal::ixr) $ \w -> passx   ixv (runReaderTX (getVal::ixr) m w)

{-
instance (Index ixr, MonadReaderX ixr r m, Monoid w) => 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)
-}