"""Run s1-llm-auto-router locally on CPU. Needs: pip install onnxruntime tokenizers numpy huggingface_hub"""
import json
import sys

import numpy as np
import onnxruntime as ort
from huggingface_hub import snapshot_download
from tokenizers import Tokenizer

CATEGORIES = ["coding", "agentic", "math", "knowledge", "long_context", "tool_use", "design", "summarisation", "general"]
DIFFICULTY = ["trivial", "easy", "moderate", "hard", "frontier"]
STAKES = ["negligible", "low", "medium", "high"]
FLAGS = ["needs_tools", "needs_vision", "needs_long_context", "follow_up"]

root = snapshot_download("decisionmodels/s1-llm-auto-router")
temps = json.load(open(f"{root}/meta.json"))["temperatures"]
tok = Tokenizer.from_file(f"{root}/tokenizer/tokenizer.json")
tok.enable_truncation(512, strategy="only_second")
sess = ort.InferenceSession(f"{root}/model.gq.onnx", providers=["CPUExecutionProvider"])


def softmax(x):
    e = np.exp(x - x.max())
    return e / e.sum()


def classify(request: str, context: str = "(new conversation)") -> dict:
    enc = tok.encode(context[:600], request)  # context: short summary, first 600 characters
    ids = np.asarray([enc.ids], dtype=np.int64)
    cat, dif, stk, flg = sess.run(None, {"input_ids": ids, "attention_mask": np.ones_like(ids)})
    pc = softmax(cat[0] / temps["category"])
    pd = softmax(dif[0] / temps["difficulty"])
    ps = softmax(stk[0] / temps["stakes"])
    pf = 1 / (1 + np.exp(-flg[0] / temps["flags"]))
    out = {
        "category": CATEGORIES[int(pc.argmax())],
        "category_confidence": round(float(pc.max()), 3),
        "difficulty": DIFFICULTY[int(pd.argmax())],
        "difficulty_expected": round(float((np.arange(5) * pd).sum()), 2),
        "stakes": STAKES[int(ps.argmax())],
        "stakes_expected": round(float((np.arange(4) * ps).sum()), 2),
    }
    out.update({name: round(float(p), 3) for name, p in zip(FLAGS, pf)})  # probability the flag is true
    return out


if __name__ == "__main__":
    text = " ".join(sys.argv[1:]) or "Refactor this ABAP report to use CDS views"
    print(json.dumps(classify(text), indent=2))