"""Edit STATE, QUESTION, and OPTIONS, then run: python run.py."""
from __future__ import annotations

import json
from pathlib import Path

import torch
from safetensors.torch import load_file
from transformers import AutoTokenizer

from data_utils import make_collate
from decision_model import DecisionModel, ModelConfig

# ----- Edit only these values -------------------------------------------------
STATE = "I was charged twice for the same purchase."
QUESTION = "Which issue is this?"
OPTIONS = ["delivery problem", "duplicate charge", "wrong item"]
# -----------------------------------------------------------------------------


def load_decision_maker(model_dir: str | Path):
    model_dir = Path(model_dir)
    config = json.loads((model_dir / "model_config.json").read_text(encoding="utf-8"))
    tokenizer = AutoTokenizer.from_pretrained(model_dir / "tokenizer")
    model = DecisionModel(ModelConfig(**config["model_config"]))
    model.resize_token_embeddings(len(tokenizer))
    model.load_state_dict(load_file(model_dir / "model.safetensors"))
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    return model.to(device).eval(), tokenizer, config, device


def decide(model, tokenizer, config, device, state: str, question: str, options: list[str]):
    row = {"id": "request", "state": state, "question": question, "options": options, "label": 0}
    batch = make_collate(tokenizer, config["max_length"])([row])
    tensors = {key: value.to(device) for key, value in batch.items() if isinstance(value, torch.Tensor) and key != "labels"}
    temperature = float(config.get("temperature", 1.0))
    with torch.no_grad():
        probabilities = (model(**tensors)[0] / temperature).softmax(-1).cpu()
    distribution = {option: round(probabilities[index].item(), 6) for index, option in enumerate(options)}
    return {"choice": options[int(probabilities.argmax())], "temperature": temperature, "probabilities": distribution}


if __name__ == "__main__":
    root = Path(__file__).parent
    model, tokenizer, config, device = load_decision_maker(root)
    print(json.dumps(decide(model, tokenizer, config, device, STATE, QUESTION, OPTIONS), indent=2))