"""One-time evaluation on the fresh cats/dogs test (600 images). Settings were fixed on validation before.

A: learned features K 800 with contrast division, standardised, ridge lam 300 (no spikes; upper reference).
B: K 800 without contrast division (folds into neuron weights), alpha 0.25, ridge lam 300.
C: fully spiking, exact step-by-step simulation (spiking_pipeline.py): layer-1 threshold 108.08, readout ridge
   lam 1e6 on exact layer-1 spike counts, readout neuron threshold 4 × median per-step input (unit from validation).

usage: ~/anvenv/bin/python final_fresh_test.py
"""
import json
import sys

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

root = "/home/int/run/cats-dogs"
fresh = np.load(f"{root}/data-v3/fresh_test.npz")
xf, yf = fresh["rgb"], fresh["labels"].astype(int)
tr, va = np.load(f"{root}/data-v3/train.npz"), np.load(f"{root}/data-v3/val.npz")
yt = np.r_[tr["labels"], tr["labels"]].astype(int)
yv = va["labels"].astype(int)
res = {}


def ridge(f, y, lam):
    f = np.hstack([f, np.ones((len(f), 1), np.float32)])
    return np.linalg.solve(f.T @ f + lam * np.eye(f.shape[1]), f.T @ (np.eye(2)[y] * 2 - 1))


def acc(w, f, y):
    f = np.hstack([f, np.ones((len(f), 1), np.float32)])
    return 100 * ((f @ w).argmax(1) == y).mean()


# A
d = np.load(f"{root}/v3_k800.npz")
model = {k: d[k] for k in ("mean", "zca", "d")}
kf.CONTRAST_DIVIDE = True
mu, sd = d["Ft"].mean(0), d["Ft"].std(0) + 1e-3
w = ridge((d["Ft"] - mu) / sd, yt, 300)
res["A"] = {"val": acc(w, (d["Fv"] - mu) / sd, yv), "fresh": acc(w, (kf.encode(xf, model) - mu) / sd, yf)}
print("A", res["A"], flush=True)

# B
d = np.load(f"{root}/v3_k800_lin.npz")
model = {k: d[k] for k in ("mean", "zca", "d")}
kf.CONTRAST_DIVIDE = False
ft = np.vstack([kf.encode(tr["rgb"], model, alpha=0.25), kf.encode(tr["rgb"][:, :, :, ::-1], model, alpha=0.25)])
mu, sd = ft.mean(0), ft.std(0) + 1e-3
w = ridge((ft - mu) / sd, yt, 300)
res["B"] = {"val": acc(w, (kf.encode(va["rgb"], model, alpha=0.25) - mu) / sd, yv),
            "fresh": acc(w, (kf.encode(xf, model, alpha=0.25) - mu) / sd, yf)}
print("B", res["B"], flush=True)

# C
s = np.load(f"{root}/spiking_v3_k800.npz")
ct, sv = s["ct"], s["sv"]
f = lambda c: np.hstack([c.astype(np.float32), np.ones((len(c), 1), np.float32)])
sol = np.linalg.solve(f(ct).T @ f(ct) + 1e6 * np.eye(ct.shape[1] + 1), f(ct).T @ (s["yt"] * 2 - 1.0))
rq = np.rint(sol / np.abs(sol).max() * 255).astype(np.int64)
bias = int(np.rint(rq[-1] / sp.ACTIVE))
unit = float(np.median(np.abs(sv[:, :sp.ACTIVE].reshape(-1, sv.shape[-1]) @ rq[:-1] + bias))) + 1
thr = 4 * unit
sf, tf = sp.layer1_counts(xf, s["wq"], s["bq"], float(s["thr1"]), per_step=True)
res["C"] = {"val": 100 * (sp.readout_neuron(sv, rq[:-1], bias, thr) == yv).mean(),
            "fresh": 100 * (sp.readout_neuron(sf, rq[:-1], bias, thr) == yf).mean(),
            "layer1_spikes_per_image": int(tf.mean()), "readout_threshold": thr}
print("C", res["C"], flush=True)
json.dump(res, open(f"{root}/final_fresh_test.json", "w"), indent=1)