"""Choose inference threshold on validation only using the existing AltAI simulator.

Weights come from real KNP STDP training. This script performs no weight training.
The chosen configuration must be checked with real KNP before reporting test results.
"""
import argparse
import csv
import json
from pathlib import Path
import sys

import numpy as np

sys.path.insert(0, '/mnt/c/KNP/viz')
import altai_sim

altai_sim.PIX = 2048
altai_sim.CLASSES = 2
ap = argparse.ArgumentParser()
ap.add_argument('run', type=Path)
a = ap.parse_args()
data = Path('/home/int/run/cats-dogs/data')
w, original_thr = altai_sim.load_dump(a.run)
images = np.fromfile(data / 'val.bin', dtype=np.uint8).reshape(-1, 2048)
labels = np.loadtxt(data / 'val.target', dtype=int)
thresholds = np.array([10, 20, 40, 60, 80, 100, 150, 200, 255], dtype=float)
pot = np.zeros((len(thresholds), w.shape[1]))
prev_input = np.zeros_like(pot, dtype=bool)
prev_current = np.zeros(w.shape[1])
states = np.zeros(2048, dtype=np.float32)
correct = np.zeros(len(thresholds), dtype=int)
unanswered = correct.copy()
predictions = np.empty((len(images), len(thresholds)), dtype=int)
for i, (img, label) in enumerate(zip(images, labels)):
    currents = altai_sim.raster(img, states) @ w
    votes = np.zeros((len(thresholds), 2), dtype=int)
    for t in range(15):
        votes += prev_input.reshape(len(thresholds), 2, -1).any(axis=2)
        pot = np.round(pot) + prev_current
        pot -= 255 * (prev_input.sum(axis=1, keepdims=True) - prev_input)
        fired = pot >= thresholds[:, None]
        pot[fired] = 0
        pot = np.maximum(pot, 0)
        prev_input, prev_current = fired, currents[t]
    pred = np.where(votes.max(axis=1) > 0, votes.argmax(axis=1), -1)
    predictions[i] = pred
    correct += pred == label
    unanswered += pred == -1
best = int(np.argmax(correct))  # stable tie break: lower threshold
rows = [{'threshold': int(t), 'correct': int(c), 'unanswered': int(u),
         'wrong': int(len(labels) - c - u), 'accuracy': float(c / len(labels))}
        for t, c, u in zip(thresholds, correct, unanswered)]
result = {'method': 'validation-only threshold sweep; offline simulation, not KNP measurements',
          'selected_threshold': int(thresholds[best]), 'rows': rows}
(a.run / 'validation-sweep.json').write_text(json.dumps(result, indent=2) + '\n')
np.save(a.run / 'validation-predictions.npy', predictions[:, best])
print(json.dumps(result), flush=True)