module Main where

import Prelude

import Control.MonadPlus (class MonadPlus, class Alternative, class Alt, class Plus, (<|>), empty)
import Control.Monad.Trans.Class (class MonadTrans, lift)
import Control.Monad.Reader.Class (class MonadReader, local, class MonadAsk, ask)
import Control.Monad.State.Class (class MonadState, state)
import Control.Monad.State (runStateT)
import Control.Monad.State.Trans (StateT(..))
import Data.Tuple.Nested (type (/\), (/\))
import Data.Maybe (Maybe(..))

class (MonadPlus m) <= MonadLogic m where
    msplit :: forall a. m a -> m (Maybe (a /\ (m a)))
    interleave :: forall a. m a -> m a -> m a

fbind :: forall a b m. MonadLogic m => m a -> (a -> m b) -> m b
fbind m f = do
    r <- msplit m
    case r of
        Nothing -> empty
        Just (a /\ m') -> interleave (f a) (m' `fbind` f)

reflect :: forall a m. MonadLogic m => Maybe (a /\ m a) -> m a
reflect Nothing = empty
reflect (Just (a /\ ma)) = pure a <|> ma

infixl 1 fbind as >>-

type FK :: Type -> Type
type FK ans = ans

type SK :: Type -> Type -> Type
type SK ans a = a -> FK ans -> ans

newtype SFKT :: (Type -> Type) -> Type -> Type
newtype SFKT m a = SFKT (forall ans. SK (m ans) a -> FK (m ans) -> m ans)

unSFKT :: forall a m. SFKT m a -> forall ans. SK (m ans) a -> FK (m ans) -> m ans
unSFKT (SFKT f) = f

instance Functor (SFKT m) where
    map f m = SFKT (\sk -> unSFKT m (\a -> sk (f a)))

instance Apply m => Apply (SFKT m) where
    apply mf ma = SFKT (\sk -> unSFKT mf (\f -> unSFKT ma (\a -> sk (f a))))

instance Applicative m => Applicative (SFKT m) where
    pure a = SFKT (\sk fk -> sk a fk)

instance Bind m => Bind (SFKT m) where
    bind m f = SFKT (\sk -> unSFKT m (\a -> unSFKT (f a) sk))

instance Monad m => Monad (SFKT m)

instance Alt (SFKT m) where
    alt m1 m2 = SFKT (\sk fk -> unSFKT m1 sk (unSFKT m2 sk fk))

instance Plus (SFKT m) where
    empty = SFKT (\_ fk -> fk)

instance Applicative m => Alternative (SFKT m)

instance Monad m => MonadPlus (SFKT m)

instance MonadTrans SFKT where
    lift m = SFKT (\sk fk -> m >>= (\a -> sk a fk))

instance Monad m => MonadLogic (SFKT m) where
    msplit ma = lift (unSFKT ma (\a fk -> pure (Just (a /\ (lift fk >>= reflect)))) (pure Nothing))
    interleave m1 m2 = do
        r <- msplit m1
        case r of
            Nothing -> m2
            Just (a /\ m1') -> pure a <|> interleave m2 m1'

instance MonadAsk r m => MonadAsk r (SFKT m) where
    ask = lift ask

instance MonadReader r m => MonadReader r (SFKT m) where
    local f m = SFKT (\sk -> unSFKT m (\a -> local f <<< sk a))

instance MonadState s m => MonadState s (SFKT m) where
    state f = lift $ state f

instance MonadLogic m => MonadLogic (StateT s m) where
    msplit sm = StateT $ \s -> do
        r <- msplit (runStateT sm s)
        case r of
            Nothing -> pure (Nothing /\ s)
            Just ((a /\ s') /\ m) -> pure (Just (a /\ (StateT $ const m)) /\ s')
    interleave m1 m2 = StateT $ \s -> runStateT m1 s <|> runStateT m2 s