"""Resume vehicles inference-in-KNP experiment on a verified CPU or CUDA backend; never evaluates test."""
import argparse
import fcntl
import json
import os
from pathlib import Path
import shutil
import subprocess
import sys
import time


def main():
    p=argparse.ArgumentParser()
    p.add_argument('--root',type=Path,default=Path.home()/'run/altai-objects/native-cuda')
    p.add_argument('--model',type=Path,default=Path.home()/'run/altai-objects/model_k200')
    p.add_argument('--data',type=Path,default=Path.home()/'run/cifar/vehicles')
    p.add_argument('--build',type=Path,default=Path(__file__).resolve().parents[1]/'build')
    p.add_argument('--backend',choices=['cpu','cuda'],default='cuda')
    p.add_argument('--jobs',type=int,default=2)
    p.add_argument('--chunk',type=int,default=1000)
    a=p.parse_args()
    a.root.mkdir(parents=True,exist_ok=True)
    lock=(a.root/'.pipeline.lock').open('w')
    fcntl.flock(lock,fcntl.LOCK_EX|fcntl.LOCK_NB)
    scripts=Path(__file__).resolve().parent
    model=a.root/'model_readout'
    counts=a.root/'counts'
    evaluation=a.root/'validation'
    status=a.root/'pipeline.json'
    started=time.time()
    phase='starting'
    def report(done=False,failed=False,error=None):
        temporary=status.with_suffix('.tmp')
        temporary.write_text(json.dumps({'phase':phase,'done':done,'failed':failed,'error':error,'pid':os.getpid(),
                                        'started_unix':started,'updated_unix':time.time(),'backend':a.backend,
                                        'counts':str(counts),'model_readout':str(model),'evaluation':str(evaluation)},indent=2)+'\n')
        temporary.replace(status)
        print(phase,flush=True)
    def run(name,arguments):
        nonlocal phase
        phase=name;report()
        with (a.root/f'{name}.log').open('w') as log:
            subprocess.run([sys.executable,*map(str,arguments)],stdout=log,stderr=subprocess.STDOUT,check=True,
                           env=dict(os.environ,OPENBLAS_NUM_THREADS='8',OMP_NUM_THREADS='8'))
    try:
        run('collecting', [scripts/'run_counts.py',a.model,counts,a.chunk,'--data',a.data,'--build',a.build,'--jobs',a.jobs,'--backend',a.backend])
        if not model.exists():
            shutil.copytree(a.model,model)
        run('training_readout',[scripts/'train_readout.py',model,counts,'--data',a.data])
        benchmark=a.root/f'benchmark_readout_{int(time.time())}'
        # A full-network CPU/CUDA gate after fitting readout, before large validation.
        run('checking_readout',[scripts/'benchmark_cuda.py',model,benchmark,'--data',a.data,'--build',a.build,'--counts','8','32'])
        run('validating',[scripts/'evaluate_knp.py',model,evaluation,'--data',a.data,'--build',a.build,'--jobs',a.jobs,
                          '--chunk',a.chunk,'--backend',a.backend])
        phase='completed';report(done=True)
    except Exception as error:
        report(failed=True,error=str(error))
        raise


if __name__=='__main__':
    main()