#!/usr/bin/env python3
"""Fit eight action-value forests with real pytoolz.wsrf, simulate and publish receipts."""
import sys,json,argparse,hashlib,inspect,platform,subprocess,importlib.metadata
from pathlib import Path
import numpy as np
import pytoolz as pt
ROOT=Path(__file__).resolve().parents[2];CACHE=ROOT/'tools/academy/.cache';CACHE.mkdir(exist_ok=True)
IDS=['corner-crew','sudoku-lanterns','prime-rail-express','lantern-tactics','stormpost-squadron','clockwork-conservatory','echo-grove','burrow-survey']
def save(p,v):
 p=Path(p);p.parent.mkdir(parents=True,exist_ok=True);p.write_text(json.dumps(v,separators=(',',':'),allow_nan=False)+'\n')
def node(script,*args):return subprocess.check_output(['node',str(ROOT/script),*map(str,args)],text=True)
def samples(id,split,n,**kw):
 seed_split=kw.pop('seed_split',split);ix=pt.indexrng.Index('academy-eight-v2',backend='blake2b');seeds=[f'{seed_split}-{i}-{int(w)}' for i,w in enumerate(ix.words(id+'/'+seed_split,n))]
 cfg=CACHE/f'{id}-{split}-request.json';out=CACHE/f'{id}-{split}.json';save(cfg,dict(id=id,seeds=seeds,output=str(out),**kw));print(node('tools/academy/data.cjs',cfg).strip(),flush=True);return json.loads(out.read_text())
def arrays(rows):return np.array([r['x'] for r in rows]),np.array([r['zone'] for r in rows]),np.array([r['y'] for r in rows])
def fit(rows,trees,depth):
 x,z,y=arrays(rows);print(f'Fitting {len(x)} rows, {trees} trees/zone, depth {depth}',flush=True)
 return pt.wsrf.WSRFRegressor(n_trees_per_zone=trees,max_depth=depth,min_samples_leaf=2,max_features=None,random_state=2026,n_jobs=1,adaptive_complexity=False,zone_prior_weight=0,calibrate=False).fit(x,y,zones=z)
def export(m,id,features):
 zones={}
 for z,tt in m.zone_trees_.items():
  trees=[]
  for tree in tt:
   t=tree.tree_;nodes=[]
   for i in range(t.node_count):nodes.append(['l',float(t.value[i].ravel()[0])] if t.children_left[i]<0 else ['b',int(t.feature[i]),float(t.threshold[i]),int(t.children_left[i]),int(t.children_right[i])])
   trees.append(nodes)
  zones[str(z)]=trees
 return dict(format='wsrf-regression-v1',game=id,features=features,zones=zones)
def train(id):
 print('\nTRAINING',id,flush=True);folder=ROOT/'games'/id;config=json.loads(node('tools/academy/config.cjs',id));features=config['featuresNames']
 rows=samples(id,'train',64,collect=True)['rows'];v=samples(id,'validation',16,collect=True);vx,vz,vy=arrays(v['rows'])
 early=fit(rows,8,10);e=export(early,id,features);earlypath=CACHE/f'{id}-early.json';save(earlypath,e)
 early_play=samples(id,'validation-early',16,seed_split='validation',model=str(earlypath))['summary']
 more=samples(id,'train-recovery',16,collect=True,guided=True,model=str(earlypath))['rows'];rows+=more
 final=fit(rows,24,18);f=export(final,id,features);finalpath=CACHE/f'{id}-final.json';save(finalpath,f)
 final_play=samples(id,'validation-final',16,seed_split='validation',model=str(finalpath))['summary']
 # Selection prioritizes task success, then task-specific quality, then pace.
 def key(s):return(s['nonloss'],s['mean_quality'] or 0,-s['mean_turns'])
 selected,selected_export,selectedpath=(final,f,finalpath) if key(final_play)>=key(early_play) else (early,e,earlypath)
 test=samples(id,'test',64,model=str(selectedpath),collect=True,paths=id=='sudoku-lanterns')
 guided=samples(id,'test-guided',32,model=str(selectedpath),guided=True)
 baseline=samples(id,'test',64);random=samples(id,'test',64,random=True);early_test=samples(id,'test',64,model=str(earlypath))
 hx,hz,hy=arrays(test['rows']);pred=selected.predict(hx,zones=hz);mae=float(np.mean(np.abs(pred-hy)));r2=float(selected.score(hx,hy,zones=hz))
 # Independent visited states plus float32 neighbors at learned split boundaries.
 fx=list(hx[:2500]);fz=list(hz[:2500])
 for z,tt in selected.zone_trees_.items():
  refs=hx[hz==z]
  if not len(refs):continue
  for tree in tt:
   for i,(feature,threshold) in enumerate(zip(tree.tree_.feature,tree.tree_.threshold)):
    if feature<0:continue
    center=np.float32(threshold)
    for value in [np.nextafter(center,np.float32(-np.inf)),center,np.nextafter(center,np.float32(np.inf))]:
     row=refs[i%len(refs)].copy();row[feature]=value;fx.append(row);fz.append(z)
 fx=np.array(fx);fz=np.array(fz);fy=selected.predict(fx,zones=fz);fixtures=[dict(x=x.tolist(),zone=str(z),y=float(y)) for x,z,y in zip(fx,fz,fy)]
 fixturepath=CACHE/f'{id}-fixtures.json';save(fixturepath,fixtures);parity=json.loads(node('tools/academy/validate.cjs',selectedpath,fixturepath))
 threshold=45 if id=='burrow-survey' else 48 if id=='lantern-tactics' else 58
 assert test['summary']['nonloss']>=threshold,(id,test['summary'])
 (folder/'training').mkdir(exist_ok=True);(folder/'forest.json').write_bytes(selectedpath.read_bytes());(folder/'early-forest.json').write_bytes(earlypath.read_bytes());(folder/'training/validation.json').write_bytes(fixturepath.read_bytes())
 source_paths=['shared/academy-core.js','shared/academy-forest.js','shared/academy-data.json',f'games/{id}/model.js','tools/academy/data.cjs','tools/academy/train.py','tools/academy/prepare.py']+config.get('extraSources',[])
 p=Path(inspect.getfile(pt.wsrf.WSRFRegressor)).resolve()
 report=dict(game=id,library='pytoolz.wsrf.WSRFRegressor',learning='Supervised action-value regression; targets are documented simulation rollouts or verified game-specific scores.',training_rows=len(rows) if selected is final else len(rows)-len(more),candidate_training_rows=len(rows),early_rows=len(rows)-len(more),features=features,selected='final' if selected is final else 'early',parameters=selected.get_params(),held_out_action_value_mae=mae,held_out_action_value_r2=r2,validation_early=early_play,validation_final=final_play,test={k:v for k,v in test.items() if k not in ['rows','paths']},guided_test={k:v for k,v in guided.items() if k!='rows'},teacher_baseline={k:v for k,v in baseline.items() if k!='rows'},random_baseline={k:v for k,v in random.items() if k!='rows'},early_test={k:v for k,v in early_test.items() if k!='rows'},export_validation=parity,versions={**{p:importlib.metadata.version(p) for p in ['numpy','pandas','scikit-learn']},'python':platform.python_version(),'wsrf':pt.load('wsrf').__version__},wsrf_source={'path':str(p),'sha256':hashlib.sha256(p.read_bytes()).hexdigest()},arcade_source_sha256={p:hashlib.sha256((ROOT/p).read_bytes()).hexdigest() for p in source_paths},forest_sha256=hashlib.sha256((folder/'forest.json').read_bytes()).hexdigest(),trees=sum(len(t) for t in selected_export['zones'].values()))
 save(folder/'training/report.json',report);save(folder/'stats.json',dict(rows=report['training_rows'],trees=report['trees'],wins=test['summary']['wins'],nonloss=test['summary']['nonloss'],days=64,parity=parity,mae=mae))
 if id=='sudoku-lanterns':
  e=test['paths'][0];sp=pt.wsrf_spn.SolvePath(puzzle=''.join(map(str,e['initial'])),moves=[pt.wsrf_spn.Move.from_wsrf_dict(dict(row=m['cell']//9,col=m['cell']%9,digit=m['digit'],technique=m['tech'],reasoning=m['proof']),step=i+1) for i,m in enumerate(e['path'])]);text=pt.wsrf_spn.write(sp);(folder/'training/example.spn').write_text(text);assert pt.wsrf_spn.parse(text).total_moves==len(e['path']);browser_path=pt.wsrf_spn.parse(e['spn']);assert browser_path.total_moves==len(e['path']);save(folder/'training/spn-report.json',dict(library='pytoolz.wsrf_spn',moves=len(e['path']),python_roundtrip=True,browser_notation_parsed=True))
 print('PUBLISHED',id,test['summary'],'MAE',mae,flush=True)
if __name__=='__main__':
 p=argparse.ArgumentParser();p.add_argument('--game',choices=['all',*IDS],default='all');a=p.parse_args()
 for id in IDS if a.game=='all' else [a.game]:train(id)