def print_dptable(V):
    print "    ",
    for i in range(len(V)): print "%7d" % i,
    print
 
    for y in V[0].keys():
        print "%.5s: " % y,
        for t in range(len(V)):
            print "%.7s" % ("%f" % V[t][y]),
        print
 
def viterbi(obs, states, start_p, trans_p, emit_p):
    V = [{}]
    path = {}
 
    # Initialize base cases (t == 0)
    for y in states:
        V[0][y] = start_p[y] * emit_p[y][obs[0]]
        path[y] = [y]
 
    # Run Viterbi for t > 0
    for t in range(1,len(obs)):
        V.append({})
        newpath = {}
 
        for y in states:
            (prob, state) = max([(V[t-1][y0] * trans_p[y0][y] * emit_p[y][obs[t]], y0) for y0 in states])
            V[t][y] = prob
            newpath[y] = path[state] + [y]
 
        # Don't need to remember the old paths
        path = newpath
 
    print_dptable(V)
    (prob, state) = max([(V[len(obs) - 1][y], y) for y in states])
    return  path[state]

def HMM(obs, states, start_p, trans_p, emit_p):
    L=[0]*len(obs)
    p=[0]*len(obs)
    state=[0]*len(obs)
    P=[{}]
    alpha=[{}]
    for y in states:
        alpha[0][y]=start_p[y]*emit_p[y][obs[0]]
    for t in range(1,len(obs)):
        alpha.append({})
        for y in states:
            alpha[t][y]=sum(alpha[t-1][y0]*trans_p[y0][y]*emit_p[y][obs[t]] for y0 in states)
    beta=[{}]
    for y in states:
        beta[-1][y]=1
    for t in range(1,len(obs)):
        beta=[{}]+beta
        for y in states:
            beta[0][y]=sum(beta[1][y0]*trans_p[y][y0]*emit_p[y][obs[len(obs)-1-t]] for y0 in states)
    for t in range(len(obs)):
        L[t]=sum(alpha[t][y]*beta[t][y] for y in states)
        for y in states:
            P[t][y]=alpha[t][y]*beta[t][y]/float(L[t])
        if t!=len(obs)-1:
            P.append({})

    print_dptable(P)
    for t in range(len(obs)):
        (p[t], state[t])=max([(P[t][y0],y0)for y0 in states])
    return  state
        
        
    
    

def example():
    return viterbi(observations,
                   states,
                   start_probability,
                   transition_probability,
                   emission_probability)


states = ('Healthy', 'Fever')
observations = ('normal', 'cold', 'dizzy','cold','dizzy','cold','normal','normal')
start_probability = {'Healthy': 0.6, 'Fever': 0.4}
transition_probability = {
   'Healthy' : {'Healthy': 0.7, 'Fever': 0.3},
   'Fever' : {'Healthy': 0.1, 'Fever': 0.9},
   }
emission_probability = {
   'Healthy' : {'normal': 0.6, 'cold': 0.3, 'dizzy': 0.1},
   'Fever' : {'normal': 0.1, 'cold': 0.2, 'dizzy': 0.7},
   }

print 'Viterbi path:'
print example()
print '----------------------'
print 'HMM nodes:'
print HMM(observations,
                   states,
                   start_probability,
                   transition_probability,
                   emission_probability)