#!/usr/bin/env python3
from pathlib import Path
import struct,subprocess,os,json,math
import argparse
ap=argparse.ArgumentParser(description='Native signed-velocity boundary regression, including full native stack/return.')
ap.add_argument('rom',type=Path);ap.add_argument('output',type=Path);ap.add_argument('--baseline',type=Path);args=ap.parse_args()
W=args.output.resolve();W.mkdir(parents=True,exist_ok=True);P=Path(__file__).resolve().parents[2];cases=[(power,x,y) for power in [192,370,393,448,463,476,481,482,491,500,600] for x,y in [(0,200),(60,200),(-60,200),(200,0),(0,-200),(60,-200),(-60,-200),(-200,1)]]
source='''.text
.globl _start
_start:
 move.w #0x2700,%sr
 movea.l #0xFFF800,%sp
 lea 0xFFF300,%a3
 lea 0xFFF100,%a2
 move.b #10,0xFFC331
 move.b #13,0x66(%a3)
'''
for power,x,y in cases:
 dist=math.isqrt(x*x+y*y) or 1
 source+=f''' move.b #{128 if y>=0 else 0},0x62(%a3)
 move.w #{power},0xFFBEDE
 bsr fire_{len(source)}
 move.w 0xFFB772,(%a2)+
 move.w 0xFFB774,(%a2)+
 bra next_{len(source)}
fire_{len(source)}:
 movem.l %d0-%d7/%a0-%a3,-(%sp)
 move.w #-1,-(%sp)
 clr.w -(%sp)
 move.w #{y},-(%sp)
 move.w #{x},-(%sp)
 move.w #{dist},%d3
 jmp 0x2FF400
next_{len(source)}:
'''
source+=' move.l #0x56454C4F,0xFFF080\nend:\n bra.s end\n'
rows=[]
for name,rom in ([('baseline',args.baseline)] if args.baseline else [])+[('fixed',args.rom)]:
 d=W/('velocity-'+name);d.mkdir(exist_ok=True);(d/'driver.s').write_text(source);subprocess.run([str(P/'tools/m68kasm.sh'),str(d/'driver.s'),'2D9000',str(d/'code.bin')],check=True,stdout=subprocess.DEVNULL);code=(d/'code.bin').read_bytes();r=bytearray(rom.read_bytes());assert all(x==255 for x in r[0x2d9000:0x2d9000+len(code)]);r[0x2d9000:0x2d9000+len(code)]=code;struct.pack_into('>I',r,4,0x2d9000);r[0x2ff800+26*10+13]=0;(d/'driver.bin').write_bytes(r);(d/'cmd.txt').write_text(f'1000 dump fff080 4 done\n1000 dump fff100 {4*len(cases)} result\n1001 exit\n');env={k:v for k,v in os.environ.items() if k not in ['GPGX_SRM','GPGX_HOOK','GPGX_WAV']};env['GPGX_DEV']='513';subprocess.run(['python3',str(P/'tools/experimental/test_core.py'),str(d/'driver.bin'),str(d/'cmd.txt'),str(d)],env=env,check=True,stdout=subprocess.DEVNULL);data={l.split()[1]:bytes.fromhex(' '.join(l.split()[2:])) for l in (d/'cmd.out').read_text().splitlines()};assert data['done']==b'VELO';out=[]
 for i,(power,x,y) in enumerate(cases):
  dist=math.isqrt(x*x+y*y) or 1;expected=[max(-32767,min(32767,int(power*68*a/dist))) for a in [x,y]];actual=list(struct.unpack_from('>hh',data['result'],i*4));out.append(dict(power=power,axis=[x,y],expected=expected,actual=actual,passed=actual==expected))
 low=[x for x in out if x['power']<=481];high=[x for x in out if x['power']>481];print(name,'low',sum(x['passed'] for x in low),'/',len(low),'high',sum(x['passed'] for x in high),'/',len(high),flush=True)
 if name=='fixed':assert all(x['passed'] for x in out),[x for x in out if not x['passed']]
 rows.append(dict(name=name,cases=out))
(W/'velocity-results.json').write_text(json.dumps(rows,indent=2))