FunctionProbabilityViterbi
// This source code is subject to the terms of the Mozilla Public License 2.0 at https://mozilla.org/MPL/2.0/
// © RicardoSantos

//@version=5

// @description The Viterbi Algorithm calculates the most likely sequence of hidden states *(called Viterbi path)*
//      that results in a sequence of observed events.
library("FunctionProbabilityViterbi")

import RicardoSantos/DebugConsole/13 as console
console.Console logger = console.new()
logger.table.set_position(position.bottom_right)
logger.table.cell_set_width(0, 0, 95)

// viterbi () {
// reference:
//      https://en.wikibooks.org/wiki/Algorithm_Implementation/Viterbi_algorithm
//      https://viterbischool.usc.edu/news/2017/03/viterbi-algorithm-demystified/
//      https://github.com/WuLC/ViterbiAlgorithm/blob/master/Viterbi.py
//      https://github.com/pandeydivesh15/AI_lab-codes/blob/master/HMM_Viterbi/hmm_model.py
//      https://www.pythonpool.com/viterbi-algorithm-python/
//      https://www.cs.princeton.edu/courses/archive/fall15/cos402/assignments/programs/viterbi/
//      https://www.freecodecamp.org/news/a-deep-dive-into-part-of-speech-tagging-using-viterbi-algorithm-17c8de32e8bc/
//      https://www.adeveloperdiary.com/data-science/machine-learning/implement-viterbi-algorithm-in-hidden-markov-model-using-python-and-r/

// @function Calculate most probable path in a Markov model.
// @param observations array<int>    . Observation states data.
// @param transitions  matrix<float> . Transition probability table, (HxH, H:Hidden states).
// @param emissions    matrix<float> . Emission probability table, (OxH, O:Observed states).
// @param initial_distribution array<float> . Initial probability distribution for the hidden states.
// @returns array<int>. Most probable path.
export viterbi (
 array<int>    observations         ,
 matrix<float> transitions          ,
 matrix<float> emissions            , 
 array<float>  initial_distribution ) =>
    // 
    int _n_obs    = observations.size()
    int _n_states = transitions.rows()
    matrix<float> _omega = matrix.new<float>(_n_obs, _n_states, 0.0)
    matrix<int>   _prev  = matrix.new<int>(_n_obs - 1, _n_states, 0)
    //
	for [_i, _e] in initial_distribution
		_omega.set(0, _i, math.log(_e * emissions.get(_i, observations.get(0))))
    //
    for _t = 1 to _n_obs - 1
        for _j = 0 to _n_states - 1
            // Forward Probability:
			_probability = _omega.row(_t - 1)
			for [_i, _p] in _probability
				_probability.set(_i, _p + math.log(transitions.get(_i, _j)) + math.log(emissions.get(_j, observations.get(_t))))
 
            // most probable state given previous state at time t:
			_maxp = _probability.max()
			_prev.set(_t - 1, _j, _probability.indexof(_maxp))
            // probability of the most probable state:
			_omega.set(_t, _j, _maxp)
    //
    // Path Array
    array<int> _S = array.new<int>(_n_obs, 0)
    // Find the most probable last hidden state
	array<float> _o = _omega.row(_n_obs - 1)
	float _maxo = _o.max()
	int _last_state = _o.indexof(_maxo)
    //
    _S.set(0, _last_state)
    int _backtrack_index = 1
    for _i = _n_obs - 2 to 0
        int _ps = _prev.get(_i, _last_state)
        _S.set(_backtrack_index, _ps)
        _last_state := _ps
        _backtrack_index += 1
    //
    // Flip the path array since we were backtracking
    _S.reverse()
	_S

// TEST: 20230318 RS: match with 2 diferent sources.
VT = matrix.new<float>(0, 2, 0.0)
matrix.add_row(VT, 0, array.from(0.7, 0.3))
matrix.add_row(VT, 1, array.from(0.4, 0.6))
VE = matrix.new<float>(0, 3, 0.0)
matrix.add_row(VE, 0, array.from(0.5, 0.4, 0.1))
matrix.add_row(VE, 1, array.from(0.1, 0.3, 0.6))
// matrix.add_row(VE, 2, array.from(0.1, 0.6))
v_init_p = array.from(0.6, 0.4)
v_obs = array.from(0, 1, 2)
// 
// v_MC = hidden.new(
//  name = 'Viterbi test', 
//  states_hidden = array.from(
// 	 state.new('healthy', 0, array.from(targetNode.new(0, 0.7), targetNode.new(1, 0.3))), 
//  	 state.new('fever'  , 1, array.from(targetNode.new(0, 0.4), targetNode.new(1, 0.6)))
// 	  ),
//  states_obs = array.from(
// 	 state.new('normal', 0, array.from(targetNode.new(0, 0.5), targetNode.new(1, 0.1))),
// 	 state.new('cold'  , 1, array.from(targetNode.new(0, 0.4), targetNode.new(1, 0.3))),
// 	 state.new('dizzy' , 2, array.from(targetNode.new(0, 0.1), targetNode.new(1, 0.6)))
//  ),
//  size_hidden = 2,
//  size_obs = 3,
//  transitions = VT,
//  emissions = VE
//  )
if barstate.islast
    new_obs = array.from(1, 2, 0, 2, 1, 0, 1)
    _str = ''
    for _o in new_obs
        v_obs.push(_o)
        vpath = viterbi(v_obs, VT, VE, v_init_p)
        _str += '\n' + str.format('Most likely path for observed states: {0} => {1}', str.tostring(v_obs, '0'), str.tostring(vpath, '0'))
	logger.queue_one(_str)

// }

logger.update()