"""Suggested inference code for v-jev: a state and typed questions in, probabilities out, in one forward pass.

    from inference import VJev
    jev = VJev("base")   # or "high-base"; the tokenizer is read from the repo root
    answers = jev.ask(state, questions)

`state` is text or any JSON value. `questions` and `answers` follow the Jev request shape (POST /v1/systemone):
  noul    {"type": "noul", "instructions": ...}                           -> {"type": "noul", "noul": p(yes)}
  choice  {"type": "choice", "instructions": ..., "criteria": {id: text}}  -> {"type": "choice", "probabilities": {id: p}}
  score   {"type": "score", "instructions": ..., "criteria": [levels]}     -> {"type": "score", "probabilities": {"0": p, ...}}
A noul may also take "criteria": {"true": text, "false": text}. A score may take its levels as an object
{id: text}, lowest first; its answer is then keyed by those ids.
Each answer also carries "na": the model's probability that the question cannot be answered from the state.
It is read from its own logit with a sigmoid, as it was trained (a separate yes/no target), not from the
softmax over the answers.

Nothing is generated. All questions of one request go into one prompt with one <ans> placeholder each, and
the model's output at each placeholder is read over the answer tokens that the question allows.

A choice may have up to 256 options. One question holds at most 64, so a longer list is asked in rounds, in
extra passes: the options are shuffled into groups of 63 plus "none of these", each group's winner goes on,
and this repeats until at most 63 are left for a final question, which also offers "none of these". The
answer has the finalists' probabilities, renormalised to sum to 1 (options that did not reach the final get
0), and "none": the final question's probability of "none of these". When every group picked "none of
these", each group's best real option goes on and the answer also says "all_groups_none": true.

The state and the question texts are data: any added token of the tokenizer written out in them is broken up
before the prompt is built, so it stays text. That covers the answer tokens ("<ans>", "<opt_3>", ...) and the
backbone's own ("<think>", "<tool_call>", "<|im_end|>", ...). Every added token starts with "<" and none
contains a space, so a space after the "<" breaks it, and the inserted space cannot complete a new one.
The state is otherwise inserted as is: text in it that imitates the prompt's "Questions:" or "Answers:" lines
cannot add answer positions, but it can sway the answers.
Needs torch and transformers (5.x, with Qwen3.5 support).
"""
import json
import os
import random
import re
import warnings

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

OPT = [f"<opt_{i}>" for i in range(1, 65)]  # choice: up to 64 options
LVL = [f"<lvl_{i}>" for i in range(10)]  # score: up to 10 levels, lowest first
NOUL = ["<no>", "<yes>"]
ANS, NA = "<ans>", "<na>"
MAX_OPTIONS, GROUP, NONE = 256, 63, "none of these"  # long lists: groups of 63 plus "none of these"
PER_PROMPT = 8  # group questions packed into one prompt, fewer when the prompt would be too long
TRAINED_TOKENS = 4608  # the longest prompt of training and the benchmarks; longer ones are untested


class TooLong(ValueError):
    """The prompt is over max_tokens."""


def tokens_and_spec(key, q):
    """The answer tokens of one question, and the question as the model sees it."""
    if not isinstance(q, dict) or not isinstance(q.get("instructions"), str):
        raise ValueError(f'question "{key}": needs "instructions" as text')
    kind, c = q.get("type"), q.get("criteria")
    spec = {"type": kind, "instructions": q["instructions"]}
    if kind == "noul":
        if c is not None:
            if not isinstance(c, dict) or not {"true", "false"} <= set(c):
                raise ValueError(f'question "{key}": a noul takes no "criteria", or an object with "true" and "false"')
            spec["criteria"] = {"<no>": c["false"], "<yes>": c["true"]}
        return NOUL, spec
    if kind == "choice":
        if not isinstance(c, dict) or not 2 <= len(c) <= MAX_OPTIONS:
            raise ValueError(f'question "{key}": a choice needs "criteria" as an object of 2 to {MAX_OPTIONS} options')
        if len(c) > len(OPT) and NONE in c:
            raise ValueError(f'question "{key}": with more than {len(OPT)} options, none may be named "{NONE}"; the rounds add it')
        spec["criteria"] = {t: f"{name}: {text}" if text else name for t, (name, text) in zip(OPT, c.items())}
        return OPT[: len(c)], spec
    if kind == "score":
        levels = list(c.values()) if isinstance(c, dict) else c
        if not isinstance(levels, list) or not 2 <= len(levels) <= len(LVL):
            raise ValueError(f'question "{key}": a score needs "criteria" as a list (or object) of 2 to {len(LVL)} levels, lowest first')
        spec["criteria"] = dict(zip(LVL, [str(x) for x in levels]))
        return LVL[: len(levels)], spec
    raise ValueError(f'question "{key}": "type" must be noul, choice or score')


class VJev:
    def __init__(self, checkpoint="base", repo=None, device=None, dtype=torch.bfloat16, max_tokens=None):
        """`checkpoint`: "base" or "high-base", the subfolder that holds the weights and their config.
        `repo`: the repo folder or Hugging Face repo ID; the default is the folder of this file. The tokenizer is
        read from its root.
        `device`: the GPU if there is one, else the CPU.
        `max_tokens`: longer requests are refused. The default is the backbone's own context length; memory
        usually runs out first. Training and the benchmarks used prompts of at most 4,608 tokens."""
        repo = repo or os.path.dirname(os.path.abspath(__file__))
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        self.tok = AutoTokenizer.from_pretrained(repo)
        self.special = re.compile("|".join(sorted(map(re.escape, self.tok.get_added_vocab()), key=len, reverse=True)))
        # Loaded tensor by tensor straight to the device, so there is one copy of the weights, not two.
        self.lm = AutoModelForCausalLM.from_pretrained(repo, subfolder=checkpoint, dtype=dtype, device_map=self.device).eval()
        self.max_tokens = max_tokens or self.lm.config.max_position_embeddings
        self.ans_id = self.tok.convert_tokens_to_ids(ANS)
        self.token_id = {t: self.tok.convert_tokens_to_ids(t) for t in OPT + LVL + NOUL + [NA]}

    def plain(self, text):
        """Text from the caller, with every added token in it broken up ("<ans>" -> "< ans>", "<|im_end|>" -> "< |im_end|>")."""
        return self.special.sub(lambda m: m[0][0] + " " + m[0][1:], text)

    def prompt(self, state, specs):
        """One chat turn: state and questions from the user, the answer lines from the model, thinking off."""
        lines = ["State:", self.plain(state if isinstance(state, str) else json.dumps(state, ensure_ascii=False)), "", "Questions:"]
        for i, s in enumerate(specs, 1):
            s = {**s, "instructions": self.plain(s["instructions"])}
            if "criteria" in s:
                s["criteria"] = {t: self.plain(str(text)) for t, text in s["criteria"].items()}
            lines.append(json.dumps({f"q{i}": s}, ensure_ascii=False))
        head = self.tok.apply_chat_template([{"role": "user", "content": "\n".join(lines)}], tokenize=False, add_generation_prompt=True, enable_thinking=False)
        return head + "Answers:\n" + "\n".join(f"q{i}: {ANS}" for i in range(1, len(specs) + 1))

    @torch.no_grad()
    def read(self, state, questions):
        """One prompt, one pass: per question, the probabilities of its answers in order, and its "na"."""
        tokens, specs = zip(*(tokens_and_spec(k, q) for k, q in questions.items()))
        ids = self.tok(self.prompt(state, specs), return_tensors="pt", add_special_tokens=False)["input_ids"].to(self.device)
        if ids.shape[1] > self.max_tokens:
            raise TooLong(f"the request is {ids.shape[1]} tokens; the limit is {self.max_tokens} (max_tokens). Send fewer questions per request or a shorter state.")
        if ids.shape[1] > TRAINED_TOKENS:
            warnings.warn(f"the request is {ids.shape[1]} tokens; training and the benchmarks used at most {TRAINED_TOKENS}, so the answers are untested at this length", stacklevel=2)
        hidden = self.lm.model(input_ids=ids, use_cache=False).last_hidden_state[0]  # one pass, nothing to cache
        at = hidden[ids[0] == self.ans_id].float()  # one row per question
        if at.shape[0] != len(specs):  # answers would be read from the wrong rows
            raise RuntimeError(f"the prompt has {at.shape[0]} answer positions for {len(specs)} questions")
        out_rows = self.lm.get_output_embeddings().weight
        out = {}
        for k, key in enumerate(questions):
            z = at[k] @ out_rows[[self.token_id[t] for t in tokens[k]] + [self.token_id[NA]]].float().T
            out[key] = (torch.softmax(z[:-1], -1).tolist(), torch.sigmoid(z[-1]).item())
        return out

    def read_many(self, state, asks):
        """Several questions about one state, PER_PROMPT to a prompt, fewer when a prompt would be too long."""
        got, items, size = {}, list(asks.items()), PER_PROMPT
        while items:
            part = dict(items[:size])
            try:
                got.update(self.read(state, part))
            except TooLong:
                if size == 1:
                    raise TooLong(f"one group of a long choice ({GROUP} options) does not fit in max_tokens ({self.max_tokens}) with this state; shorten the state or the option texts, or raise max_tokens") from None
                size = max(1, size // 2)
                continue
            items = items[size:]
        return got

    def rounds(self, state, q):
        """A choice with more than 64 options: (finalists' probabilities, the final "none of these" probability,
        "na" of the final question, whether every group of some round picked "none of these")."""
        rng = random.Random(0)  # the grouping is random but the same on every call
        left, c, all_none = list(q["criteria"]), q["criteria"], False
        while len(left) > GROUP:
            rng.shuffle(left)
            groups = [left[s : s + GROUP] for s in range(0, len(left), GROUP)]
            asks = {str(g): {"type": "choice", "instructions": q["instructions"], "criteria": {**{o: c[o] for o in group}, NONE: None}} for g, group in enumerate(groups)}
            got = self.read_many(state, asks)
            best = [max(range(len(group) + 1), key=lambda i: got[str(g)][0][i]) for g, group in enumerate(groups)]  # index len(group) is "none of these"
            left = [group[i] for group, i in zip(groups, best) if i < len(group)]
            if not left:  # no group found a fit: go on with each group's best real option, and say so
                all_none = True
                left = [group[max(range(len(group)), key=lambda i: got[str(g)][0][i])] for g, group in enumerate(groups)]
        p, na = self.read_many(state, {"final": {"type": "choice", "instructions": q["instructions"], "criteria": {**{o: c[o] for o in left}, NONE: None}}})["final"]
        real = sum(p[:-1])
        return {o: x / real for o, x in zip(left, p)}, p[-1], na, all_none

    def ask(self, state, questions):
        for k in questions:
            tokens_and_spec(k, questions[k])  # validate everything before the first pass
        long = {k for k, q in questions.items() if q["type"] == "choice" and len(q["criteria"]) > len(OPT)}
        short = {k: q for k, q in questions.items() if k not in long}
        got = self.read(state, short) if short else {}
        answers = {}
        for key, q in questions.items():
            if key in long:
                finalists, none, na, all_none = self.rounds(state, q)
                answers[key] = {"type": "choice", "probabilities": {o: finalists.get(o, 0.0) for o in q["criteria"]}, "none": none, "na": na}
                if all_none:
                    answers[key]["all_groups_none"] = True
                continue
            p, na = got[key]
            if q["type"] == "noul":
                answers[key] = {"type": "noul", "noul": p[1], "na": na}
            else:
                names = list(q["criteria"]) if isinstance(q["criteria"], dict) else [str(i) for i in range(len(p))]
                answers[key] = {"type": q["type"], "probabilities": dict(zip(names, p)), "na": na}
        return answers


if __name__ == "__main__":
    import sys

    jev = VJev(sys.argv[1] if len(sys.argv) > 1 else "base")
    state = "Customer writes: Our nightly export has failed three nights in a row and month-end close is on Friday. If this is not fixed by Thursday we will have to look at other vendors."
    questions = {
        "export_failing": {"type": "noul", "instructions": "Is a scheduled data export currently failing?"},
        "refund_asked": {"type": "noul", "instructions": "Does the customer ask for a refund?"},
        "team": {"type": "choice", "instructions": "Which team should own the issue?",
                 "criteria": {"engineering": "Fixes broken product functionality", "billing": "Handles invoices, charges and refunds", "sales": "Handles plans, upgrades and renewals"}},
        "churn_risk": {"type": "score", "instructions": "How likely is the customer to leave?",
                       "criteria": ["No sign of leaving", "Mild dissatisfaction", "Mentions alternatives", "States an intent to leave"]},
    }
    print(json.dumps(jev.ask(state, questions), indent=2))