#!/usr/bin/env python3
"""Paired defensive spacing and actual controller-driven rush/shot scenarios.
Both sides/ends. Native AI, collisions, clock, keeper and puck run unchanged
through each observation. No reset of positions/velocities after setup.
"""
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor
import argparse,json,os,subprocess,statistics
ap=argparse.ArgumentParser(description=__doc__);ap.add_argument('rom',type=Path);ap.add_argument('output',type=Path);ap.add_argument('--baseline',type=Path,required=True);ap.add_argument('--state',type=Path,required=True);a=ap.parse_args();here=Path(__file__).resolve().parent;out=a.output.resolve();out.mkdir(parents=True,exist_ok=True)
def run(profile,rom,name,s):
 p=out/(profile+'-'+name);p.mkdir(exist_ok=True);(p/'spec.json').write_text(json.dumps(s));end=s.get('end_frame',308);(p/'commands.txt').write_text(f'100 load {a.state.resolve()}\n148 snap\n188 snap\n248 snap\n{end+1} exit\n');env={k:v for k,v in os.environ.items() if not k.startswith('GPGX_')};env.update(GPGX_DEV='513',GPGX_HOOK=str(here/'ai_gap_gameplay_hook.py'))
 with (p/'log.txt').open('w') as log:subprocess.run(['nice','-n','19','python3',str(here/'test_core.py'),str(rom.resolve()),str(p/'commands.txt'),str(p)],env=env,stdout=log,stderr=log,check=True,timeout=30)
 rows=json.loads((p/'trajectory.json').read_text());assert len(rows)==end-108
 assert 0<rows[-1]['clock']<rows[0]['clock'] and all(x['sprite_bytes']<=0x280 for x in rows),('native play stopped or sprite overflow',s)
 if s['moving']:
  assert min(x['carrier_xy'][1] for x in rows[:60])<rows[0]['carrier_xy'][1]-12,('fixture did not rush toward defended goal',s)
 return rows

def pair(s):
 name='t{team}-o{orient}-s{sign}-m{moving}-d{difficulty}-shot{shot_at}'.format(**s);old=run('old',a.baseline,name,s);new=run('new',a.rom,name,s)
 stop=60 if s['shot_at'] else 90
 pairs=[(b,n) for b,n in zip(old,new) if 15<=b['frame']<=stop and b['owner']==n['owner']==b['carrier'] and b['jobs'][0]==n['jobs'][0]==2]
 assert len(pairs)>=12,('insufficient native paired defensive frames',s,len(pairs))
 row=dict(spec=s,samples=len(pairs))
 for key in ('gap','cover'):
  row['baseline_'+key]=statistics.mean(b[key] for b,n in pairs);row['fixed_'+key]=statistics.mean(n[key] for b,n in pairs)
 row['gap_improvement']=row['baseline_gap']-row['fixed_gap'];row['cover_improvement']=row['baseline_cover']-row['fixed_cover']
 if s['difficulty']!=1:
  assert abs(row['gap_improvement'])<5 and abs(row['cover_improvement'])<5,('Normal/Easy spacing changed beyond native timing tolerance',row)
 else:
  assert row['gap_improvement']>-10 and row['cover_improvement']>-5,('defensive pair regressed',row)
 if s['shot_at']:
  for profile,rows in [('baseline',old),('fixed',new)]:
   side=1-s['team'];row[profile+'_goals']=max(x['score'][side] for x in rows)-rows[0]['score'][side]
   row[profile+'_shot']=any(x['shotplayer']==x['carrier'] for x in rows)
   row[profile+'_recovery']=any(s['team']*6<=x['owner']<s['team']*6+6 for x in rows if x['frame']<s['shot_at'])
 return row
spacing=[dict(team=t,orient=o,sign=s,moving=m,difficulty=1,depth=-140,shot_at=0) for t in (0,1) for o in (-1,1) for s in (-1,1) for m in (False,True)]
spacing += [dict(team=0,orient=1,sign=s,moving=m,difficulty=d,depth=-140,shot_at=0) for d in (0,2) for s in (-1,1) for m in (False,True)]
pressure=[dict(team=t,orient=o,sign=s,moving=True,difficulty=1,depth=-140,shot_at=shot,end_frame=708) for t in (0,1) for o in (-1,1) for s in (-1,1) for shot in (50,70,90)]
with ThreadPoolExecutor(max_workers=3) as pool:results=list(pool.map(pair,spacing+pressure))
hard=[r for r in results if r['spec']['difficulty']==1 and not r['spec']['shot_at']];shots=[r for r in results if r['spec']['shot_at']]
summary=dict(spacing_cases=len(spacing),pressure_cases=len(shots),hard_improved_cases=sum(r['gap_improvement']>0 for r in hard),mean_gap_improvement=statistics.mean(r['gap_improvement'] for r in hard),worst_gap_improvement=min(r['gap_improvement'] for r in hard),baseline_goals=sum(r['baseline_goals'] for r in shots),fixed_goals=sum(r['fixed_goals'] for r in shots),baseline_shots=sum(r['baseline_shot'] for r in shots),fixed_shots=sum(r['fixed_shot'] for r in shots))
assert summary['hard_improved_cases']>=14 and summary['mean_gap_improvement']>20,('Hard containment did not improve',summary)
assert summary['baseline_goals']>=2 and summary['fixed_goals']<=summary['baseline_goals'],('scripted rush scoring regressed or fixture uninformative',summary)
# C is attempted in every rush; native pressure can stop a release by winning
# the puck first. Count observed releases rather than crediting all 24 as shots.
assert summary['baseline_shots']>=8 and summary['fixed_shots']>=4,('native shot paths not exercised',summary)
assert summary['fixed_shots']<=summary['baseline_shots'],('pressure allowed more scripted attacker releases',summary)
(out/'results.json').write_text(json.dumps(results,indent=2)+'\n');(out/'summary.json').write_text(json.dumps(summary,indent=2)+'\n');print('PASS',len(results),'paired native spacing/rush scenarios;',json.dumps(summary))