#!/usr/bin/env python3
"""Train the six watch-and-guide characters with actual pytoolz.wsrf.
Existing Miso and flight models are never loaded for training or overwritten.
"""
import argparse
import hashlib
import json
from pathlib import Path
import numpy as np
import pytoolz as pt
from train import ROOT, CACHE, save, node, arrays, fit, export, accuracy, publish

IDS = ['mosslight-garden', 'firefly-ferry', 'little-lantern-works',
       'starlight-observatory', 'tidepool-rescue', 'cloudline-bakery']
NAMESPACE = 'watch-and-guide-v1'


def request(id, tag, **kwargs):
    out, config = CACHE / f'{id}-{tag}.json', CACHE / f'{id}-{tag}-request.json'
    save(config, dict(id=id, output=str(out), **kwargs))
    print(node('tools/wsrf/watch-data.cjs', config).strip(), flush=True)
    return json.loads(out.read_text())


def examples(id, split, n):
    ix = pt.indexrng.Index(NAMESPACE, backend='blake2b')
    u = np.asarray(ix.words(f'{id}/{split}', n*13), dtype=np.float64).reshape(n, 13) / 2**32
    return request(id, split, values=u.tolist())['rows']


def episodes(id, split, n, **kwargs):
    ix = pt.indexrng.Index(NAMESPACE, backend='blake2b')
    seeds = [f'{split}-{int(w)}' for w in ix.words(f'{id}/{split}', n)]
    return request(id, split, seeds=seeds, **kwargs)


def train(id):
    print(f'\nTraining {id} with pytoolz.wsrf', flush=True)
    config = json.loads(node('tools/wsrf/watch-config.cjs'))[id]
    rows = examples(id, 'train-states', 32000)
    rows += episodes(id, 'train-days', 256, collect=True)['rows']
    rows += episodes(id, 'train-guided', 128, collect=True, guided=True)['rows']
    vx, vz, vy = arrays(examples(id, 'validation-states', 6000))
    candidates, best, best_key = [], None, None
    path = CACHE / f'{id}-candidate.json'
    for trees, depth in [(16, 14), (32, 18)]:
        x, z, y = arrays(rows)
        m = fit(x, z, y, trees, depth)
        e = export(m, config['features'], list(range(len(config['actions']))), id)
        save(path, e)
        recovered = episodes(id, 'train-recovery', 64, model=str(path), collect=True, guided=True)
        augmented = rows + recovered['rows']
        ax, az, ay = arrays(augmented)
        m = fit(ax, az, ay, trees, depth)
        e = export(m, config['features'], list(range(len(config['actions']))), id)
        save(path, e)
        played = episodes(id, 'validation-days', 32, model=str(path))['summary']
        guided = episodes(id, 'validation-guided', 32, model=str(path), guided=True)['summary']
        acc, _ = accuracy(m, vx, vz, vy)
        candidate = dict(trees_per_zone=trees, max_depth=depth, state_accuracy=acc,
                         play=played, guided=guided, training_states=len(ax), bytes=path.stat().st_size)
        candidates.append(candidate)
        print('  candidate:', candidate, flush=True)
        key = (played['wins']+guided['wins'], -played['mean_time']-guided['mean_time'], acc, -path.stat().st_size)
        if best_key is None or key > best_key:
            best_key, best = key, (m, e, len(ax))
    m, e, trained = best
    save(path, e)
    test = episodes(id, 'test-days', 64, model=str(path), collect=True)
    guided = episodes(id, 'test-guided', 64, model=str(path), collect=True, guided=True)
    reference = episodes(id, 'test-days', 64)
    baseline = episodes(id, 'test-days', 64, random=True)
    held = examples(id, 'test-states', 10000)
    tx, tz, ty = arrays(held)
    acc, pred = accuracy(m, tx, tz, ty)
    assert test['summary']['wins'] >= 62, test['summary']
    assert guided['summary']['wins'] >= 62, guided['summary']
    assert acc >= .80, acc
    # Include states actually visited by the selected model, independent
    # synthetic states, and every tree boundary in Python/browser fixtures.
    px, pz, py = arrays(held + test['rows'] + guided['rows'])
    confusion = np.zeros((len(config['actions']), len(config['actions'])), dtype=int)
    for a, b in zip(ty, pred):
        confusion[a,b] += 1
    report = dict(dataset_namespace=NAMESPACE, training_rows=trained, initial_synthetic_rows=32000,
                  validation_rows=len(vx), held_out_rows=len(tx), candidates=candidates,
                  held_out_teacher_accuracy=acc, confusion_matrix_actual_rows=confusion.tolist(),
                  test={k:v for k,v in test.items() if k!='rows'}, guided_test={k:v for k,v in guided.items() if k!='rows'},
                  teacher_reference=reference, random_baseline=baseline,
                  library='pytoolz.wsrf.WSRFClassifier', task_labels=config['actions'],
                  features=config['features'], guide_inputs='one-hot priority, weather, supplies',
                  teacher='tools/wsrf/watch-data.cjs', learning='supervised imitation with on-policy teacher relabeling')
    report = publish(id, m, e, report, px, pz)
    source_files=['shared/watch-worlds.js','shared/wsrf.js','tools/wsrf/train-watch.py','tools/wsrf/watch-data.cjs','tools/wsrf/train.py']
    report['arcade_source_sha256']={f:hashlib.sha256((ROOT/f).read_bytes()).hexdigest() for f in source_files}
    save(ROOT/'games'/id/'training/report.json',report)
    save(ROOT/'games'/id/'stats.json',dict(examples=trained,accuracy=acc,trees=report['trees'],test_wins=test['summary']['wins'],test_days=64,guided_wins=guided['summary']['wins'],parity=report['export_validation']))


if __name__ == '__main__':
    parser=argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--game',choices=['all',*IDS],default='all')
    args=parser.parse_args();CACHE.mkdir(parents=True,exist_ok=True)
    for id in IDS if args.game=='all' else [args.game]:
        train(id)