"""Compare CPU backends on the same small AltAI/STDP validation task."""
import argparse
import json
import os
from pathlib import Path
import signal
import subprocess
import sys
import time

ap = argparse.ArgumentParser()
ap.add_argument('--prefix', default='backend_bench')
ap.add_argument('--skip-stock', action='store_true')
a = ap.parse_args()
root = Path('/home/int/run/cats-dogs')
rows = []
cases = [('single', 'single', 1), ('stock24', 'multi', 24), ('tuned2', 'tuned', 2), ('tuned4', 'tuned', 4)]
for label, backend, threads in cases:
    if a.skip_stock and backend == 'multi':
        continue
    name = f'{a.prefix}_{label}'
    command = [sys.executable, '/mnt/c/KNP/cats-dogs/run.py', name,
               '--data-dir', str(root / 'pilot-features'), '--dataset', 'sparse_hog_validation',
               '--train', '800', '--count', '200', '--threshold', '4000', '--quant', 'weights',
               '--backend', backend, '--threads', str(threads)]
    started = time.monotonic()
    try:
        process = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                                   text=True, start_new_session=True)
        try:
            stdout, stderr = process.communicate(timeout=45)
        except subprocess.TimeoutExpired:
            os.killpg(process.pid, signal.SIGTERM)
            process.communicate()
            raise
        if process.returncode:
            raise RuntimeError(stderr[-2000:])
        row = {'backend': label, **json.loads((root / name / 'metrics.json').read_text())}
        ref = root / f'{a.prefix}_single'
        row['same_weights_and_result'] = all((ref / f).read_bytes() == (root / name / f).read_bytes()
                                              for f in ('float_weights.csv', 'quant_weights.csv', 'result.csv'))
    except subprocess.TimeoutExpired:
        row = {'backend': label, 'timeout_seconds': 45}
    except Exception as error:
        row = {'backend': label, 'error': str(error)}
    row['wall_seconds'] = round(time.monotonic() - started, 3)
    rows.append(row)
    print(json.dumps(row), flush=True)
(root / f'acceleration/{a.prefix}-benchmark.json').write_text(json.dumps(rows, indent=2) + '\n')