{-# LANGUAGE RankNTypes, FlexibleInstances, DeriveFunctor, TypeOperators #-}

module Control.Monad.PlusMonad ((::+)(..), Lifting, Dist(..), leftMap, rightMap, leftSum, rightSum, inl, inr, sym, commute, File, runFile, readLine, mapPlus, refl) where

import qualified Control.Monad.State.Strict as Strict
import Control.Monad.State
import Control.Monad.Reader
import Control.Monad.Writer hiding (Sum)
import Control.Monad.Error
import Control.Monad.List
import Control.Monad.Identity
import Control.Monad.Morph
import Control.Applicative
import Control.Exception
import Data.Functor.Sum
import Data.Functor.Compose
import System.IO

-- | The following construction on two monads is a monad whenever the two monads have extended distributive laws, defined below.
data (m ::+ n) t = ExtendedComposition !(forall x. (Monad x) => Lifting (Sum m n) x (m (n t))) | Rtn t

-- | Lifting along an arbitrary morphism.
type Lifting m x t = (forall u. m u -> x u) -> x t

-- | An extended distributive law allows to permute two layers in a composition.
--
-- Laws are:
--
-- join . T dist = dist . join :: TTS -> TST
--
-- TS join . dist . dist = dist :: TS -> TST
class Dist n where
	dist :: (Applicative m) => n (m t) -> n (m (n t))

instance Dist (StateT s Identity) where
	dist m = do
		n <- m
		s <- get
		return (fmap (\x -> put s >> return x) n)

instance (Monoid s) => Dist (WriterT s Identity) where
	dist m =
		let (n, w) = runWriter m in
			return (fmap (\x -> tell w >> return x) n)

instance Dist [] where
	dist ls = return (sequenceA ls)

-- I/O is equipped with a trivial distributive law.
instance Dist IO where
	dist m = fmap (fmap return) m

instance Dist Identity where
	dist m = fmap (fmap return) m

instance Dist Maybe where
	dist m = fmap (fmap return) m

sumMap :: (m t -> x u) -> (n t -> y u) -> Sum m n t -> Sum x y u
sumMap f _ (InL x) = InL (f x)
sumMap _ g (InR x) = InR (g x)

-- | Left and right maps...
leftMap :: (Monad x) => (forall u. m u -> n u) -> (m ::+ x) t -> (n ::+ x) t
leftMap f (ExtendedComposition g) = ExtendedComposition (\h -> liftM f (g (h . sumMap f id)))
leftMap _ (Rtn x) = Rtn x

rightMap :: (Monad x) => (forall u. m u -> n u) -> (x ::+ m) t -> (x ::+ n) t
rightMap f (ExtendedComposition g) = ExtendedComposition (\h -> liftM (liftM f) (g (h . sumMap id f)))
rightMap _ (Rtn x) = Rtn x

-- | And sums...
leftSum (ExtendedComposition f) (ExtendedComposition g) = ExtendedComposition (\h -> liftM2 mplus (f h) (g h))

rightSum (ExtendedComposition f) (ExtendedComposition g) = ExtendedComposition (\h -> liftM2 (liftM2 mplus) (f h) (g h))

distributive1 x = fmap getCompose (dist (fmap Compose x))

instance (Dist m, Dist n, Monad m, Monad n) => Dist (m ::+ n) where
	dist (ExtendedComposition f) = ExtendedComposition (\g -> (
		fmap (fmap (fmap (fmap (\m -> ExtendedComposition (\_ -> return m)))))
		. fmap (fmap distributive1)
		. fmap distributive1)
		(f g))
	dist (Rtn x) = return (fmap return x)

instance (Functor m, Functor n) => Functor (m ::+ n) where
	fmap f (ExtendedComposition g) = ExtendedComposition (\h -> fmap (fmap (fmap f)) (g h))
	fmap f (Rtn x) = Rtn (f x)

distributive x = (fmap (getCompose . fmap join) . dist . fmap Compose) x

instance (Dist m, Dist n, Monad m, Monad n) => Monad (m ::+ n) where
	return x = ExtendedComposition (\_ -> return (return (return x)))
	ExtendedComposition h >>= f = ExtendedComposition (\g -> (
		join
		. join
		. fmap join
		. fmap (g . InL)
		. fmap (fmap (g . InR))
		. fmap (fmap distributive)
		. fmap distributive
		. fmap (fmap (fmap (\x -> case f x of
			ExtendedComposition i -> i g
			Rtn x -> return (return (return x))))))
		(h g))
	Rtn x >>= f = f x
	fail s = ExtendedComposition (\_ -> fail s)

instance (Dist m, Dist n, Monad m, Monad n) => Applicative (m ::+ n) where
	pure = return
	(<*>) = ap

instance (Dist m, Dist n, MonadPlus m, MonadPlus n) => MonadPlus (m ::+ n) where
	mzero = ExtendedComposition (\_ -> return (return mzero))
	mplus = rightSum

instance (Dist m, Dist n, MonadPlus m, MonadPlus n) => Alternative (m ::+ n) where
	empty = mzero
	(<|>) = mplus

instance (Monad m) => MonadTrans ((::+) m) where
	lift = inr

instance (Monad m) => MFunctor ((::+) m) where
	hoist = rightMap

instance (Dist m, Dist n, Monad m, Monad n, MonadIO n) => MonadIO (m ::+ n) where
	liftIO = inr . liftIO

-- | Injections into the 'ExtendedComposition' type.
inl m = ExtendedComposition (\_ -> return (liftM return m))

inr m = ExtendedComposition (\_ -> return (return m))

symSum (InL x) = x
symSum (InR x) = x

-- | If you have an 'ExtendedComposition' over a monad, you can extract the underlying action.
sym (ExtendedComposition f) = join (join (f symSum))
sym (Rtn x) = return x

mirror (InL x) = InR x
mirror (InR x) = InL x

-- | 'ExtendedComposition' is commutative, provided the distributive laws for 'm' and 'n' witness an isomorphism.
commute :: (Dist m, Monad n) => (m ::+ n) t -> (n ::+ m) t
commute (ExtendedComposition f) = ExtendedComposition (\g -> join (fmap (g . InR . dist) (f (g . mirror))))
commute (Rtn x) = Rtn x

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

-- | Example of an IO-performing ADT.
newtype File t = File (StateT Handle IO t) deriving Functor

runFile (File m) path = do
	hdl <- openFile path ReadMode
	finally (evalStateT m hdl) (hClose hdl)

readLine = File (do
	hdl <- get
	lift (hGetLine hdl))

instance Monad File where
	return = File . return
	File m >>= f = File (m >>= \x -> case f x of File m -> m)
	fail = File . fail

instance Applicative File where
	pure = return
	(<*>) = ap

instance Dist File where
	dist m = do
		n <- m
		s <- File get
		return (fmap (\x -> File (put s) >> return x) n)

mapPlus :: (Monad m, Monad n1) => (forall u. m u -> m1 u) -> (forall u. n u -> n1 u) -> (m ::+ n) t -> (m1 ::+ n1) t
mapPlus f g = leftMap f . rightMap g

refl :: (MonadPlus m) => (m ::+ m) t -> m t
refl = sym