"""Decision Index engine, vLLM path, one instance: direct slot pass for every question, learned router, stock-model rationale only for the escalated questions, adapted slot pass.

Same decisions as decision_engine_v6r:RoutedSlotEngine (Hugging Face) but served by vLLM (CUDA graphs on, prefix cache on). vLLM does not return hidden states, so the model runner's
`sample(hidden_states, input_batch, ...)` is wrapped (single-process engine core, VLLM_ENABLE_V1_MULTIPROCESSING=0): for every request the final-norm hidden states at the wanted slot
positions are kept, and the answer-slot head (frozen LM-head rows of the answer tokens + trained residual, head.pt) and the router head (CKPT/router-direct-brief.npz) run on them in torch.
The rationale is generated by the stock model (adapter off, greedy, at most `brief_cap` tokens) and cleaned like the training rationales (dense/contexts.py).

  VLLM_ENABLE_V1_MULTIPROCESSING=0 VLLM_ALLOW_INSECURE_SERIALIZATION=1 /mnt/datasets/venv-vllm/bin/python -m decision_index run --engine decision_engine_vr:VllmRoutedEngine \
      --option model=Qwen/Qwen3.5-9B --option run=RUN --option ckpt=final [--option mode=routed|direct|brief] [--option tau=0.05] --suite-dir ... --rows ... --out ...
Options: model, run, ckpt, mode (routed | direct | brief), tau (override; inf = never escalate), brief_cap (64), max_tokens (prompt limit, 32768), gpu_util (0.88), max_loras lora (converted adapter dir; made next to the checkpoint when absent).
Per-option (parallel slot) questions are always answered by the direct pass (v6b recipe). Temperature: none.
"""
import json
import os
import shutil
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
os.environ.setdefault("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
os.environ.setdefault("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")

import numpy as np  # noqa: E402
import torch  # noqa: E402
import torch.nn.functional as F  # noqa: E402

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


def _lg(p):
    p = np.clip(p, 1e-4, 1 - 1e-4)
    return np.log(p / (1 - p))


class VllmRoutedEngine(Engine):
    name = "dense-slot-routed-vllm"
    latency = ("In-process request wall time (vLLM, one instance, CUDA graphs): the direct slot pass of all questions of the request in one batch, the router head, and for the escalated "
               "questions a batched stock-model rationale (adapter off, at most 64 tokens) followed by the adapted slot pass. Excludes model loading.")

    def __init__(self, model="Qwen/Qwen3.5-9B", run=None, ckpt="final", mode="routed", tau=None, brief_cap=64, max_tokens=32768, gpu_util=0.88, lora=None, **options):
        super().__init__(**options)
        if not run:
            raise ValueError("VllmRoutedEngine needs a trained run (option run=...)")
        from transformers import AutoTokenizer
        from vllm import LLM, SamplingParams
        from vllm.lora.request import LoRARequest
        self.mode = str(mode)
        if self.mode not in ("routed", "direct", "brief"):
            raise ValueError("mode must be routed, direct or brief")
        ck = ckpt if os.path.isdir(ckpt) else os.path.join(run, ckpt)
        self.tok = AutoTokenizer.from_pretrained(model)
        if self.tok.pad_token_id is None:
            self.tok.pad_token_id = self.tok.eos_token_id
        self.codec = AS.SlotCodec(self.tok, state_max=None, option_max=None, template_kwargs={"enable_thinking": False}, mask_id=self.tok.pad_token_id or 0, strip_label_keys=True, autoregressive=True)
        self.brief_cap = int(brief_cap)
        self.ctx = Contexts(self.tok, self.codec, brief_cap=self.brief_cap)
        self.eos = self.tok.convert_tokens_to_ids("<|im_end|>")
        self.model_id = model
        lora_dir = lora or os.path.join(ck, "lora-vllm")
        if not os.path.isfile(os.path.join(lora_dir, "adapter_model.safetensors")):
            self._convert_lora(os.path.join(ck, "lora"), lora_dir)
        self.limit = int(max_tokens)
        self.llm = LLM(model, dtype="bfloat16", max_model_len=self.limit + 256, max_num_batched_tokens=self.limit + 256, gpu_memory_utilization=float(gpu_util), enable_prefix_caching=True,
                       enable_lora=True, max_lora_rank=32, max_loras=1)
        self.eng = self.llm.llm_engine
        core = self.eng.engine_core
        core = core.engine_core if hasattr(core, "engine_core") else core
        self.runner = core.model_executor.driver_worker.worker.model_runner
        self.lora = LoRARequest("slot", 1, lora_dir)
        self.sp_slot = SamplingParams(temperature=0, max_tokens=1)
        self.sp_slot_nc = SamplingParams(temperature=0, max_tokens=1, skip_reading_prefix_cache=True)
        self.sp_gen = SamplingParams(temperature=0, max_tokens=self.brief_cap, stop_token_ids=[self.eos])
        self.cap, self.n_req = {}, 0
        orig = self.runner.sample
        cap = self.cap

        def wrapped(hidden_states, input_batch, *a, **k):
            for b, rid in enumerate(input_batch.req_ids):
                w = cap.get(rid.rsplit("-", 1)[0])
                if w is None:
                    continue
                start, comp, n = int(input_batch.query_start_loc_np[b]), int(input_batch.num_computed_tokens_np[b]), int(input_batch.num_scheduled_tokens[b])
                for p in w["pos"]:
                    if comp <= p < comp + n and p not in w["got"]:
                        w["got"][p] = hidden_states[start + p - comp].detach().clone()
            return orig(hidden_states, input_batch, *a, **k)
        self.runner.sample = wrapped
        # answer-slot head: frozen LM-head rows of the answer tokens (+ softcap) + trained residual (head.pt)
        m = self.runner.get_model() if hasattr(self.runner, "get_model") else self.runner.model
        w = next(p for n, p in m.named_parameters() if n.endswith("lm_head.weight")) if any(n.endswith("lm_head.weight") for n, _ in m.named_parameters()) else m.get_input_embeddings().weight
        self.cand_w = w[torch.tensor(self.codec.cand, device=w.device)].detach().clone()
        self.softcap = getattr(m.config.get_text_config(), "final_logit_softcapping", None) if hasattr(m, "config") else None
        self.head = AS.SlotResidual(self.cand_w.size(1), len(self.codec.cand)).to("cuda").eval()
        self.head.load_state_dict(torch.load(os.path.join(ck, "head.pt"), map_location="cuda"))
        # router head
        blob = np.load(os.path.join(ck, "router-direct-brief.npz"))
        self.knots = blob["knots"]
        self.tau = float(blob["tau"]) if tau in (None, "") else float(tau)
        self.W = {p: {k: blob[f"{p}_{k}"] for k in ("mh", "sh", "mc", "sc", "w")} for p in ("w", "f")}
        self.provenance = {"kind": "dense-answer-slot", "repo": model, "run": run, "checkpoint": ck, "mode": self.mode, "tau": self.tau, "layout": AS.LAYOUT, "engine": "vLLM, hidden states from a wrapped model-runner sample()",
                           "context_limit_tokens": self.limit, "temperature": "none",
                           "policy": "Direct answer-slot pass for every question; a learned head escalates a question to a stock-model rationale of at most %d tokens before the slot is read. Prompts over the limit are refused as unsupported, never truncated." % self.brief_cap}

    @staticmethod
    def _convert_lora(src, dst):
        from safetensors.torch import load_file, save_file
        os.makedirs(dst, exist_ok=True)
        sd = load_file(os.path.join(src, "adapter_model.safetensors"))
        new = {k.replace("base_model.model.model.layers.", "base_model.model.model.language_model.layers.", 1): v for k, v in sd.items()}
        assert len(new) == len(sd) and any("language_model" in k for k in new)
        save_file(new, os.path.join(dst, "adapter_model.safetensors"), metadata={"format": "pt"})
        shutil.copy(os.path.join(src, "adapter_config.json"), os.path.join(dst, "adapter_config.json"))

    # ------------------------------------------------------------------ rendering
    def _item(self, state, q):
        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 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}}

    # ------------------------------------------------------------------ vLLM plumbing
    def _add(self, ids, sp, lora=None):
        self.n_req += 1
        rid = f"r{self.n_req}"
        self.eng.add_request(rid, {"prompt_token_ids": ids}, sp, lora_request=lora)
        return rid

    def _drain(self, rids):
        wait, done = set(rids), {}
        while len(done) < len(wait):
            for o in self.eng.step():
                if o.request_id in wait and o.finished:
                    done[o.request_id] = o
        return done

    def _states(self, jobs):
        """jobs: list of (ids, slot positions) -> list of (S, H) bf16 hidden states at the slot positions (adapter on). A slot position that sits in the prefix cache is
        not recomputed, so such a job is run again with the prefix-cache read skipped."""
        res, todo = [None] * len(jobs), list(range(len(jobs)))
        for attempt in (0, 1):
            rids = {}
            for i in todo:
                ids, pos = jobs[i]
                self.n_req += 1
                rid = f"s{self.n_req}"
                self.cap[rid] = {"pos": set(pos), "got": {}}
                self.eng.add_request(rid, {"prompt_token_ids": ids}, self.sp_slot_nc if attempt else self.sp_slot, lora_request=self.lora)
                rids[rid] = i
            self._drain(rids)
            todo = []
            for rid, i in rids.items():
                got, pos = self.cap.pop(rid)["got"], jobs[i][1]
                if all(p in got for p in pos):
                    res[i] = torch.stack([got[p] for p in pos])
                else:
                    todo.append(i)
            if not todo:
                return res
        raise RuntimeError("slot hidden state not captured")

    @torch.inference_mode()
    def _readout(self, hs, r):
        """Option logits (K,) and pooled slot state (H,) from slot hidden states hs (S, H), as DensePipeline.readout."""
        dev = "cuda"
        t = lambda x: torch.tensor([x], device=dev)
        k = len(r["target"])
        with torch.amp.autocast("cuda", dtype=torch.bfloat16):
            h = hs[None]
            lin = F.linear(h, self.cand_w).float()
            if self.softcap:
                lin = self.softcap * torch.tanh(lin / self.softcap)
            z = lin + self.head(h).float()
            zs = z.gather(1, t(r["slot_of_opt"]).unsqueeze(-1).expand(-1, -1, z.size(-1)))
            pos = zs.gather(2, t(r["pos_idx"]).unsqueeze(-1)).squeeze(-1)
            neg = zs.gather(2, t(r["neg_idx"]).clamp(min=0).unsqueeze(-1)).squeeze(-1)
            logits = (pos - neg * (t(r["neg_idx"]) >= 0)).masked_fill(~torch.ones(1, k, dtype=torch.bool, device=dev), -1e4)[0]
        return logits.float(), hs.float().mean(0)

    def _score(self, pooled, conf, k):
        x = _lg(np.array([conf], dtype=np.float64))
        C = np.concatenate([x[:, None], np.maximum(0, x[:, None] - self.knots[None]), np.log([[float(k)]]), np.log([[float(k)]]) * x[:, None]], 1)
        h = pooled.double().cpu().numpy()[None]
        f = lambda M: float(np.clip((np.concatenate([(h - M["mh"]) / M["sh"], (C - M["mc"]) / M["sc"], np.ones((1, 1))], 1) @ M["w"])[0], 0, 1))
        return f(self.W["w"]) * f(self.W["f"])

    @staticmethod
    def _slot_ids(r):
        """Prompt ids and slot positions; a single slot is the last position, so everything after it is dropped (causal)."""
        sp = list(r["slot_pos"])
        return (r["ids"][:sp[0] + 1] if len(sp) == 1 else r["ids"]), sp

    # ------------------------------------------------------------------ request
    def __call__(self, state, questions):
        qids = list(questions)
        items = {qid: self._item(state, questions[qid]) for qid in qids}
        direct = {}
        for qid in qids:
            keys, item = items[qid]
            r = self.codec.render_ar(item)
            if len(r["ids"]) + self.brief_cap > self.limit:
                raise Unsupported(f"prompt of {len(r['ids'])} tokens exceeds the {self.limit}-token limit")
            direct[qid] = r
        jobs = [self._slot_ids(direct[qid]) for qid in qids]
        hs = self._states(jobs)
        answers, raw, esc, probs_of = {}, {}, [], {}
        for qid, h in zip(qids, hs):
            keys, item = items[qid]
            r = direct[qid]
            logits, pooled = self._readout(h, r)
            p = torch.softmax(logits, -1).cpu().numpy().astype(np.float64)
            p = p / p.sum()
            q = questions[qid]
            score = None
            routable = q["type"] == "choice" and AS.slot_mode(r["qtype"], len(keys)) != "parallel"
            go = False
            if routable and self.mode == "brief":
                go = True
            elif routable and self.mode == "routed":
                score = self._score(pooled, float(p.max()), len(keys))
                go = score >= self.tau
            if go:
                esc.append(qid)
                raw[qid] = {"routed": "brief", "router_score": None if score is None else round(score, 5), "direct_conf": round(float(p.max()), 4)}
                continue
            probs_of[qid] = dict(zip(keys, (float(x) for x in p)))
            raw[qid] = {"routed": "direct", "router_score": None if score is None else round(score, 5), "prompt_tokens": len(r["ids"]), "direct_conf": round(float(p.max()), 4)}
        if esc:
            gen_prompts = [self.ctx.gen_prompt(items[qid][1]) for qid in esc]
            for qid, gp in zip(esc, gen_prompts):
                if len(gp) + self.brief_cap > self.limit:
                    raise Unsupported(f"prompt of {len(gp)} tokens exceeds the {self.limit}-token limit")
            rids = {self._add(gp, self.sp_gen): qid for gp, qid in zip(gen_prompts, esc)}
            outs = self._drain(rids)
            rendered = {}
            for rid, qid in rids.items():
                g = list(outs[rid].outputs[0].token_ids)
                if self.eos in g:
                    g = g[:g.index(self.eos)]
                rat = self.ctx.clean_brief(g)
                keys, item = items[qid]
                r = self.ctx.render({**item, "rationale": {"brief": {"tokens": rat, "n": len(rat)}}}, "brief")
                if len(r["ids"]) > self.limit:
                    raise Unsupported(f"prompt of {len(r['ids'])} tokens exceeds the {self.limit}-token limit")
                rendered[qid] = (r, rat)
            order = list(rendered)
            hs2 = self._states([self._slot_ids(rendered[qid][0]) for qid in order])
            for qid, h in zip(order, hs2):
                r, rat = rendered[qid]
                keys = items[qid][0]
                logits, _ = self._readout(h, r)
                p = torch.softmax(logits, -1).cpu().numpy().astype(np.float64)
                p = p / p.sum()
                probs_of[qid] = dict(zip(keys, (float(x) for x in p)))
                raw[qid].update({"prompt_tokens": len(r["ids"]), "rationale_tokens": len(rat), "rationale": self.tok.decode(rat)})
        for qid in qids:
            probs = probs_of[qid]
            answers[qid] = {"type": "choice", "choice": max(probs, key=probs.get), "probabilities": probs} if questions[qid]["type"] == "choice" else {"type": "noul", "noul": probs["true"]}
        return {"model": self.model_id, "answers": answers, "usage": {"input_tokens": sum(v["prompt_tokens"] for v in raw.values())}}, raw