"""Run identical native KNP networks on CPU/CUDA and compare binary detector/class outputs."""
import argparse
import json
import os
from pathlib import Path
import re
import subprocess
import time
from run_counts import digest

p=argparse.ArgumentParser()
p.add_argument('model',type=Path)
p.add_argument('out',type=Path)
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('--counts',type=int,nargs='+',default=[8,32])
a=p.parse_args()
a.out.mkdir(parents=True,exist_ok=False)
exe=a.build/'bin/altai_infer'
env=dict(os.environ,LD_LIBRARY_PATH=str(a.build/'bin')+':'+os.environ.get('LD_LIBRARY_PATH',''))
rows=[]
for count in a.counts:
    pair={}
    for backend in ['cpu','cuda']:
        library=a.build/'bin'/('libknp-cpu-single-threaded-backend.so' if backend=='cpu' else 'libknp-altai-cuda-backend.so')
        prefix=a.out/f'{backend}_{count}'
        command=[str(exe),str(a.data/'val_images.bin'),'0',str(count),str(a.model),str(library),str(prefix)]
        start=time.monotonic()
        with prefix.with_suffix('.log').open('w') as log:
            subprocess.run(command,env=env,stdout=log,stderr=subprocess.STDOUT,check=True)
        wall=time.monotonic()-start
        log=prefix.with_suffix('.log').read_text()
        match=re.search(r'images \d+\.\.\d+, ([0-9.eE+-]+) s',log)
        pair[backend]={'wall_seconds':wall,'execution_seconds':float(match.group(1)) if match else None,
                       'command':command,'backend_sha256':digest(library)}
    files=['_counts.bin']
    if (a.out/f'cpu_{count}_classes.csv').exists(): files.append('_classes.csv')
    same={suffix:(a.out/f'cpu_{count}{suffix}').read_bytes()==(a.out/f'cuda_{count}{suffix}').read_bytes() for suffix in files}
    row={'count':count,'timing':pair,'byte_identical':same}
    rows.append(row)
    (a.out/'benchmark.json').write_text(json.dumps(rows,indent=2)+'\n')
    print(json.dumps(row),flush=True)
    if not all(same.values()): raise RuntimeError('CPU/CUDA output mismatch')