import numpy as np

A = {
    '+': {
        '+': 0.8,
        '-': 0.2
    },
    
    '-':{
        '+': 0.1,
        '-': 0.9
    }
}

B = {
    '+': {
        'A': 0.2,
        'T': 0.2,
        'C': 0.3,
        'G': 0.3
    },
    
    '-': {
        'A': 0.4,
        'T': 0.4,
        'C': 0.1,
        'G': 0.1
    }
}

P = {
    '+': 0.5,
    '-': 0.5
}

class HMM:
    def __init__(self, l = None):        
        if l != None:
            A, B, P = l
            self.A = A
            self.B = B
            self.P = P
        
    def a(self, q_1, q):
        return self.A[q_1][q]
    
    def b(self, q, x):
        return self.B[q][x]
    
    def pi(self, q):
        return self.P[q]
    
    def state_num(self, q):
        return list(self.A.keys()).index(q)
    
    def num_state(self, num):
        return list(self.A.keys())[num]
    
    def viterbi(self, X):
        T = len(X)
        N = len(A)

        v_matrix = [[0 for _ in range(T)] for _ in range(N)]
        backtrack_matrix = [[-1 for _ in range(T)] for _ in range(N)]
    
        for t in range(T):
            x = X[t]
            
            if t == 0:
                for i in range(N):
                    q = self.num_state(i)
                    
                    transition_prob = self.pi(q)
                    emission_prob = self.b(q, x)
                    
                    prob = transition_prob * emission_prob
                    
                    v_matrix[i][t] = prob
                    
            else:
                for i in range(N):
                    
                    max_prob = float('-inf')
                    max_prop_state = -1
                    
                    q = self.num_state(i)
                    
                    emission_prob = self.b(q, x)
                    
                    for j in range(N):
                        q_1 = self.num_state(j)
                        
                        prev_prob = v_matrix[j][t - 1]
                        
                        transition_prob = self.a(q_1, q)
                        
                        prob = transition_prob * emission_prob * prev_prob
                        
                        if prob > max_prob:
                            max_prob = prob
                            max_prop_state = j
                            
                    v_matrix[i][t] = max_prob
                    backtrack_matrix[i][t] = max_prop_state
                    
        
        # Rekonstrukcija puta
        last_index = np.argmax(np.array(v_matrix)[:,t - 1])
        
        path = ''
        t = T - 1
        
        while last_index != -1:
            last_state = self.num_state(last_index)
            path += last_state
            
            last_index = backtrack_matrix[last_index][t]
            t -= 1
            
        return ''.join(list(reversed(path)))

l = (A, B, P)
hmm = HMM(l)

X = 'CGCAGCGCTTGCGAAAAAAAAATTAAGTAAAAAAACCTGAGGA'

res = hmm.viterbi(X)

print(X)
print(res)

CGCAGCGCTTGCGAAAAAAAAATTAAGTAAAAAAACCTGAGGA
+++++++++++++----------------------++++++++