# Copyright 2026 djduecaster. Original governor/head adapter modification, Apache-2.0.
# Inherited model credited in NOTICE; see LICENSE.
from pathlib import Path
import sys,hashlib
import onnx,numpy as np
from onnx import helper as h,numpy_helper as nh
def export(base,path):
 p=[1,1,.2,0,0,0,0,-.075,1]
 p=np.array(p,np.float32);m=onnx.load(base);inp=m.graph.input[0].name;output=m.graph.output[0].name;nodes=[];init=[]
 def c(name,value):
  init.append(nh.from_array(np.array(value),name));return name
 def n(op,inputs,out,**kw):nodes.append(h.make_node(op,inputs,[out],**kw));return out
 def gather(i):return n('Gather',[inp,c('idx'+str(i),np.array([i],np.int64))],'slot'+str(i),axis=1)
 vx=gather(48);yaw=gather(50);gx=gather(3);gz=gather(2)
 green=n('Cast',[n('Not',[n('Equal',[vx,c('zero',np.float32(0))],'eq')],'neq')],'green',to=1)
 y2=n('Mul',[yaw,yaw],'y2');den=n('Add',[c('one',np.float32(1)),n('Mul',[c('k',p[2]),y2],'ky2')],'den')
 eff=n('Div',[n('Mul',[vx,c('vxscale',p[0])],'vxs'),den],'eff');vnew=n('Add',[vx,n('Sub',[eff,vx],'diff')],'vnew')
 ynew=n('Mul',[yaw,c('yscale',p[1])],'ynew');gxnew=n('Add',[gx,n('Mul',[green,c('gxbias',p[3])],'gxd')],'gxnew');gznew=n('Mul',[gz,c('gzscale',p[8])],'gznew')
 obs=inp
 for i,val in [(48,vnew),(50,ynew),(3,gxnew),(2,gznew)]:
  obs=n('ScatterElements',[obs,c('scatter'+str(i),np.array([[i]],np.int64)),val],'obs'+str(i),axis=1)
 old=list(m.graph.node)
 for node in old:
  for i,v in enumerate(node.input):
   if v==inp:node.input[i]=obs
  for i,v in enumerate(node.output):
   if v==output:node.output[i]='base_output'
 bias=np.zeros((1,14),np.float32);bias[0,[2,3,4,5,6,11,12,13]]=[p[4],p[5],p[6],p[7],p[7],-p[4],-p[5],-p[6]]
 delta=n('Mul',[green,c('bias',bias)],'act_delta')
 # final add must occur after base nodes
 final=h.make_node('Add',['base_output',delta],[output]);del m.graph.node[:];m.graph.node.extend(nodes+old+[final]);m.graph.initializer.extend(init)
 onnx.checker.check_model(m);onnx.save(m,path)

if __name__=='__main__':
 if len(sys.argv)!=3:raise SystemExit('Usage: python export_selected.py UPSTREAM_alpha_walking.onnx OUTPUT.onnx')
 base=Path(sys.argv[1]);out=Path(sys.argv[2])
 assert hashlib.sha256(base.read_bytes()).hexdigest()=='e36332d383997d51401897734cd3e79cf5038406feddb18b4d57ecfb141daa6c'
 export(base,out)
 assert hashlib.sha256(out.read_bytes()).hexdigest()=='e9300f63f8374029ea21c732f78f27a84b54926baabe0d240d13cd27a669f44a'
 print('Exact selected policy bytes reproduced')