implementation module DTMC

import StdArray
import StdBool
import StdClass
import StdFile
from StdFunc import o
import StdInt
import StdList
import StdString
import StdTuple

import Data.Error
from Data.Func import $, mapSt, seqSt
import Data.List
import qualified Data.Map as M
from Data.Map import :: Map, alter
import Data.Maybe
import System.File
import System.FilePath
import Text

paths :: !*DTMC -> *([[Int]], !*DTMC)
paths dtmc=:{nr_states,states}
# (inits,states) = getInitStates (nr_states-1) states
# (ps,states) = mapSt (paths` []) inits states
= (flatten ps, {dtmc & states=states})
where
	getInitStates :: !Int !*{Maybe State} -> *(![Int], !*{Maybe State})
	getInitStates i ss
	| i < 0 = ([], ss)
	# (s,ss) = ss![i]
	| isNothing s = ([], ss)
	# s = fromJust s
	| s.init
		# (inits,ss) = getInitStates (i-1) ss
		= ([i:inits],ss)
	= getInitStates (i-1) ss

	paths` :: ![Int] !Int !*{Maybe State} -> *(![[Int]], !*{Maybe State})
	paths` seen cur ss
	| isMember cur seen = ([], ss)
	# (s,ss) = ss![cur]
	| isNothing s = ([], ss)
	# s = fromJust s
	| isAbsorbing s = ([[cur]], ss)
	# succ = 'M'.keys s.transitions
	# (pss,ss) = mapSt (paths` [cur:seen]) succ ss
	= ([[cur:p] \\ ps <- pss, p <- ps], ss)

stateElimination :: !*DTMC -> *DTMC
stateElimination dtmc=:{nr_states}
# (ps,dtmc) = paths dtmc
= case [p \\ p=:[_:_:_:_] <- ps] of
	[]    -> {dtmc & states=prune (nr_states-1) dtmc.states}
	[p:_] -> stateElimination {dtmc & states=uncurry elim (penultimate_trans p) dtmc.states}
		with
			penultimate_trans [x,y,_] = (x,y)
			penultimate_trans [_:xs] = penultimate_trans xs
where
	prune :: !Int !*{Maybe State} -> *{Maybe State}
	prune i ss
	| i < 0 = ss
	# (Just s,ss) = ss![i]
	| s.init || isAbsorbing s = prune (i-1) ss
	| otherwise               = prune (i-1) $ remove i (nr_states-1) ss
	where
		remove :: !Int !Int !*{Maybe State} -> *{Maybe State}
		remove i j ss
		| j < 0 = {ss & [i]=Nothing}
		# (s,ss) = ss![j]
		| isNothing s = remove i (j-1) ss
		# s = fromJust s
		= remove i (j-1) ss

	elim :: !Int !Int !*{Maybe State} -> *{Maybe State}
	elim s1 s ss
	# (Just s1s,ss) = ss![s1]
	# prob = 'M'.get s s1s.transitions

	# (succ,ss) = ss![s]
	| isNothing succ = ss
	# succ = fromJust succ
	| succ.init || isAbsorbing succ = ss
	# succ = succ.transitions

	# self_prob = fromMaybe "0" $ 'M'.get s succ
	# succ = 'M'.toList succ
	| s1 == s
		# ss = seqSt (uncurry $ update_eq self_prob) succ ss
		# (Just s1s,ss) = ss![s1]
		# ss & [s1] = Just {s1s & transitions='M'.del s s1s.transitions}
		= ss
	| otherwise
		# ss = seqSt (uncurry $ update_ne self_prob) succ ss
		# (Just s1s,ss) = ss![s1]
		# ss & [s1] = Just {s1s & transitions='M'.del s s1s.transitions}
		= ss
	where
		update_eq :: !String !Int !String !*{Maybe State} -> *{Maybe State}
		update_eq self_prob s2 prob2 ss
		| s == s2 = ss
		# (Just s1s,ss) = ss![s1]
		# ss & [s1] = Just
			{ s1s
			& transitions = 'M'.put s2
				("(" +++ prob2 +++ ")/(1-(" +++ self_prob +++ "))")
				s1s.transitions
			}
		= ss

		update_ne :: !String !Int !String !*{Maybe State} -> *{Maybe State}
		update_ne self_prob s2 prob2 ss
		| s == s2 = ss
		# (Just s1s,ss) = ss![s1]
		# ps1s2 = fromMaybe "0" $ 'M'.get s2 s1s.transitions
		# ps1s = fromMaybe "0" $ 'M'.get s s1s.transitions
		# ss & [s1] = Just
			{ s1s
			& transitions = 'M'.put s2
				(ps1s2 +++ "+(" +++ ps1s +++ ")*(" +++ prob2 +++ ")/(1-(" +++ self_prob +++ "))")
				s1s.transitions
			}
		= ss

isAbsorbing :: !State -> Bool
isAbsorbing s = all ((==) s.state_id) [to \\ (to,_) <- 'M'.toList s.transitions]

parseDTMC :: !FilePath !*World -> *(!DTMC, !*World)
parseDTMC fp w
# (lines,w) = readFileLines fp w
= (parseDTMCFromLines [s % (0,size s-2) \\ s <- fromOk lines], w)

parseDTMCFromLines :: ![String] -> DTMC
parseDTMCFromLines lines
# [nr_states:_:lines] = tl $ dropWhile ((<>) "@nr_states") lines
# states = map parseState $ groupLines "\t" [] lines
=
	{ nr_states = toInt nr_states
	, states    = {Just s \\ s <- states}
	}
where
	parseState :: ![String] -> State
	parseState [head:actions] =
		{ state_id    = toInt $ trim $ head % (5,size head-1)
		, transitions = 'M'.unions $ map (snd o parseAction) $ groupLines "\t\t" [] actions
		, init        = endsWith "init" head
		}
	where
		parseAction :: ![String] -> (!Int, !Map Int String)
		parseAction [head:transitions] =
			( toInt (head % (8, size head-1))
			, 'M'.fromList $ map parseTransition transitions
			)
		where
			parseTransition :: !String -> (!Int, !String)
			parseTransition s = case split " : " (trim s) of
				[id:prob:_] -> (toInt id, prob)

	groupLines :: !String ![String] ![String] -> [[String]]
	groupLines head current [] = [reverse current]
	groupLines head current [line:lines]
	| startsWith head line = groupLines head [line:current] lines
	| isEmpty current      = groupLines head [line] lines
	| otherwise            = [reverse current:groupLines head [line] lines]

printDTMC :: !*DTMC -> *(!String, !*DTMC)
printDTMC dtmc=:{nr_states,states}
# (state_strings,states) = printStates 0 states
# s = join "\n"
	[ "@type: DTMC"
	, "@parameters"
	, "@reward_models"
	, "@nr_states"
	, toString nr_states
	, "@model"
	: state_strings
	]
= (s, {dtmc & states=states})
where
	printStates :: !Int !*{Maybe State} -> *(![String], !*{Maybe State})
	printStates i ss
	| i == nr_states = ([], ss)
	# (s,ss) = ss![i]
	| isNothing s = printStates (i+1) ss
	# s = fromJust s
	# repr =
		[ "state " <+ s.state_id
		, "\taction 0"
		: ["\t\t" <+ to <+ " : " <+ prob \\ (to,prob) <- 'M'.toList s.transitions]
		]
	# (reprs,ss) = printStates (i+1) ss
	= (repr ++ reprs,ss)

addVars :: *DTMC -> *DTMC
addVars dtmc = {DTMC | nr_states = dtmc.nr_states, states = addVarsS 0 dtmc.nr_states dtmc.states}
where
	addVarsS :: !Int !Int !*{Maybe State} -> *{Maybe State}
	addVarsS fi m ss
	| fi >= m = ss
	| otherwise = addVarsS (fi + 1) m (addVarsT fi 0 m ss)

	addVarsT :: !Int !Int !Int !*{Maybe State} -> *{Maybe State}
	addVarsT fi ti m ss
	| ti >= m = ss
	# (s, ss) = ss![fi]
	= case s of
		Nothing -> ss
		(Just s) -> addVarsT fi (ti + 1) m {ss & [fi] = Just {s & transitions = alter (addVar fi ti) ti s.transitions}}

	addVar :: Int Int (Maybe String) -> Maybe String
	addVar _ _ Nothing = Nothing
	addVar f t (Just fn) = Just ("((" <+ fn <+ ") + v" <+ f <+ "_" <+ t <+ ")")

Start w
# (dtmc,w) = parseDTMC "die.drn" w
# dtmc = addVars dtmc
# dtmc = stateElimination dtmc
= printDTMC dtmc