"""Request schema shared by data loading, prompting, the head and the client."""
from __future__ import annotations

import math
from dataclasses import asdict, dataclass, field
from typing import Optional

QUESTION_TYPES = ("noul", "choice", "score")
TYPE_INDEX = {t: i for i, t in enumerate(QUESTION_TYPES)}
MAX_CHOICE_OPTIONS = 64
MAX_SCORE_LEVELS = 10
NOUL_OPTIONS = ["true", "false"]
NOUL_DEFAULT_DESCRIPTIONS = ["yes, the statement is true", "no, the statement is false"]
ALIASES = (
    [chr(c) for c in range(ord("A"), ord("Z") + 1)]
    + [chr(c) for c in range(ord("a"), ord("z") + 1)]
    + ["A" + chr(c) for c in range(ord("A"), ord("L") + 1)]
)


@dataclass
class Question:
    qid: str
    type: str
    instructions: str
    options: list[str]
    descriptions: list[str]

    @classmethod
    def from_jev(cls, qid: str, spec: dict) -> "Question":
        qtype = spec.get("type")
        if qtype not in QUESTION_TYPES:
            raise ValueError(f"{qid}: unknown question type {qtype!r}")
        instructions = str(spec.get("instructions", "")).strip()
        criteria = spec.get("criteria")
        if qtype == "noul":
            options = list(NOUL_OPTIONS)
            descriptions = list(NOUL_DEFAULT_DESCRIPTIONS)
            if criteria is not None:
                if not isinstance(criteria, dict) or set(criteria) - set(NOUL_OPTIONS):
                    raise ValueError(f"{qid}: noul criteria must be a mapping with 'true' and 'false' keys")
                descriptions = [str(criteria.get("true") or descriptions[0]),
                                str(criteria.get("false") or descriptions[1])]
        elif qtype == "choice":
            if not isinstance(criteria, dict):
                raise ValueError(f"{qid}: choice needs a criteria mapping")
            options = [str(k) for k in criteria.keys()]
            descriptions = [str(v) if v else str(k) for k, v in criteria.items()]
            if not (2 <= len(options) <= MAX_CHOICE_OPTIONS):
                raise ValueError(f"{qid}: choice needs 2..{MAX_CHOICE_OPTIONS} options, got {len(options)}")
        else:
            if not isinstance(criteria, list):
                raise ValueError(f"{qid}: score needs a criteria list")
            if not (2 <= len(criteria) <= MAX_SCORE_LEVELS):
                raise ValueError(f"{qid}: score needs 2..{MAX_SCORE_LEVELS} levels, got {len(criteria)}")
            options = [str(i) for i in range(len(criteria))]
            descriptions = [str(c) for c in criteria]
        if len(set(options)) != len(options):
            raise ValueError(f"{qid}: duplicate options")
        return cls(qid=qid, type=qtype, instructions=instructions,
                   options=options, descriptions=descriptions)

    @property
    def k(self) -> int:
        return len(self.options)


def argmax(values) -> int:
    """Index of the first maximum, so ties go to the lowest index. Shared by evaluation and the zero-shot read."""
    return max(range(len(values)), key=lambda i: values[i])


def align_distribution(probs: dict, options: list[str]) -> list[float]:
    vals = [float(probs.get(o, 0.0)) for o in options]
    if any(v < 0 or not math.isfinite(v) for v in vals):
        raise ValueError("distribution values must be finite and non-negative")
    total = sum(vals)
    if total <= 0:
        raise ValueError("distribution has no mass on the listed options")
    return [v / total for v in vals]


@dataclass
class Example:
    id: str
    source: str
    state: str
    questions: list[Question]
    gold: dict[str, Optional[list[float]]]
    meta: dict = field(default_factory=dict)

    def to_dict(self) -> dict:
        return {
            "id": self.id, "source": self.source, "state": self.state,
            "questions": [asdict(q) for q in self.questions],
            "gold": self.gold, "meta": self.meta,
        }

    @classmethod
    def from_dict(cls, d: dict) -> "Example":
        return cls(id=d["id"], source=d["source"], state=d["state"],
                   questions=[Question(**q) for q in d["questions"]],
                   gold={k: (list(v) if v is not None else None) for k, v in d["gold"].items()},
                   meta=dict(d.get("meta", {})))