"""Builds altai_viz.html: embeds trained weights, sample test images and full-test statistics into app.html.

usage: python3 export_data.py <out.html> <images.bin> <labels.target> "<name>|<dump_dir>|<inhib>" ... [--readout r.npz]

--readout: stroke features + spiking readout network (npz with w, thr, kernels, scale from
readout_coarseA_page.npz); its test accuracy is computed here with the same code as the experiments.
"""
import base64
import json
import sys
from pathlib import Path

import numpy as np

from altai_sim import CLASSES, Sim, load_dump, load_images, predict

GALLERY_RANDOM, GALLERY_HARD = 60, 30


def b64(arr, dtype):
    return base64.b64encode(np.ascontiguousarray(arr, dtype=dtype).tobytes()).decode()


def main():
    args = sys.argv[1:]
    readout = None
    if "--readout" in args:
        i = args.index("--readout")
        readout = args[i + 1]
        del args[i:i + 2]
    out, images, labels_path, *specs = args
    total = len(open(labels_path).read().split())
    imgs, labels = load_images(images, labels_path, total - 10000, 10000)

    models, preds = [], {}
    for spec in specs:
        name, dump, inhib = spec.split("|")
        w, thr = load_dump(dump)
        sim = Sim(w, thr, float(inhib))
        p = np.array([predict(sim.run_image(img)[0]) for img in imgs])
        preds[name] = p
        confusion = np.zeros((CLASSES, CLASSES + 1), dtype=int)  # last column: no answer
        for lab, pr in zip(labels, p):
            confusion[lab, pr if pr >= 0 else CLASSES] += 1
        models.append({
            "name": name, "inhib": float(inhib), "col": len(thr) // CLASSES,
            "thr": thr.astype(int).tolist(),
            "w": b64(w.T, np.int8),  # [neuron][pixel]
            "confusion": confusion.tolist(),
            "accuracy": float((p == labels).mean()),
        })
        print(name, f"{(p == labels).mean():.4f}")

    rng = np.random.default_rng(7)
    first, last = list(preds.values())[0], list(preds.values())[-1]
    fixed = np.flatnonzero((first != labels) & (last == labels))
    broken = np.flatnonzero(last != labels)
    pick = list(rng.choice(10000, GALLERY_RANDOM, replace=False))
    pick += list(rng.choice(fixed, min(GALLERY_HARD, len(fixed)), replace=False))
    pick += list(rng.choice(broken, min(GALLERY_HARD // 2, len(broken)), replace=False))
    pick = sorted(set(int(i) for i in pick))
    extra = {}
    if readout:
        sys.path.insert(0, "/mnt/c/knp/features")
        sys.path.insert(0, "/mnt/c/knp/readout")
        import stroke_features as sf
        import spiking_readout as sr
        r = np.load(readout)
        sf.use_preset("coarse")
        sf.P["scale"] = float(r["scale"])
        p = sr.spiking_predict(sf.features(imgs.reshape(-1, 28, 28)), r["w"], float(r["thr"]))
        confusion = np.zeros((CLASSES, CLASSES + 1), dtype=int)
        for lab, pr in zip(labels, p):
            confusion[lab, pr if pr >= 0 else CLASSES] += 1
        extra["readout"] = {"name": "Штрихи + спайковый выход", "kernels": sf.KERNELS.astype(float).round(6).tolist(),
                            "scale": float(r["scale"]), "thr": float(r["thr"]), "w": r["w"].astype(int).tolist(),
                            "confusion": confusion.tolist(), "accuracy": float((p == labels).mean())}
        print("readout", f"{(p == labels).mean():.4f}")
    data = {
        **extra,
        "models": models,
        "gallery": {"index": pick, "labels": labels[pick].tolist(), "pixels": b64(imgs[pick], np.uint8)},
    }
    html = (Path(__file__).parent / "app.html").read_text(encoding="utf-8")
    Path(out).write_text(html.replace("/*DATA*/null", json.dumps(data, separators=(",", ":"))), encoding="utf-8")
    print("written", out, len(pick), "gallery images")


if __name__ == "__main__":
    main()