--
-- Support code of the figure
--  "Order algorithm for the visit's rules"
--

{-# LANGUAGE GADTs, RankNTypes, MultiParamTypeClasses #-}
module RankMonad where

import Control.Monad
import Control.Monad.State
import Control.Monad.Error
import Control.Applicative
import Data.Maybe
import Data.Map(Map)
import qualified Data.Map as Map
import Data.Set(Set)
import qualified Data.Set as Set
import Data.List(sortBy)


data Node = Node Int
  deriving (Eq, Ord, Show)

data C = C1 | C2 | C3 | C4 | C5
  deriving (Eq, Ord, Show, Enum, Bounded)

data R a where
  R :: (forall b . ((a -> S -> (Maybe b, S)) -> S -> (Maybe b, S))) -> R a

data S = S { sRank :: Int
           , sRanks :: Map Node Int
           , sClasses :: Map Node (Set C)
           }


instance Functor R where
  fmap f (R h) = R (\c -> h (c . f))

instance Monad R where
  return x = R (\c -> c x)
  R g >>= f = R (\c -> g (\a -> case f a of
                                  R h -> h c
                         ))
  fail _ = R (\_ s -> (Nothing, s))

instance MonadState S R where
  get = R (\c s -> c s s)
  put s = R (\c _ -> c () s)

instance MonadError () R where
  throwError ()   =  fail ""
  catchError m h  =  R (\ c s ->  let  f c  Nothing   s  = runR (h ()) s (g c)
                                       f c  (Just x)  s  = c x s
                                       g c  Nothing   s  = (Nothing, s)
                                       g c  (Just x)  s  = c x s
                                  in runR m s (f c))

runR :: R a -> S -> (Maybe a -> S -> b) -> b
runR (R f) s g = case f term s of (mb, s') -> g mb s'

term :: a -> S -> (Maybe a, S)
term x s = (Just x, s)

execR :: R () -> S
execR m = runR m initS (\(Just _) s -> s)

initS :: S
initS = S 1 Map.empty Map.empty

foreach :: [a] -> R a
foreach xs = R (\c s -> let f c (r0, s0) x = case c x s0 of
                              (Nothing, s') -> (r0, s')
                              (r, s')       -> (r, s')
                        in foldl (f c) (Nothing, s) xs)

guard :: Bool -> R ()
guard g = when (not g) (throwError ())

iter :: R () -> R ()
iter m = (m >> iter m) `orElse` ()

orElse :: R a -> a -> R a
orElse m r = m `catchError` (\_ -> return r)

foreachL :: [R a] -> R [a]
foreachL = foldM f [] where
  f acc m = do  acc' <- singleR m `orElse` []
                return (acc ++ acc')

singleR :: R a -> R [a]
singleR m = m >>= return . return

rank :: Node -> R ()
rank n = modify (\s -> s { sRank = sRank s + 1, sRanks = Map.insert n (sRank s) (sRanks s) } )

hasClass :: Node -> C -> R Bool
hasClass n c = (\mp -> c `Set.member` Map.findWithDefault Set.empty n mp) <$> gets sClasses

isRanked :: Node -> R Bool
isRanked n = Map.member n <$> gets sRanks

sortAsc :: (Node -> Node -> Ordering) -> [Node] -> R [Node]
sortAsc c xs = (\ranks -> sortBy (cmp ranks) xs) <$> gets sRanks where
  cmp ranks l r = compare (Map.lookup l ranks) (Map.lookup r ranks)