"""Max-product decoding and backtrace, with a greedy counterexample."""
from fractions import Fraction as F
from itertools import product
import json
PI=[F(9,10),F(1,10)]
T=[[F(9,10),F(1,10)],[F(1,100),F(99,100)]]
B=[{'R':F(2,5),'B':F(3,5),'G':F(0)},{'R':F(9,10),'B':F(1,10),'G':F(0)}]
def joint(q,obs):
    p=PI[q[0]]*B[q[0]][obs[0]]
    for t in range(1,len(obs)):p*=T[q[t-1]][q[t]]*B[q[t]][obs[t]]
    return p
def decode(obs,pi=PI,trans=T,emit=B):
    if not obs:raise ValueError('nonempty observations required')
    if any(o not in emit[0] for o in obs):raise ValueError('unknown observation')
    delta=[pi[j]*emit[j][obs[0]] for j in range(2)];rows=[delta];backs=[[None,None]]
    for o in obs[1:]:
        new=[];back=[]
        for j in range(2):
            choices=[delta[i]*trans[i][j] for i in range(2)]
            i=max(range(2),key=lambda i:choices[i]) # first index wins exact ties
            new.append(choices[i]*emit[j][o]);back.append(i)
        delta=new;rows.append(delta);backs.append(back)
    state=max(range(2),key=lambda i:delta[i]);score=delta[state]
    if score==0:return None,score,rows,backs
    path=[state]
    for t in range(len(obs)-1,0,-1):path.append(backs[t][path[-1]])
    return path[::-1],score,rows,backs
def greedy(obs):
    q=[max(range(2),key=lambda j:PI[j]*B[j][obs[0]])]
    for o in obs[1:]:q.append(max(range(2),key=lambda j:T[q[-1]][j]*B[j][o]))
    return q
if __name__=='__main__':
    obs=['R']*3;path,score,rows,backs=decode(obs);g=greedy(obs)
    all_paths=list(product(range(2),repeat=3));best=max(all_paths,key=lambda q:joint(q,obs))
    assert path==list(best)==[1,1,1] and g==[0,0,0]
    assert score==F(7144929,100000000) and joint(g,obs)==F(729,15625)
    tie_emit=[{'R':F(1)},{'R':F(1)}];tie_trans=[[F(1,2)]*2 for _ in range(2)]
    assert decode(obs,[F(1,2)]*2,tie_trans,tie_emit)[0]==[0,0,0]
    assert decode(['G'])[0] is None and decode(['G'])[1]==0
    print(json.dumps({'observations':obs,'delta_rows':[[str(v) for v in r] for r in rows],'backpointers':backs,'MAP_path':path,'joint_score':str(score),'greedy_path':g,'greedy_joint':str(joint(g,obs)),'all_eight_path_joints':[{'path':q,'joint':str(joint(q,obs))} for q in all_paths],'impossible_path':None,'parameters_learned':False},indent=2))