import sys
from BA10H import print_matrices, HMMParameterEstimation
from BA10C import Viterbi


def Viterbi_Learning(x, init_transition_matrix, init_emission_matrix, alphabet, all_states, max_iterations):
    transition_matrix = init_transition_matrix
    emission_matrix = init_emission_matrix

    for iteration in range(max_iterations):
        # Step 1: Determine path using current parameters
        path = Viterbi(x, all_states, transition_matrix, emission_matrix)

        # Step 2: Determine new parameters using path
        transition_matrix, emission_matrix = HMMParameterEstimation(x, path, alphabet, all_states)

    return transition_matrix, emission_matrix


if __name__ == "__main__":
    '''
    Given: A sequence of emitted symbols x = x1 ... xn in an alphabet A, generated by a k-state HMM with unknown 
    transition and emission probabilities, initial Transition and Emission matrices and a number of iterations i.
    Return: A matrix of transition probabilities Transition and a matrix of emission probabilities Emission that 
    maximizes Pr(x, π) over all possible transition and emission matrices and over all hidden paths π.
    '''

    tmp = sys.stdin.read().splitlines()

    j = tmp[0]
    x = tmp[2]
    alphabet = tmp[4].split()
    all_states = tmp[6].split()

    init_transition_matrix = {}
    init_emission_matrix = {}

    # initial transition matrix
    col_syms = tmp[8].split('\t')[1:]
    transition_end = 8 + len(all_states)

    for i in range(9, transition_end + 1):
        current_line = tmp[i].rstrip().split('\t')
        row_sym = current_line[0]
        init_transition_matrix[row_sym] = {}
        for j in range(1, len(current_line)):
            init_transition_matrix[row_sym][col_syms[j - 1]] = float(current_line[j])

    # emission matrix
    col_syms = tmp[transition_end + 2].split('\t')[1:]

    for i in range(transition_end + 3, len(tmp)):
        current_line = tmp[i].rstrip().split('\t')
        row_sym = current_line[0]
        init_emission_matrix[row_sym] = {}
        for j in range(1, len(current_line)):
            init_emission_matrix[row_sym][col_syms[j - 1]] = float(current_line[j])

    transition_matrix, emission_matrix = Viterbi_Learning(x, init_transition_matrix, init_emission_matrix, alphabet,
                                                          all_states, j)

    print_matrices(transition_matrix, emission_matrix)