/* File:      Viterbi.P or .pl (tested on XSB Prolog, SWI prolog, Sicstus, YAP, LPA and tuProlog)
* 
% The forward Viterbi algorithm computes the probability of a sequence of observed events and the most likely sequence of hidden states (denoted the Viterbi path) that result in the sequence of observed events (in the context of hidden Markov models).
* Author: Paul Fodor <pfodor@cs.stonybrook.edu>
* Stony Brook University, 2008
*/

% forward_viterbi(+Observations, +States, +Start_probabilities, +Transition_probabilities, +Emission_probabilities, -Prob, -Viterbi_path, -Viterbi_prob)
% forward_viterbi(['walk', 'shop', 'clean'], ['Ranny', 'Sunny'], [0.6, 0.4], [[0.7, 0.3],[0.4,0.6]], [[0.1, 0.4, 0.5], [0.6, 0.3, 0.1]], Prob, Viterbi_path, Viterbi_prob) will return: Prob = 0.03361, Viterbi_path = [Sunny, Rainy, Rainy, Rainy], Viterbi_prob=0.0094 
forward_viterbi(Obs, States, Sp, Tp, Ep, Prob, V_path, V_prob) :-
  size(Obs, ObsLength),
  size(States, StatesLength),
  StatesLength1 is StatesLength - 1,
  initializeT(StatesLength1, Sp, [], T),
  viterbiMain(0, ObsLength, T, StatesLength, Tp, Ep, NewT),
  Total is 0.0,
  Argmax = [],
  Valmax is 0.0,
  viterbiFinalMain(0, StatesLength, Total, Argmax, Valmax, Prob, V_path_inverse, V_prob, NewT),
  reverse(V_path_inverse, V_path).

% initializeT(+State, +StartProbabilities, +TemporaryT, -FinalT)
% e.g. initializeT(2, [0.6, 0.4], [], FinalT) will return FinalT=
initializeT(-1, _SP, T, T).
initializeT(State, SP, TemporaryT, FinalT) :-
  State >= 0,
  NextState is State - 1,
  listElementSelect(SP, [State], SPState),
  initializeT(NextState, SP, [tnode(SPState, [State], SPState)|TemporaryT], FinalT).

% viterbiMain(+Output, +ObsLength, +T, +StatesLength, +Tp, +Ep, -NewT)
viterbiMain(ObsLength, ObsLength, T, _StatesLength, _Tp, _Ep, T).
viterbiMain(Output, ObsLength, T, StatesLength, Tp, Ep, NewT) :-
  Output < ObsLength,
  viterbiInternalMain(0, StatesLength, Output, T, Tp, Ep, U),
  NextOutput is Output + 1,
  viterbiMain(NextOutput, ObsLength, U, StatesLength, Tp, Ep, NewT).
  
% viterbiInternalMain(+NextState, +StatesLength, +Output, +T, +Tp, +Ep, -U)
viterbiInternalMain(StatesLength, StatesLength, _Output, _T,_Tp, _Ep, []).
viterbiInternalMain(NextState, StatesLength, Output, T, Tp, Ep, [H|RestU]) :-
  Total is 0.0,
  Argmax = [],
  Valmax is 0.0,
  viterbiInternalInternalMain(0, StatesLength, Total, Argmax, Valmax, NextState, ReturnTotal, ReturnArgmax, ReturnValmax, Output, T, Tp, Ep),
  H = tnode(ReturnTotal, ReturnArgmax, ReturnValmax),
  NextState2 is NextState + 1,
  viterbiInternalMain(NextState2, StatesLength, Output, T, Tp, Ep, RestU).
  
% viterbiInternalInternalMain(+State, +StatesLength, +Total, +Argmax, +Valmax, +NextState, -NextTotal, -NextArgmax, -NextValmax, +Output, +T, +Tp, +Ep)
viterbiInternalInternalMain(StatesLength, StatesLength, Total, Argmax, Valmax, _NextState, Total, Argmax, Valmax, _Output, _T, _Tp, _Ep).
viterbiInternalInternalMain(State, StatesLength, Total, Argmax, Valmax, NextState, NextTotal, NextArgmax, NextValmax, Output, T, Tp, Ep) :-  
  listElementSelect(T, [State], tnode(Prob, V_path, V_prob)),
  listElementSelect(Ep, [State, Output], P1),
  listElementSelect(Tp, [State, NextState], P2),
  P is P1 * P2,
  Prob2 is Prob * P,
  V_prob2 is V_prob * P,
  Total2 is Total + Prob2,
  (V_prob2 > Valmax
  ->
    Argmax2 = [NextState|V_path], Valmax2 is V_prob2
  ; Argmax2 = Argmax, Valmax2 is Valmax
  ),
  State2 is State + 1,
  viterbiInternalInternalMain(State2, StatesLength, Total2, Argmax2, Valmax2, NextState, NextTotal, NextArgmax, NextValmax, Output, T, Tp, Ep).

% viterbiFinalMain(+State, +StatesLength, +Total, +Argmax, +Valmax, -NextTotal, -NextArgmax, -NextValmax, +Output, +T)
viterbiFinalMain(StatesLength, StatesLength, Total, Argmax, Valmax, Total, Argmax, Valmax, _T).
viterbiFinalMain(State, StatesLength, Total, Argmax, Valmax, NextTotal, NextArgmax, NextValmax, T) :-  
  listElementSelect(T, [State], tnode(Prob, V_path, V_prob)),
  Total2 is Total + Prob,
  (V_prob > Valmax
  ->
    Argmax2 = V_path, Valmax2 is V_prob
  ; Argmax2 = Argmax, Valmax2 is Valmax
  ),
  State2 is State + 1,
  viterbiFinalMain(State2, StatesLength, Total2, Argmax2, Valmax2, NextTotal, NextArgmax, NextValmax, T).

% listElementSelect(+List, +IndexList, -Element) - indexes start from 0 like arrays in C
% e.g. listElementSelect([['00','01','02'], ['10','11','12']], [0,2], E) will return E = 02
listElementSelect(E, [], E).
listElementSelect([Head|_RestList], [0|T], E) :-
  listElementSelect(Head, T, E).
listElementSelect([_Head|RestList], [H|T], E) :-
  H \= 0,
  H1 is H - 1,
  listElementSelect(RestList, [H1|T], E).

%writePath(States,Viterbi_path)
writePath(_States,[]).
writePath(States,[H|T]) :-
  listElementSelect(States, [H], E),
  write(E), write(' '),
  writePath(States, T).

size(List, Size) :- size(List, 0, Size).
size([], N, N).
size([_H|T], N, M) :-
  N1 is N + 1,
  size(T, N1, M).

reverse(List, Inverse) :- 
  reverse(List, [], Inverse).
reverse([], Inverse, Inverse).
reverse([H|T], SoFar, Inverse) :-
  reverse(T, [H|SoFar], Inverse).

?- Observations = ['walk', 'shop', 'clean'], States = ['Ranny', 'Sunny'], Start_probabilities = [0.6, 0.4], Transmission_probabilities = [[0.7, 0.3],[0.4,0.6]], Emission_probabilities = [[0.1, 0.4, 0.5], [0.6, 0.3, 0.1]],forward_viterbi(Observations, States, Start_probabilities, Transmission_probabilities, Emission_probabilities, Prob, Viterbi_path, Viterbi_prob), write('The probability of the sequence of observed events: '),  write(Observations), write(' is: '), write(Prob), nl, write('The most likely sequence of hidden states is: ['), writePath(States,Viterbi_path), write('] and has a probability of: '), write(Viterbi_prob), nl.