import json, subprocess, os, re
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor, as_completed
root=Path('artifacts/mixed-pairing-20260911')
s=json.load(open('artifacts/exchange-frontier-materialized-20260911/transfers.json'))
prepared=[]
for p in s['phases']:
 tx=[0]*s['tile_count']; rx=tx.copy(); legal=[]
 for i,t in enumerate(p['transfers']):
  tx[t['source']]+=t['words']
  for d in t['destinations']:rx[d['tile']]+=t['words']
  ds={d['tile'] for d in t['destinations']}
  if t['words']%2==0 and t['words']>=2 and all(a%8==0 for a in t['source_addresses']) and all(d['address']%8==0 and d['tile']^1 in ds for d in t['destinations']):legal.append(i)
 if legal:
  prepared.append((p,legal,tx,rx))
print('Eligible phases:',[(p['phase'],len(ls)) for p,ls,_,_ in prepared],flush=True)
# Largest payloads among phases with Repeat weight addresses, plus the largest
# remaining representative phases; retain the input's exact addresses/timings.
prepared.sort(key=lambda x:sum(x[0]['transfers'][i]['words'] for i in x[1]),reverse=True)
prepared=prepared[:6]
jobs=[]
for p,legal,tx,rx in prepared:
 variants={
  'ordinary':[], 'all':legal,
  'idle-partner':[i for i in legal if tx[p['transfers'][i]['source']^1]==0],
  'dominant':[i for i in legal if (tx[p['transfers'][i]['source']],-(p['transfers'][i]['source'])) > (tx[p['transfers'][i]['source']^1],-(p['transfers'][i]['source']^1))],
  'receive-bound':[i for i in legal if max(rx[d['tile']] for d in p['transfers'][i]['destinations'])>2*tx[p['transfers'][i]['source']^1]],
 }
 seen={}
 for name,indices in variants.items():
  key=tuple(indices)
  if key in seen:continue
  seen[key]=name
  chosen=set(indices)
  phase={**p,'transfers':[{**t,'width':'Paired64' if i in chosen else 'Word32'} for i,t in enumerate(p['transfers'])]}
  path=root/f"phase{p['phase']}-{name}.json"
  path.write_text(json.dumps({**s,'phases':[phase]},separators=(',',':')))
  jobs.append((p['phase'],name,len(indices),path))
def run(job):
 phase,name,count,path=job; log=path.with_suffix('.log')
 cmd=['target/release/ipu-exchange-schedule-bench',str(path),'--stream-words','1024','--balance-streams']
 with log.open('w') as f:r=subprocess.run(cmd,stdout=f,stderr=subprocess.STDOUT,env={**os.environ,'RAYON_NUM_THREADS':'1'})
 text=log.read_text(); fields=dict(re.findall(r'(horizonCycles|endpointLowerBoundCycles|maximumRowWords|rowWords)=(\d+)',text))
 result={'phase':phase,'variant':name,'paired':count,'exit':r.returncode,**{k:int(v) for k,v in fields.items()}}
 print(result,flush=True);return result
with ThreadPoolExecutor(max_workers=6) as pool:
 results=list(pool.map(run,jobs))
(root/'results.json').write_text(json.dumps(results,indent=2)+'\n')
