"""Test-time shift voting in the offline AltAI simulator.

Each image is fed several times with small shifts; votes of all copies are summed.
Also measures robustness: the test images themselves are shifted / thickened first.

usage: python tta_sim.py <dump_dir> <images.bin> <labels.target> --first 30000 --test 2000
"""
import argparse
import sys

import numpy as np

sys.path.insert(0, "/mnt/c/knp/viz")
import altai_sim as s  # noqa: E402

CROSS = [(0, 0), (0, 1), (0, -1), (1, 0), (-1, 0)]
SQUARE = [(dy, dx) for dy in (-1, 0, 1) for dx in (-1, 0, 1)]


def shift(a, dy, dx):
    return np.roll(np.roll(a, dy, -2), dx, -1)


def grow(a):
    return np.maximum.reduce([shift(a, dy, dx) for dy, dx in CROSS])


def thin(a):
    return np.minimum.reduce([shift(a, dy, dx) for dy, dx in CROSS])


def evaluate(w, thr, imgs, labels, offsets, inhib):
    sim = s.Sim(w, thr, inhib)
    correct = silent = 0
    for img, lab in zip(imgs, labels):
        votes = sum(sim.run_image(shift(img, dy, dx).reshape(-1))[0] for dy, dx in offsets)
        p = s.predict(votes)
        correct += p == lab
        silent += p < 0
    return 100 * correct / len(labels), 100 * silent / len(labels)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("dump")
    ap.add_argument("images")
    ap.add_argument("labels")
    ap.add_argument("--first", type=int, default=30000)
    ap.add_argument("--test", type=int, default=2000)
    ap.add_argument("--inhib", type=float, default=255)
    a = ap.parse_args()
    w, thr = s.load_dump(a.dump)
    imgs, labels = s.load_images(a.images, a.labels, a.first, a.test)
    imgs = imgs.reshape(-1, 28, 28).astype(np.int32)
    variants = {"as is": imgs, "shift 2px": shift(imgs, 0, 2), "thicker 1px": grow(imgs), "thinner 1px": thin(imgs)}
    for vote_name, offsets in [("1 copy", [(0, 0)]), ("5 copies (cross)", CROSS), ("9 copies (3x3)", SQUARE)]:
        row = []
        for name, x in variants.items():
            acc, sil = evaluate(w, thr, x.clip(0, 255).astype(np.uint8), labels, offsets, a.inhib)
            row.append(f"{name} {acc:5.1f}%")
        print(f"{vote_name:17s} | " + " | ".join(row), flush=True)


if __name__ == "__main__":
    main()