% The hidden Markov model with Viterbi optimization

% The start state
start(q1).

% The final states
final(q8).

transition(q1, o1, q2, 0.5, 0.6).
transition(q1, o2, q2, 0.5, 0.3).
transition(q1, k, q2, 0.5, 0.1).
transition(q2, o1, q2, 0.75, 0.6).
transition(q2, o2, q2, 0.75, 0.3).
transition(q2, k, q2, 0.75, 0.1).
transition(q2, o1, q3, 0.25, 0.0).
transition(q2, o2, q3, 0.25, 0.0).
transition(q2, k, q3, 0.25, 1.0).
transition(q3, o1, q3, 0.7, 0.0).
transition(q3, o2, q3, 0.7, 0.0).
transition(q3, k, q3, 0.7, 1.0).
transition(q3, o1, q4, 0.3, 0.65).
transition(q3, o2, q4, 0.3, 0.25).
transition(q3, k, q4, 0.3, 0.1).
transition(q4, o1, q4, 0.65, 0.65).
transition(q4, o2, q4, 0.65, 0.25).
transition(q4, k, q4, 0.65, 0.1).
transition(q1, o1, q5, 0.5, 0.3).
transition(q1, o2, q5, 0.5, 0.7).
transition(q1, k, q5, 0.5, 0.0).
transition(q5, o1, q5, 0.8, 0.3).
transition(q5, o2, q5, 0.8, 0.7).
transition(q5, k, q5, 0.8, 0.0).
transition(q5, o1, q6, 0.2, 0.0).
transition(q5, o2, q6, 0.2, 0.0).
transition(q5, k, q6, 0.2, 1.0).
transition(q6, o1, q6, 0.8, 0.0).
transition(q6, o2, q6, 0.8, 0.0).
transition(q6, k, q6, 0.8, 1.0).
transition(q6, o1, q7, 0.2, 0.65).
transition(q6, o2, q7, 0.2, 0.25).
transition(q6, k, q7, 0.2, 0.1).
transition(q7, o1, q7, 0.7, 0.65).
transition(q7, o2, q7, 0.7, 0.25).
transition(q7, k, q7, 0.7, 0.1).

transition(q4, e, q8, 0.35, 1.0).
transition(q7, e, q8, 0.3, 1.0).	

search(Observations, Paths) :-
	start(StartState),
	viterbi(Observations, [-(1.0, [StartState])], Paths).

viterbi([], Paths, OptimalPaths) :-
	findall(-(Prob, [State | Path]), 
		(member(-(Prob, [State | Path]), Paths), final(State)),
		OptimalPaths).
viterbi([Observation | Observations], Paths, Result) :-
	extend_paths(Observation, Paths, NewPaths),
	keysort(NewPaths, SortedPaths),
	discard_paths(SortedPaths, OptimalPaths),
	viterbi(Observations, OptimalPaths, Result).

% extend_paths(+Observation, +Paths, -NewPaths)
% Given an observation, it extends all the paths in the list
extend_paths(_, [], []).
extend_paths(Observ, [Path | Paths], NewPaths):-
	findall(NewPath, extend_path(Observ, Path, NewPath), FirstNewPaths),
	extend_paths(Observ, Paths, RestNewPaths),
	append(FirstNewPaths, RestNewPaths, NewPaths).

%extend_path(+Observation, +Path, -NewPath)
% Given an observation, it extends one path
extend_path(Observation, -(Prob, [State | Path]), -(NextProb, [NextState, State | Path])) :-
	    transition(State, Observation, NextState, ProbTrans, ProbObserve),
	    ProbObserve =\= 0.0,
	    NextProb is Prob * ProbTrans * ProbObserve.

% discard_paths(+SortedKeyList, -OptimalList)
% Discards the non-optimal paths.
discard_paths([], []).
% There is another path leading to a same state with a
% higher probability. We discard it.
discard_paths([-(_, [State | _]) | SortedPaths], OptimalPaths) :-
	member(-(_, [State | _]), SortedPaths),
	discard_paths(SortedPaths, OptimalPaths).
discard_paths([-(Prob, [State | Path]) | SortedPaths], [-(Prob, [State | Path]) | OptimalPaths]) :-
	\+ member(-(_, [State | _]), SortedPaths),
	discard_paths(SortedPaths, OptimalPaths).



path(Paths) :-
	search([o1, o2, o1, k, k, o1, o1, e], Paths).