"""Owned known-model stochastic MDP, synchronous exact Bellman sweeps and policy oracle."""
from fractions import Fraction as F
from itertools import product
import json
def model():
    return {'A':{'quit':[(F(1),'T',F(1))],'try':[(F(1,2),'B',F(0)),(F(1,2),'T',F(4))]},'B':{'finish':[(F(1),'T',F(3))],'wait':[(F(1),'B',F(1))]},'T':{}}
def validate(M,gamma):
    if not 0<=gamma<1:raise ValueError('discounted continuing-model subset requires gamma<1')
    for actions in M.values():
        for transitions in actions.values():
            if sum(p for p,s,r in transitions)!=1 or any(p<0 or s not in M for p,s,r in transitions):raise ValueError('valid known transition probabilities required')
def qvalues(M,V,s,gamma):return {a:sum(p*(r+gamma*V[nxt]) for p,nxt,r in ts) for a,ts in M[s].items()}
def backup(M,V,gamma,reverse=False):
    out={}
    for s in list(M)[::-1] if reverse else M:
        qs=qvalues(M,V,s,gamma);out[s]=max(qs.values(),default=F(0))
    return out
def run(M,gamma,sweeps):
    validate(M,gamma);V={s:F(0) for s in M};trace=[]
    for k in range(sweeps):
        new=backup(M,V,gamma);assert new==backup(M,V,gamma,True)
        trace.append({'sweep':k+1,'V_before':V,'Q_from_previous_snapshot':{s:qvalues(M,V,s,gamma) for s in M},'V_after':new});V=new
    BV=backup(M,V,gamma);res=max(abs(BV[s]-V[s]) for s in M);policy={s:(max(qvalues(M,V,s,gamma),key=qvalues(M,V,s,gamma).get) if M[s] else None) for s in M}
    return {'sweeps':sweeps,'trace':trace,'reported_V':V,'fresh_Bellman_backup':BV,'Bellman_residual_infinity':res,'greedy_policy_from_reported_V':policy,'exact_fixed_point_in_this_run':res==0}
def demo():
    M=model();g=F(1,2);full=run(M,g,2);short=run(M,g,1)
    assert full['reported_V']=={'A':F(11,4),'B':F(3),'T':0} and full['Bellman_residual_infinity']==0
    assert short['reported_V']=={'A':2,'B':3,'T':0} and short['Bellman_residual_infinity']==F(3,4)
    oracle=[]
    for a,b in product(['quit','try'],['finish','wait']):
        VB=F(3) if b=='finish' else F(1)/(1-g);VA=F(1) if a=='quit' else F(2)+g*F(1,2)*VB
        oracle.append({'policy':{'A':a,'B':b,'T':None},'independently_solved_V':{'A':VA,'B':VB,'T':F(0)}})
    assert max(x['independently_solved_V']['A'] for x in oracle)==F(11,4) and max(x['independently_solved_V']['B'] for x in oracle)==3
    assert full['greedy_policy_from_reported_V']=={'A':'try','B':'finish','T':None}
    return {'authored_known_MDP':M,'reward_convention':'per transition, includes terminal-entry reward exactly once; terminal has no actions and V0','gamma':g,'two_synchronous_sweeps':full,'one_sweep_not_fixed':short,'four_stationary_deterministic_policy_oracle':oracle,'experience_learning_or_Gymnasium_executed':False}
if __name__=='__main__':print(json.dumps(demo(),default=str,indent=2))