import json
import sys

import numpy as np
import onnxruntime as ort
import sentencepiece as spm
from huggingface_hub import hf_hub_download

REPO = "Prosperitytech/lighthouse-keeper"

cfg = json.load(open(hf_hub_download(REPO, "config.json")))
sp = spm.SentencePieceProcessor(model_file=hf_hub_download(REPO, "keeper.spm"))
session = ort.InferenceSession(hf_hub_download(REPO, "keeper.onnx"), providers=["CPUExecutionProvider"])


def score(text):
    ids = [cfg["bos_id"]] + sp.encode(text)[: cfg["max_tokens"] - 1]
    probs = session.run(None, {"input_ids": np.array([ids], dtype=np.int64)})[0][0]
    return {cat: dict(zip(cfg["levels"], map(float, p))) for cat, p in zip(cfg["categories"], probs)}


def flag(text, min_risk=cfg["suggested_min_risk"]):
    priority = cfg["tie_break_order"]
    best = None
    for cat, p in score(text).items():
        level = max(p, key=p.get)
        risk = 1 - p["none"]
        if level != "none" and risk >= min_risk:
            key = (cfg["levels"].index(level), p[level] + 0.06 * (1 - priority.index(cat) / len(priority)))
            if best is None or key > best[1]:
                best = (cat, key, p[level])
    return None if best is None else {"category": best[0], "level": cfg["levels"][best[1][0]], "confidence": round(best[2], 3)}


if __name__ == "__main__":
    print(flag(" ".join(sys.argv[1:]) or "send me your BVN make I help you collect the money"))