"""Load an exported adapter and answer a typed request without text generation."""

import json
from pathlib import Path

import torch

from .data import tokenizer_for
from .model import collate, load_model


class Predictor:
    def __init__(self, export, device="cuda"):
        self.export = Path(export)
        saved = json.loads((self.export / "themis_config.json").read_text())
        self.cfg = saved["config"]
        self.codec = tokenizer_for(self.cfg, self.export / "tokenizer")
        if self.codec.config() != saved["codec"]:
            raise ValueError("Saved tokenizer and option codes disagree")
        self.model = load_model(self.cfg, self.codec, device, adapter=self.export, trainable=False)
        self.model.eval()
        path = self.export / "calibration.json"
        self.calibration = (
            json.loads(path.read_text())
            if path.exists()
            else {"completed": False, "temperatures": {}}
        )

    @torch.no_grad()
    def predict(self, state, questions):
        encoded = self.codec.encode(state, questions)
        device = next(self.model.parameters()).device
        batch = collate([encoded], self.codec.pad, device)
        with torch.autocast(device.type, dtype=torch.bfloat16, enabled=device.type == "cuda"):
            logits = self.model(batch).float().cpu()
        answers = {}
        for index, (name, keys, typ) in enumerate(
            zip(encoded["question_names"], encoded["option_keys"], encoded["types"], strict=True)
        ):
            t = self.calibration["temperatures"].get(typ, 1.0)
            p = (logits[index, : len(keys)] / t).softmax(-1)
            best = int(p.argmax())
            answer = keys[best] if typ == "choice" else best == 0 if typ == "noul" else best
            item = {
                "type": typ,
                "answer": answer,
                "probabilities": dict(zip(keys, p.tolist(), strict=True)),
                "confidence": float(p.max()),
            }
            if typ == "score":
                item["expected_score"] = float((p * torch.arange(len(p))).sum())
            answers[name] = item
        return {
            "model": "themis-v0",
            "answers": answers,
            "calibrated": bool(self.calibration["completed"]),
            "input_tokens": len(encoded["input_ids"]),
        }