#!/usr/bin/env python3
"""Benchmark per-tick training cost vs readout vocab size (decides 8k vs 16k)."""
import sys
import time

sys.path.insert(0, "/home/z/my-project/download/spira")

import numpy as np

from spira.spiking import SpikingCore

rng = np.random.default_rng(0)
T = 3000

for V in (4097, 8193, 16385):
    core = SpikingCore(d_in=1024, k_in=12, n_res=512, n_out=V, seed=1)
    neurons = rng.integers(0, 1024, size=(T, 12)).astype(np.int32)
    cols = rng.integers(0, V - 1, size=T).astype(np.int64)
    core.train_pass(neurons[:200], cols[:200])          # warm-up
    t0 = time.perf_counter()
    core.train_pass(neurons, cols)
    dt = time.perf_counter() - t0
    print(f"n_out={V:>6}: {T/dt:7.0f} ticks/s  ({dt/T*1000:.2f} ms/tick) "
          f"-> 250k ticks = {250_000/(T/dt)/60:.1f} min")