"""Arbiter: load the model from this repository and answer typed decisions.

    from arbiter import Arbiter
    from questions import OUTLOOK

    arbiter = Arbiter.from_pretrained("ternary/Arbiter-8B")
    answer = arbiter.decide(state="...article or 10-K passage...", question=OUTLOOK, focus="Meta Platforms")
    print(answer["choice"], answer["probabilities"])

One forward pass per decision; the answer carries a probability for every option, calibrated with the
per-type temperatures in calibration.json. Every fallback (state truncated to fit, uncalibrated
temperature) is counted in `arbiter.fallback_counts`.
"""

from __future__ import annotations

import json
from collections import Counter
from pathlib import Path
from typing import Any

from model import QUESTION_TYPES, jev_answer, load_decision_model, option_logits, probabilities, question_type


class Arbiter:
    def __init__(self, model, tokenizer, temperatures: dict[str, float], max_length: int, counter: Counter):
        self.model, self.tokenizer = model, tokenizer
        self.temperatures, self.max_length = temperatures, max_length
        self.fallback_counts = counter

    @classmethod
    def from_pretrained(cls, path_or_repo: str = "ternary/Arbiter-8B", adapter: bool = False,
                        quantization: str | None = None, max_memory: str | None = None) -> "Arbiter":
        """Load from a local folder or a Hugging Face repo id.
        Default: the full model (LoRA merged into the weights, fp16/bf16, no base download).
        adapter=True: the LoRA in adapter/ on Qwen/Qwen3-8B, 4-bit as trained (6.5 GB GPU memory measured).
        quantization overrides the default ('none' or '4bit')."""
        from safetensors.torch import load_file

        path = Path(path_or_repo)
        if not path.exists():
            from huggingface_hub import snapshot_download
            path = Path(snapshot_download(path_or_repo))
        meta = json.loads((path / "metadata.json").read_text())
        head_config = json.loads((path / "head_config.json").read_text())
        counter: Counter = Counter()
        if (path / "config.json").exists() and not adapter:
            # Full model: the LoRA is merged into the weights in this folder; no base download, no peft.
            model_id, adapter_dir, default_quantization = str(path), None, "none"
            calibration = path / "calibration.json"
        else:
            # Adapter on the base model, 4-bit as trained.
            model_id, adapter_dir, default_quantization = meta["base_model"], path / "adapter", meta["quantization"]
            calibration = path / "adapter" / "calibration.json"
            if not calibration.exists():
                calibration = path / "calibration.json"
        model, tokenizer, letter_ids, _ = load_decision_model(
            model_id, quantization=quantization or default_quantization, max_memory=max_memory,
            adapter_dir=adapter_dir, counter=counter)
        if letter_ids != head_config["letter_ids"]:
            raise ValueError("this tokenizer's letter ids differ from the ones the head was trained with")
        model.head.load_state_dict(load_file(str(path / "decision_head.safetensors")))
        model.head.to(next(model.backbone.parameters()).device)
        temperatures = {kind: 1.0 for kind in QUESTION_TYPES}
        if calibration.exists():
            temperatures.update(json.loads(calibration.read_text())["temperatures"])
        else:
            counter["uncalibrated_temperature_1"] += 1
        return cls(model, tokenizer, temperatures, meta["max_length"], counter)

    def decide(self, state: str, question: dict[str, Any], focus: str | None = None) -> dict[str, Any]:
        """question: {"type": "choice"|"noul"|"score", "instructions": str, "criteria": {option_id: definition}}.
        noul criteria are {"true": ..., "false": ...}; score criteria are {"0": lowest, "1": ..., ...}."""
        labels = ["true", "false"] if question["type"] == "noul" else list(question["criteria"])
        row = {"state": state, "question": question, "labels": labels,
               "focus": {"surface": focus} if focus else {}}
        logits, truncated = option_logits(self.model, self.tokenizer, row, self.max_length, self.fallback_counts)
        probs = probabilities(logits, self.temperatures[question_type(row)])
        answer = jev_answer(row, probs)
        answer["probabilities"] = probs
        answer["state_truncated"] = truncated
        return answer