"""trigger_gate.py — runtime recall gate (TRIGGER_V3 head, wiring stage).

Scores the current question in the production layout [BOS][SP(kept)][raw window]
[question] with the question-position pre-RoPE queries (layers 8/14/20) and a
4608-dim logistic head. fires() == True means "the answer needs something that
was evicted from this session's window" — only then does the app pay for an
archive retrieve + verbatim re-injection.

Threshold note: the gate npz carries a TPR-first threshold (-2.6: heldout TPR
1.000, FPR 0.043) — a false fire costs one redundant retrieve (today's per-turn
status quo), a miss costs the recall itself.
"""
import numpy as np
import torch

LAYERS = (8, 14, 20)


class TriggerGate:
    def __init__(self, llm, tok, pooler, head_path, thresh=None):
        z = np.load(head_path)
        self.coef = z["coef"].astype(np.float32).ravel()
        self.b = float(np.ravel(z["intercept"])[0])
        self.mean = z["mean"].astype(np.float32)
        self.scale = z["scale"].astype(np.float32)
        if thresh is not None:
            self.thresh = float(thresh)
        elif "thresh" in z.files:
            self.thresh = float(z["thresh"])
        else:
            self.thresh = 0.0
        self.llm, self.tok, self.pooler = llm, tok, pooler
        self.dev = next(llm.parameters()).device
        self.embT = llm.get_input_embeddings()
        self.last_score = 0.0

    @torch.no_grad()
    def score(self, kept_ids, window_ids, q_text):
        from transformers.cache_utils import DynamicCache
        pre = self.tok.encode(f"<｜end▁of▁sentence｜><｜User｜>{q_text}<｜Assistant｜>"
                              f"<think>\n\n</think>\n\n", add_special_tokens=False)
        cache = DynamicCache()
        bos = self.tok.bos_token_id
        self.llm.model(input_ids=torch.tensor([[bos]], device=self.dev),
                       past_key_values=cache, use_cache=True)
        parts = []
        if kept_ids:
            sp = self.pooler(self.embT(torch.tensor([kept_ids], device=self.dev))
                             .float()).to(self.embT.weight.dtype)
            parts.append(sp)
        if window_ids:
            parts.append(self.embT(torch.tensor([window_ids], device=self.dev)))
        if parts:
            self.llm.model(inputs_embeds=torch.cat(parts, 1), past_key_values=cache,
                           use_cache=True)
        outs, hooks = {}, []
        for li in LAYERS:
            attn = self.llm.model.layers[li].self_attn

            def mk(li):
                def fn(_m, _i, o):
                    outs[li] = o.detach()
                return fn
            hooks.append(attn.q_proj.register_forward_hook(mk(li)))
        try:
            self.llm.model(inputs_embeds=self.embT(torch.tensor([pre], device=self.dev)),
                           past_key_values=cache, use_cache=True)
        finally:
            for h in hooks:
                h.remove()
        f = np.concatenate([outs[li][0].mean(0).float().cpu().numpy().ravel()
                            for li in LAYERS])
        self.last_score = float(((f - self.mean) / self.scale) @ self.coef + self.b)
        return self.last_score

    def fires(self, kept_ids, window_ids, q_text):
        return self.score(kept_ids, window_ids, q_text) > self.thresh