#!/usr/bin/env python3
"""Score text questions with this MLX checkpoint.

The vision tower is stored at full precision in ``model.safetensors`` and is not
used here. Load requires ``mlx``, ``mlx-lm`` and ``transformers``.
"""

from __future__ import annotations

import json
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))

import d3_mlx  # noqa: E402


EXAMPLE = {
    "state": "The order arrived damaged yesterday. The customer has a receipt and asks for a replacement today.",
    "questions": {
        "route": {
            "type": "choice",
            "instructions": "Which team should handle this request?",
            "criteria": {
                "returns": "Refunds, replacements and damaged deliveries",
                "billing": "Payments, invoices and charges",
                "technical": "Product setup and faults",
            },
        },
        "receipt": {"type": "noul", "instructions": "Does the customer have a receipt?"},
        "urgency": {
            "type": "score",
            "instructions": "How urgent is this request?",
            "criteria": ["Routine", "Soon", "Today"],
        },
    },
}


def option_labels(question: dict) -> list[str]:
    rendered = d3_mlx.normalize_question(question)
    if rendered["type"] == "noul":
        return ["false", "true"]
    return list(rendered["criteria"])


def predict(model, readout, tokenizer, codes, state, questions: dict) -> dict:
    answers = {}
    for key, question in questions.items():
        text, n_options = d3_mlx.render_prompt(tokenizer, state, question, codes)
        token_ids = tokenizer.encode(text, add_special_tokens=False)
        probabilities = d3_mlx.probs_from_ids(model, readout, token_ids, n_options)
        labels = option_labels(question)
        pairs = {label: float(value) for label, value in zip(labels, probabilities)}
        best = max(pairs, key=pairs.get)
        answers[key] = {"best": best, "probabilities": pairs}
    return answers


def main() -> None:
    from transformers import AutoTokenizer

    root = Path(__file__).resolve().parent
    model, _config, decision, readout = d3_mlx.load_model(root)
    tokenizer = AutoTokenizer.from_pretrained(str(root))
    payload = EXAMPLE
    if len(sys.argv) > 1:
        payload = json.loads(Path(sys.argv[1]).read_text(encoding="utf-8"))
    answers = predict(
        model,
        readout,
        tokenizer,
        decision["codes"],
        payload["state"],
        payload["questions"],
    )
    print(json.dumps(answers, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()