{-# LANGUAGE Arrows #-}

module Data.FSM
  (FSM, State, state, content, addTransition,
   fromList, checkStates, Problem, reactMachine,
   reactMachineMult, reactMachineHist, fsmToDot)

where

import FRP.Yampa
import qualified Data.Map as M
import Data.List (nub, (\\), foldl')
import Data.Monoid
import Data.Function (on)


-- *************************************************************************
--
-- Datatypes for States and FSMs
--
-- *************************************************************************

type SId = Int

data Problem t = NonUniqueState SId
                 | MissingTarget SId t
 deriving (Show)

-- State consists of an unique Id, a content describing the state
-- (possibly an enum), functions for generating messages on
-- leaving the current state or entering a new state, resp., and
-- a transitiion table from a given state to its neighboring states
data State a t perception messages = State {
    stateId :: SId,
    content :: a,
    while, onEnter, onExit :: perception -> messages,
    transition :: M.Map t Int
}
instance (Show t, Show a) => Show (State a t p m)
-- where show s = "State " ++ show (content s) ++ ", ("
--                 ++ concatMap (\(t,y) -> ' ' : (show t ++ "->" ++ show y)) (M.toList (transition s))
--                 ++ " )"
 where show s = "{" ++ show (stateId s) ++ "}"

instance Eq (State a t p m)
 where (==) = (==) `on` stateId

-- a FSM is defined as a Map of State Ids to its corresponding State data
type FSM a t p m = M.Map SId (State a t p m)


-- *************************************************************************
--
-- Basic functions for creating States and FSMs and for running transitions
-- on a given FSM
--
-- *************************************************************************

-- Construction of new State; new States instance start with an empty
-- transition table
state :: Int -> a -> (p -> m)-> (p -> m)-> (p -> m)-> State a t p m
state id' a' while' onEnter' onExit' = State id' a' while' onEnter' onExit' M.empty

-- adds a new transition to a given State instance
addTransition :: Ord t => t -> Int -> State a t p m -> State a t p m
addTransition t transStateId s =
    let newTrans = M.insert t transStateId (transition s)
    in s {transition = newTrans }

-- Creates a new FSM from a list of previously generated states, returns
-- either a Right FSM or - in case of duplicate states or transitions to
-- non-existent states - a list of detected problems in the Leftt
fromList :: [State a t p m] -> Either [Problem t] (FSM a t p m)
fromList ss =
   case checkStates ss of
      [] -> Right $ foldl (\m s -> M.insert (stateId s) s m) M.empty ss
      ps -> Left ps

-- Given a FSM, a current state and a transition, the function computes the
-- next state (if transition is applicable) or Nothing (if transition is not
-- applicable); if transition was applicable, the messages from leaving the
-- old state and entering the new state will be collected
runTrans :: (Ord t, Monoid m) => FSM a t p m -> State a t p m -> t -> p -> Maybe (State a t p m, m)
runTrans fsm currentState trans perception = do
    newId <- M.lookup trans (transition currentState)
    let Just newState = M.lookup newId fsm
    return (newState, onExit currentState perception `mappend` onEnter newState perception)


-- Same as runTrans, but takes a list of transistions instead of a single transistion, returns
-- an additional Bool that indicates whether any transition occured
runTransMult :: (Ord t, Monoid m) => FSM a t p m -> State a t p m -> [t] -> p -> (State a t p m, m, Bool)
runTransMult _ currentState [] _ =
   (currentState, mempty, False)
runTransMult fsm currentState (trans:transs) perception =
        case runTrans fsm currentState trans perception of
            Just (s1, m1) -> let (s2, m2, _) = runTransMult fsm s1 transs perception in (s2, m1 `mappend` m2, True)
            Nothing -> let (s2, m2, occ) = runTransMult fsm currentState transs perception in (s2, mempty `mappend` m2, occ)


-- *************************************************************************
--
-- Create dot-output for rendering with graphviz. To do so: Capture the
-- output of fsmToDot in a file, then use dot to convert
--
-- Example of dot-format:
--
-- digraph simple_hierarchy {
--     B [label="The boss"]      // node B
--     E [label="The employee"]  // node E
--     B->E [label="commands", fontcolor=darkgreen] // edge B->E
-- }
-- *************************************************************************

fsmToDot :: (Show a, Show t) => FSM a t p m -> String
fsmToDot fsm =
   "digraph fsm\n{\n" ++
   (foldl (++) "" $ map statesToDot (M.elems fsm)) ++
   "}\n"
   where
       formatTrans source k a result =
           result ++ show source ++ "->" ++ show a ++ " [label=\"" ++ show k ++ "\"]\n"
       statesToDot s =
           let ts = M.map (\dest -> content ((M.!) fsm dest)) $ transition s
               curr = content s
           in
               M.foldrWithKey (formatTrans curr) "" ts -- (M.fromList [(5,"a"), (3,"b")])


-- *************************************************************************
--
-- Sanity checks on a set of FSM states, checks for duplicate states and
-- transitions with missing target states
--
-- *************************************************************************

checkStates :: [State a e p m] -> [Problem e]
checkStates states =
    map NonUniqueState (checkDoubles states) ++
    map (\(i,e,_) -> MissingTarget i e) (checkTransitions states)

checkTransitions :: [State a e p m] -> [(Int, e, Int)]
checkTransitions = foldl' checkProblem [] . allElems

checkProblem :: [(Int, e, Int)] -> (State a e p m, [State a e p m]) -> [(Int, e, Int)]
checkProblem ps (s, ts) = ps ++ missingStates (stateId s) (transition s) ts

missingStates :: Int -> M.Map e Int -> [State a e p m] -> [(Int, e, Int)]
missingStates s trans ss =
    let sIds = map stateId ss
        ts = M.toList trans
    in concatMap (\(e, sId) -> if elem sId sIds then [] else [(s, e, sId)]) ts

allElems :: [a] -> [(a, [a])]
allElems xs =
    let ixs = zip [0..length xs - 1] xs
    in map (\(i,_) -> ((xs!!i), xs)) ixs

checkDoubles :: [State a e p m] -> [Int]
checkDoubles states =
    let ids = map stateId states
    in ids \\ nub ids


-- *************************************************************************
--
-- Reactivity
--
-- *************************************************************************

-- Wrapper function for easy access to runTrans by accumHold; if no transition
-- applies, the wrapper returns the old state and no messages
-- additionaly, the function pushes through a state parameter that will be attached
-- to the new state
runMachine :: (Ord t, Monoid m) =>
    FSM a t (s, p) m -> (State a t (s, p) m, (m, s)) -> ((t, s), p) -> (State a t (s, p) m, (m, s))
runMachine fsm curr ((trans, stateParam), perc) =
  let (curr', (_, currParam)) = curr
  in
    case runTrans fsm curr' trans (stateParam, perc) of
         Nothing  -> (curr', (mempty, currParam))
         Just (newState, messages) -> (newState, (messages, stateParam))

-- Reactive FSM transition: yields the (time-varying) current state and
-- for every transition (defined by an Event containing the transition
-- and the perception for the onExit / onEnter functions) a Monoid of
-- the collected messages from the onExit / onEnter functions
-- additionaly, the function pushes through a state parameter that will be attached
-- to the new state
reactTransition :: (Ord t, Monoid m) =>
    FSM a t (s, p) m -> State a t (s, p) m -> s -> SF (Event ((t, s), p)) (State a t (s, p) m, Event (m, s))
reactTransition fsm init' initParam =
    accumBy (runMachine fsm) (init', (mempty, initParam)) >>>
    (fst ^<< hold (init', (mempty, initParam))) &&& arr (snd . splitE)

-- Yields the (time-varying) current state and the messages (that originate from
-- either the onExit / onEnter-function on a transition or the while-function
-- if no transition occured)
-- The event that yields a state transition (Event (t, s)) consist of the actual
-- transition t that advances the FSM, and a state parameter s that will be attached
-- to the new state. E.g. a event could be "Event (RunTo, Position 10 20)", telling
-- a player in state "Waiting" to change to state "RunTo", and attaching the additional
-- information "Position 10 20" to the new state
reactMachine :: (Ord t, Monoid m) =>
    FSM a t (s, p) m -> State a t (s, p) m -> s -> SF (p, Event (t, s)) ((State a t (s, p) m, s), m)
reactMachine fsm initState initParam = proc (perception, ets) -> do
    (state', result) <- reactTransition fsm initState initParam -< attach ets perception
    param <- hold initParam -< snd (splitE result)
    returnA -< ((state', param), if isEvent result then fst (fromEvent result) else while state' (param, perception))


-- *************************************************************************
--
-- Similar to reactMachine, but puts a whole list of transition through
-- the FSM at a given point in time
--
-- *************************************************************************

rMM :: (Ord t, Monoid m) =>
    FSM a t (s, p) m -> (State a t (s, p) m, s, Event [(t, s)], p) -> (State a t (s, p) m, (m, s))
rMM _ (s0, sp0, Event [], _) = (s0, (mempty, sp0))
rMM fsm (s0, sp0, Event ((t, s):tss), perc) =
    let (s1, (m1, sp1)) = case runTrans fsm s0 t (s, perc) of
                 Nothing -> (s0, (mempty, sp0))
                 Just (s', m') -> (s', (m', s))
        (s2, (m2, sp2)) =  rMM fsm (s1, sp1, Event tss, perc)
    in (s2, (m1 `mappend` m2, sp2))
rMM _ (s0, sp0, _, _) = (s0, (mempty, sp0))

reactMachineMult :: (Ord t, Monoid m, Eq m) =>
    FSM a t (s, p) m -> State a t (s, p) m -> s -> SF (p, Event [(t, s)]) ((State a t (s, p) m, s), m)
reactMachineMult fsm initState initParam = proc (perception, tss) -> do
    rec
       (s2, sp2) <- iPre (initState, initParam) -< (s1, sp1)
       (s1, (ms, sp1)) <- arr (rMM fsm) -< (s2, sp2, tss, perception)
-- Achtung, hier auch die while-Events aufsammeln!!!
    returnA -< ((s1, sp1), ms `mappend` (while s1 (sp1, perception)))

-- instance (Show e) => Show (Event e) where
--  show NoEvent = "NoEvent"
--  show (Event e) = "Event " ++ show e


-- *************************************************************************
--
-- Similar to reactMachineMult, but also provides the state parameter from
-- the previous state to the message generators. (Special case, used by
-- the key parser to determine the duration between keydown and keyup states
--
-- *************************************************************************

runMachineHist :: (Ord t, Monoid m) =>
    FSM a t (s, (p,s)) m -> (State a t (s, (p,s)) m, (m, s)) -> (([t], s), (s,p)) -> (State a t (s, (p,s)) m, (m, s))
runMachineHist fsm (curr, (_,currParam)) ((transs, stateParam), (oldParam, perc)) =
    let (state', messages, transOccured) = runTransMult fsm curr transs (stateParam, (perc, oldParam))
    in (state', (messages, if transOccured then stateParam else currParam))

reactHistTransitions :: (Ord t, Monoid m) =>
    FSM a t (s, (p,s)) m -> State a t (s, (p,s)) m -> s -> SF (Event (([t], s), (s, p))) (State a t (s, (p,s)) m, Event (m, s))
reactHistTransitions fsm init' initParam =
    accumBy (runMachineHist fsm) (init', (mempty, initParam)) >>>
    (fst ^<< hold (init', (mempty, initParam))) &&& arr (snd . splitE)

-- This is a bit more complicated than it's single input counterpart: the onEnter etc. functions
-- are not only provided with the stateParam s and the perception p, but also with the state
-- param of the previous state, so the input to the onEnter etc. functions is (s, (p,s)) and
-- not (s, p) as in reactMachine.
reactMachineHist :: (Ord t, Monoid m)
    => FSM a t (s, (p,s)) m -> State a t (s, (p,s)) m -> s -> SF (p, Event ([t], s)) ((State a t (s, (p,s)) m, s), m)
reactMachineHist fsm initState initParam = proc (perception, ets) -> do
    rec
      (state, result) <- reactHistTransitions fsm initState initParam -< attach ets (oldParam, perception)
      param <- hold initParam -< snd (splitE result)
      oldParam <- iPre initParam -< param
    returnA -< ((state, param), if isEvent result then fst (fromEvent result) else while state (param, (perception, oldParam)))