"""Decision Index engine for the v6 checkpoints (dense/train_dense.py --contexts): every question gets a short stock-model rationale, then the slot is read with the adapter ON.

Per request: (1) the rationale prompts of ALL questions of the request are built (dense/contexts.py: system line + question prompt + "Rationale:\\n") and generated together in
length-sorted batches with the adapter DISABLED (greedy, at most `brief_cap` tokens, stop <|im_end|>), (2) the generated text gets the same clean-up as the training rationales
(trailing bare-answer line dropped), (3) every question is rendered in the `brief` layout (rationale then "\\n\\nAnswer: ", per-option items put it before the slot block) and the
slot is read with the adapter ON: option logits = frozen LM-head rows of the answer tokens + trained residual, softmax over options (optional fitted temperature).
This is the accuracy path (Hugging Face generate); the latency path is tools/vllm_requests_0f.py (vLLM) with the same layout.

  python -m decision_index run --engine decision_engine_v6:BriefSlotEngine --option model=Qwen/Qwen3.5-9B --option run=RUN --option ckpt=best|final --suite-dir ... --rows ... --out ...
Options: model, run, ckpt (best|final|dir), context (brief|direct: direct = the v5 behaviour, no rationale), brief_cap (64), gen_tokens_budget (padded tokens per generate batch, 24000),
temperature (auto|none|...: same files as decision_engine_dense), max_tokens (32768).
"""
import os
import sys
import time

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)

import numpy as np  # noqa: E402
import torch  # noqa: E402

import answer_slot as AS  # noqa: E402
from common import QTYPES  # noqa: E402
from decision_engine_dense import DenseSlotEngine, _flag  # noqa: E402
try:
    from decision_index.engines.base import Unsupported  # noqa: E402
except ImportError:   # Decision Index kit not installed
    from engine_base import Unsupported  # noqa: E402
from dense.contexts import Contexts  # noqa: E402


class BriefSlotEngine(DenseSlotEngine):
    name = "dense-slot-brief"
    latency = ("In-process request wall time including prompt construction, the batched stock-model rationale generation of all questions of the request (Hugging Face generate, "
               "adapter disabled) and one adapted slot pass per question; excludes model loading.")

    def __init__(self, model="Qwen/Qwen3.5-9B", run=None, ckpt="best", context="brief", parallel_context="brief", brief_cap=64, gen_tokens_budget=24000, temperature="auto", max_tokens=32768, **options):
        if not run:
            raise ValueError("BriefSlotEngine needs a trained run (option run=...)")
        self.temp_name = f"temperature-{context}"   # fitted by tools/fit_temperature_v6.py
        super().__init__(model=model, run=run, ckpt=ckpt, quant="none", merge=False, temperature=temperature, prefix_cache=False, max_tokens=max_tokens, **options)
        self.context = str(context)
        self.parallel_context = str(parallel_context)   # context of per-option ('parallel' slot) questions: brief | direct (v6b: trained without a rationale for these families)
        if self.context not in ("brief", "direct"):
            raise ValueError("context must be brief or direct")
        self.brief_cap, self.gen_budget = int(brief_cap), int(gen_tokens_budget)
        self.ctx = Contexts(self.tok, self.codec, brief_cap=self.brief_cap)
        self.peft = self.pipe.peft
        self.eos = self.tok.convert_tokens_to_ids("<|im_end|>")
        self.provenance.update({"context": self.context, "brief_cap": self.brief_cap, "generation": "stock model (adapter disabled), greedy, Hugging Face generate",
                                "policy": self.provenance["policy"] + " Before the slot is read, a stock-model rationale of at most %d tokens is generated for every question." % self.brief_cap})

    # ------------------------------------------------------------------ generation (adapter off)
    @torch.inference_mode()
    def _generate(self, prompts):
        """Greedy rationales for a list of token-id prompts; length-sorted batches under a padded-token budget."""
        out = [None] * len(prompts)
        order = sorted(range(len(prompts)), key=lambda i: len(prompts[i]))
        s = 0
        pad = self.tok.pad_token_id
        while s < len(order):
            e = s + 1
            while e < len(order) and (e - s + 1) * len(prompts[order[e]]) <= self.gen_budget:
                e += 1
            idx = order[s:e]
            L = max(len(prompts[i]) for i in idx)
            ids = torch.full((len(idx), L), pad, dtype=torch.long)
            att = torch.zeros((len(idx), L), dtype=torch.long)
            for j, i in enumerate(idx):
                p = prompts[i]
                ids[j, L - len(p):] = torch.tensor(p)
                att[j, L - len(p):] = 1
            base = self.peft.get_base_model()
            with self.peft.disable_adapter(), torch.amp.autocast("cuda", dtype=torch.bfloat16):
                gen = base.generate(input_ids=ids.cuda(), attention_mask=att.cuda(), max_new_tokens=self.brief_cap, do_sample=False, eos_token_id=self.eos, pad_token_id=pad, use_cache=True)
            for j, i in enumerate(idx):
                g = gen[j, L:].tolist()
                if self.eos in g:
                    g = g[:g.index(self.eos)]
                while g and g[-1] == pad:
                    g.pop()
                out[i] = self.ctx.clean_brief(g)
            s = e
        return out

    def _item(self, state, q):
        """The codec item of a question (as DenseSlotEngine._render, but returns the item)."""
        if q["type"] not in ("choice", "noul"):
            raise Unsupported("Unsupported question type " + str(q["type"]))
        keys = list(q["criteria"]) if q["type"] == "choice" else ["false", "true"]
        if not keys:
            raise Unsupported("question without options")
        crit = q["criteria"] if q["type"] == "choice" else (q.get("criteria") or {})
        k = len(keys)
        text = "" if state in (None, "", {}, []) else (state if isinstance(state, str) else __import__("json").dumps(state, ensure_ascii=False, separators=(",", ":")))
        return keys, {"qtype": QTYPES[q["type"]], "target": [1.0 / k] * k, "label": 0, "state": text, "source": "index",
                      "qdef": {"type": q["type"], "instructions": q.get("instructions", ""), "criteria": crit}}

    def __call__(self, state, questions):
        if self.context == "direct":
            return super().__call__(state, questions)
        answers, raw = {}, {}
        items = {qid: self._item(state, q) for qid, q in questions.items()}
        if self.parallel_context == "direct":   # per-option questions skip the rationale: they are answered by the plain slot path
            par = {qid for qid, (_, it) in items.items() if AS.slot_mode(it["qtype"], len(it["target"])) == "parallel"}
            if par:
                resp_p, raw_p = DenseSlotEngine.__call__(self, state, {qid: questions[qid] for qid in par})
                rest = {qid: q for qid, q in questions.items() if qid not in par}
                resp_r, raw_r = self.__call__(state, rest) if rest else ({"model": self.model_id, "answers": {}, "usage": {"input_tokens": 0}}, {})
                resp_p["answers"].update(resp_r["answers"]); resp_p["usage"]["input_tokens"] += resp_r["usage"]["input_tokens"]
                return resp_p, {**raw_p, **raw_r}
        qids = list(items)
        try:
            gen_prompts = [self.ctx.gen_prompt(items[qid][1]) for qid in qids]
            for qid, p in zip(qids, gen_prompts):
                if len(p) + self.brief_cap > self.limit:
                    raise Unsupported(f"prompt of {len(p)} tokens exceeds the {self.limit}-token limit")
            t0 = time.perf_counter()
            rats = self._generate(gen_prompts)
            t_gen = time.perf_counter() - t0
            for qid, rat in zip(qids, rats):
                keys, item = items[qid]
                ex = {**item, "rationale": {"brief": {"tokens": rat, "n": len(rat)}}}
                r = self.ctx.render(ex, "brief")
                if len(r["ids"]) > self.limit:
                    raise Unsupported(f"prompt of {len(r['ids'])} tokens exceeds the {self.limit}-token limit")
                logits, pooled = self._logits(r, None)
                t = self._temperature(logits, pooled, r["qtype"], len(keys))
                p = torch.softmax(logits / t, -1).cpu().numpy().astype(np.float64)
                p = p / p.sum()
                probs = dict(zip(keys, (float(x) for x in p)))
                q = questions[qid]
                if q["type"] == "choice":
                    answers[qid] = {"type": "choice", "choice": max(keys, key=lambda x: probs[x]), "probabilities": probs}
                else:
                    answers[qid] = {"type": "noul", "noul": probs["true"]}
                raw[qid] = {"prompt_tokens": len(r["ids"]), "rationale_tokens": len(rat), "rationale": self.tok.decode(rat), "temperature": round(t, 4), "generation_s": round(t_gen, 3)}
        except torch.OutOfMemoryError:
            torch.cuda.empty_cache()
            raise Unsupported("out of memory on a long prompt")
        return {"model": self.model_id, "answers": answers, "usage": {"input_tokens": sum(v["prompt_tokens"] for v in raw.values())}}, raw