"""Authored contextual two-action Bernoulli simulation. Observed chosen reward only; no future bootstrap."""
from fractions import Fraction as F
import random,json
ACTIONS=('A','B');CONTEXTS=('red','blue');EPS=F(1,5)
RATES={'red':{'A':F(4,5),'B':F(1,5)},'blue':{'A':F(1,5),'B':F(4,5)}}
def demo():
    policy=random.Random(7);environment=random.Random(11)
    counts={c:{a:0 for a in ACTIONS} for c in CONTEXTS};successes={c:{a:0 for a in ACTIONS} for c in CONTEXTS};trace=[]
    for t in range(60):
        if t<4:
            context,action=[('red','A'),('red','B'),('blue','A'),('blue','B')][t];reason='forced_warmup';propensity=F(1);greedy=None
        else:
            context=CONTEXTS[(t-4)%2]
            means={a:F(successes[context][a],counts[context][a]) for a in ACTIONS}
            greedy=min(ACTIONS,key=lambda a:(-means[a],a))
            if policy.random()<float(EPS):action=policy.choice(ACTIONS);reason='exploration'
            else:action=greedy;reason='greedy_branch'
            propensity=(1-EPS if action==greedy else F(0))+EPS/len(ACTIONS)
        before={a:{'count':counts[context][a],'successes':successes[context][a]} for a in ACTIONS}
        reward=int(environment.random()<float(RATES[context][action]));counts[context][action]+=1;successes[context][action]+=reward
        after={a:{'count':counts[context][a],'successes':successes[context][a]} for a in ACTIONS}
        other='B' if action=='A' else 'A';assert before[other]==after[other]
        trace.append({'step':t+1,'context':context,'action':action,'reason':reason,'greedy_action_before_selection':greedy,'selection_propensity_conditional_on_history':str(propensity),'observed_selected_reward':reward,'before':before,'after':after,'expected_pseudo_regret_from_known_simulator':str(max(RATES[context].values())-RATES[context][action]),'unchosen_reward_observed':False})
    assert sum(sum(v.values()) for v in counts.values())==60
    assert all(r['selection_propensity_conditional_on_history'] in ['1/10','9/10'] for r in trace[4:])
    return {'simulation_only':True,'epsilon':str(EPS),'policy_seed':7,'environment_seed':11,'warmup_steps':4,'environment_rates':{c:{a:str(p) for a,p in v.items()} for c,v in RATES.items()},'trace':trace,'final_counts':counts,'final_successes':successes,'total_observed_reward':sum(r['observed_selected_reward'] for r in trace),'expected_pseudo_regret':str(sum((F(r['expected_pseudo_regret_from_known_simulator']) for r in trace),F(0))),'real_users_clicks_sales_or_actions':False,'VW_executed':False}
if __name__=='__main__':print(json.dumps(demo(),ensure_ascii=False,indent=2))