"""Download the Hugging Face checkpoint and run one typed Choice prediction."""

from __future__ import annotations

import sys
from pathlib import Path

import torch
from huggingface_hub import snapshot_download


REPO_ID = "Syigen/rgcm"


def main() -> None:
    repo_path = Path(snapshot_download(REPO_ID))
    # The model repo ships this small runtime package beside the checkpoint.
    sys.path.insert(0, str(repo_path))

    from concept_memory import ConceptMemoryModel
    from concept_memory.types import Question

    checkpoint = repo_path / "model.pt"
    payload = torch.load(checkpoint, map_location="cpu", weights_only=False)

    model = ConceptMemoryModel.from_pretrained(
        payload["encoder"],
        decision_size=payload["decision_size"],
        max_length=payload["max_length"],
    )
    model.load_state_dict(payload["model_state_dict"])
    model.eval()

    question = Question(
        id="intent",
        type="choice",
        instructions="Which team should handle this customer request?",
        criteria=[
            {"id": "billing", "description": "Payment or subscription issues"},
            {"id": "technical", "description": "Bugs or integration problems"},
            {"id": "sales", "description": "Pricing or account questions"},
        ],
    )

    with torch.inference_mode():
        result = model.predict(
            "Stripe integration keeps failing after several attempts.", [question]
        )[0]

    criteria = question.normalized_criteria()
    predicted = criteria[result.prediction.item()]
    print(f"Prediction: {predicted.id} - {predicted.description}")
    print("Probabilities:", result.probabilities.tolist())


if __name__ == "__main__":
    main()