{-# LANGUAGE KindSignatures, FlexibleContexts #-}

import Control.Applicative
import Control.Monad
import Control.Monad.Trans.Either

{-|
    An EitherLike is a Kleisli category over the flipped monad, with:

> returnE >|> f = f
> f >|> returnE = f
> (f >|> g) >|> h = f >|> (g >|> h)

    Two functors define the interaction between the ordinary and flipped monad.

    (fmap k .) defines a functor from one Kleisli category of the flipped monad
    to another Kleisli category of the flipped monad:

> forall k . (fmap k .) returnE = returnE
> forall k . (fmap k .) (f >|> g) = (fmap k .) f >|> (fmap k .) g

    (fmapE k .) defines a functor from one Kleisli category of the ordinary
    monad to another Kleisli category of the ordinary monad.

> forall k . (fmapE k .) return = return
> forall k . (fmapE k .) (f >=> g) = (fmapE k .) f >=> (fmapE k .) g

-}
class EitherLike e where
    returnE :: (Monad m) =>  a -> e a m r
    (>|>)   :: (Monad m) => (a -> e b m r) -> (b -> e c m r) -> (a -> e c m r)

instance EitherLike EitherT where
    returnE = left
    f >|> g = \x -> EitherT $ do
        y <- runEitherT (f x)
        runEitherT $ case y of
            Left  z -> g z
            Right r -> right r

throwE :: (EitherLike e, Monad m) => a -> e a m r
throwE = returnE

catchE :: (EitherLike e, Monad m) => e a m r -> (a -> e b m r) -> e b m r
m `catchE` h = ((\() -> m) >|> h) ()

{-|
    Two functors define how an EitherLike may be lifted.

    (liftE .) defines a functor from the Kleisli category of the ordinary lower
    monad to the Kleisli category of the ordinary transformed monad:

> (liftE .) return = return
> (liftE .) (f >=> g) = (liftE .) f >=> (liftE .) g

    (liftE .) defines a functor from the Kleisli category of the flipped lower
    monad to the Kleisli category of the flipped transformed monad:

> (liftE .) returnE = returnE
> (liftE .) (f >|> g) = (liftE .) f >|> (liftE .) g
-}
class EitherTrans t where
    liftE :: (Monad (e l m), EitherLike e) => e l m r -> t e l m r

newtype StateE s e l (m :: * -> *) r = StateE { runStateE :: s -> e l m (r, s) }

instance (Monad (e l m)) => Functor (StateE s e l m) where
    fmap = liftM

instance (Monad (e l m)) => Applicative (StateE s e l m) where
    pure  = return
    (<*>) = ap

instance (Monad (e l m)) => Monad (StateE s e l m) where
    return r = StateE $ \s -> return (r, s)
    m >>= f  = StateE $ \s -> do
        (a, s') <- runStateE m s
        runStateE (f a) s'

instance (EitherLike e) => EitherLike (StateE s e) where
    returnE l = StateE $ \s -> returnE l
    f >|> g = \a -> StateE $ \s ->
                 (((`runStateE` s) . f) >|> ((`runStateE` s) . g)) a

instance EitherTrans (StateE s) where
    liftE m = StateE $ \s -> liftM (\r -> (r, s)) m

get :: (Monad (e l m)) => StateE s e l m s
get = StateE $ \s -> return (s, s)

put :: (Monad (e l m)) => s -> StateE s e l m ()
put s = StateE $ \_ -> return ((), s)

{- Proof that StateE satisfies the EitherLike laws:

-- Note that this is not a required law, but I'm using it as a helper proof
Identity law for fmap = ((`runStateE` s) .), id = returnE:
(`runStateE` s) . returnE = returnE

Proof:
(`runStateE` s) . returnE
= \a -> runStateE (returnE a) s
= \a -> runStateE (StateE $ \s -> returnE a) s
= \a -> (\s -> returnE a) s
= \a -> returnE a
= returnE

Left identity law for (.) = (>|>), id = returnE:
returnE >|> g = g

Proof:
returnE >|> g
= \a -> StateE $ \s -> (((`runStateE` s) . returnE) >|> ((`runStateE` s) . g)) a
= \a -> StateE $ \s -> (returnE >|> ((`runStateE` s) . g)) a
= \a -> StateE $ \s -> ((`runStateE` s) . g) a
= \a -> StateE $ \s -> runStateE (g a) s
= \a -> StateE $ runStateE (g a)
= \a -> g a
= g

Right identity law for (.) = (>|>), id = returnE:
f >|> returnE = f

Proof:
f >|> returnE
= \a -> StateE $ \s -> (((`runStateE` s) . f) >|> ((`runStateE` s) . returnE)) a
= \a -> StateE $ \s -> (((`runStateE` s) . f) >|> returnE) a
= \a -> StateE $ \s -> ((`runStateE` s) . f) a
= \a -> StateE $ \s -> runStateE (f a) s
= \a -> StateE $ runStateE (f a)
= \a -> f a
= f

-- Note that this is not a required law, but I'm using it as a helper proof
Composition law for fmap = ((`runStateE` s) .), (.) = (>|>):
(`runStateE` s) . (f >|> g) = ((`runStateE` s) . f) >|> ((`runStateE` s) . g)

Proof:
(`runStateE` s) . (f >|> g)
= (`runStateE` s) . (\b -> StateE $ \s ->
    ((`runStateE` s) . f) >|> ((`runStateE` s) . g) b )
= \b -> (`runStateE` s) $ StateE $ \s ->
    ((`runStateE` s) . f) >|> ((`runStateE` s) . g) b )
= \b -> ((`runStateE` s) . f) >|> ((`runStateE` s) . g) b
= ((`runStateE` s) . f) >|> ((`runStateE` s) . g)

Associativity law for (.) = (>|>):
f >|> (g >|> h) = (f >|> g) >|> h

Proof:
f >|> (g >|> h)
= \a -> StateE $ \s ->
    ( ((`runStateE` s) . f)
  >|> ((`runStateE` s) . (g >|> h)) ) a
= \a -> StateE $ \s ->
    ( ((`runStateE` s) . f)
  >|> ((`runStateE` s) . (g >|> h)) ) a
= \a -> StateE $ \s ->
      ( ((`runStateE` s) . f)
  >|> ( ((`runStateE` s) . g)
  >|>   ((`runStateE` s) . h) ) ) a
= \a -> StateE $ \s ->
    ( ( ((`runStateE` s) . f)
  >|>   ((`runStateE` s) . g) )
  >|>   ((`runStateE` s) . h) ) a
= \a -> StateE $ \s ->
    ( ((`runStateE` s) . (f >|> g))
  >|> ((`runStateE` s) . h)
= (f >|> g) >|> h

-- STILL MISSING: The functor law proofs
-}

{- Proof that this satisfies the EitherTrans laws:

The monad transformer laws are just the same proofs as for StateT, since it is
implemented identically.

Identity law for fmap = (liftE .) and id = returnE
(liftE .) returnE = returnE

Proof:
liftE . returnE
= \l -> StateE $ \s -> liftM (\r -> (r, s)) (returnE l)
= \l -> StateE $ \s -> returnE l
= returnE

Composition law for fmap = (liftE .) and (.) = (>|>)
(liftE .) f >|> (liftE .) g = (liftE .) (f >|> g)

Proof:
liftE . f >|> liftE . g
=    (\a -> StateE $ \s -> liftM (\r -> (r, s)) (f a))
 >|> (\b -> StateE $ \s -> liftM (\r -> (r, s)) (g b))
= \a -> StateE $ \s ->
   ( ((`runStateE` s) . (\a -> StateE $ \s -> liftM (\r -> (r, s)) (f a)))
 >|> ((`runStateE` s) . (\b -> StateE $ \s -> liftM (\r -> (r, s)) (g b))) ) a
= \a -> StateE $ \s ->
   ( (\a -> liftM (\r -> (r, s)) (f a))
 >|> (\b -> liftM (\r -> (r, s)) (g b)) ) a
= \a -> StateE $ \s ->
   ( (liftM (\r -> (r, s)) . f)
 >|> (liftM (\r -> (r, s)) . g) ) a
= \a -> StateE $ \s -> liftM (\r -> (r, s)) . (f >|> g)
= liftE . (f >|> g)
-}