"""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>" ...
"""
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():
    out, images, labels_path, *specs = sys.argv[1:]
    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))
    data = {
        "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()