{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE UndecidableInstances #-}

 -- whether indexes match or not.
{-# LANGUAGE IncoherentInstances #-}
{-# LANGUAGE OverlappingInstances #-}

-- for type coercions to reduce a RWSX constraint to constraints ReaderX, WriterX, StateX.
{-# LANGUAGE TypeFamilies #-}

{-
Module      :  Control.Monad.RWSX.Lazy
Copyright   :  (c) Mark Snyder 2012.
License     :  BSD-style
Maintainer  :  Mark Snyder, msnyde14@gmu.edu
Stability   :  experimental
Portability :  non-portable (multi-param classes, functional dependencies)
-}


module Control.Monad.RWSX.Lazy (
    RWSX(..),
    mkRWSX,
    runRWSX,
    evalRWSX,
    execRWSX,
    mapRWSX,
    withRWSX,
    RWSTX(..),
    mkRWSTX,
    runRWSTX,
    evalRWSTX,
    execRWSTX,
    mapRWSTX,
    withRWSTX,
    module Control.Monad,
    module Control.Monad.Fix,
    module Control.Monad.Trans,
    module Data.Monoid,
    module Control.Monad.RWSX.Class,
  ) where

import Control.Monad
import Control.Monad.Cont.Class
import Control.Monad.Error.Class
import Control.Monad.Fix

import Control.Monad.RWS.Class
import Control.Monad.Reader.Class
import Control.Monad.Writer.Class
import Control.Monad.State.Class
import Control.Monad.Trans
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
import Control.Monad.ReaderX.Class()
import Control.Monad.StateX.Class()
import Control.Monad.WriterX.Class()
import Control.Monad.RWSX.Class()
import Control.Monad.RWSX.Class

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

newtype RWSX ix r w s a = RWSX { runRWSX' :: r -> s -> (a, s, w) }
mkRWSX :: (Index ix) => ix -> (r->s->(a,s,w)) -> RWSX ix r w s a
mkRWSX _ v = RWSX v
runRWSX :: (Index ix) => ix -> RWSX ix r w s a -> r -> s -> (a,s,w)
runRWSX _ m r s = runRWSX' m r s

evalRWSX :: (Index ix) => ix -> RWSX ix r w s a -> r -> s -> (a, w)
evalRWSX (_::ix) m r s = let
    (a, _, w) = runRWSX' m r s
    in (a, w)

execRWSX :: (Index ix) => ix -> RWSX ix r w s a -> r -> s -> (s, w)
execRWSX (_::ix) m r s = let
    (_, s', w) = runRWSX' m r s
    in (s', w)

mapRWSX :: (Index ix) => ix -> ((a, s, w) -> (b, s, w')) -> RWSX ix r w s a -> RWSX ix r w' s b
mapRWSX (ixv::ix) f m = mkRWSX ixv $ \r s -> f (runRWSX' m r s)

withRWSX :: (Index ix) => ix -> (r' -> s -> (r, s)) -> RWSX ix r w s a -> RWSX ix r' w s a
withRWSX (ixv::ix) f m = mkRWSX ixv $ \r s -> uncurry (runRWSX' m) (f r s)

instance (Index ix) => Functor (RWSX ix r w s) where
    fmap f m = mkRWSX (getVal::ix) $ \r s -> let
        (a, s', w) = runRWSX' m r s
        in (f a, s', w)

instance (Monoid w, Index ix) => Monad (RWSX ix r w s) where
    return a = mkRWSX (getVal::ix) $ \_ s -> (a, s, mempty)
    m >>= k  = mkRWSX (getVal::ix) $ \r s -> let
        (a, s',  w)  = runRWSX' m r s
        (b, s'', w') = runRWSX' (k a) r s'
        in (b, s'', w `mappend` w')


instance (Monoid w, Index ix) => MonadFix (RWSX ix r w s) where
    mfix f = mkRWSX (getVal::ix) $ \r s -> let (a, s', w) = runRWSX' (f a) r s in (a, s', w)

-- providing the three monads' behaviors
instance (Monoid w, Index ix) => MonadReaderX ix r (RWSX ix r w s) where
    askx (ixv::ix)       = mkRWSX ixv $ \r s -> (r, s, mempty)
    localx (ixv::ix) f m = mkRWSX ixv $ \r s -> runRWSX' m (f r) s

instance (Monoid w, Index ix) => MonadWriterX ix w (RWSX ix r w s) where
    tellx (ixv::ix)   w = mkRWSX ixv $ \_ s -> ((), s, w)
    listenx (ixv::ix) m = mkRWSX ixv $ \r s -> let
        (a, s', w) = runRWSX' m r s
        in ((a, w), s', w)
    passx (ixv::ix)   m = mkRWSX ixv $ \r s -> let
        ((a, f), s', w) = runRWSX' m r s
        in (a, s', f w)

instance (Monoid w, Index ix) => MonadStateX ix s (RWSX ix r w s) where
    getx (ixv::ix)   = mkRWSX ixv $ \_ s -> (s, s, mempty)
    putx (ixv::ix) s = mkRWSX ixv $ \_ _ -> ((), s, mempty)

instance (Monoid w, Index ix) => MonadRWSX ix r w s (RWSX ix r w s) where

-- ---------------------------------------------------------------------------
-- Our parameterizable RWS monad, with an inner monad

newtype RWSTX ix r w s m a = RWSTX { runRWSTX' :: r -> s -> m (a, s, w) }
mkRWSTX :: (Index ix) => ix -> (r->s->m(a,s,w)) -> RWSTX ix r w s m a
mkRWSTX _ m = RWSTX m
runRWSTX :: (Index ix) => ix -> RWSTX ix r w s m a -> (r -> s -> m (a,s,w))
runRWSTX _ = runRWSTX'

evalRWSTX :: (Monad m, Index ix) => ix -> RWSTX ix r w s m a -> r -> s -> m (a, w)
evalRWSTX (_::ix) m r s = do
    ~(a, _, w) <- runRWSTX' m r s
    return (a, w)

execRWSTX :: (Monad m, Index ix) => ix -> RWSTX ix r w s m a -> r -> s -> m (s, w)
execRWSTX (_::ix) m r s = do
    ~(_, s', w) <- runRWSTX' m r s
    return (s', w)

mapRWSTX :: (Index ix) => ix -> (m (a, s, w) -> n (b, s, w')) -> RWSTX ix r w s m a -> RWSTX ix r w' s n b
mapRWSTX (ixv::ix) f m = mkRWSTX ixv $ \r s -> f (runRWSTX' m r s)

withRWSTX :: (Index ix) => ix -> (r' -> s -> (r, s)) -> RWSTX ix r w s m a -> RWSTX ix r' w s m a
withRWSTX (ixv::ix) f m = mkRWSTX ixv $ \r s -> uncurry (runRWSTX' m) (f r s)

instance (Monad m, Index ix) => Functor (RWSTX ix r w s m) where
    fmap f m = mkRWSTX (getVal::ix) $ \r s -> do
        ~(a, s', w) <- runRWSTX' m r s
        return (f a, s', w)

instance (Monoid w, Monad m, Index ix) => Monad (RWSTX ix r w s m) where
    return a = mkRWSTX (getVal::ix) $ \_ s -> return (a, s, mempty)
    m >>= k  = mkRWSTX (getVal::ix) $ \r s -> do
        ~(a, s', w)  <- runRWSTX' m r s
        ~(b, s'',w') <- runRWSTX' (k a) r s'
        return (b, s'', w `mappend` w')
    fail msg = mkRWSTX (getVal::ix) $ \_ _ -> fail msg

instance (Monoid w, MonadPlus m, Index ix) => MonadPlus (RWSTX ix r w s m) where
    mzero       = mkRWSTX (getVal::ix) $ \_ _ -> mzero
    m `mplus` n = mkRWSTX (getVal::ix) $ \r s -> runRWSTX' m r s `mplus` runRWSTX' n r s

instance (Monoid w, MonadFix m, Index ix) => MonadFix (RWSTX ix r w s m) where
    mfix f = mkRWSTX (getVal::ix) $ \r s -> mfix $ \ ~(a, _, _) -> runRWSTX' (f a) r s


-- providing the three monads' behaviors
-- ReaderX
instance (Monoid w, Monad m, Index ix) => MonadReaderX ix r (RWSTX ix r w s m) where
    askx   (ixv::ix)     = mkRWSTX ixv $ \r s -> return (r, s, mempty)
    localx (ixv::ix) f m = mkRWSTX ixv $ \r s -> runRWSTX' m (f r) s
-- WriterX
instance (Monoid w, Monad m, Index ix) => MonadWriterX ix w (RWSTX ix r w s m) where
    tellx (ixv::ix)   w = mkRWSTX ixv $ \_ s -> return ((),s,w)
    listenx (ixv::ix) m = mkRWSTX ixv $ \r s -> do
        ~(a, s', w) <- runRWSTX' m r s
        return ((a, w), s', w)
    passx (ixv::ix)   m = mkRWSTX ixv $ \r s -> do
        ~((a, f), s', w) <- runRWSTX' m r s
        return (a, s', f w)
-- StateX
instance (Monoid w, Monad m, Index ix) => MonadStateX ix s (RWSTX ix r w s m) where
    getx (ixv::ix)   = mkRWSTX ixv $ \_ s -> return (s, s, mempty)
    putx (ixv::ix) s = mkRWSTX ixv $ \_ _ -> return ((), s, mempty)

-- RWSX same index.
instance (Monad m, Monoid w1, Monoid w2, Index ix
         , r1~r2, s1~s2,w1~w2
         )
    => MonadRWSX ix r1 w1 s1 (RWSTX ix r2 w2 s2 m) where

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


instance (Monoid w, Index ix) => MonadTrans (RWSTX ix r w s) where
    lift m = mkRWSTX (getVal::ix) $ \_ s -> do
        a <- m
        return (a, s, mempty)

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

instance (Monoid w, MonadCont m, Index ix) => MonadCont (RWSTX ix r w s m) where
    callCC f = mkRWSTX (getVal::ix) $ \r s ->
        callCC $ \c ->
        runRWSTX' (f (\a -> mkRWSTX (getVal::ix) $ \_ s' -> c (a, s', mempty))) r s


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

-- Error
instance (Monoid w, MonadError e m, Index ix) => MonadError e (RWSTX ix r w s m) where
    throwError       = lift . throwError
    m `catchError` h = mkRWSTX (getVal::ix) $ \r s -> runRWSTX' m r s
        `catchError` \e -> runRWSTX' (h e) r s

-- ErrorX
instance (Monoid w, MonadErrorX ixe e m, Index ix, Index ixe) => MonadErrorX ixe e (RWSTX ix r w s m) where
    throwErrorx (ixv::ixe)     = lift . throwErrorx ixv
    catchErrorx (ixv::ixe) m h = mkRWSTX (getVal::ix) $ \r s -> 
                                catchErrorx
                                (ixv)
                                (runRWSTX' m r s)
                                (\e -> runRWSTX' (h e) r s)

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

-- interactions with the basic MTL monads is simple...

-- Reader
instance (Monoid w2, Monad m, Index ix2, MonadReader r1 m) => MonadReader r1 (RWSTX ix2 r2 w2 s2 m) where
    ask       = mkRWSTX (getVal::ix2) $ \_{-r2-} s2 -> ask    >>= \r1 -> return (r1, s2, mempty)
    local f m = mkRWSTX (getVal::ix2) $ \r2 s2 -> local  f (runRWSTX' m r2 s2)


-- Writer
instance (Monoid w1, Monoid w2, Monad m, Index ix2, MonadWriter w1 m) => MonadWriter w1 (RWSTX ix2 r2 w2 s2 m) where
    tell   w1 = mkRWSTX (getVal::ix2) $ \_ s2 -> tell w1 >> return ((),s2,mempty)
    listen  m = mkRWSTX (getVal::ix2) $ \r2 s2 -> do
        ~((a,s2',w2'),w1) <- listen (runRWSTX' m r2 s2)
        return ((a,w1),s2',w2')
    pass m = mkRWSTX (getVal::ix2) $ \r2 s2 -> do
        ~(a,b,c) <- runRWSTX' m r2 s2
        ~a' <- pass $ return a
        return (a',b,c)

-- State
instance (Monoid w2, Monad m, Index ix2, MonadState s1 m) => MonadState s1 (RWSTX ix2 r2 w2 s2 m) where
    get    = mkRWSTX (getVal::ix2) $ \_ (s::s2) -> get >>= (\v1 -> return (v1, s, mempty))
    put s1 = mkRWSTX (getVal::ix2) $ \_ (s::s2) -> put s1 >>  return ((), s, mempty)

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

-- RWS
instance (Monoid w1, Monoid w2, Index ix2
         , MonadRWS  r1 w1 s1 m
         )
    => MonadRWS      r1 w1 s1 (RWSTX ix2 r2 w2 s2 m) where









-- interactions with other MTLX monads must be handled with different indexes.

-- RWSX (different index).
instance (Monoid w1, Monoid w2, Index ix1, Index ix2
         ,MonadRWSX ix1 r1 w1 s1 m
         )
    => MonadRWSX ix1 r1 w1 s1 (RWSTX ix2 r2 w2 s2 m) where

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

-- ReaderX, WriterX, and StateX with different indexes.

-- ReaderX
instance (Monoid w2, Monad m, Index ix1, Index ix2, MonadReaderX ix1 r1 m) => MonadReaderX ix1 r1 (RWSTX ix2 r2 w2 s2 m) where
    askx   (_::ix1)     = mkRWSTX (getVal::ix2) $ \_{-r2-} s2 -> askx   (getVal::ix1) >>= \r1 -> return (r1, s2, mempty)
    localx (_::ix1) f m = mkRWSTX (getVal::ix2) $ \r2 s2 -> localx (getVal::ix1) f (runRWSTX' m r2 s2)

-- StateX
instance (Monoid w2, Monad m, Index ix1, Index ix2, MonadStateX ix1 s1 m) => MonadStateX ix1 s1 (RWSTX ix2 r2 w2 s2 m) where
    getx (_::ix1)    = mkRWSTX (getVal::ix2) $ \_ (s::s2) -> getx (getVal::ix1) >>= (\v1 -> return (v1, s, mempty))
    putx (_::ix1) s1 = mkRWSTX (getVal::ix2) $ \_ (s::s2) -> putx (getVal::ix1) s1 >>  return ((), s, mempty)

-- WriterX
instance (Monoid w1, Monoid w2, Monad m, Index ix1, Index ix2, MonadWriterX ix1 w1 m) => MonadWriterX ix1 w1 (RWSTX ix2 r2 w2 s2 m) where
    tellx (_::ix1)  w1 = mkRWSTX (getVal::ix2) $ \_ s2 -> tellx (getVal::ix1) w1 >> return ((),s2,mempty)
    listenx (_::ix1) m = mkRWSTX (getVal::ix2) $ \r2 s2 -> do
        ~((a,s2',w2'),w1) <- listenx (getVal::ix1) (runRWSTX' m r2 s2)
        return ((a,w1),s2',w2')
    passx (_::ix1) m = mkRWSTX (getVal::ix2) $ \r2 s2 -> do
        ~(a,b,c) <- runRWSTX' m r2 s2
        ~a' <- passx (getVal::ix1) $ return a
        return (a',b,c)