MarkovChain
// 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 Generic Markov Chain type functions.
// ---
// A Markov chain or Markov process is a stochastic model describing a sequence of possible events in which the 
// 	probability of each event depends only on the state attained in the previous event.
// ---
// reference:
//		Understanding Markov Chains, Examples and Applications. Second Edition. Book by Nicolas Privault.
//		https://en.wikipedia.org/wiki/Markov_chain
//		https://www.geeksforgeeks.org/finding-the-probability-of-a-state-at-a-given-time-in-a-markov-chain-set-2/
//		https://towardsdatascience.com/brief-introduction-to-markov-chains-2c8cab9c98ab
//		https://github.com/mxgmn/MarkovJunior
//		https://stats.stackexchange.com/questions/36099/estimating-markov-transition-probabilities-from-sequence-data
//		https://timeseriesreasoning.com/contents/hidden-markov-models/
//		https://www.ris-ai.com/markov-chain
//		https://github.com/coin-or/jMarkov/blob/master/src/jmarkov/MarkovProcess.java
//		https://gist.github.com/mschauer/4c81a0529220b21fdf819e097f570f06
//		https://github.com/rasmusab/bayes.js/blob/master/mcmc.js
//		https://gist.github.com/sathomas/cf526d6495811a8ca779946ef5558702
//		https://writings.stephenwolfram.com/2022/06/games-and-puzzles-as-multicomputational-systems/
//		https://kevingal.com/blog/boardgame.html
//		https://towardsdatascience.com/brief-introduction-to-markov-chains-2c8cab9c98ab
//		https://spedygiorgio.github.io/markovchain/reference/index.html
//		https://github.com/alexsosn/MarslandMLAlgo/blob/4277b24db88c4cb70d6b249921c5d21bc8f86eb4/Ch16/HMM.py
//		https://www.projectrhea.org/rhea/index.php/Introduction_to_Hidden_Markov_Chains
library(title='MarkovChain')



//#region	-> Notes:

// Table 7.1 Summary of Markov chain properties
// Property                     Definition
// Absorbing (state)            Pi,i = 1
// Recurrent (state)            P(Tri < ∞ | X0 = i) = 1
// Transient (state)            P(Tri < ∞ | X0 = i) < 1
// Positive recurrent (state)   Recurrent and IE[Tri | X0 = i] < ∞
// Null recurrent (state)       Recurrent and IE[Tri | X0 = i] = ∞
// Aperiodic (state or chain)   Period(s) = 1
// Ergodic (state or chain)     Positive recurrent and aperiodic
// Irreducible (chain)          All states communicate
// Regular (chain)              All coefficients of Pn are > 0 for some n ≥ 1
// Stationary distribution π    Obtained from solving π = πP

// validity errors: 
// - states transition not probability (range 0>1)
// - transition matrix not square.
// - transition matrix probability row/col do not adds to 1

//#endregion



//#region	-> Helper Functions:

import RicardoSantos/WIPTensor/2 as Tensor
import RicardoSantos/FunctionBaumWelch/1 as fb
import RicardoSantos/FunctionProbabilityViterbi/1 as vit

// matrix_pow () {
// reference: https://numpy.org/doc/stable/reference/generated/numpy.linalg.matrix_power.html

// @function Raise a square matrix to the provided exponent power.
// @param mat      matrix<float> . Source square matrix.
// @param exponent int           . Exponent value.
// @returns matrix<float>. matrix with same shape as `A`, 
matrix_pow (matrix<float> mat, int exponent) =>
	int _dim = matrix.rows(mat)
	matrix<float> _M = matrix.copy(mat)
	matrix<float> _ID = matrix.new<float>(_dim, _dim, 0.0)
	for _i = 0 to _dim - 1
		matrix.set(_ID, _i, _i, 1.0)
	if matrix.columns(mat) != _dim
		runtime.error('Matrix is not square')
	int _p = exponent
	while true
		if _p > 0
			if (_p % 2) == 1
				_ID := matrix.mult(_ID, _M)
			_M := matrix.mult(_M, _M)
			_p /= 2
			int(na)
		else
			break
			int(na)
	_M

// }

//#endregion



//#region	-> Structure:

// Node {

// @type Target node.
// @field index       . Key index of the node.
// @field probability . Probability rate of activation.
export type Node
	int   index
	float probability

//}
// state {

// @type State reference.
// @field name         . Name of the state.
// @field index        . Key index of the state.
// @field target_nodes . List of index references and probabilities to target states.
export type state
	string      name
	int			index
	array<Node> target_nodes

//}
// MC {

// @type Markov Chain reference object.
// @field name        . Name of the chain.
// @field states      . List of state nodes and its name, index, targets and transition probabilities.
// @field size        . Number of unique states
// @field transitions . Transition matrix
export type MC
	string        name = 'Unnamed Markov Chain'
	array<state>  states
	int           size
	matrix<float> transitions

//}
// HMC {

// @type Hidden Markov Chain reference object.
// @field name          . Name of thehidden chain.
// @field states_hidden . List of state nodes and its name, index, targets and transition probabilities.
// @field states_obs    . List of state nodes and its name, index, targets and transition probabilities.
// @field transitions   . Transition matrix
// @field emissions     . Emission matrix
export type HMC
	string        name = 'Unnamed Markov Chain' 
	array<state>  states_hidden                 
	array<state>  states_obs                    
	matrix<float> transitions                   
	matrix<float> emissions    
	array<float>  initial_distribution

// }

//#endregion



//#region	-> Functions:

// to_string                  () {

// @function Translate a Markov Chain object to a string format.
// @param this `MC` . Markov Chain object.
// @returns string
export method to_string (MC this) =>
	string _str = this.name + '\n'
	for _state in this.states
		string _str1 = ''
		for _target_node in _state.target_nodes
			_str1 += str.format(' ({0,number,integer}: {1,number,percent})', _target_node.index, _target_node.probability)
		_str += str.format('{0,number,integer} {1}: ({2} )\n', _state.index, _state.name, _str1)
	_str += str.format('\n{0}', str.tostring(this.transitions))
	_str

// }
// to_table                   () {

// @function Translate a Markov Chain object to a table.
// @param this MC, Markov Chain object.
// @returns table
export method to_table (MC this, string position = position.bottom_right, color text_color = #e2e2e2, string text_size = size.auto) =>
	table _T = table.new(position, this.size + 2, this.size + 3)
	table.merge_cells(_T, 0, 0, this.size + 1, 0)
	table.cell(_T, 0, 0, this.name, text_color=text_color, text_size=text_size)
	for [_i, _state] in this.states
		table.cell(_T, 0, 2 + _i, str.format('{0,number,integer}: \"{1}\"', _state.index, _state.name), text_color=text_color, text_size=text_size)
		table.cell(_T, 1 + _i, 1, str.format('{0,number,integer}: \"{1}\"', _state.index, _state.name), text_color=text_color, text_size=text_size)
		for [_j, _target_node] in _state.target_nodes
			table.cell(_T, 1 + _j, 2 + _i, str.format('{0,number,percent}', _target_node.probability), text_color=text_color, text_size=text_size)
	_T

// }
// create_transition_matrix   () {

// @function Creates a transition matrix from a Markov Chain object.
// @param this MC, Markov Chain object.
// @returns matrix<float>
export method create_transition_matrix (MC this) =>
	matrix<float> _T = matrix.new<float>(this.size, this.size, 0.0)
	for _state in this.states
		for _tnode in _state.target_nodes
			matrix.set(_T, _state.index, _tnode.index, _tnode.probability)
	_T

//}
// generate_transition_matrix () {

// @function Generate the transitions matrix property of the Markov Chain.
// @param this MC, Markov Chain object.
// @returns void, updates the Markov Chain object MC.
export method generate_transition_matrix (MC this) =>
	matrix<float> _T = this.create_transition_matrix()
	this.transitions := _T
	this

// }
// new_chain                  () {

// @function Create a new Markov Chain object from a list of provided states.
// @param states array<state>, List of states and its properties.
// @param name   string      , Name of the Markov Chain Object, optional.
// @returns MC, Markov Chain object.
export new_chain (array<state> states, string name = 'Unnamed Markov Chain') =>
	MC _MC = MC.new(name, states, array.size(states))
	_MC.generate_transition_matrix()
	_MC

//}
// from_data                  () {

// @function Create a new Markov Chain object from a list of provided states.
// @param data array<string>, Formated array with classes or states to extract Markov Chain.
// @param name string       , Name of the Markov Chain Object, optional.
// @returns MC, Markov Chain object.
export from_data (array<string> data, string name = 'Unnamed Markov Chain') =>
	// identify unique states:
	array<string> _unique = array.new<string>(0)
	for _e in data
		if not array.includes(_unique, _e)
			array.push(_unique, _e)
	// generate transition counts for state to state:
	int _size_u = array.size(_unique)
	matrix<int> _transition_counters = matrix.new<int>(_size_u, _size_u, 0)
	int _previous = 0
	int _total_count = array.size(data)
	for _i = 1 to _total_count-1
		int _current = array.indexof(_unique, array.get(data, _i))
		int _counter = matrix.get(_transition_counters, _previous, _current) + 1
		matrix.set(_transition_counters, _previous, _current, _counter)
		_previous := _current
	// generate probabilities based on counters:
	matrix<float> _transition_probability = matrix.new<float>(_size_u, _size_u, 0.0)
	for _i = 0 to _size_u - 1
		_total_count := array.sum(matrix.row(_transition_counters, _i))
		for _j = 0 to _size_u - 1
			matrix.set(_transition_probability, _i, _j, float(matrix.get(_transition_counters, _i, _j) / _total_count))
	// generate chain:
	array<state> _states = array.new<state>(_size_u)
	for _i = 0 to _size_u - 1
		string _name = array.get(_unique, _i)
		array<Node> _tnodes = array.new<Node>(0)
		for _j = 0 to _size_u - 1
			float _p = matrix.get(_transition_probability, _i, _j)
			if _p != 0.0
				array.push(_tnodes, Node.new(_j, _p))
		array.set(_states, _i, state.new(_name, _i, _tnodes))
	MC _MC = new_chain(_states, name)

//}
// probability_at_step        () {
// reference: 
//		https://www.geeksforgeeks.org/finding-the-probability-of-a-state-at-a-given-time-in-a-markov-chain-set-2/

// @function Probability of reaching target_state at target_step after starting from start_state
// @param this MC, Markov Chain object.
// @param target_step       int          , Number of steps to find probability.
// @returns 
// - `matrix<float>`: Probability.
export method probability_at_step (MC this, int target_step) =>
    matrix_pow(this.transitions, target_step)

//}
// state_at_step              () {
// reference: 
//		https://www.geeksforgeeks.org/finding-the-probability-of-a-state-at-a-given-time-in-a-markov-chain-set-2/

// @function Probability of reaching target_state at target_step after starting from start_state
// @param this MC, Markov Chain object.
// @param start_state       int          , State at which to start.
// @param target_state      int          , State to find probability.
// @param target_step       int          , Number of steps to find probability.
// @returns
// - `float`: Probability.
export method state_at_step (MC this, int start_state, int target_state, int target_step) =>
    matrix<float> _T = matrix_pow(this.transitions, target_step)
    matrix.get(_T, target_state - 1, start_state - 1)

//}
// 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 this			Hidden markov chain object (states, transitions, emissions and parameters).
// > - transitions			matrix<float>   Transition matrix(T) where T(i,j) stores the transition probability of transitioning from State(i) to State(j)
// > - emissions			matrix<float>   Emission matrix(E) where E(i, j) stores the probability of observing Observation(j) from State(i)
// > - initial_distribution array<float> Initial probability distribution.
// @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 method forward (HMC this, array<int> obs) =>
	fb.forward(this.initial_distribution, this.transitions, this.emissions, obs)

// }
// 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 this		hidden		Hidden markov chain object (states, transitions, emissions and parameters).
// > - transitions	matrix<float>   Transition matrix(T) where T(i,j) stores the transition probability of transitioning from State(i) to State(j)
// > - emissions	matrix<float>   Emission matrix(E) where E(i, j) stores the probability of observing Observation(j) from State(i)
// @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 method backward (HMC this, array<int> obs) =>
	fb.backward(this.transitions, this.emissions, obs)

// }
// viterbi () {

// @function Calculate most probable path in a Markov model.
// Given a Hidden Markov chain (A,B,P) and an observation sequence X, find out the best hidden state sequence S. 
// @param HMC			Hidden markov chain object (states, transitions, emissions and parameters).
// > - transitions			matrix<float>   Transition matrix(T) where T(i,j) stores the transition probability of transitioning from State(i) to State(j)
// > - emissions			matrix<float>   Emission matrix(E) where E(i, j) stores the probability of observing Observation(j) from State(i)
// > - initial_distribution array<float> Initial probability distribution.
// @param observations         array<int>      Observation space.
// @returns 
// - `array<int>`. Array with the most probable sequence of hidden states.
export method viterbi (HMC this, array<int> observations) =>
	vit.viterbi(observations, this.transitions, this.emissions, this.initial_distribution)

// 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(Node.new(0, 0.7), Node.new(1, 0.3))), 
// //  	 state.new('fever'  , 1, array.from(Node.new(0, 0.4), Node.new(1, 0.6)))
// // 	  ),
// //  states_obs = array.from(
// // 	 state.new('normal', 0, array.from(Node.new(0, 0.5), Node.new(1, 0.1))),
// // 	 state.new('cold'  , 1, array.from(Node.new(0, 0.4), Node.new(1, 0.3))),
// // 	 state.new('dizzy' , 2, array.from(Node.new(0, 0.1), Node.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.tostring(vpath)
// 	label.new(bar_index, 0.0, str.format('{0}', _str))

// }
// baumwelch () {

// @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 this			Hidden markov chain object (states, transitions, emissions and parameters).
// > - transitions			matrix<float>   Transition matrix(T) where T(i,j) stores the transition probability of transitioning from State(i) to State(j)
// > - emissions			matrix<float>   Emission matrix(E) where E(i, j) stores the probability of observing Observation(j) from State(i)
// > - initial_distribution array<float> Initial probability distribution.
// @param observations 	List of observed 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 method baumwelch (HMC this, array<int> observations) =>
	fb.baumwelch(observations, this.initial_distribution, this.transitions, this.emissions)

// }

//#endregion



// _MC = MC.new(array.from(
// 	 state.new(  'up', 0, array.from(Node.new(0, 0.1), Node.new(1, 0.1), Node.new(2, 0.1))),
// 	 state.new('down', 1, array.from(Node.new(0, 0.1), Node.new(2, 0.1), Node.new(3, 0.1))),
// 	 state.new('fast', 2, array.from(Node.new(1, 0.1), Node.new(3, 0.1), Node.new(5, 0.1))),
// 	 state.new('doub', 3, array.from(Node.new(1, 0.1), Node.new(2, 0.1), Node.new(5, 0.1))),
// 	 state.new('trip', 4, array.from(Node.new(2, 0.1), Node.new(3, 0.1), Node.new(4, 0.1))),
// 	 state.new('slow', 5, array.from(Node.new(0, 0.1), Node.new(4, 0.1), Node.new(5, 0.1)))
// 	 ), 6)

atr = ta.atr(100)
diff = ta.change(close)
classify = switch
	diff > atr => "> +1"
	diff < -atr => "< -1"
	=> "> 0 <"
var array<string> data = array.new<string>()
if not na(atr)
	array.push(data, classify)

plot(classify=='> +1'?1:(classify=='< -1'?-1:0))
if barstate.islast
	_MC = from_data(data, 'Returns to ATR Odds MC')
	// label.new(bar_index, 0.0, to_string(MC) + '\n' + str.tostring(MC.transitions))
	to_table(_MC)

