import os
os.environ['OPENBLAS_NUM_THREADS']='4';os.environ['OMP_NUM_THREADS']='4'
import json,argparse
from pathlib import Path
import numpy as np
from lib.numpy_policy import Dense
from lib.commander_demonstrations import load_demonstrations
parser=argparse.ArgumentParser();parser.add_argument('--demonstrations',help='Exported human commander JSON; only explicitly trainable records are imported');args=parser.parse_args()
root=Path(__file__).resolve().parent.parent;out=root/'.local/battle-v3';rng=np.random.default_rng(82823)
base=json.loads((root/'neural-model/commander-v1.json').read_text());old=Dense(base);model={**base,'name':'commander-macro-v3','inputDim':116,'outputDim':41,'inputMean':base['inputMean']+[0.]*58,'inputStd':base['inputStd']+[1.]*58,'observationSchema':'field-orders.commander-observation.v3','actionSchema':'field-orders.macro-action.v1'}
w=old.w[0];first=np.zeros((116,w.shape[1]),'f');first[:58]=w;last=np.zeros((old.w[-1].shape[0],41),'f');last[:,:22]=old.w[-1];b=np.zeros(41,'f');b[:22]=old.b[-1];model['layers']=[{'weight':first.T.flatten().tolist(),'bias':old.b[0].tolist()},base['layers'][1],{'weight':last.T.flatten().tolist(),'bias':b.tolist()}];net=Dense(model)
x=np.fromfile(out/'macro-x.f32',dtype='f').reshape(-1,116);y=np.fromfile(out/'macro-y.f32',dtype='f').reshape(-1,8);split=int(len(x)*.8);train=np.arange(split);test=np.arange(split,len(x));human_report=None
if args.demonstrations:
 hx,hy,human_report=load_demonstrations(args.demonstrations)
 if hx:
  # Human records are training only; preserve the original held-out teacher rows.
  train=np.concatenate([train,np.arange(len(x),len(x)+len(hx))]);x=np.concatenate([x,np.asarray(hx,dtype='f')]);y=np.concatenate([y,np.asarray(hy,dtype='f')])
heads=[(0,11),(11,4),(15,7),(22,9),(31,5),(38,3)];weights=[1,.3,.3,.7,.7,.2]
def lossgrad(raw,target):
 grad=np.zeros_like(raw);loss=0
 for h,((start,width),weight) in enumerate(zip(heads,weights)):
  logits=raw[:,start:start+width];p=np.exp(logits-logits.max(1,keepdims=True));p/=p.sum(1,keepdims=True);label=target[:,h].astype(int);loss-=weight*np.log(p[np.arange(len(p)),label]+1e-8).mean();p[np.arange(len(p)),label]-=1;grad[:,start:start+width]=p*weight/len(p)
 predicted=np.tanh(raw[:,36:38]);mask=(target[:,4]>0)[:,None];err=(predicted-target[:,6:8])*mask;loss+=float((err**2).sum(1).mean());grad[:,36:38]=2*err*(1-predicted**2)/len(raw);return float(loss),grad
best=float('inf');saved=None
for epoch in range(100):
 for batch in np.array_split(rng.permutation(train),max(1,len(train)//128)):
  pred,cache=net.forward(x[batch],True);loss,grad=lossgrad(pred,y[batch]);net.update(net.gradients(cache,grad),.0012)
 val,_=lossgrad(net.forward(x[test]),y[test])
 if val<best:best=val;saved=[p.copy() for p in net.p]
 if epoch%20==0:print(json.dumps({'epoch':epoch,'heldoutLoss':val}),flush=True)
for p,v in zip(net.p,saved):p[:]=v
pred=net.forward(x[test]);accuracy={name:float((pred[:,start:start+width].argmax(1)==y[test,i]).mean()) for i,(name,(start,width)) in enumerate(zip(['command','formation','strategy','group','task','ttl'],heads))}
net.meta['training']='offline macro demonstrations before PPO';net.export(out/'commander-macro-bc.json','commander-macro-v3',{'stage':'macro-behavior-cloning','seed':82823,'samples':split});m=json.loads((out/'commander-macro-bc.json').read_text());m['training']='offline macro demonstrations before PPO';(out/'commander-macro-bc.json').write_text(json.dumps(m,separators=(',',':')))
soldier=json.loads((root/'neural-model/soldier-ppo-v2.json').read_text());old=Dense(soldier);first=np.zeros((523,old.w[0].shape[1]),'f');first[:126]=old.w[0];soldier.update({'name':'soldier-temporal-v3','inputDim':523,'inputMean':soldier['inputMean']+[0.]*397,'inputStd':soldier['inputStd']+[1.]*397,'observationSchema':'field-orders.temporal-observation.v3','temporalOffsets':[.5,1.5,3]});soldier['layers'][0]['weight']=first.T.flatten().tolist();(out/'soldier-temporal-initial.json').write_text(json.dumps(soldier,separators=(',',':')))
report={'method':'offline macro demonstrations, not live rules','trainingSamples':len(train),'humanDemonstrations':human_report,'heldoutSamples':len(test),'heldoutLoss':best,'headAccuracy':accuracy};(out/'macro-bc-report.json').write_text(json.dumps(report,indent=2));print(json.dumps(report),flush=True)