import json, math, random
from pathlib import Path
import numpy as np
SEED=2516; np.random.seed(SEED); random.seed(SEED)
# State-qualified event labels make B a valid deterministic event-reachability map.
# States: 0=start, 1=dead-end, 2=progress, 3=goal.
# Events: wait0,risky0,progress0,wait1,wait2,progress2,wait3.
N,M=4,7
T={(0,0):0,(0,1):1,(0,2):2,(1,3):1,(2,4):2,(2,5):3,(3,6):3}
C=np.zeros((N,M),int); A=np.zeros((N,N),int); B=np.zeros((N,M),int)
for (s,e),ns in T.items(): C[s,e]=1; A[ns,s]=1; B[ns,e]=1

def enabled(s): return np.flatnonzero(C[s]).tolist()
def update(s,e): return (A@np.eye(N,dtype=int)[:,s])*(B@np.eye(M,dtype=int)[:,e])
def transition(s,e):
 y=update(s,e); assert y.sum()==1,(s,e,y); return int(np.flatnonzero(y)[0])
def logits(s,bias):
 z=np.full(M,-100.); z[0]=-2.; z[1]=bias; z[2]=0.
 if s==1:z[3]=0.
 if s==2:z[4]=-2.; z[5]=0.
 if s==3:z[6]=0.
 return z
def softmax(z):
 q=np.exp(z-np.max(z)); return q/q.sum()
def raw_action(s,b,rng): return int(rng.choice(M,p=softmax(logits(s,b))))
def shield_action(s,b,rng):
 es=enabled(s); return int(rng.choice(es,p=softmax(logits(s,b)[es])))
def plan_action(s,b,H,gamma=1.):
 best=None
 def rec(cur,d,score,first):
  nonlocal best
  if d==H:
   cand=(score+5.*(cur==3),first)
   if best is None or cand>best: best=cand
   return
  for e in enabled(cur): rec(transition(cur,e),d+1,score+gamma**d*logits(cur,b)[e],e if first is None else first)
 rec(s,0,0.,None); return best[1]
def exact_certificate(): return all(update(s,e).sum()==1 for s,e in T)
def run(b,mode,H=1,episodes=4000,max_steps=8):
 rng=np.random.default_rng(SEED+int(round((b+10)*100))); viol=inter=first=done=total=0
 for _ in range(episodes):
  s=0
  for t in range(max_steps):
   raw=raw_action(s,b,rng); e=raw if mode=='raw' else shield_action(s,b,rng) if mode=='shield' else plan_action(s,b,H)
   if t==0 and e!=raw:first+=1
   if e not in enabled(s):viol+=1;break
   inter+=e!=raw;total+=1;s=transition(s,e)
   if s==3:done+=1;break
 return {'violation_rate':viol/episodes,'completion_rate':done/episodes,'intervention_rate':inter/max(total,1),'first_intervention_rate':first/episodes}
def main():
 assert exact_certificate()
 rows=[]
 for b in [-2,-1,0,1,2,3,4,5]:
  r=run(b,'shield'); rr=run(b,'raw'); p=math.exp(b)/(math.exp(-2)+math.exp(b)+1)
  rows.append({'bias':b,'predicted_start_mask_probability':p,'observed_first_intervention':r['first_intervention_rate'],'shield_forbidden':r['violation_rate'],'raw_forbidden':rr['violation_rate'],'shield_completion':r['completion_rate']})
 hs=[{'H':h,**run(4.,'plan',h,1000)} for h in [1,2,3,4]]
 out={'seed':SEED,'printed_ABt_certificate_dimensionally_ill_typed':True,'exact_enabled_pair_certificate':exact_certificate(),'transition_formula_ok':exact_certificate(),'prediction_zero_forbidden':max(x['shield_forbidden'] for x in rows)==0.,'prediction_first_intervention_matches_softmax':max(abs(x['observed_first_intervention']-x['predicted_start_mask_probability']) for x in rows)<.03,'horizon_sweep':hs,'bias_sweep':rows,'matrices':{'A':A.tolist(),'B':B.tolist(),'C':C.tolist()}}
 Path('results.json').write_text(json.dumps(out,indent=2));print(json.dumps(out,indent=2))
if __name__=='__main__':main()