#!/usr/bin/env python3
"""Check the card against float32 sentence-transformers references shipped in refs/.

    python3 selftest.py            # all items in refs/items.json
Prints cosine to the reference for every item, retrieval top-1 per group, and timings.
"""

import json
import time
from pathlib import Path

import numpy as np

from eg2_host import Eg2

HERE = Path(__file__).resolve().parent


def main():
    import argparse

    ap = argparse.ArgumentParser()
    ap.add_argument("--skip", default="", help="comma list of encoders not to load/test: vision,audio")
    skip = set(filter(None, ap.parse_args().skip.split(",")))
    refs = HERE / "refs"
    items = json.loads((refs / "items.json").read_text(encoding="utf-8"))
    ref = np.load(refs / "ref.npy")
    keep = [i for i, it in enumerate(items)
            if not ("vision" in skip and it.get("images")) and not ("audio" in skip and it.get("audios"))
            and not (it.get("group") == "image" and "vision" in skip) and not (it.get("group") == "audio" and "audio" in skip)]
    items, ref = [items[i] for i in keep], ref[keep]
    eg = Eg2(HERE / "assets", load=tuple(k for k in ("text", "vision", "audio") if k not in skip))
    try:
        got, times = [], {}
        for it in items:
            t0 = time.perf_counter()
            v = eg.encode(text=it.get("text", ""), prompt=it.get("prompt"),
                          images=[refs / "files" / f for f in it.get("images", [])],
                          audios=[refs / "files" / f for f in it.get("audios", [])])
            times.setdefault(it["kind"], []).append((time.perf_counter() - t0) * 1000)
            got.append(v)
        got = np.stack(got)
        cos = (got * ref).sum(1)
        print(f"{'kind':10s} {'n':>3s} {'cos min':>8s} {'cos mean':>9s} {'ms avg':>8s}")
        for kind in times:
            sel = [i for i, it in enumerate(items) if it["kind"] == kind]
            print(f"{kind:10s} {len(sel):3d} {cos[sel].min():8.5f} {cos[sel].mean():9.5f} {np.mean(times[kind]):8.1f}")
        for g in sorted({it["group"] for it in items if it.get("group")}):
            q = [i for i, it in enumerate(items) if it.get("group") == g and it.get("role") == "query"]
            d = [i for i, it in enumerate(items) if it.get("group") == g and it.get("role") == "doc"]
            want = [items[i]["answer"] for i in q]
            for name, e in (("reference", ref), ("card", got)):
                top = (e[q] @ e[d].T).argmax(1)
                acc = np.mean([items[d[t]]["id"] == w for t, w in zip(top, want)])
                print(f"retrieval {g:14s} {name:9s} top-1 {acc:.2f}")
        ok = cos.min() > 0.99
        print("PASS" if ok else "FAIL", f"(min cosine {cos.min():.5f})")
    finally:
        eg.close()


if __name__ == "__main__":
    main()