"""Decisions, questions and the ``/v1/systemone`` wire format.

A *decision* is a state (text, JSON or a list) plus optional images and one or more typed
questions. Each question has an ordered list of labels:

* ``choice``: labels are the option names; ``definitions`` holds each option's criterion.
* ``score``: labels are the level descriptions, index = level (2 to 10 levels).
* ``noul``: labels are always ``["no", "yes"]``; ``definitions`` holds the
  ``[false, true]`` criteria when given.

Requests carry ``state``, ``questions`` and optionally ``images`` (``data:`` URLs); responses
carry one answer per question with ``probabilities`` (choice, score) or ``noul`` (the yes
probability).
"""

from __future__ import annotations

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

QUESTION_TYPES = ("choice", "score", "noul")
NOUL_LABELS = ["no", "yes"]
MAX_CHOICE_OPTIONS = 255
MIN_SCORE_LEVELS, MAX_SCORE_LEVELS = 2, 10
MAX_IMAGES = 8


class SchemaError(ValueError):
    """A request that does not follow the schema."""


def normalise(values: list[float]) -> list[float]:
    total = float(sum(values))
    return [float(v) / total for v in values]


@dataclass
class Question:
    type: str
    instructions: str
    labels: list[str]
    definitions: list[str | None] | None = None

    def validate(self) -> None:
        if self.type not in QUESTION_TYPES:
            raise SchemaError(f"unknown question type {self.type!r}")
        if not isinstance(self.instructions, str):
            raise SchemaError("instructions must be a string")
        n = len(self.labels)
        if self.type == "noul":
            if self.labels != NOUL_LABELS:
                raise SchemaError(f"noul labels must be {NOUL_LABELS}, got {self.labels}")
        elif self.type == "choice":
            if not 2 <= n <= MAX_CHOICE_OPTIONS:
                raise SchemaError(f"choice needs 2-{MAX_CHOICE_OPTIONS} options, got {n}")
            if len(set(self.labels)) != n:
                raise SchemaError("choice labels must be unique")
        elif not MIN_SCORE_LEVELS <= n <= MAX_SCORE_LEVELS:
            raise SchemaError(f"score needs {MIN_SCORE_LEVELS}-{MAX_SCORE_LEVELS} levels, got {n}")
        if any(not isinstance(label, str) or not label for label in self.labels):
            raise SchemaError("labels must be non-empty strings")
        if self.definitions is not None and len(self.definitions) != n:
            raise SchemaError("definitions must align with labels")

    @classmethod
    def from_wire(cls, q: dict[str, Any]) -> "Question":
        qtype = q.get("type")
        instructions = q.get("instructions", "") or ""
        criteria = q.get("criteria")
        if qtype == "choice":
            if isinstance(criteria, dict):
                labels = [str(k) for k in criteria]
                definitions = [None if v is None else str(v) for v in criteria.values()]
            elif isinstance(criteria, list):
                labels, definitions = [str(c) for c in criteria], None
            else:
                raise SchemaError("choice question needs criteria")
            question = cls("choice", instructions, labels, definitions)
        elif qtype == "score":
            if not isinstance(criteria, list):
                raise SchemaError("score question needs a list of levels")
            question = cls("score", instructions, [str(c) for c in criteria], None)
        elif qtype == "noul":
            definitions = None
            if isinstance(criteria, dict):
                definitions = [criteria.get("false"), criteria.get("true")]
            question = cls("noul", instructions, list(NOUL_LABELS), definitions)
        else:
            raise SchemaError(f"unknown question type {qtype!r}")
        question.validate()
        return question


@dataclass
class Decision:
    """One state with its typed questions."""

    id: str
    state: Any
    questions: dict[str, Question]
    images: list[dict[str, Any]] = field(default_factory=list)

    def validate(self) -> None:
        if not isinstance(self.state, (str, dict, list)):
            raise SchemaError("state must be a string, object or list")
        if not self.questions:
            raise SchemaError("a decision needs at least one question")
        for name, q in self.questions.items():
            if not name:
                raise SchemaError("question names must be non-empty")
            q.validate()
        if len(self.images) > MAX_IMAGES:
            raise SchemaError(f"at most {MAX_IMAGES} images per request")


def choice_confidence(probs: list[float]) -> float:
    """(p_max - 1/n) / (1 - 1/n): 0 for a uniform distribution, 1 for a certain one."""
    n = len(probs)
    return (max(probs) - 1.0 / n) / (1.0 - 1.0 / n)


def score_confidence(probs: list[float]) -> float:
    """1 - (probability-weighted distance from the peak level) / (same for a uniform distribution)."""
    n = len(probs)
    peak = max(range(n), key=lambda i: probs[i])
    spread = sum(p * abs(i - peak) for i, p in enumerate(probs))
    uniform = sum(abs(i - peak) for i in range(n)) / n
    return 1.0 - spread / uniform if uniform else 1.0


def make_answer(question: Question, probs: list[float]) -> dict[str, Any]:
    """Format a distribution over ``question.labels`` as a ``/v1/systemone`` answer."""
    probs = normalise(probs)
    if question.type == "noul":
        return {"type": "noul", "noul": probs[1]}
    if question.type == "choice":
        best = max(range(len(probs)), key=lambda i: probs[i])
        return {"type": "choice", "choice": question.labels[best],
                "probabilities": dict(zip(question.labels, probs)),
                "confidence": choice_confidence(probs)}
    expected = sum(i * p for i, p in enumerate(probs))
    return {"type": "score", "score": expected, "confidence": score_confidence(probs),
            "probabilities": {str(i): p for i, p in enumerate(probs)}}