"""Arthur's input layout and output calibration, without torch.

Copied verbatim from arthur_v0_8_0.ipynb (sha256 fb6836d8c9534e9e…): LAYOUT, QTYPES and the token constants and
assemble() from its cell 17, softmax() and bucket() from its cell 20. normalize_question() and state_text() are
FalconDec's (LightDec v1.0.2 falcondec_modeling.py, _normalize_question and _state_text), because Arthur was
trained on FalconDec_V2's data and its option texts. Do not edit these copies: change the notebook and re-copy.
"""
from __future__ import annotations

import json

import numpy as np

LAYOUT = {"max_len": 1024, "long_max_len": 2048, "long_opts_threshold": 24, "head_max_len": 192, "max_tok_per_opt": 24}
QTYPES = {"choice": 0, "noul": 1, "score": 2}
PAD, OPT, SEP, CLS, VOCAB, BPT, STRIDE = 0, 257, 258, 259, 260, 4, 4


def assemble(d):
    """[CLS] question [SEP] [OPT] option ... [SEP] state [SEP] in byte ids; each [OPT] starts a 4-byte pooling window."""
    n = len(d["options"])
    max_tok = int(d.get("seq_len", LAYOUT["max_len"]))
    eff = BPT * (max_tok if n <= LAYOUT["long_opts_threshold"] else max(max_tok, LAYOUT["long_max_len"]))
    q = list(str(d.get("question", "")).encode("utf-8"))[:96 * BPT]
    opt_budget = BPT * int(d.get("opt_budget", LAYOUT["max_tok_per_opt"]))
    head = min(eff - 64 * BPT, max(LAYOUT["head_max_len"] * BPT, len(q) + 2 + n * (opt_budget + STRIDE)))
    per = max(2 * BPT, min(opt_budget, (head - len(q) - 2) // max(n, 1) - STRIDE))
    ids, windows = [CLS] + [b + 1 for b in q] + [SEP], []
    for o in d["options"]:
        ids += [PAD] * ((-len(ids)) % STRIDE)
        windows.append(len(ids) // STRIDE)
        ids += [OPT] + [b + 1 for b in list(str(o).encode("utf-8"))[:per]]
    ids.append(SEP)
    room = eff - len(ids) - 1
    if room > 0 and d.get("state"):
        ids += [b + 1 for b in list(str(d["state"]).encode("utf-8"))[:room]] + [SEP]
    return ids[:eff], windows


def softmax(z, t=1.0):
    e = np.exp((np.asarray(z, dtype=np.float64) - np.max(z)) / t)
    return e / e.sum()


def bucket(k):
    return 0 if k <= 2 else 1 if k <= 5 else 2 if k <= 12 else 3


def normalize_question(q):
    qtype = q.get("type", "choice")
    text = q.get("question") or q.get("instructions") or ""
    if qtype == "noul":
        lab = q.get("labels") or {}
        return qtype, text, [True, False], [str(lab.get("true", "Yes")), str(lab.get("false", "No"))]
    crit = q.get("criteria", q.get("options"))
    if qtype == "score":
        opts = [str(c) for c in crit]
        return qtype, text, list(range(len(opts))), opts
    if isinstance(crit, dict):
        return qtype, text, list(crit), [f"{k}: {v}" if v else str(k) for k, v in crit.items()]
    return qtype, text, list(crit), [str(c) for c in crit]


def state_text(state):
    if state is None:
        return ""
    return state if isinstance(state, str) else json.dumps(state, ensure_ascii=False)


def probabilities(logits, qtype: str, temperatures) -> np.ndarray:
    """Calibrated probabilities, as the notebook's report(): softmax at temperatures[qtype][bucket(n_options)]."""
    return softmax(logits, temperatures[QTYPES[qtype]][bucket(len(logits))])