"""Run real KNP training/inference and preserve logs, settings and count-based metrics."""
import argparse
import csv
import io
import json
import os
from pathlib import Path
import subprocess
import time

ap = argparse.ArgumentParser()
ap.add_argument('name')
ap.add_argument('--dataset', default='validation')
ap.add_argument('--train', type=int, default=8000)
ap.add_argument('--count', type=int, default=2000)
ap.add_argument('--threshold', type=int, default=60)
ap.add_argument('--data-dir', type=Path, default=Path('/home/int/run/cats-dogs/data'))
ap.add_argument('--columns', type=int, default=20)
ap.add_argument('--seed', type=int, default=42)
ap.add_argument('--jitter', type=float, default=20)
ap.add_argument('--quant', choices=['default', 'weights'], default='default')
ap.add_argument('--backend', choices=['single', 'multi', 'tuned', 'fast', 'delta'], default='single')
ap.add_argument('--threads', type=int, default=2)
a = ap.parse_args()
if a.backend == 'multi':
    ap.error('Stock multi-threaded KNP has no AltAI types; use --backend tuned for the experimental AltAI port')
if not 1 <= a.threads <= 256:
    ap.error('--threads must be between 1 and 256')
root = Path('/home/int/run/cats-dogs')
out = root / a.name
out.mkdir(exist_ok=False)
env = dict(os.environ, LD_LIBRARY_PATH='/home/int/knp-build/bin:/home/int/deps/lib',
           KNP_COL_JITTER=str(a.jitter), KNP_COL_SEED=str(a.seed), KNP_COL_SIZE=str(a.columns),
           KNP_MNIST_THR=str(a.threshold), KNP_MNIST_ADD_THR_SCALE='0',
           KNP_MNIST_INFER_INHIB='255', KNP_MNIST_DUMP=str(out))
if a.quant == 'weights':
    env['KNP_MNIST_QUANT'] = 'weights'
else:
    env.pop('KNP_MNIST_QUANT', None)
env['KNP_CPU_THREADS'] = str(a.threads)
backends = {'single': '/home/int/knp-build/bin/libknp-cpu-single-threaded-backend.so',
            'multi': '/home/int/knp-build/bin/libknp-cpu-multi-threaded-backend.so',
            'tuned': str(root / 'acceleration/libknp-cpu-tuned-backend.so'),
            'fast': str(root / 'fast-cpu/libknp-cpu-fast-backend.so'),
            'delta': '/home/int/run/delta-cpu/libknp-cpu-delta-backend.so'}
command = [str(root / 'build/cats_dogs'), '-m', 'altai',
           '--images', str(a.data_dir / (a.dataset + '.bin')),
           '--labels', str(a.data_dir / (a.dataset + '.target')),
           '-t', str(a.train), '-i', str(a.count), '-l', 'error',
           '--training_backend', backends[a.backend]]
settings = {'arguments': {k: str(v) if isinstance(v, Path) else v for k, v in vars(a).items()}, 'command': command,
            'environment': {k: v for k, v in env.items() if k.startswith('KNP_')},
            'started_utc': time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime())}
(out / 'settings.json').write_text(json.dumps(settings, indent=2) + '\n')
start = time.monotonic()
with (out / 'run.log').open('w') as log:
    result = subprocess.run(command, input='\n', text=True, env=env, stdout=log, stderr=subprocess.STDOUT)
text = (out / 'run.log').read_text()
if result.returncode:
    raise RuntimeError(f'KNP exited {result.returncode}: {text[-2000:]}')
lines = text[text.index('CLASS,TOTAL_VOTES'):].splitlines()[:3]
(out / 'result.csv').write_text('\n'.join(lines) + '\n')
rows = list(csv.DictReader(io.StringIO('\n'.join(lines))))
metrics = {'elapsed_seconds': round(time.monotonic() - start, 2), 'classes': {}}
for row in rows:
    label = 'cat' if row['CLASS'] == '0' else 'dog'
    metrics['classes'][label] = {'correct': int(row['TRUE_POSITIVES']),
                                  'unanswered': int(row['FALSE_NEGATIVES']),
                                  'wrong': int(row['FALSE_POSITIVES'])}
for key in ('correct', 'unanswered', 'wrong'):
    metrics[key] = sum(row[key] for row in metrics['classes'].values())
assert metrics['correct'] + metrics['unanswered'] + metrics['wrong'] == a.count
metrics['accuracy'] = metrics['correct'] / a.count
(out / 'metrics.json').write_text(json.dumps(metrics, indent=2) + '\n')
print(json.dumps(metrics), flush=True)