#!/usr/bin/env python3
import json
import sys
from pathlib import Path

import numpy as np
import onnxruntime as ort

REPO_ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(REPO_ROOT))
from features import build_features  # noqa: E402

TASK_LABELS = [
    "general_chat", "writing", "translation", "summarization", "research",
    "coding", "mathematics", "document_analysis", "high_stakes",
]
COMPLEXITY_LABELS = ["low", "medium", "high"]


def distribution(labels, values):
    return {label: float(value) for label, value in zip(labels, values, strict=True)}


def classify(text: str) -> dict:
    model_path = REPO_ROOT / "model.onnx"
    session = ort.InferenceSession(str(model_path), providers=["CPUExecutionProvider"])
    features = build_features(text).astype(np.float32).reshape(1, -1)
    outputs = session.run(None, {"features": features})
    task_probs = outputs[0][0]
    complexity_probs = outputs[1][0]
    task_index = int(np.argmax(task_probs))
    complexity_index = int(np.argmax(complexity_probs))
    return {
        "task": {
            "label": TASK_LABELS[task_index],
            "confidence": float(task_probs[task_index]),
            "probabilities": distribution(TASK_LABELS, task_probs),
        },
        "complexity": {
            "label": COMPLEXITY_LABELS[complexity_index],
            "confidence": float(complexity_probs[complexity_index]),
            "probabilities": distribution(COMPLEXITY_LABELS, complexity_probs),
        },
    }


def main() -> None:
    if len(sys.argv) < 2:
        raise SystemExit('Usage: python examples/python/classify.py "your prompt"')
    print(json.dumps(classify(" ".join(sys.argv[1:])), indent=2))


if __name__ == "__main__":
    main()