{-# LANGUAGE RankNTypes, GADTs, EmptyDataDecls, MultiParamTypeClasses, TypeFamilies #-}
module Labyrinth where

import Data.BreadthFirstCata.Cata

import Control.Monad
import Control.Monad.State
import Control.Monad.Identity
import Control.Applicative
import Data.Array
import Data.Set hiding (empty)
import qualified Data.Set(empty)


{-
-- Example of finding a shortest simple path through a labyrinth.
-- This is an example of a monad with non-deterministic choice
-- and failure, with specifically as feature a lazy result if the
-- choice is committed.

-- This file is actually about the implementation of this monad
-- using Data.BreadthFirstCata.Cata.

type Lab a = NonT (State LabState) a
data LabState
  = LS { ls_lab   :: Array Pos Bool
       , ls_end   :: Pos
       , ls_cur   :: Pos
       , ls_trail :: Set Pos
       , ls_path  :: (Path -> Path)
       }

data Dir  = North | East | South | West deriving Enum
type Pos  = (Int,Int)
data Path = Forward Dir Path | End

data NonT m a
instance Monad m => Monad (NonT m) where
  return = undefined
  (>>=) = undefined

instance Functor (NonT m)
instance Applicative (NonT m)
instance Alternative (NonT m)
instance MonadState LabState (NonT (State LabState))

search :: Lab ()
search =     finished
        <|> non [ do move d
                     search
                | d <- [North .. West] ]

finished :: Lab ()
finished = do p <- gets ls_cur
              e <- gets ls_end
              if p == e
               then return ()
               else fail "not at the end yet"

non :: [Lab a] -> Lab a
non []     = empty
non (x:xs) = x <|> non xs

move :: Dir -> Lab ()
move d = do p  <- gets ls_cur
            let p' = forward d p
            a  <- gets ls_lab
            if a ! p'
             then do ps <- gets ls_trail
                     if p' `member` ps
                      then fail "already been there"
                      else modify (\s -> s { ls_cur   = p'
                                           , ls_trail = insert p' (ls_trail s)
                                           , ls_path  = Forward d . ls_path s })
             else fail "hit against the wall"

forward :: Dir -> Pos -> Pos
forward North (x,y) = (x,y+1)
forward East  (x,y) = (x+1,y)
forward South (x,y) = (x,y-1)
forward West  (x,y) = (x-1,y)

run :: Lab () -> Path
run m = ls_path sFin End
  where
    sIn = LS { ls_lab   = undefined
             , ls_cur   = (0,0)
             , ls_end   = (100,100)
             , ls_trail = Data.Set.empty
             , ls_path  = id
             }
    sFin = execState (evalNonT m) sIn
    
    evalNonT :: NonT m a -> m a
    evalNonT = undefined
-}




-- Here the actual contents start

data NonT :: (* -> *) -> * -> * where
  Return :: a -> NonT m a
  Bind   :: NonT m b -> (b -> NonT m a) -> NonT m a
  Alt    :: NonT m a -> NonT m a -> NonT m a
  Fail   :: String -> NonT m a
  Work   :: NonT m a -> NonT m a

data Info
  = SomeWork
  | Failed String

data NonAlgT m a
  = NonAlgT
    { alg_return :: a -> Sem Info (NonT m a)
    , alg_bind   :: forall b . Sem Info (NonT m b) -> (b -> Sem Info (NonT m a)) -> Sem Info (NonT m a)
    , alg_alt    :: Sem Info (NonT m a) -> Sem Info (NonT m a) -> Sem Info (NonT m a)
    , alg_fail   :: String -> Sem Info (NonT m a)
    , alg_work   :: Sem Info (NonT m a) -> Sem Info (NonT m a)
    }

foldNonT :: (forall b . NonAlgT m b) -> NonT m a -> Sem Info (NonT m a)
foldNonT alg (Return x) = alg_return alg x
foldNonT alg (Bind m f) = alg_bind alg (foldNonT alg m) (foldNonT alg . f)
foldNonT alg (Alt l r)  = alg_alt alg (foldNonT alg l) (foldNonT alg r)
foldNonT alg (Fail s)   = alg_fail alg s
foldNonT alg (Work m)   = alg_work alg (foldNonT alg m)

nonAlgT :: NonAlgT m a
nonAlgT = NonAlgT sem_Return sem_Bind sem_Alt sem_Fail sem_Work

-- Directly program with the algebra

newtype N m a = N (Sem Info (NonT m a))

instance Monad (N m) where
  return      = N . sem_Return
  (N m) >>= f = N (sem_Bind m ((\(N m') -> m') . f))
  fail        = N . sem_Fail

instance Functor (N m) where
  fmap f n = return f <*> n

instance Applicative (N m) where
  pure x  = return x
  p <*> q = p >>= \f -> (q >>= (return . f))

instance Alternative (N m) where
  empty = N (sem_Fail "empty")
  (N l) <|> (N r) = N (sem_Alt l r)

work :: N m a -> N m a
work (N m) = N (sem_Work m)


-- feed an |a|, finally get a |b|
data Rem m a b where
  Done :: Rem m a a
  More :: (a -> Sem Info (NonT m c)) -> Rem m c b -> Rem m a b

data DictM m where
  DictM :: Monad m => DictM m

data instance Inh (NonT m a) c = Inh { cont :: Rem m a c, dict :: DictM m }
data instance Syn (NonT m a) c = Syn { res :: c } -- transformer part not working yet


-- smart constructor that lifts the contents of a Sem that
-- expect a |Done| as cont-input, to a Sem that allows any
-- cont.
sem' :: CtxSem a Info (NonT m a) -> Sem Info (NonT m a)
sem' f
  = Sem (\inh ->
      let kInh  = Inh { cont = Done, dict = dict inh }
          kComp = invokeWithCtx f kInh
      in resume kComp $ \kSyn ->
           let rInh = Inh { cont = cont inh, dict = dict inh }
               rSem  = sem_Rem (res kSyn)
               rComp = invoke rSem rInh
           in resume rComp $ \rSyn ->
                final $ Syn { res = res rSyn })

sem_Rem :: a -> Sem Info (NonT m a)
sem_Rem v
  = Sem (\inh ->
      case cont inh of
        Done     -> final $ Syn { res = v }
        More f r ->
          let fInh  = Inh { cont = r, dict = dict inh }
              fComp = invoke (f v) fInh
          in resume fComp $ \fSyn ->
               final $ Syn { res = res fSyn })

-- normal cases
sem_Return :: a -> Sem Info (NonT m a)
sem_Return v
  = sem' $ CtxSem $ \inh ->
      final $ Syn { res = v }

sem_Fail :: String -> Sem Info (NonT m a)
sem_Fail s
  = sem' (CtxSem (\inh ->
      resume (inject $ Failed s) $ \_ ->
        final $ Syn { res = error ("error: " ++ s) }))

sem_Bind :: Sem Info (NonT m b) -> (b -> Sem Info (NonT m a)) -> Sem Info (NonT m a)
sem_Bind m f
  = Sem (\inh ->
      let mInh  = Inh { cont = More f (cont inh), dict = dict inh }
          mComp = invoke m mInh
      in resume mComp $ \mSyn ->
           final $ Syn { res = res mSyn })

sem_Work :: Sem Info (NonT m a) -> Sem Info (NonT m a)
sem_Work k
  = Sem (\inh ->
      resume (inject SomeWork) $ \_ ->
        let kInh = Inh { cont = cont inh, dict = dict inh }
            kComp = invoke k kInh
        in resume kComp $ \kSyn ->
             final $ Syn { res = res kSyn })

sem_Alt :: Sem Info (NonT m a) -> Sem Info (NonT m a) -> Sem Info (NonT m a)
sem_Alt l r
  = Sem (\inh ->
      let lInh  = Inh { cont = cont inh, dict = dict inh }
          lComp = invoke l lInh
          
          rInh  = Inh { cont = cont inh, dict = dict inh }
          rComp = invoke r rInh
          
          kComp = best lComp rComp
      in resume kComp $ \kSyn ->
           final $ Syn { res = res kSyn })

best :: Comp Info (NonT m a) s -> Comp Info (NonT m a) s -> Comp Info (NonT m a) s
best l r = best' (oneStep l) (oneStep r)
  where
    best' (Fin synL) _ = final synL
    best' _ (Fin synR) = final synR
    best' (Step infoL remL) (Step infoR remR)
      = case infoR of
          Failed _ -> info infoL remL
          SomeWork -> case infoL of
                        Failed _ -> info infoR remR
                        SomeWork -> info SomeWork (best remL remR)


-- wrapper

evalNonT :: Monad m => N m a -> m a
evalNonT (N m) = return (res syn)
  where inh = Inh { cont = Done, dict = DictM }
        sem = invoke m inh
        syn = lazyEval sem

evalNon :: N Identity a -> a
evalNon = runIdentity . evalNonT


-- test cases. Open in GHCI and run: ex1 .. ex6

m1 :: N Identity String
m1 = return "Hello, World!"
ex1 = evalNon m1 

-- returns False, because the left alternative fails
m2 :: N Identity Bool
m2 = (do fail "bah"
         return True)
    <|> return False
ex2 = evalNon m2

-- returns True, because the right alternative fails
m3 :: N Identity Bool
m3 = return True
    <|> (do fail "bah"
            return False)
ex3 = evalNon m3

-- fail with the right one, because that one did a bit more work
m4 :: N Identity Bool
m4 = fail "left" <|> work (fail "right")
ex4 = evalNon m4

-- returns the list already before evaluating the failure
m5 :: N Identity [Bool]
m5 = do x <- fail "x"
        return [x]
ex5 = evalNon m5

-- picks the left one, because both did an equal amount of work
m6 :: N Identity Bool
m6 = (work $ work $ work $ work $ return False) <|> (work $ work $ work $ work $ return True)
ex6 = evalNon m6