FunctionBaumWelch
// 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 Baum-Welch Algorithm, also known as Forward-Backward Algorithm, uses the well known EM algorithm 
// 	  to find the maximum likelihood estimate of the parameters of a hidden Markov model given a set of observed 
// 	  feature vectors. 
// ---
// ### Function List:
// > `forward (array<float> pi, matrix<float> a, matrix<float> b, array<int> obs)`
// > `forward (array<float> pi, matrix<float> a, matrix<float> b, array<int> obs, bool scaling)`
// > `backward (matrix<float> a, matrix<float> b, array<int> obs)`
// > `backward (matrix<float> a, matrix<float> b, array<int> obs, array<float> c)`
// > `baumwelch (array<int> observations, int nstates)`
// > `baumwelch (array<int> observations, array<float> pi, matrix<float> a, matrix<float> b)`
// ---
// ### Reference:
// > https://en.wikipedia.org/wiki/Baum%E2%80%93Welch_algorithm
// > https://github.com/alexsosn/MarslandMLAlgo/blob/4277b24db88c4cb70d6b249921c5d21bc8f86eb4/Ch16/HMM.py
// > https://en.wikipedia.org/wiki/Forward_algorithm
// > https://www.rdocumentation.org/packages/HMM/versions/1.0.1/topics/forward
// > https://www.rdocumentation.org/packages/HMM/versions/1.0.1/topics/backward
library("FunctionBaumWelch")


//#region	-> Helpers:

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

// array_random () {

array_random(int size, float min, float max) =>
	array<float> _r = array.new<float>(size)
	for _i = 0 to size - 1
		_r.set(_i, math.random(min, max))
	_r

array_random(int size, float max) =>
	array<float> _r = array.new<float>(size)
	for _i = 0 to size - 1
		_r.set(_i, math.random(max))
	_r

array_random(int size) =>
	array<float> _r = array.new<float>(size)
	for _i = 0 to size - 1
		_r.set(_i, math.random())
	_r

// }
// matrix_random () {

matrix_random (int rows, int columns, float min, float max) =>
	matrix<float> _r = matrix.new<float>(rows, columns)
	for _i = 0 to rows - 1
		for _j = 0 to columns - 1
			_r.set(_i, _j, math.random(min, max))
	_r

matrix_random (int rows, int columns, float max) =>
	matrix<float> _r = matrix.new<float>(rows, columns)
	for _i = 0 to rows - 1
		for _j = 0 to columns - 1
			_r.set(_i, _j, math.random(max))
	_r

matrix_random (int rows, int columns) =>
	matrix<float> _r = matrix.new<float>(rows, columns)
	for _i = 0 to rows - 1
		for _j = 0 to columns - 1
			_r.set(_i, _j, math.random())
	_r

// }

//#endregion



// forward () {
// reference:
//		https://github.com/alexsosn/MarslandMLAlgo/blob/4277b24db88c4cb70d6b249921c5d21bc8f86eb4/Ch16/HMM.py
//		https://www.rdocumentation.org/packages/HMM/versions/1.0.1/topics/forward

// @function Computes forward probabilities for state `X` up to observation at time `k`, is defined as the
// 		probability of observing sequence of observations `e_1 ... e_k` and that the state at time `k` is `X`.
// @param pi		Initial probabilities.
// @param a   		Transmissions, hidden transition matrix a or alpha = transition probability matrix of changing 
//						states given a state matrix is size (M x M) where M is number of states.
// @param b   		Emissions, matrix of observation probabilities b or beta = observation probabilities. Given 
// 						state matrix is size (M x O) where M is number of states and O is number of different 
//						possible observations.
// @param obs 		List with actual state observation data.
// @returns 
// - `matrix<float> _alpha`: Forward probabilities. The probabilities are given on a logarithmic scale (natural logarithm). The first 
// dimension refers to the state and the second dimension to time.
export forward (array<float> pi, matrix<float> a, matrix<float> b, array<int> obs) =>
	int _n_states = matrix.rows(b)
	int _size_t = array.size(obs)
	matrix<float> _alpha = matrix.new<float>(_n_states, _size_t, 0.0)
	int _obs0 = array.get(obs, 0)
	for _i = 0 to _n_states - 1
		matrix.set(_alpha, _i, 0, array.get(pi, _i) * matrix.get(b, _i, _obs0))
	//
	for _t = 1 to _size_t - 1
		for s = 0 to _n_states - 1
			float _alpha_sum = 0.0
			for _i = 0 to _n_states - 1
				_alpha_sum += matrix.get(_alpha, _i, _t - 1) * matrix.get(a, _i, s)
			matrix.set(_alpha, s, _t, matrix.get(b, s, array.get(obs, _t)) * _alpha_sum)
	_alpha

// pi = array.from(0.25, 0.25, 0.25, 0.25)
// a = matrix.new<float>(0, 4)
// a.add_row(0, array.from(0.05, 0.70, 0.05, 0.20))
// a.add_row(1, array.from(0.10, 0.40, 0.30, 0.20))
// a.add_row(2, array.from(0.10, 0.60, 0.05, 0.25))
// a.add_row(3, array.from(0.25, 0.30, 0.40, 0.05))
// b = matrix.new<float>(0, 0)
// b.add_row(0, array.from(0.3,0.4,0.2,0.1))
// b.add_row(1, array.from(0.2,0.1,0.2,0.5))
// b.add_row(2, array.from(0.4,0.2,0.1,0.3))
// b.add_row(3, array.from(0.3,0.05,0.3,0.35))
// array<int> obs = array.from(3,1,1,3,0,3,3,3,1,1,0,2,2)
// if barstate.islast
// 	fw = forward(pi, a, b, obs)
// 	label.new(bar_index, 0.0, str.tostring(fw))



// @function Computes forward probabilities for state `X` up to observation at time `k`, is defined as the
//		probability of observing sequence of observations `e_1 ... e_k` and that the state at time `k` is `X`.
// @param pi		Initial probabilities.
// @param a   		Transmissions, hidden transition matrix a or alpha = transition probability matrix of changing 
//						states given a state matrix is size (M x M) where M is number of states.
// @param b   		Emissions, matrix of observation probabilities b or beta = observation probabilities. Given 
// 						state matrix is size (M x O) where M is number of states and O is number of different 
//						possible observations.
// @param obs 		List with actual state observation data.
// @param scaling 	Normalize `alpha` scale.
// @returns 
// - #### Tuple with:
// > - `matrix<float> _alpha`: Forward probabilities. The probabilities are given on a logarithmic scale (natural logarithm). The first 
// dimension refers to the state and the second dimension to time.
// > - `array<float> _c`: Array with normalization scale.
export forward (array<float> pi, matrix<float> a, matrix<float> b, array<int> obs, bool scaling) =>
	_alpha = forward(pi, a, b, obs)
	int _n_states = matrix.rows(b)
	int _size_t   = array.size(obs)
	//	
	array<float> _c = array.new<float>(_size_t, 1.0)
	if scaling
		for _t = 0 to _size_t - 1
			_alpha_sum = array.sum(matrix.col(_alpha, _t))
			array.set(_c, _t, _alpha_sum)
			for _i = 0 to _n_states - 1
				matrix.set(_alpha, _i, _t, matrix.get(_alpha, _i, _t) / _alpha_sum)
	[_alpha, _c]

// }
// backward () {
// reference:
//		https://github.com/alexsosn/MarslandMLAlgo/blob/4277b24db88c4cb70d6b249921c5d21bc8f86eb4/Ch16/HMM.py
//		https://www.rdocumentation.org/packages/HMM/versions/1.0.1/topics/backward

// @function Computes backward probabilities for state `X` and observation at time `k`, is defined as the probability of observing the sequence of observations `e_k+1, ... , e_n` under the condition that the state at time `k` is `X`.
// @param a 	Transmissions, hidden transition matrix a or alpha = transition probability matrix of changing states
//    				given a state matrix is size (M x M) where M is number of states
// @param b 	Emissions, matrix of observation probabilities b or beta = observation probabilities. given state 
//    				matrix is size (M x O) where M is number of states and O is number of different possible observations
// @param obs 	Array with actual state observation data.
// @returns 
// - `matrix<float> _beta`: Backward probabilities. The probabilities are given on a logarithmic scale (natural logarithm). The first dimension refers to the state and the second dimension to time.
export backward (matrix<float> a, matrix<float> b, array<int> obs) =>
	int _n_states = matrix.rows(b)
	int _size_t = array.size(obs)
	matrix<float> _beta = matrix.new<float>(_n_states, _size_t, 0.0)
	for _i = 0 to _n_states - 1
		matrix.set(_beta, _i, _size_t - 1, 1.0) // aLast
	//
	for _t = _size_t - 2 to 0
		int _obs_t1 = array.get(obs, _t + 1)
		for s = 0 to _n_states - 1
			float _beta_sum = 0.0
			for _i = 0 to _n_states - 1
				_beta_sum += matrix.get(b, _i, _obs_t1) * matrix.get(_beta, _i, _t + 1) * matrix.get(a, s, _i)
			matrix.set(_beta, s, _t, _beta_sum)
	_beta



// @function Computes backward probabilities for state `X` and observation at time `k`, is defined as the probability of observing the sequence of observations `e_k+1, ... , e_n` under the condition that the state at time `k` is `X`.
// @param a 	Transmissions, hidden transition matrix a or alpha = transition probability matrix of changing states
//    				given a state matrix is size (M x M) where M is number of states
// @param b 	Emissions, matrix of observation probabilities b or beta = observation probabilities. given state 
//    				matrix is size (M x O) where M is number of states and O is number of different possible observations
// @param obs 	Array with actual state observation data.
// @param c 	Array with Normalization scaling coefficients.
// @returns 
// - `matrix<float> _beta`: Backward probabilities. The probabilities are given on a logarithmic scale (natural logarithm). The first dimension refers to the state and the second dimension to time.
export backward (matrix<float> a, matrix<float> b, array<int> obs, array<float> c) =>
	int _n_states = matrix.rows(b)
	int _size_t = array.size(obs)
	_beta = backward(a, b, obs)
	for _t = 0 to _size_t - 1
		for _i = 0 to _n_states - 1
			matrix.set(_beta, _i, _t, matrix.get(_beta, _i, _t) / array.get(c, _t))
	_beta

// baumwelch () {
// AKA: Forward-Backward Algorithm
// reference:
//		https://en.wikipedia.org/wiki/Baum%E2%80%93Welch_algorithm
//		https://github.com/alexsosn/MarslandMLAlgo/blob/4277b24db88c4cb70d6b249921c5d21bc8f86eb4/Ch16/HMM.py

// @function **(Random Initialization)** Baum–Welch algorithm is a special case of the expectation–maximization algorithm used to find the
//				unknown parameters of a hidden Markov model (HMM). It makes use of the forward-backward algorithm 
//				to compute the statistics for the expectation step. 
// @param observations 	List of observed states.
// @param nstate 		Number of hidden states.
// @returns 
// - #### Tuple with:
// > - `array<float> _pi`: Initial probability distribution.
// > - `matrix<float> _a`: Transition probability matrix.
// > - `matrix<float> _b`: Emission probability matrix.
// ---
// requires: `import RicardoSantos/WIPTensor/2 as Tensor`
export baumwelch (array<int> observations, int nstates) =>
	int _nobs   = observations.size()
	int _obsmax = observations.max()
	// 	Initialise pi, a, b randomly
	array<float>  _pi = array.new<float>(nstates, 1.0 / nstates)
	matrix<float> _a  = matrix.new<float>(nstates, nstates)
	matrix<float> _b  = matrix.new<float>(nstates, _obsmax + 1)
	for _i = 0 to nstates - 1
		for _j = 0 to nstates - 1
			_a.set(_i, _j, math.random())
		for _j = 0 to _obsmax
			_b.set(_i, _j, math.random())
	//
	float  _tol    = 1.0e-5
	float  _error  = _tol + 1.0
	int    _maxits = 100
	int    _nits   = 0
	while (_error > _tol) and (_nits < _maxits)
		_nits += 1
		array<float>  _oldpi = _pi.copy()
		matrix<float> _olda  = _a.copy()
		matrix<float> _oldb  = _b.copy()
		// E step
		matrix<float> _alpha = forward(_pi, _a, _b, observations)
		matrix<float> _beta  = backward(_a, _b, observations)
		Tensor.Tensor _xi = Tensor.new(nstates, nstates, _nobs, 0.0)
		for _t = 0 to _nobs - 2
			for _i = 0 to nstates - 1
				for _j = 0 to nstates - 1
					_xi.set(_i, _j, _t, _alpha.get(_i, _t) * _a.get(_i, _j) * _b.get(_j, observations.get(_t + 1)) * _beta.get(_j, _t + 1))
			float _sum = 0.0
			for _i = 0 to nstates - 1
				for _j = 0 to nstates - 1
					_sum += _xi.get(_i, _j, _t)
			for _i = 0 to nstates - 1
				for _j = 0 to nstates - 1
					_xi.set(_i, _j, _t, nz(_xi.get(_i, _j, _t) / _sum))
		for _i = 0 to nstates - 1
			for _j = 0 to nstates - 1
				_xi.set(_i, _j, _nobs - 1, _alpha.get(_i, _nobs - 1) * _a.get(_i, _j))
		// The last step has no b, beta in
		float _sum = 0.0
		for _i = 0 to nstates - 1
			for _j = 0 to nstates - 1
				_sum += _xi.get(_i, _j, _nobs - 1)
		for _i = 0 to nstates - 1
			for _j = 0 to nstates - 1
				_xi.set(_i, _j, _nobs - 1, nz(_xi.get(_i, _j, _nobs - 1) / _sum))
		// M step
		for _i = 0 to nstates - 1
			float _sum2 = 0.0
			for _j = 0 to nstates - 1
				_sum2 += _xi.get(_i, _j, 0)
			_pi.set(_i, _sum2)
			for _j = 0 to nstates - 1
				float _numerator_a   = _xi.get_vector_xy(_i, _j).sum()
				float _denominator_a = 0.0
				for _k = 0 to nstates - 1
					_denominator_a += _xi.get_vector_xy(_i, _k).sum()
				_a.set(_i, _j, nz(_numerator_a / _denominator_a))
			for _k = 0 to _obsmax
				array<int> _found = array.new<int>()
				for [_idx, _o] in observations
					if _k == _o
						_found.push(_idx)
				float _numerator_b   = 0.0
				float _denominator_b = 0.0
				for _j = 0 to nstates - 1
					for _f in _found
						_numerator_b += _xi.get(_i, _j, _f)
					for _s = 0 to nstates - 1
						_denominator_b += _xi.get_vector_xy(_i, _s).sum()
				_b.set(_i, _k, nz(_numerator_b / _denominator_b)) 
		//
		float _max_e = 0.0
		for _i = 0 to nstates - 1
			for _j = 0 to nstates - 1
				_max_e := math.max(_max_e, math.abs(_a.get(_i, _j) - _olda.get(_i, _j)))
			_error := _max_e
			_max_e := 0.0
			for _o = 0 to _obsmax
				_max_e := math.max(_max_e, math.abs(_b.get(_i, _o) - _oldb.get(_i, _o)))
			_error += _max_e
	[_pi, _a, _b]



// @function Baum–Welch algorithm is a special case of the expectation–maximization algorithm used to find the
//				unknown parameters of a hidden Markov model (HMM). It makes use of the forward-backward algorithm 
//				to compute the statistics for the expectation step. 
// @param observations 	List of observed states.
// @param pi 			Initial probability distribution.
// @param a 			Transmissions, hidden transition matrix a or alpha = transition probability matrix of changing states
//    						given a state matrix is size (M x M) where M is number of states
// @param b 			Emissions, matrix of observation probabilities b or beta = observation probabilities. given state 
//    						matrix is size (M x O) where M is number of states and O is number of different possible observations
// @returns 
// - #### Tuple with:
// > - `array<float> _pi`: Initial probability distribution.
// > - `matrix<float> _a`: Transition probability matrix.
// > - `matrix<float> _b`: Emission probability matrix.
// ---
// requires: `import RicardoSantos/WIPTensor/2 as Tensor`
export baumwelch (array<int> observations, array<float> pi, matrix<float> a, matrix<float> b) =>
	int _nobs    = observations.size()	, int _obsmax   = observations.max()
	int _nstates = a.rows()            	, int _nstates1 = _nstates - 1
	// 	Initialise pi, a, b
	array<float>  _pi = pi.copy()
	matrix<float> _a  = a.copy()
	matrix<float> _b  = b.copy()
	//
	float  _tol    = 1.0e-5
	float  _error  = _tol + 1.0
	int    _maxits = 100
	int    _nits   = 0
	while (_error > _tol) and (_nits < _maxits)
		_nits += 1
		array<float>  _oldpi = _pi.copy()
		matrix<float> _olda  = _a.copy()
		matrix<float> _oldb  = _b.copy()
		// E step
		matrix<float> _alpha = forward(_pi, _a, _b, observations)
		matrix<float> _beta  = backward(_a, _b, observations)
		Tensor.Tensor _xi    = Tensor.new(_nstates, _nstates, _nobs, 0.0)
		for _t = 0 to _nobs - 2
			for _i = 0 to _nstates1
				for _j = 0 to _nstates1
					_xi.set(_i, _j, _t, _alpha.get(_i, _t) * _a.get(_i, _j) * _b.get(_j, observations.get(_t + 1)) * _beta.get(_j, _t + 1))
			float _sum = 0.0
			for _i = 0 to _nstates1
				for _j = 0 to _nstates1
					_sum += _xi.get(_i, _j, _t)
			for _i = 0 to _nstates1
				for _j = 0 to _nstates1
					_xi.set(_i, _j, _t, nz(_xi.get(_i, _j, _t) / _sum))
		for _i = 0 to _nstates1
			for _j = 0 to _nstates1
				_xi.set(_i, _j, _nobs - 1, _alpha.get(_i, _nobs - 1) * _a.get(_i, _j))
		// The last step has no b, beta in
		float _sum = 0.0
		for _i = 0 to _nstates1
			for _j = 0 to _nstates1
				_sum += _xi.get(_i, _j, _nobs - 1)
		for _i = 0 to _nstates1
			for _j = 0 to _nstates1
				_xi.set(_i, _j, _nobs - 1, nz(_xi.get(_i, _j, _nobs - 1) / _sum))
		// M step
		for _i = 0 to _nstates1
			float _sum2 = 0.0
			for _j = 0 to _nstates1
				_sum2 += _xi.get(_i, _j, 0)
			_pi.set(_i, _sum2)
			for _j = 0 to _nstates1
				float _numerator_a   = _xi.get_vector_xy(_i, _j).sum()
				float _denominator_a = 0.0
				for _k = 0 to _nstates1
					_denominator_a += _xi.get_vector_xy(_i, _k).sum()
				_a.set(_i, _j, nz(_numerator_a / _denominator_a))
			for _k = 0 to _obsmax
				array<int> _found = array.new<int>()
				for [_idx, _o] in observations
					if _k == _o
						_found.push(_idx)
				float _numerator_b   = 0.0
				float _denominator_b = 0.0
				for _j = 0 to _nstates1
					for _f in _found
						_numerator_b += _xi.get(_i, _j, _f)
					for _s = 0 to _nstates1
						_denominator_b += _xi.get_vector_xy(_i, _s).sum()
				_b.set(_i, _k, nz(_numerator_b / _denominator_b)) 
		//
		float _max_e = 0.0
		for _i = 0 to _nstates1
			for _j = 0 to _nstates1
				_max_e := math.max(_max_e, math.abs(_a.get(_i, _j) - _olda.get(_i, _j)))
			_error := _max_e
			_max_e := 0.0
			for _o = 0 to _obsmax
				_max_e := math.max(_max_e, math.abs(_b.get(_i, _o) - _oldb.get(_i, _o)))
			_error += _max_e
	[_pi, _a, _b]

array<int> obs = array.from(1, 0, 2, 1, 0, 2, 1, 1, 0, 2, 0, 0, 1)
ns = input.int(2)
log_line = input.int(1)
if barstate.islast
	[pi, a, b] = baumwelch(obs, ns)

	logger.queue_one(str.format('Initial probability: {0}\n\nTransition probability:\n{1}\n\nEmission probability:\n{2}', str.tostring(pi,'0.000'), str.tostring(a,'0.000'), str.tostring(b,'0.000')))


// def evenings():
// 	pi = np.array([0.25, 0.25, 0.25, 0.25])
// 	a = np.array([[0.05,0.7, 0.05, 0.2],[0.1,0.4,0.3,0.2],[0.1,0.6,0.05,0.25],[0.25,0.3,0.4,0.05]])
// 	b = np.array([[0.3,0.4,0.2,0.1],[0.2,0.1,0.2,0.5],[0.4,0.2,0.1,0.3],[0.3,0.05,0.3,0.35]])
	
// 	obs = np.array([3,1,1,3,0,3,3,3,1,1,0,2,2])
// 	print Viterbi(pi,a,b,obs)[0]
// 	alpha,c = HMMfwd(pi,a,b,obs)
// 	print np.sum(alpha[:,-1])

// def test():
// 	np.random.seed(4)
// 	pi = np.array([0.25,0.25,0.25,0.25])
// 	aLast = np.array([0.25,0.25,0.25,0.25])
// 	#a = np.array([[.7,.3],[.4,.6]] )
// 	a = np.array([[.4,.3,.1,.2],[.6,.05,.1,.25],[.7,.05,.05,.2],[.3,.4,.25,.05]])
// 	#b = np.array([[.2,.4,.4],[.5,.4,.1]] )
// 	b = np.array([[.2,.1,.2,.5],[.4,.2,.1,.3],[.3,.4,.2,.1],[.3,.05,.3,.35]])
// 	obs = np.array([0,0,3,1,1,2,1,3])
// 	#obs = np.array([2,0,2])
// 	HMMfwd(pi,a,b,obs)
// 	Viterbi(pi,a,b,obs)
// 	print BaumWelch(obs,4)
	
// def biased_coins():
// 	a = np.array([[0.4,0.6],[0.9,0.1]])
// 	b = np.array([[0.49,0.51],[0.85,0.15]])
// 	pi = np.array([0.5,0.5])

// 	obs = np.array([0,1,1,0,1,1,0,0,1,1,0,1,1,1,0,0,1,0,0,1,1,0,1,1,1,1,0,1,0,0,1,0,1,0,0,1,1,1,0])	
// 	print Viterbi(pi,a,b,obs)[0]

// 	print BaumWelch(obs,2)

// }

logger.update()