"""Tiny deterministic environment and returns; policies are hand-written."""
from fractions import Fraction as F
import json
class Line:
    def reset(self):self.position=0;self.steps=0;self.done=False;return self.position
    def step(self,action):
        if self.done:raise ValueError('reset after episode end')
        if action not in ('L','R'):raise ValueError('L or R required')
        self.position=max(0,min(3,self.position+(1 if action=='R' else -1)));self.steps+=1
        terminated=self.position==3;truncated=self.steps>=4 and not terminated
        reward=5 if terminated else -1;self.done=terminated or truncated
        return self.position,reward,terminated,truncated
def rollout(actions):
    env=Line();state=env.reset();rows=[]
    for action in actions:
        new,reward,terminated,truncated=env.step(action)
        rows.append({'state':state,'action':action,'next_state':new,'reward':reward,'terminated':terminated,'truncated':truncated});state=new
        if terminated or truncated:break
    discounted=sum((F(9,10)**i)*row['reward'] for i,row in enumerate(rows))
    return {'trace':rows,'reward_sum':sum(r['reward'] for r in rows),'gamma':'9/10','discounted_return':str(discounted)}
if __name__=='__main__':
    goal=rollout('RRR');limit=rollout('RLRR')
    assert goal['reward_sum']==3 and goal['discounted_return']=='43/20'
    assert goal['trace'][-1]['terminated'] and not goal['trace'][-1]['truncated']
    assert limit['reward_sum']==-4 and limit['trace'][-1]['next_state']==2 and limit['trace'][-1]['truncated'] and not limit['trace'][-1]['terminated']
    env=Line();env.reset();env.step('R');env.step('R');env.step('R')
    try:env.step('R')
    except ValueError:pass
    else:raise AssertionError('step after end accepted')
    print(json.dumps({'goal_trace':goal,'time_limit_trace':limit,'policy_learned':False,'Gymnasium_run':False},indent=2))