implementation module DTMC

import StdArray
import StdBool
import StdClass
import StdFile
from StdFunc import o
import StdInt
import StdList
from StdMisc import abort
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.CommandLine
import System.File
import System.FilePath
import Text
import Text.GenParse

import Expression
import Z3

derive gParse Op, (,,)

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:_]
		# (unzs,states) = uncurry elim (penultimate_trans p) dtmc.states
		# dtmc & states = states
		# dtmc & unzero = unzs ++ dtmc.unzero
		= stateElimination dtmc
		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 ss
	where
		remove :: !Int !*{Maybe State} -> *{Maybe State}
		remove i ss = {ss & [i]=Nothing}

	elim :: !Int !Int !*{Maybe State} -> *(![String], !*{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.0" $ 'M'.get s succ
	# succ = 'M'.toList succ

	# (unzs,ss) = mapSt (uncurry $ if (s1==s) update_eq update_ne self_prob) succ ss
	# (Just s1s,ss) = ss![s1]
	# ss & [s1] = Just {s1s & transitions='M'.del s s1s.transitions}
	= (catMaybes unzs,ss)
	where
		update_eq :: !String !Int !String !*{Maybe State} -> *(!Maybe String, !*{Maybe State})
		update_eq self_prob s2 prob2 ss
		| s == s2 = (Nothing, ss)
		# (Just s1s,ss) = ss![s1]
		# ss & [s1] = Just
			{ s1s
			& transitions = 'M'.put s2
				("(" +++ prob2 +++ ")/" +++ denominator)
				s1s.transitions
			}
		= (Just $ "(not (= 0 " +++ parenthesise (parseInfix denominator) +++ "))", ss)
		where denominator = "(1.0-(" +++ self_prob +++ "))"

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

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
	, parameters = []
	, unzero     = []
	, 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,parameters,states}
# (state_strings,states) = printStates 0 states
# s = join "\n" $
	[ "@type: DTMC"
	, "@parameters"
	] ++ if (isEmpty parameters) [""] parameters ++
	[ "@reward_models\n"
	, "@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=:{nr_states} = addVarsS (nr_states-1) dtmc
where
	addVarsS :: !Int !*DTMC -> *DTMC
	addVarsS i dtmc=:{states}
	| i < 0 = dtmc
	# (s,states) = states![i]
	# dtmc & states = states
	| isNothing s = dtmc
	# s = fromJust s
	# dtmc = seqSt (uncurry $ addVar i) ('M'.toList s.transitions) dtmc
	= addVarsS (i-1) dtmc

	addVar :: !Int !Int !String !*DTMC -> *DTMC
	addVar f t prob dtmc=:{parameters}
	# (s,dtmc) = dtmc!states.[f]
	# s = fromJust s
	# s & transitions = 'M'.put t ("((" <+ prob <+ ") + " <+ var <+ ")") s.transitions
	# dtmc & states.[f] = Just s
	# dtmc & parameters = [var:parameters]
	= dtmc
	where
		var = "v" <+ f <+ "x" <+ t

repairConfig :: !CostFunction !*DTMC -> *(!RepairConfig, !*DTMC)
repairConfig cf dtmc=:{nr_states,parameters,states}
# (ss,states) = collectStates (nr_states-1) states
# transs = concatMap collectTransitions ss
# dtmc & states = states
# config =
	{ totality_restrictions   = map totality ss
	, transition_restrictions = concatMap probabilityBounds transs
	, minimalisation_goal     = minimise cf parameters
	}
= (config, dtmc)
where
	collectStates :: !Int !*{Maybe State} -> *(![State], !*{Maybe State})
	collectStates i ss
	| i < 0 = ([], ss)
	# (s,ss) = ss![i]
	| isNothing s = collectStates (i-1) ss
	# s = fromJust s
	# (rest,ss) = collectStates (i-1) ss
	= ([s:rest],ss)

	collectTransitions :: !State -> [(!Int,!Int,!String)]
	collectTransitions s = [(s.state_id, to, prob) \\ (to,prob) <- 'M'.toList s.transitions]

	totality :: !State -> String
	totality s = "(= 1.0 " <+ parenthesise (parseInfix expr) <+ ")"
	where
		expr = "(" <+ (join ")+(" $ 'M'.elems s.transitions) <+ ")"

	probabilityBounds :: !(!Int, !Int, !String) -> [String]
	probabilityBounds (_,_,prob) =
		[ "(<= 0.0 " <+ parenthesise (parseInfix prob) <+ ")"
		, "(>= 1.0 " <+ parenthesise (parseInfix prob) <+ ")"
		]

	minimise :: !CostFunction ![String] -> String
	minimise cf params = foldr1 (\a b -> "(+ " <+ a <+ " " <+ b <+ ")") $ case cf of
		NumberModifiedTransitions ->
			["(if (= 0.0 " <+ p <+ ") 0.0 1.0)" \\ p <- params]
		SumSquaredErrors ->
			["(* " <+ p <+ " " <+ p <+ ")" \\ p <- params]

findInitial :: !*DTMC -> (!State, !*DTMC)
findInitial dtmc=:{nr_states,states}
# (init,states) = findInitial` (nr_states-1) states
# dtmc & states = states
= (init, dtmc)
where
	findInitial` :: !Int !*{Maybe State} -> *(!State, !*{Maybe State})
	findInitial` i states
	| i < 0 = abort "No initial state"
	# (state, states) = states![i]
	| isJust state && (fromJust state).init = (fromJust state, states)
	| otherwise = findInitial` (i - 1) states

instance toString Op
where
	toString LT = "<"
	toString EQ = "="
	toString GT = ">"

assertProperties :: !Properties !*DTMC !Z3 !*World -> (*DTMC, *World)
assertProperties [] dtmc _ w = (dtmc,w)
assertProperties [prop:props] dtmc z3 w
# (dtmc, w) = assertProperty prop dtmc z3 w
= assertProperties props dtmc z3 w

assertProperty :: !Property !*DTMC !Z3 !*World -> (*DTMC, *World)
assertProperty (op, prob, target) dtmc z3 w
# (init, dtmc) = findInitial dtmc
# formula = 'M'.get target init.transitions
| isNothing formula = abort "Cannot reach property target from initial state\n"
# formula = fromJust formula
# formula = "(" <+ op <+ " " <+ prob <+ " " <+ parenthesise (parseInfix formula) <+ ")"
# w = addAssert z3 formula w
= (dtmc, w)

z3AddVars :: !*DTMC !Z3 !*World -> (!*DTMC, !*World)
z3AddVars dtmc z3 w
# (params,dtmc) = dtmc!parameters
# w = seqSt (addVariable z3 "Real") params w
= (dtmc,w)

parenthesise :: !a -> String | toString a
parenthesise x = let s = toString x in if (s.[0] == '(') s ("(" +++ s +++ ")")

parseProperty :: !String -> Property
parseProperty s = case parseString s of
	Nothing -> abort $ "Could not parse '" +++ s +++ "' as a property\n"
	Just p  -> p

Start w
# ([prog:args],w) = getCommandLine w
| length args <> 2 = "Usage: " +++ prog +++ " DRN PROPERTIES\n"
# [drn_in,props:_] = args

# (dtmc,w) = parseDTMC drn_in w
# (props,w) = readFileLines props w
# props = map (parseProperty o trim) $ fromOk props

# dtmc = addVars dtmc
# (config,dtmc) = repairConfig SumSquaredErrors dtmc
# dtmc = stateElimination dtmc

# (z3,w) = startZ3 w
# (dtmc,w) = z3AddVars dtmc z3 w
# (dtmc,w) = assertProperties props dtmc z3 w
# w = seqSt (addAssert z3) config.totality_restrictions w
# w = seqSt (addAssert z3) config.transition_restrictions w
# w = seqSt (addAssert z3) dtmc.unzero w
# w = addVariable z3 "Real" "cost" w
# w = addAssert z3 ("(= cost " +++ config.minimalisation_goal +++ ")") w
//# w = addMinimize z3 "cost" w // NOTE: z3 fails with this, but you can approach a bound by asserting an upper bound on cost

# (sat,w) = checkSat z3 w
| not (sat=:Unsat)
	# (model,w) = getModel z3 w
	= model
= "Could not satisfy the property\n"