"""Calibration and the answer objects of the hosted API.

Calibration (calibration.json):

    choice, score   p = softmax(z / T(K)),  T(K) = t0 + t1 * ln K, K clamped to the fitted range
    truth           p_yes = sigmoid((z_yes - z_no) / T + b)
    coord           p = softmax(z), a point or box question's reading, never calibrated (places.py)

A parameter set exists per input kind ("text", "image"); a missing one is the identity.

The answers, rounded as the hosted API rounds them:

    choice  {"type", "choice", "confidence", "probabilities": {option: p}}      options in alphabetical order
    score   {"type", "score", "confidence", "legend": {"0": level}, "probabilities": {"0": p}}
    truth   {"type", "truth"}
    number  {"type", "number", "expected", "confidence", "probabilities": {value: p}}
    point   {"type", "x", "y", "present", "confidence", "image", "width", "height"}         places.py
    box     {"type", "x1", "y1", "x2", "y2", "present", "confidence", "image", "width", "height"}
            with `max`: {"type", "count", "present", "complete", "points" | "boxes": [...], "image", "width", "height"}

`choice`, `score`, `number` and the confidences come from the unrounded probabilities. A probability list is rounded
to 2 decimals entry by entry, then its largest entries move by 0.01 until it sums to 1 (`round_probabilities`).
`exact_sum` (the rounding of /v1/systemone, where benchmarks check the sum): the probabilities are normalised first, and
the list sums to exactly 1.00 in hundredths, every value a clean 2-decimal number.
"""

from __future__ import annotations

import math
from decimal import Decimal, localcontext
from typing import Sequence

T_MIN, T_MAX = 0.05, 100.0
UNIT = 0.01


# ---------------------------------------------------------------------------
# Calibration
# ---------------------------------------------------------------------------


def _temperature(params: dict, k: int) -> float:
    if params.get("k_min") is not None:
        k = min(max(k, params["k_min"]), params["k_max"])
    t = float(params.get("t0", 1.0)) + float(params.get("t1", 0.0)) * math.log(k)
    return min(T_MAX, max(T_MIN, t))


def calibrated(logits: Sequence[float], kind: str, calibration: dict, has_images: bool) -> list[float]:
    """Calibrated class probabilities from a question's code logits (class order: options or levels as given,
    [yes, no] for truth)."""
    if kind == "coord":
        return softmax(logits)
    params = (calibration.get("image" if has_images else "text") or {}).get("choice" if kind == "number" else kind) or {}
    z = [float(x) for x in logits]
    if kind == "truth":
        x = (z[0] - z[1]) / float(params.get("temperature", 1.0)) + float(params.get("bias", 0.0))
        p = 1.0 / (1.0 + math.exp(-x)) if x >= 0 else math.exp(x) / (1.0 + math.exp(x))
        return [p, 1.0 - p]
    if len(z) == 1:
        return [1.0]
    t = _temperature(params, len(z))
    scaled = [x / t for x in z]
    top = max(scaled)
    log_total = math.log(sum(math.exp(x - top) for x in scaled))
    return [math.exp(x - top - log_total) for x in scaled]


def softmax(logits: Sequence[float]) -> list[float]:
    """A plain softmax (the largest score subtracted first): a coordinate row's probabilities."""
    z = [float(x) for x in logits]
    if not z or not all(math.isfinite(x) for x in z):
        raise ValueError("scores must be a list of finite numbers")
    top = max(z)
    e = [math.exp(x - top) for x in z]
    total = math.fsum(e)
    return [x / total for x in e]


# ---------------------------------------------------------------------------
# Rounding
# ---------------------------------------------------------------------------


def _units(x: float) -> int:
    """x in hundredths, halves up (the epsilon absorbs binary error: 0.805 rounds to 0.81)."""
    return int(math.floor(x * 100.0 + 0.5 + 1e-9))


def round2(x: float) -> float:
    value = _units(x) / 100.0
    return 0.0 if value == 0 else value


def normalised(probs: Sequence[float]) -> list[float]:
    """The probabilities divided by their float64 sum (a sum within 1e-9 of 1); a list that sums to 0 as it is."""
    p = [float(x) for x in probs]
    total = math.fsum(p)
    if total <= 0 or not math.isfinite(total):
        return p
    return [x / total for x in p]


def round_probabilities(probs: Sequence[float], exact_sum: bool = False) -> list[float]:
    """Round to 2 decimals, then move the |n| largest entries by 0.01, n = int((1 - fsum(rounded)) / 0.01).
    A moved entry is not rounded again (0.81 + 0.01 is 0.8200000000000001), as the hosted API returns it.

    `exact_sum` (pass normalised probabilities): n is 100 less the sum of the rounded hundredths, and every value is the
    clean 2-decimal number, so the hundredths sum to exactly 100."""
    p = [float(x) for x in probs]
    units = [_units(x) for x in p]
    rounded = [u / 100.0 for u in units]
    if exact_sum:
        left = abs(100 - sum(units))
        step = 1 if sum(units) < 100 else -1
        order = sorted(range(len(p)), key=lambda i: (-p[i], i))
        while left:  # every entry moves once, the largest first; a list that needs more moves round again
            progressed = False
            for i in order:
                if not left:
                    break
                if units[i] + step < 0:
                    continue
                units[i] += step
                left -= 1
                progressed = True
            if not progressed:
                raise ValueError("probabilities cannot be rounded to a list that sums to 1")
        return [0.0 if u == 0 else u / 100.0 for u in units]
    n = int((1.0 - math.fsum(rounded)) / UNIT)
    if n:
        step = 1 if n > 0 else -1
        for i in sorted(range(len(p)), key=lambda i: (-p[i], i))[: abs(n)]:
            rounded[i] = rounded[i] + UNIT if step > 0 else rounded[i] - UNIT
    return [0.0 if x == 0 else x for x in rounded]


def _argmax_first(probs: Sequence[float]) -> int:
    best = 0
    for i, p in enumerate(probs):
        if p > probs[best]:
            best = i
    return best


# ---------------------------------------------------------------------------
# Answers
# ---------------------------------------------------------------------------


def choice_answer(keys: Sequence[str], probs: Sequence[float], exact_sum: bool = False) -> dict:
    if exact_sum:
        probs = normalised(probs)
    best = _argmax_first(probs)
    k = len(probs)
    confidence = 1.0 if k == 1 else min(1.0, max(0.0, (k * probs[best] - 1.0) / (k - 1.0)))
    shown = round_probabilities(probs, exact_sum)
    return {"type": "choice", "choice": keys[best], "confidence": round2(confidence), "probabilities": dict(sorted(zip(keys, shown)))}


def score_answer(levels: Sequence, probs: Sequence[float], exact_sum: bool = False) -> dict:
    if exact_sum:
        probs = normalised(probs)
    k = len(probs)
    if k == 1:
        confidence = 1.0
    else:
        mode = _argmax_first(probs)
        spread = sum(p * abs(i - mode) for i, p in enumerate(probs))
        confidence = min(1.0, max(0.0, 1.0 - spread / ((k * k // 4) / k)))
    return {
        "type": "score",
        "score": round2(sum(i * p for i, p in enumerate(probs))),
        "confidence": round2(confidence),
        "legend": {str(i): level for i, level in enumerate(levels)},
        "probabilities": {str(i): x for i, x in enumerate(round_probabilities(probs, exact_sum))},
    }


def truth_answer(probs: Sequence[float]) -> dict:
    return {"type": "truth", "truth": round2(probs[0])}


def _expected_text(value: float, decimals: int) -> str:
    """Halves up to `decimals` (at most 15 significant digits), with a decimal point, never -0.0."""
    with localcontext() as ctx:
        ctx.prec = 200
        exact = Decimal(value)
        whole = abs(int(exact))
        decimals = min(decimals, 15 - (len(str(whole)) if whole else 0))
        units = (exact * Decimal(10) ** decimals + Decimal("0.5") + Decimal("1e-9")).to_integral_value(rounding="ROUND_FLOOR")
        text = format(units * Decimal(10) ** -decimals, "f")
    if "." in text:
        text = text.rstrip("0")
        text = text + "0" if text.endswith(".") else text
    else:
        text += ".0"
    return "0.0" if text == "-0.0" else text


def _number(text: str) -> int | float:
    return float(text) if "." in text else int(text)


def number_answer(grid, probs: Sequence[float]) -> dict:
    """The most probable value, the probability-weighted mean (`expected`), the probability of the value
    (`confidence`) and every value above 0.001 to 3 decimals (rounded entries that sum above 1 lose 0.001 each,
    the largest first, the answer's own last)."""
    best = _argmax_first(probs)
    texts = grid.texts()
    expected = math.fsum(float(text) * p for text, p in zip(texts, probs))
    listed = [i for i in range(grid.n) if probs[i] > 0.001 or i == best]
    millis = {i: int(math.floor(probs[i] * 1000.0 + 0.5 + 1e-9)) for i in listed}
    excess = sum(millis.values()) - 1000
    order = sorted((i for i in listed if i != best), key=lambda i: (-probs[i], i))
    while excess > 0:
        lowered = False
        for i in order:
            if excess <= 0:
                break
            if millis[i] > 0:
                millis[i] -= 1
                excess -= 1
                lowered = True
        if not lowered:
            millis[best] -= 1
            excess -= 1
    return {
        "type": "number",
        "number": _number(texts[best]),
        "expected": float(_expected_text(expected, max(2, grid.d))),
        "confidence": round2(probs[best]),
        "probabilities": {texts[i]: millis[i] / 1000.0 for i in listed},
    }