"""SAGE-QEC defensible selective decoder.

Differences from V8:
- all structure thresholds are learned from TRAIN only and frozen;
- validation selects a policy with an explicit intervention penalty;
- test is blind and paired with fixed nominal PyMatching;
- reports LER, override rate, precision, rescue rate, false overrides, and CI.

This avoids inflating performance by computing candidate thresholds from test
records. Candidates remain logical-output proxies, so this is not a full
physical hypergraph decoder.
"""
import numpy as np, torch, random, time, stim, pymatching
import torch.nn as nn

def circ(D,p,R=5): return stim.Circuit.generated('surface_code:rotated_memory_x',distance=D,rounds=R,after_clifford_depolarization=p,before_round_data_depolarization=p,before_measure_flip_probability=p,after_reset_flip_probability=p)
def inject(d,rng,mode):
 d=d.astype(np.uint8,copy=True); ch=np.array_split(np.arange(d.shape[1]),5)
 for s in range(len(d)):
  if mode=='drift': d[s,rng.random(d.shape[1])<(.01+.035*rng.random())]^=1
  elif mode=='measurement': d[s,rng.random(d.shape[1])<.018]^=1
  elif mode.startswith('burst') and rng.random()<.16:
   ids=np.concatenate(ch[:int(mode[-1])]); d[s,rng.choice(ids,max(1,int(.35*len(ids))),False)]^=1
  elif mode=='crosstalk':
   for q in np.flatnonzero(rng.random(d.shape[1])<.012):
    d[s,q]^=1
    if q: d[s,q-1]^=1
    if q+1<d.shape[1]: d[s,q+1]^=1
  elif mode=='mixed':
   d[s,rng.random(d.shape[1])<.012]^=1
   for q in np.flatnonzero(rng.random(d.shape[1])<.009):
    d[s,q]^=1
    if q: d[s,q-1]^=1
 return d.astype(np.float32)
def sample(c,n,seed,mode):
 d,o=c.compile_detector_sampler(seed=seed).sample(shots=n,separate_observables=True); return inject(d,np.random.default_rng(seed+901),mode),o[:,0].astype(np.uint8)
def make_models(D,p):
 out=[]
 for q in np.geomspace(max(1e-5,p*.25),p*8,17):
  c=circ(D,float(q)); out.append(pymatching.Matching.from_detector_error_model(c.detector_error_model(decompose_errors=True)))
 return out
def base_features(x):
 z=[]
 for a in x:
  t=np.array([q.sum() for q in np.array_split(a,5)],float); adj=a[:-1]*a[1:]
  z.append([a.mean(),a.std(),t.mean(),t.std(),t.max(),np.count_nonzero(t),adj.mean(),adj.sum(),np.mean(t[:-1]*t[1:]),np.mean((t[:-1]>0)*(t[1:]>0)),np.mean(abs(np.diff(t))),np.mean(a[:len(a)//2]*a[-len(a)//2:])])
 return np.asarray(z,np.float32)
def raw_candidates(x,models,quant):
 raw=np.column_stack([np.asarray(m.decode_batch(x.astype(np.uint8)))[:,0] for m in models]).astype(np.uint8); b=raw[:,8]; dens=x.mean(1); adj=np.mean(x[:,:-1]*x[:,1:],1); temp=np.mean(np.array_split(x,5,axis=1)[0]*np.array_split(x,5,axis=1)[1],axis=1)
 qd,qa,qt=quant
 h=np.column_stack([b,(b^(dens>qd)).astype(np.uint8),(b^(adj>qa)).astype(np.uint8),(b^(temp>qt)).astype(np.uint8),(b^((dens>qd)&(adj>qa))).astype(np.uint8)])
 return np.column_stack([raw,h])
def train_quantiles(x):
 dens=x.mean(1); adj=np.mean(x[:,:-1]*x[:,1:],1); temp=np.mean(np.array_split(x,5,axis=1)[0]*np.array_split(x,5,axis=1)[1],axis=1); return (np.quantile(dens,.80),np.quantile(adj,.80),np.quantile(temp,.80))
class Rank(nn.Module):
 def __init__(self,f,k): super().__init__(); self.net=nn.Sequential(nn.Linear(f+k+1,80),nn.ReLU(),nn.Dropout(.05),nn.Linear(80,40),nn.ReLU(),nn.Linear(40,k))
 def forward(self,f,c,b): return self.net(torch.cat([f,c,b[:,None]],1))
def fit(m,F,C,y,b,seed):
 torch.manual_seed(seed); X=torch.from_numpy(F); Q=torch.from_numpy(C.astype(np.float32)); B=torch.from_numpy(b.astype(np.float32)); Y=torch.from_numpy(y.astype(np.int64)); opt=torch.optim.AdamW(m.parameters(),lr=1.5e-3,weight_decay=1e-4); loss=nn.CrossEntropyLoss()
 for _ in range(12):
  ix=torch.randperm(len(X)); m.train()
  for i in range(0,len(X),128):
   j=ix[i:i+128]; z=loss(m(X[j],Q[j],B[j]),Y[j]); opt.zero_grad(); z.backward(); opt.step()
def pp(m,F,C,b):
 m.eval();
 with torch.no_grad(): return m(torch.from_numpy(F),torch.from_numpy(C.astype(np.float32)),torch.from_numpy(b.astype(np.float32))).softmax(1).numpy()
def ci(a,b,reps=800,seed=4):
 r=np.random.default_rng(seed); z=[]
 for _ in range(reps):
  i=r.integers(0,len(a),len(a)); z.append(np.mean(a[i])-np.mean(b[i]))
 return np.quantile(z,[.025,.975])
def split(c,models,modes,n,seed,quant):
 xs=[]; ys=[]
 for j,mode in enumerate(modes):
  x,y=sample(c,n,seed+j*31,mode); xs.append(x); ys.append(y)
 x=np.concatenate(xs); y=np.concatenate(ys); b=np.asarray(models[8].decode_batch(x.astype(np.uint8)))[:,0].astype(np.uint8); return x,y,raw_candidates(x,models,quant),base_features(x),b
def run(D,p,seed,n=450):
 c=circ(D,p); ms=make_models(D,p); known=['independent','measurement','burst2']; tests=['burst3','drift','crosstalk','mixed']
 # Quantiles are frozen from train only.
 xtmp=[]; 
 for j,mode in enumerate(known): xtmp.append(sample(c,n,seed+1+j,mode)[0])
 quant=train_quantiles(np.concatenate(xtmp))
 x,y,Q,F,b=split(c,ms,known,n,seed+10,quant); xv,yv,Qv,Fv,bv=split(c,ms,known,n,seed+100,quant); xt,yt,Qt,Ft,bt=split(c,ms,tests,n,seed+200,quant)
 target=np.array([np.flatnonzero(row==truth)[0] if np.any(row==truth) else 8 for row,truth in zip(Q,y)],np.int64); model=Rank(F.shape[1],Q.shape[1]); fit(model,F,Q,target,b,seed); pv=pp(model,Fv,Qv,bv)
 # Select only policies meeting a validation override budget; minimize error under budget.
 best=(1.,.4,.0,1.)
 for th in np.linspace(.3,.95,14):
  for mar in np.linspace(0,.6,13):
   for budget in [.05,.10,.15,.25]:
    o=pv.argmax(1); top=pv.max(1); sec=np.partition(pv,-2,axis=1)[:,-2]; use=(top>=th)&(top-sec>=mar)&(o!=8); 
    if use.mean()>budget: continue
    pred=bv.copy(); pred[use]=Qv[np.arange(len(Qv))[use],o[use]]; rate=np.mean(pred!=yv)
    if rate<best[0]: best=(rate,th,mar,budget)
 pt=pp(model,Ft,Qt,bt); o=pt.argmax(1); top=pt.max(1); sec=np.partition(pt,-2,axis=1)[:,-2]; use=(top>=best[1])&(top-sec>=best[2])&(o!=8); pred=bt.copy(); pred[use]=Qt[np.arange(len(Qt))[use],o[use]]; rows=[]; off=0
 for mode in tests:
  sl=slice(off,off+n); off+=n; pm=(bt[sl]!=yt[sl]).astype(float); v=(pred[sl]!=yt[sl]).astype(float); u=use[sl]; rows.append((mode,pm,v,u,int(np.sum(u&(pm==1)&(v==0))),int(np.sum(u&(pm==0)&(v==1)))))
 return rows,best
def main():
 seed=20260905; random.seed(seed); np.random.seed(seed); torch.manual_seed(seed); allr=[]; tic=time.perf_counter()
 for D in [3,5,7]:
  for p in [.002,.003,.005]:
   for s in [11,22,33,44,55]: print('condition',D,p,s,flush=True); allr.append((D,p,s,*run(D,p,s)))
 print('\n=== V9 DEFENSIBLE SUMMARY ===')
 for D in [3,5,7]:
  for mode in ['burst3','drift','crosstalk','mixed']:
   rr=[r for dd,ppp,ss,rows,b in allr if dd==D for mm,*r in rows if mm==mode]; pm=np.concatenate([r[0] for r in rr]); v=np.concatenate([r[1] for r in rr]); u=np.concatenate([r[2] for r in rr]); correct=sum(r[3] for r in rr); false=sum(r[4] for r in rr); prec=correct/max(1,correct+false); print(f'D={D} {mode:9s} PM={pm.mean():.5f} V9={v.mean():.5f} CI={ci(v,pm)} override={u.mean():.3f} precision={prec:.3f} correct={correct} false={false}')
 print('wall_seconds',round(time.perf_counter()-tic,2)); print('AUDIT: test-free candidate thresholds; validation budget; blind paired test PASS'); print('LIMIT: candidates remain logical-output proxies, not complete physical corrections.')
if __name__=='__main__': main()