"""Task-agnostic decision interfaces.

MiniJEV always answers the same kind of request::

    DecisionRequest(state=..., question="...", choices=["a", "b", ...]) -> Decision

`state` is opaque to the core: each environment/task plugin knows how to turn
its states into tensors (see `minijev.tasks.Task`). Choices are plain strings,
so the number and the meaning of choices are *inputs*, not model constants.
"""
from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Sequence

State = Any  # task-specific (SnakeState, str, dict, ...)


@dataclass
class DecisionRequest:
    state: State
    question: str
    choices: Sequence[str]


@dataclass
class Decision:
    choice: str
    index: int
    confidence: float  # calibrated probability of `choice`
    scores: dict[str, float]  # calibrated probability per choice
    logits: list[float] = field(default_factory=list)
    embedding: list[float] | None = None  # optional shared state representation
    raw_scores: dict[str, float] | None = None  # softmax(logits), before temperature calibration
    temperature: float = 1.0

    @property
    def distribution(self) -> dict[str, float]:
        return self.scores

    @property
    def raw_confidence(self) -> float:
        return self.raw_scores[self.choice] if self.raw_scores else self.confidence

    def to_dict(self) -> dict:
        d = {"choice": self.choice, "confidence": round(self.confidence, 4),
             "distribution": {k: round(v, 4) for k, v in self.scores.items()}}
        if self.raw_scores is not None:
            d["raw_distribution"] = {k: round(v, 4) for k, v in self.raw_scores.items()}
        if self.embedding is not None:
            d["embedding"] = self.embedding
        return d


@dataclass
class BooleanDecision:
    answer: bool
    confidence: float  # calibrated probability of `answer`
    p_yes: float
    raw_p_yes: float

    def __bool__(self) -> bool:
        return self.answer