#!/usr/bin/env python3
"""Fit and evaluate eight new forests through the installed pytoolz/WSRF library."""
import sys,json,argparse,gzip,hashlib,inspect,platform,subprocess,importlib.metadata
from pathlib import Path
import numpy as np
import pytoolz as pt
ROOT=Path(__file__).resolve().parents[2]
sys.path.insert(0,str(ROOT/'tools/academy'))
from train import fit,export,arrays,save
CACHE=ROOT/'tools/curiosity/.cache';CACHE.mkdir(exist_ok=True)
IDS=['cagelight-sudoku','inequality-garden','prime-pair-pavilion','prime-gap-atlas','polyrhythm-postcards','chord-quilt','square-patchwork','index-isles']
NAMESPACE='curiosity-eight-v1'
def node(script,*args):return subprocess.check_output(['node',str(ROOT/script),*map(str,args)],cwd=ROOT,text=True)
def samples(id,split,n,**kw):
    seed_split=kw.pop('seed_split',split);ix=pt.indexrng.Index(NAMESPACE,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/curiosity/data.cjs',cfg).strip(),flush=True);return json.loads(out.read_text())
def train(id):
    print('TRAINING',id,flush=True);folder=ROOT/'games'/id
    features=json.loads(subprocess.check_output(['node','-e',f"console.log(JSON.stringify(require('./games/{id}/model.js').featuresNames))"],cwd=ROOT,text=True))
    rows=samples(id,'train',96,collect=True,explore=True)['rows'];early_count=len(rows)
    early=fit(rows,8,10);ep=CACHE/f'{id}-early.json';save(ep,export(early,id,features))
    ev=samples(id,'validation-early',24,seed_split='validation',model=str(ep))['summary']
    rows+=samples(id,'recovery',24,collect=True,guided=True,model=str(ep))['rows']
    final=fit(rows,24,16);fp=CACHE/f'{id}-final.json';save(fp,export(final,id,features))
    fv=samples(id,'validation-final',24,seed_split='validation',model=str(fp))['summary']
    key=lambda r:(r['wins'],r['mean_quality'] or r['mean_score'],-r['mean_turns'])
    chosen,mp,label=(final,fp,'final') if key(fv)>=key(ev) else (early,ep,'early')
    test=samples(id,'test',64,model=str(mp),collect=True);guided=samples(id,'guided',32,model=str(mp),guided=True)
    teacher=samples(id,'test-reference',64,seed_split='test');random=samples(id,'test-random',64,seed_split='test',random=True);et=samples(id,'test-early',64,seed_split='test',model=str(ep))
    hx,hz,hy=arrays(test['rows']);pred=chosen.predict(hx,zones=hz);fx=list(hx[:1500]);fz=list(hz[:1500])
    for zone,trees in chosen.zone_trees_.items():
        refs=hx[hz==zone]
        assert len(refs),'Missing held-out zone'
        for tree in trees:
            for i,(feature,t) in enumerate(zip(tree.tree_.feature,tree.tree_.threshold)):
                if feature<0:continue
                center=np.float32(t)
                for value in [np.nextafter(center,np.float32(-np.inf)),center,np.nextafter(center,np.float32(np.inf))]:
                    x=refs[i%len(refs)].copy();x[feature]=value;fx.append(x);fz.append(zone)
    fy=chosen.predict(np.array(fx),zones=np.array(fz));fixtures=[dict(x=x.tolist(),zone=str(z),y=float(y)) for x,z,y in zip(fx,fz,fy)]
    fixture=CACHE/f'{id}-fixtures.json';save(fixture,fixtures);parity=json.loads(node('tools/academy/validate.cjs',mp,fixture))
    (folder/'training').mkdir(exist_ok=True);(folder/'forest.json').write_bytes(mp.read_bytes());(folder/'early-forest.json').write_bytes(ep.read_bytes());raw=fixture.read_bytes();(folder/'training/validation.json.gz').write_bytes(gzip.compress(raw,compresslevel=9,mtime=0))
    save(folder/'training/fixture-package.json',dict(file='validation.json.gz',cases=len(fixtures),uncompressed_sha256=hashlib.sha256(raw).hexdigest()))
    paths=[f'games/{id}/model.js','tools/curiosity/data.cjs','tools/curiosity/train.py','tools/curiosity/prepare.py','tools/curiosity/prepare-puzzles.cjs','tools/curiosity/load.cjs','shared/curiosity/logic.js','shared/curiosity/number-world.js','shared/curiosity/data.json','shared/academy-core.js','shared/academy-forest.js','shared/pytoolz-lite.js','tools/academy/train.py']
    estimator=Path(inspect.getfile(pt.wsrf.WSRFRegressor)).resolve();clean=lambda x:{k:v for k,v in x.items() if k!='rows'}
    report=dict(game=id,library='pytoolz.wsrf.WSRFRegressor',learning='Supervised regression of game-specific action scores; prime dispatch and travel use four-step reference continuations.',namespace=NAMESPACE,training_rows=len(rows) if label=='final' else early_count,candidate_training_rows=len(rows),early_rows=early_count,features=features,selected=label,parameters=chosen.get_params(),validation_early=ev,validation_final=fv,test=clean(test),guided_test=clean(guided),teacher_baseline=clean(teacher),random_baseline=clean(random),early_test=clean(et),held_out_action_value_mae=float(np.mean(np.abs(pred-hy))),held_out_action_value_r2=float(chosen.score(hx,hy,zones=hz)),export_validation=parity,trees=sum(len(t) for t in chosen.zone_trees_.values()),versions={**{p:importlib.metadata.version(p) for p in ['numpy','pandas','scikit-learn']},'python':platform.python_version(),'wsrf':pt.load('wsrf').__version__},wsrf_source=dict(path=str(estimator),sha256=hashlib.sha256(estimator.read_bytes()).hexdigest()),arcade_source_sha256={p:hashlib.sha256((ROOT/p).read_bytes()).hexdigest() for p in paths},forest_sha256=hashlib.sha256((folder/'forest.json').read_bytes()).hexdigest())
    save(folder/'training/report.json',report);save(folder/'stats.json',dict(rows=report['training_rows'],trees=report['trees'],wins=test['summary']['wins'],days=64,parity=parity))
    print('PUBLISHED',id,label,test['summary'],guided['summary'],flush=True)
if __name__=='__main__':
    p=argparse.ArgumentParser();p.add_argument('--game',choices=['all',*IDS],default='all');args=p.parse_args()
    for id in IDS if args.game=='all' else [args.game]:train(id)