from __future__ import annotations

from typing import Any, Sequence

import torch
from transformers import AutoModel, BitsAndBytesConfig, MistralCommonBackend


class MinistralDecisionInference:
    def __init__(self, model, tokenizer):
        self.model = model.eval()
        self.tokenizer = tokenizer
        self.device = model.get_input_embeddings().weight.device
        self.temperature = float(model.decision_temperature)

    @classmethod
    def from_pretrained(cls, model_id: str = ".", *, quantized: bool = True):
        quantization = None
        if quantized:
            quantization = BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_quant_type="nf4",
                bnb_4bit_use_double_quant=True,
                bnb_4bit_compute_dtype=torch.float16,
            )
        model = AutoModel.from_pretrained(
            model_id,
            trust_remote_code=True,
            quantization_config=quantization,
            device_map={"": 0} if torch.cuda.is_available() else "cpu",
            low_cpu_mem_usage=True,
        )
        tokenizer = MistralCommonBackend.from_pretrained(model_id)
        return cls(model, tokenizer)

    @staticmethod
    def _pair_text(state: str, question: str, option: dict[str, str]) -> str:
        criteria = option.get("criteria", "").strip()
        criterion_text = f"\nCandidate criteria: {criteria}" if criteria else ""
        return (
            "Choose the best answer for this decision.\n"
            f"Task: {question}\n"
            f"Context: {state}\n"
            f"Candidate answer: {option['label']}{criterion_text}\n"
            "Does this candidate correctly answer the task?"
        )

    def score(self, state: str, question: str, options: Sequence[dict[str, str]], *, abstain_below: float | None = None) -> dict[str, Any]:
        if len(options) < 2:
            raise ValueError("Provide at least two options")
        ids = [option["id"] for option in options]
        if len(ids) != len(set(ids)):
            raise ValueError("Option ids must be unique")
        rows = []
        for option in options:
            encoded = self.tokenizer.encode(self._pair_text(state, question, option), return_tensors="pt")
            input_ids = encoded.input_ids if hasattr(encoded, "input_ids") else encoded
            rows.append(input_ids[0] if input_ids.ndim == 2 else input_ids)
        pad_id = getattr(self.tokenizer, "pad_token_id", None)
        if pad_id is None:
            pad_id = getattr(self.tokenizer, "eos_token_id", 0) or 0
        scores = []
        for start in range(0, len(rows), 4):
            batch = rows[start:start + 4]
            lengths = torch.tensor([row.numel() for row in batch], dtype=torch.long)
            width = int(lengths.max().item())
            input_ids = torch.full((len(batch), width), int(pad_id), dtype=torch.long)
            mask = torch.zeros((len(batch), width), dtype=torch.long)
            for index, row in enumerate(batch):
                input_ids[index, :row.numel()] = row
                mask[index, :row.numel()] = 1
            input_ids, mask = input_ids.to(self.device), mask.to(self.device)
            with torch.inference_mode():
                scores.extend(self.model.score_candidate_pairs(input_ids, mask).float().cpu().tolist())
        probabilities = torch.softmax(torch.tensor(scores) / self.temperature, dim=0).tolist()
        best = max(range(len(options)), key=probabilities.__getitem__)
        confidence = probabilities[best]
        abstain = abstain_below is not None and confidence < abstain_below
        return {
            "choice": None if abstain else ids[best],
            "confidence": confidence,
            "probabilities": dict(zip(ids, probabilities)),
            "scores": dict(zip(ids, scores)),
            "abstain": abstain,
            "calibrated_temperature": self.temperature,
        }