"""Robustness of the spiking readout: distort raw validation pixels, compute stroke features, run the readout.

usage: python robustness.py <raw.bin> <raw.target> <feat_train.bin> <feat_train.target> --train N [--preset coarse --scale 0.3842]
"""
import argparse
import sys

import numpy as np

sys.path.insert(0, "/mnt/c/knp/features")
sys.path.insert(0, "/mnt/c/knp/augment")
import stroke_features as sf  # noqa: E402
import spiking_readout as sr  # noqa: E402
from tta_sim import grow, shift, thin  # noqa: E402


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("raw")
    ap.add_argument("raw_labels")
    ap.add_argument("feat")
    ap.add_argument("feat_labels")
    ap.add_argument("--train", type=int, required=True)
    ap.add_argument("--eval-first", type=int, default=30000)
    ap.add_argument("--eval", type=int, default=10000)
    ap.add_argument("--preset", default="coarse")
    ap.add_argument("--scale", type=float, default=0.3842)
    ap.add_argument("--thr", type=float, default=16)
    ap.add_argument("--lam", type=float, default=10.0)
    ap.add_argument("--save")
    a = ap.parse_args()
    sf.use_preset(a.preset)
    sf.P["scale"] = a.scale
    width = 784
    xf = np.fromfile(a.feat, dtype=np.uint8).reshape(-1, width)[:a.train]
    yf = np.array([int(v) for v in open(a.feat_labels).read().split()])[:a.train]
    w = sr.train_ridge(xf, yf, 10, a.lam)
    wq, scale = sr.quantize(w)
    raw = np.fromfile(a.raw, dtype=np.uint8).reshape(-1, 28, 28)[a.eval_first:a.eval_first + a.eval].astype(np.int32)
    y = np.array([int(v) for v in open(a.raw_labels).read().split()])[a.eval_first:a.eval_first + a.eval]
    row = []
    for name, v in [("as is", raw), ("shift 2px", shift(raw, 0, 2)), ("thicker 1px", grow(raw)), ("thinner 1px", thin(raw))]:
        p = sr.spiking_predict(sf.features(v.clip(0, 255).astype(np.uint8)), wq, a.thr)
        row.append(f"{name} {100 * (p == y).mean():5.1f}%")
    print(" | ".join(row))
    if a.save:
        np.savez(a.save, w=wq, thr=a.thr, scale=scale, preset=a.preset, feat_scale=a.scale)


if __name__ == "__main__":
    main()