"""Decision Index engine with a learned router between the two v6 levels: every question gets the direct slot pass first; the router head (tools/router_fit_final.py,
CKPT/router-direct-brief.npz: P(direct wrong) x P(brief right | direct wrong) from the direct pass pooled slot hidden state + confidence) decides per question whether the brief
rationale path (BriefSlotEngine) is run instead. Per-option (parallel) and non-choice questions stay direct. Escalated questions of a request are batched through the brief path together.

  python -m decision_index run --engine decision_engine_v6r:RoutedSlotEngine --option model=Qwen/Qwen3.5-9B --option run=RUN --option ckpt=final [--option tau=0.02] ...
Options: those of BriefSlotEngine (context is fixed to brief for the escalated path), tau (override of the stored threshold; `inf` = never escalate, `-inf` = always).
"""
import os
import sys

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 decision_engine_dense import DenseSlotEngine  # noqa: E402
from decision_engine_v6 import BriefSlotEngine  # noqa: E402


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


class RoutedSlotEngine(BriefSlotEngine):
    name = "dense-slot-routed"
    latency = ("In-process request wall time (Hugging Face path): direct slot pass for every question, router head, and the brief path (batched rationale generation + adapted slot pass) "
               "only for the questions the router escalates. Not the board latency path; see tools/vllm_requests_0f.py for decode timing.")

    def __init__(self, tau=None, **options):
        options["context"] = "brief"
        options.setdefault("temperature", "none")
        super().__init__(**options)
        run, ck = options["run"], options.get("ckpt", "best")
        blob = np.load(os.path.join(ck if os.path.isdir(ck) else os.path.join(run, 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.update({"router": f"direct->brief head, tau {self.tau}"})

    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"])

    def __call__(self, state, questions):
        answers, raw, esc = {}, {}, {}
        tokens = 0
        try:
            for qid, q in questions.items():
                keys, r = self._render(state, q)
                logits, pooled = self._logits(r, None)
                p = torch.softmax(logits, -1).cpu().numpy().astype(np.float64)
                p = p / p.sum()
                score = None
                if q["type"] == "choice" and AS.slot_mode(r["qtype"], len(keys)) != "parallel":
                    score = self._score(pooled, float(p.max()), len(keys))
                    if score >= self.tau:
                        esc[qid] = q
                        raw[qid] = {"routed": "brief", "router_score": round(score, 5), "direct_conf": round(float(p.max()), 4)}
                        continue
                probs = dict(zip(keys, (float(x) for x in p)))
                answers[qid] = {"type": "choice", "choice": max(keys, key=lambda x: probs[x]), "probabilities": probs} if q["type"] == "choice" else {"type": "noul", "noul": probs["true"]}
                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)}
                tokens += len(r["ids"])
        except torch.OutOfMemoryError:
            torch.cuda.empty_cache()
            try:
                from decision_index.engines.base import Unsupported
            except ImportError:
                from engine_base import Unsupported
            raise Unsupported("out of memory on a long prompt")
        if esc:
            resp_b, raw_b = BriefSlotEngine.__call__(self, state, esc)
            answers.update(resp_b["answers"]); tokens += resp_b["usage"]["input_tokens"]
            for qid, v in raw_b.items():
                raw[qid].update(v)
        return {"model": self.model_id, "answers": {qid: answers[qid] for qid in questions}, "usage": {"input_tokens": tokens}}, raw