"""Authored depth-three probability tree: exact beam pruning, EOS, exhaustive oracle and seeded sampling."""
from fractions import Fraction as F
import random,json
EOS='<eos>';CAP=3
TREE={():[('oak',F(2,5)),('pine',F(7,20)),('elm',F(1,4))],('oak',):[('a',F(3,5)),(EOS,F(2,5))],('oak','a'):[(EOS,F(1,2)),('b',F(1,2))],('pine',):[('a',F(1,2)),('b',F(1,2))],('pine','a'):[(EOS,F(1))],('pine','b'):[(EOS,F(1,2)),('d',F(1,2))],('elm',):[('a',F(1))],('elm','a'):[(EOS,F(1))]}
def serial(state):
    tokens,p=state;return {'tokens':list(tokens),'joint_probability':str(p),'completed':bool(tokens and tokens[-1]==EOS),'length_cap_reached':len(tokens)==CAP and tokens[-1]!=EOS}
def greedy():
    tokens=();p=F(1)
    for _ in range(CAP):
        token,q=min(TREE[tokens],key=lambda pair:(-pair[1],pair[0]));tokens+=(token,);p*=q
        if token==EOS:break
    return serial((tokens,p))
def beam(width):
    current=[((),F(1))];trace=[]
    for step in range(1,CAP+1):
        expanded=[]
        for tokens,p in current:
            if tokens and tokens[-1]==EOS:expanded.append((tokens,p))
            else:expanded.extend((tokens+(token,),p*q) for token,q in TREE[tokens])
        expanded.sort(key=lambda row:(-row[1],row[0]));current=expanded[:width]
        trace.append({'step':step,'expanded':[serial(r) for r in expanded],'kept':[serial(r) for r in current],'pruned':[serial(r) for r in expanded[width:]]})
    complete=[r for r in current if r[0][-1]==EOS]
    return {'width':width,'trace':trace,'best_completed':serial(complete[0]) if complete else None}
def exhaustive(tokens=(),p=F(1)):
    if tokens and (tokens[-1]==EOS or len(tokens)==CAP):return [(tokens,p)]
    return [leaf for token,q in TREE[tokens] for leaf in exhaustive(tokens+(token,),p*q)]
def sample():
    rng=random.Random(11);tokens=();p=F(1);draws=[]
    for step in range(CAP):
        u=rng.random();cdf=0.0
        for token,q in TREE[tokens]:
            cdf+=float(q)
            if u<cdf:break
        draws.append({'step':step+1,'uniform_draw':u,'selected':token});tokens+=(token,);p*=q
        if token==EOS:break
    return {'seed':11,'cdf_order':'stored TREE order; not lexical sorting','draws':draws,'result':serial((tokens,p))}
def demo():
    assert all(sum((q for _,q in rows),F(0))==1 for rows in TREE.values())
    leaves=exhaustive();assert sum((p for _,p in leaves),F(0))==1
    oracle=min([r for r in leaves if r[0][-1]==EOS],key=lambda row:(-row[1],row[0]));beams=[beam(w) for w in [1,2,3]];g=greedy();s=sample()
    assert g['joint_probability']=='3/25' and beams[0]['best_completed']==g
    assert beams[1]['best_completed']['joint_probability']=='7/40'
    assert beams[2]['best_completed']==serial(oracle) and str(oracle[1])=='1/4'
    assert s['result']['tokens']==['pine','b','d'] and not s['result']['completed']
    return {'tree':[{'prefix':list(k),'conditional':[[t,str(q)] for t,q in v]} for k,v in TREE.items()],'max_tokens_including_EOS':CAP,'score':'raw joint probability, no length normalization','ties':'probability descending then lexical token tuple','greedy':g,'beams':beams,'exhaustive_leaves':[serial(r) for r in leaves],'best_completed_oracle':serial(oracle),'seeded_sampling':s,'language_model_or_transformers_executed':False}
if __name__=='__main__':print(json.dumps(demo(),ensure_ascii=False,indent=2))