"""Train the feature detectors offline and write integer weights for altai_infer (inference runs in KNP).

k-means on 6×6×3 patches of the vehicles training images (patch centring and ZCA, no contrast division), then
centring, ZCA and the prototype are folded into one weight vector per detector type on raw pixels, the offset
and alpha into a per-step bias (objects/spiking_pipeline.py); weights are scaled to integers in ±wscale.
usage: ~/anvenv/bin/python build_model.py <model_dir> [--k 200 --stride 2 --alpha 0.25 --wscale 127]
"""
import argparse
import sys
from pathlib import Path

import numpy as np

sys.path.insert(0, "/mnt/c/knp/objects")
import kmeans_features as kf  # noqa: E402
import spiking_pipeline as sp  # noqa: E402

ap = argparse.ArgumentParser()
ap.add_argument("out")
ap.add_argument("--k", type=int, default=200)
ap.add_argument("--stride", type=int, default=2)
ap.add_argument("--alpha", type=float, default=0.25)
ap.add_argument("--wscale", type=float, default=127)
ap.add_argument("--classes", type=int, default=4)
ap.add_argument("--data", default="/home/int/run/cifar/vehicles")
a = ap.parse_args()
out = Path(a.out)
out.mkdir(parents=True, exist_ok=True)
kf.CONTRAST_DIVIDE = False
xt = np.load(f"{a.data}/train.npz")["rgb"]
model = kf.learn(xt, k=a.k)
w, b = sp.fold(model, a.alpha)  # [108, K], [K] for patches in [0, 1]
scale = a.wscale / np.abs(w).max()
wq, bq = np.rint(w * scale).astype(np.int32), np.rint(b * scale).astype(np.int32)
p = kf.patches_at(xt[:300].astype(np.float32) / 255, range(0, 27, a.stride), range(0, 27, a.stride)).reshape(-1, 108)
act = p @ w + b
thr = int(round(float(np.percentile(act[act > 0], 99)) * scale))
wq.T.copy().tofile(out / "det_w.bin")
bq.tofile(out / "det_b.bin")
(out / "meta.txt").write_text(f"{a.k} {a.stride} {thr} {a.classes} 0\n")
np.savez(out / "kmeans.npz", **model)
print(f"K {a.k} stride {a.stride}: weight scale {scale:.2f}, detector threshold {thr}, bias range {bq.min()}..{bq.max()}")