"""
Forward(HMM, O)

A = [
[] (indexed on the ending state)
] (indexed on the starting state)

A[start][end] -> transition from start to end

B = [
[] (indexed on observation)
] (indexed on hidden state state)

B[state][observation] -> probability of state emitting observation

pi -> initial probabilities
"""

def forward(A, B, PI, O):
    a = [[0] * len(A) for _ in range(len(O))]

    for i in range(len(A)):
        a[0][i] = PI[i] * B[i][O[0]]
    
    for t in range(1, len(O)):
        for j in range(len(A)):
            a[t][j] = 0

            for i in range(len(A)):
                a[t][j] = a[t][j] + a[t-1][i] * A[i][j]
            
            a[t][j] = a[t][j] * B[j][O[t]]
    
    p = 0
    for i in range(len(A)):
        p += a[len(O)-1][i]
    
    return p


A = [
    [0.5, 0.5],
    [0.4, 0.6]
]
B = [
    [0.5, 0.4, 0.1],
    [0.2, 0.4, 0.4]
]
PI = [0.2, 0.8]

O = [2, 0, 2]

print(forward(A, B, PI, O))