"""Evaluate class spikes emitted by real KNP; does not implement recognition in Python."""
import argparse
from concurrent.futures import ThreadPoolExecutor, as_completed
import csv
import fcntl
import json
import os
from pathlib import Path
import subprocess
import time
from run_counts import digest


def main():
    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('--backend', choices=['cpu','cuda'], default='cpu')
    p.add_argument('--jobs', type=int)
    p.add_argument('--chunk', type=int, default=100)
    p.add_argument('--limit', type=int)
    a = p.parse_args()
    if a.jobs is None: a.jobs = 2 if a.backend == "cuda" else 12
    if not 1 <= a.jobs <= 64 or a.chunk < 1 or (a.limit is not None and a.limit < 1):
        p.error('Invalid jobs, chunk or limit')
    a.model, a.out, a.data, a.build = (v.resolve() for v in (a.model,a.out,a.data,a.build))
    a.out.mkdir(parents=True,exist_ok=True)
    lock=(a.out/'.lock').open('w')
    fcntl.flock(lock,fcntl.LOCK_EX|fcntl.LOCK_NB)
    k, stride, threshold, classes, out_threshold=map(int,(a.model/'meta.txt').read_text().split())
    if out_threshold < 1:
        raise ValueError('Train readout first')
    y=[int(v) for v in (a.data/'val_labels.txt').read_text().split()]
    if a.limit:
        y=y[:a.limit]
    n=len(y)
    exe=a.build/'bin/altai_infer'
    backend=a.build/'bin'/('libknp-altai-cuda-backend.so' if a.backend=='cuda' else 'libknp-cpu-single-threaded-backend.so')
    signature={'count':n,'chunk':a.chunk,'files':{str(v):digest(v) for v in [exe,backend,a.data/'val_images.bin',a.data/'val_labels.txt',*(a.model/name for name in ['meta.txt','det_w.bin','det_b.bin','readout_w.bin','readout_b.bin'])]}}
    config=a.out/'settings.json'
    if config.exists() and json.loads(config.read_text()) != signature:
        raise ValueError('Different model/data/build; use a new evaluation folder')
    config.write_text(json.dumps(signature,indent=2)+'\n')
    jobs=[(first,min(a.chunk,n-first)) for first in range(0,n,a.chunk)]
    def read(first,count):
        with (a.out/f'val_{first}_classes.csv').open() as f:
            rows=list(csv.DictReader(f))
        if len(rows)!=count or [int(r['index']) for r in rows] != list(range(first,first+count)):
            raise ValueError('Incomplete or misindexed class CSV')
        return [[int(r[f'c{c}']) for c in range(classes)] for r in rows]
    def valid(first,count):
        if not (a.out/f'val_{first}.done').exists():
            return False
        try:
            read(first,count)
            return True
        except (OSError,KeyError,ValueError):
            return False
    finished=sum(count for first,count in jobs if valid(first,count))
    start=time.time()
    def progress(failed=False,error=None):
        temporary=a.out/'progress.tmp'
        temporary.write_text(json.dumps({'current':finished,'total':n,'done':finished==n and not failed,'failed':failed,'error':error,'started_unix':start,'updated_unix':time.time()})+'\n')
        temporary.replace(a.out/'progress.json')
    def one(first,count):
        prefix=a.out/f'val_{first}'
        env=dict(os.environ,LD_LIBRARY_PATH=str(a.build/'bin')+':'+os.environ.get('LD_LIBRARY_PATH',''))
        with Path(str(prefix)+'.log').open('w') as f:
            subprocess.run([str(exe),str(a.data/'val_images.bin'),str(first),str(count),str(a.model),str(backend),str(prefix)],env=env,stdout=f,stderr=subprocess.STDOUT,check=True)
        read(first,count)
        (a.out/f'val_{first}.done').write_text('ok\n')
        return count
    progress()
    errors=[]
    with ThreadPoolExecutor(max_workers=a.jobs) as pool:
        futures=[pool.submit(one,first,count) for first,count in jobs if not valid(first,count)]
        for future in as_completed(futures):
            try:
                finished+=future.result()
            except Exception as error:
                errors.append(str(error))
            progress(bool(errors),'; '.join(errors) if errors else None)
    if errors:
        raise RuntimeError('; '.join(errors))
    counts=[row for first,count in jobs for row in read(first,count)]
    predictions=[max(range(classes),key=lambda c:row[c]) for row in counts]
    correct=sum(a==b for a,b in zip(predictions,y))
    metrics={'source':'KNP class spike CSV, argmax; first class breaks ties','count':n,'correct':correct,'accuracy':correct/n,
             'silent':sum(max(row)==0 for row in counts),'ties':sum(row.count(max(row))>1 for row in counts),
             'per_class':{str(c):{'count':sum(v==c for v in y),'correct':sum(t==c and pred==c for t,pred in zip(y,predictions))} for c in range(classes)}}
    (a.out/'metrics.json').write_text(json.dumps(metrics,indent=2)+'\n')
    print(json.dumps(metrics),flush=True)


if __name__=='__main__':
    main()