"""The three question types, and the record every source is converted into.

Every decision is one forward pass that scores a list of candidates. The three
types differ only in where the candidates come from:

- choice: the caller supplies them.
- score:  they are the points of a scale (`min`..`max`), in order.
- null:   they are "yes" and "no"; the answer is P(yes).

One head scores every candidate the same way, so what the model learns about
"which of these fits the state" transfers across all three.

A record, one per JSONL line:

    {"type": "choice", "state": "...", "question": "...",
     "options": ["a", "b", "c"], "label": 1,
     "source": "banking77", "licence": "cc-by-4.0"}

    {"type": "score", "state": "...", "question": "...",
     "scale": {"min": 1, "max": 5, "labels": {"1": "terrible", "5": "excellent"}},
     "label": 4, ...}                 # label is the scale value, not an index

    {"type": "null", "state": "...", "question": "The ticket is urgent.",
     "label": true, ...}
"""

from __future__ import annotations

from dataclasses import dataclass, field

TYPES = ("choice", "score", "null")
NULL_OPTIONS = ("yes", "no")


@dataclass
class Decision:
    type: str
    state: str
    question: str
    options: list[str] = field(default_factory=list)   # choice only
    scale: dict | None = None                           # score only
    label: object = None                                # index / value / bool
    source: str = ""
    licence: str = ""

    def candidates(self) -> list[str]:
        """The strings the head scores, in the order the model sees them."""
        if self.type == "choice":
            return list(self.options)
        if self.type == "null":
            return list(NULL_OPTIONS)
        lo, hi = int(self.scale["min"]), int(self.scale["max"])
        names = {str(k): v for k, v in (self.scale.get("labels") or {}).items()}
        return [f"{v} ({names[str(v)]})" if str(v) in names else str(v)
                for v in range(lo, hi + 1)]

    def scale_values(self) -> list[int]:
        return list(range(int(self.scale["min"]), int(self.scale["max"]) + 1))

    def target(self) -> int | None:
        """The index of the right candidate, or None for an unlabelled record."""
        if self.label is None:
            return None
        if self.type == "choice":
            return int(self.label)
        if self.type == "null":
            return 0 if bool(self.label) else 1
        return self.scale_values().index(int(round(float(self.label))))

    def validate(self) -> "Decision":
        if self.type not in TYPES:
            raise ValueError(f"type must be one of {TYPES}, not {self.type!r}")
        if not self.question.strip():
            raise ValueError("question is empty")
        if self.type == "choice":
            if len(self.options) < 2:
                raise ValueError("a choice needs at least two options")
            if len(set(self.options)) != len(self.options):
                raise ValueError("choice options must be distinct")
        if self.type == "score":
            if not self.scale or int(self.scale["max"]) <= int(self.scale["min"]):
                raise ValueError("a score needs a scale with max > min")
            if int(self.scale["max"]) - int(self.scale["min"]) > 20:
                raise ValueError("a scale has at most 21 points")
        t = self.target()
        if t is not None and not 0 <= t < len(self.candidates()):
            raise ValueError(f"label {self.label!r} is outside the candidates")
        return self

    @classmethod
    def from_dict(cls, d: dict) -> "Decision":
        return cls(type=d["type"], state=d.get("state", ""), question=d["question"],
                   options=list(d.get("options") or []), scale=d.get("scale"),
                   label=d.get("label"), source=d.get("source", ""),
                   licence=d.get("licence", "")).validate()

    def to_dict(self) -> dict:
        d = {"type": self.type, "state": self.state, "question": self.question,
             "label": self.label, "source": self.source, "licence": self.licence}
        if self.type == "choice":
            d["options"] = self.options
        if self.type == "score":
            d["scale"] = self.scale
        return d