"""Typed-decision core from STEAV's published System One inference implementation.

Apache-2.0. The encoder layout and decision head follow Convai Laya's
Apache-2.0 implementation. No remote model code is executed.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import math
import os
import re
import sys
import warnings
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Iterable, Sequence

import numpy as np
import torch
from torch import nn

__all__ = ["SystemOneDecisionModel", "LayaDecisionHead", "render_state", "offset_logit", "apply_prior_shift",
           "encode_decision", "load_laya_encoder", "ModelError"]

# --- Contract constants (identical to the trainer that produced these models) ---
PROMPT_TEMPLATE = "cid/system-one-encoder-prompt/laya-v1"
HEAD_SCHEMA = "cid/system-one-head/v1"
CALIBRATION_SCHEMA = "cid/system-one-calibration/v1"
RELEASE_SCHEMA = "system-one-release/v1"
LAYA_QTYPES = {"choice": 0, "score": 1, "noul": 2}
LAYA_HEAD_MAX_LEN = 192
LAYA_OPTION_TOKENS = 48
MAX_SLOTS = 128
OFFSET_CLIP = 1e-6
PRIOR_CLIP = 1e-9  # the benchmark scorer's clip before the prior shift
DEFAULT_DESCRIPTIONS = {"choice": "Which answer is correct?", "score": "Which level applies?", "noul": "Is the proposition true?"}
NOUL_SLOTS = ({"id": "false", "description": "The proposition is false."},
              {"id": "true", "description": "The proposition is true."})
NOUL_OPTION_TEXTS = ("false: no, the statement does not hold", "true: yes, the statement holds")
_NUMBER = re.compile(r"[+-]?(?:[0-9]+(?:\.[0-9]*)?|\.[0-9]+)(?:[eE][+-]?[0-9]+)?")

# --- Base model pin: Hugging Face convaiinnovations/laya, root files (unchanged since the initial release commit
# 00c37c405e3c3ad73ee070227614c89cda06b99e). Byte-identical to the base used in training except config.json's
# "architectures" entry, which AutoModel.from_config ignores.
LAYA_BASE_REPO = "convaiinnovations/laya"
LAYA_BASE_REVISION = "7b928d828b7b0e022f929d9bd2e44165aa270148"
LAYA_BASE_SHA256 = "891102d372688fc2a094dac56a384bc537b87c63f21f9f3dac0be2b7cbc8d86c"
LAYA_BASE_FILES = ("model.safetensors", "encoder/config.json", "tokenizer/tokenizer.json", "tokenizer/tokenizer_config.json")


class ModelError(ValueError):
    """Raised for invalid model folders, inputs that break the input contract, or failed integrity checks."""


# ----------------------------------------------------------------------------- input contract

def question_slots(question: dict) -> list[dict]:
    """Answer slots in scoring order: choice options, score levels, or [false, true] for noul."""
    qtype = str(question.get("type", "")).strip()
    if qtype == "choice":
        slots = question.get("options") or []
    elif qtype == "score":
        slots = question.get("levels") or []
    elif qtype == "noul":
        if question.get("options") or question.get("levels"):
            raise ModelError("noul takes no options or levels")
        slots = NOUL_SLOTS
    else:
        raise ModelError("question type must be choice, score or noul")
    if not 2 <= len(slots) <= MAX_SLOTS:
        raise ModelError("questions must carry between 2 and %d answer slots" % MAX_SLOTS)
    ids = [str(slot.get("id", "")).strip() for slot in slots]
    if any(not identity for identity in ids) or len(set(ids)) != len(ids):
        raise ModelError("answer ids must be nonempty and distinct")
    return [{"id": identity, "description": str(slot.get("description") or "")} for identity, slot in zip(ids, slots)]


def option_texts(question: dict, slots: Sequence[dict]) -> list[str]:
    """Answer texts as rendered after each [MASK] marker."""
    qtype = str(question.get("type", "")).strip()
    if qtype == "score":
        return ["level %d: %s" % (i, " ".join(slot["description"].split()) or slot["id"]) for i, slot in enumerate(slots)]
    if qtype == "noul":
        return list(NOUL_OPTION_TEXTS)
    return [(slot["id"] + ": " + " ".join(slot["description"].split())) if slot["description"].strip() else slot["id"]
            for slot in slots]


def render_state(state: Any, offset_field: str | None = None, offset_hidden: bool = False) -> str:
    """Evidence text exactly as trained.

    A dict is rendered as compact JSON with sorted keys and raw Unicode. Key order therefore decides which fields
    survive end truncation: digits < uppercase (Context_*, Tree_*) < "_" < lowercase. A str state is used as given
    (stripped). With offset_hidden, only the literal offset key is removed from the text.
    """
    if offset_hidden and offset_field and isinstance(state, dict):
        state = {key: value for key, value in state.items() if key != offset_field}
    if isinstance(state, str):
        return state.strip()
    return json.dumps(state, sort_keys=True, separators=(",", ":"), ensure_ascii=False)


def offset_logit(state: Any, question: dict, field: str | None) -> float:
    """Fixed upstream log-odds added to the "true" slot of noul questions.

    The value must be a JSON number or a numeric string in [0, 1]; booleans are rejected. It is clipped to
    [1e-6, 1 - 1e-6], so the offset lies within about +/-13.8.
    """
    if not field or str(question.get("type", "")).strip() != "noul":
        return 0.0
    if not isinstance(state, dict) or field not in state:
        raise ModelError("offset field %r missing from state" % field)
    value = state[field]
    if isinstance(value, bool) or not isinstance(value, (int, float, str)):
        raise ModelError("offset field %r must be a probability in [0, 1]" % field)
    if isinstance(value, str) and not _NUMBER.fullmatch(value.strip()):
        raise ModelError("offset field %r must be a probability in [0, 1]" % field)
    try:
        p = float(value)
    except (TypeError, ValueError, OverflowError):
        raise ModelError("offset field %r must be a probability in [0, 1]" % field) from None
    if not math.isfinite(p) or p < 0 or p > 1:
        raise ModelError("offset field %r must be a probability in [0, 1]" % field)
    p = min(max(p, OFFSET_CLIP), 1 - OFFSET_CLIP)
    return math.log(p / (1 - p))


def bucket_key(question: dict, count: int) -> str:
    """Calibration bucket: '<type>:2', ':3-4', ':5-8' or ':9+'."""
    width = "2" if count <= 2 else "3-4" if count <= 4 else "5-8" if count <= 8 else "9+"
    return str(question.get("type", "")).strip() + ":" + width


@dataclass
class Encoded:
    ids: list[int]
    positions: list[int]  # index of each answer's [MASK] marker
    qtype: int
    offset: float = 0.0
    state_tokens: int = 0  # evidence tokens before truncation
    dropped_state_tokens: int = 0  # evidence tokens cut from the END to fit max_seq_length


def encode_decision(tokenizer, state: Any, question: dict, limit: int, offset_field: str | None = None,
                    offset_hidden: bool = False) -> Encoded:
    """Laya sequence (template cid/system-one-encoder-prompt/laya-v1):

    [CLS] "<type> question: <description>" [SEP] [MASK] " <option 0>" [MASK] " <option 1>" ... [SEP] <evidence> [SEP]

    Answers are never dropped; the question text and, if needed, the answers are shortened to the answer budget;
    the evidence keeps its first tokens. Raw text is tokenized as is: special-token strings such as "[MASK]"
    inside the evidence become special tokens, as they did in training.
    """
    mask_id, cls_id, sep_id = tokenizer.mask_token_id, tokenizer.cls_token_id, tokenizer.sep_token_id
    if mask_id is None or cls_id is None or sep_id is None:
        raise ModelError("the tokenizer needs mask, cls and sep tokens")
    slots = question_slots(question)
    qtype = str(question.get("type", "")).strip()
    description = " ".join(str(question.get("description") or "").split()) or DEFAULT_DESCRIPTIONS[qtype]
    head_ids = tokenizer.encode(qtype + " question: " + description, add_special_tokens=False)
    option_ids = [[mask_id] + tokenizer.encode(" " + text, add_special_tokens=False)[:LAYA_OPTION_TOKENS]
                  for text in option_texts(question, slots)]
    head_max_len = max(LAYA_HEAD_MAX_LEN, min(8 * len(option_ids), limit // 2))
    budget = head_max_len - sum(len(o) for o in option_ids)
    if budget < 16:
        per = max(4, (head_max_len - 16) // max(1, len(option_ids)))
        option_ids = [o[:per] for o in option_ids]
        budget = head_max_len - sum(len(o) for o in option_ids)
    head_ids = head_ids[:max(8, budget)]
    ids = [cls_id] + head_ids + [sep_id]
    positions = []
    for option in option_ids:
        positions.append(len(ids))
        ids.extend(option)
    ids.append(sep_id)
    room = limit - len(ids) - 1
    if room < 8:
        raise ModelError("decision prompt leaves no room for evidence; raise max_seq_length or shorten the answers")
    evidence = tokenizer.encode(render_state(state, offset_field, offset_hidden), add_special_tokens=False)
    ids = ids + evidence[:room] + [sep_id]
    if len(ids) > limit:
        raise ModelError("decision prompt exceeds max_seq_length")
    return Encoded(ids, positions, LAYA_QTYPES[qtype], offset_logit(state, question, offset_field),
                   len(evidence), max(0, len(evidence) - room))


def apply_prior_shift(p_true: float, training_prior: float, target_prior: float) -> float:
    """Bayes prior correction of a calibrated P(true), as the benchmark scorer applied it."""
    for prior in (training_prior, target_prior):
        if not 0 < prior < 1:
            raise ModelError("priors must lie strictly between 0 and 1")
    shift = math.log(target_prior / (1 - target_prior)) - math.log(training_prior / (1 - training_prior))
    log_odds = math.log(max(p_true, PRIOR_CLIP) / max(1 - p_true, PRIOR_CLIP)) + shift
    return 1 / (1 + math.exp(-log_odds))


# ----------------------------------------------------------------------------- model parts

class LayaDecisionHead(nn.Module):
    """Laya's typed decision head (Apache-2.0 `DecisionModel` minus the encoder), run in float32.

    Parameter names match head.safetensors: head.layers.{0,1}.*, type_emb.weight, scorer.*,
    act_head.* (frozen act/escalate head, unused for scoring) and the inert `temperature` buffer (Laya's own
    per-type temperatures; candidates are calibrated by calibration.json instead).
    """

    def __init__(self, hidden: int, layers: int = 2, actions: int = 2, dropout: float = 0.1):
        super().__init__()
        layer = nn.TransformerEncoderLayer(hidden, max(1, hidden // 64), 4 * hidden, dropout, batch_first=True, norm_first=True)
        self.head = nn.TransformerEncoder(layer, layers, enable_nested_tensor=False) if layers > 0 else None
        self.type_emb = nn.Embedding(3, hidden)
        self.scorer = nn.Sequential(nn.LayerNorm(hidden), nn.Linear(hidden, hidden), nn.GELU(), nn.Linear(hidden, 1))
        self.act_head = nn.Sequential(nn.Linear(hidden + 4, 256), nn.GELU(), nn.Linear(256, actions))
        self.register_buffer("temperature", torch.ones(3))

    def score(self, hidden: torch.Tensor, attention_mask: torch.Tensor, positions: Sequence[int], qtype: int) -> torch.Tensor:
        """hidden [L, d] float32 of one padded sequence -> one logit per answer marker.

        Rows are scored one at a time, as in training, so batching never changes the head's arithmetic.
        """
        h = hidden[None, :, :] + self.type_emb(torch.tensor([qtype], device=hidden.device))[:, None, :]
        if self.head is not None:
            pad = ~attention_mask[None, :].bool()
            for layer in self.head.layers:
                h = layer(h, src_key_padding_mask=pad)
        markers = h[0, list(positions)]
        return self.scorer(markers).squeeze(-1).float()


def _sha256(path: Path) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as stream:
        for chunk in iter(lambda: stream.read(1 << 22), b""):
            h.update(chunk)
    return h.hexdigest()


def _local_dir(source: str | os.PathLike, revision: str | None, allow_patterns: Iterable[str] | None = None) -> Path:
    """A local directory as is; otherwise a Hugging Face repo id downloaded at `revision`."""
    path = Path(source).expanduser()
    if path.is_dir():
        return path
    if not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9._-]{0,95}/[A-Za-z0-9][A-Za-z0-9._-]{0,95}", str(source)):
        raise ModelError("%s is neither a local directory nor a Hugging Face repo id" % source)
    from huggingface_hub import snapshot_download
    return Path(snapshot_download(repo_id=str(source), revision=revision, allow_patterns=list(allow_patterns) if allow_patterns else None))


def base_layout(base_dir: Path) -> tuple[Path, Path, Path]:
    """(weights, encoder config dir, tokenizer dir) for the HF repo layout or the flat bundle layout used in training.

    HF convaiinnovations/laya: model.safetensors + encoder/config.json + tokenizer/ (its root config.json is a
    `model_type: laya` stub that transformers cannot load). Flat bundle: model.safetensors + config.json
    (model_type modernbert) + tokenizer files in the same directory.
    """
    weights = base_dir / "model.safetensors"
    if not weights.is_file():
        raise ModelError("base directory has no model.safetensors: %s" % base_dir)
    if (base_dir / "encoder" / "config.json").is_file():
        config_dir = base_dir / "encoder"
        tokenizer_dir = base_dir / "tokenizer" if (base_dir / "tokenizer" / "tokenizer.json").is_file() else base_dir
    elif (base_dir / "config.json").is_file():
        config_dir = tokenizer_dir = base_dir
    else:
        raise ModelError("base directory has neither encoder/config.json nor config.json: %s" % base_dir)
    return weights, config_dir, tokenizer_dir


def load_laya_encoder(base_dir: Path, dtype: torch.dtype, device: torch.device, expected_sha256: str | None = None,
                      attn_implementation: str = "sdpa", rope_dtype: torch.dtype | None = torch.bfloat16) -> nn.Module:
    """Laya's ModernBERT encoder, loaded the way the reference scorer loads it.

    from_config + strict key accounting over the `encoder.`-prefixed tensors, then ONE `.to(dtype)` for the whole
    module. The checkpoint stores float16 tensors; they are copied into float32 parameters and, on CUDA, cast to
    bfloat16. `.to(dtype)` also casts the rotary inv_freq buffers (from_pretrained(dtype=...) keeps them float32),
    so do not replace this with from_pretrained when reproducing stored scores.

    rope_dtype: The original candidates trained and evaluated on CUDA in bfloat16, i.e. with bfloat16-rounded rotary
    frequencies (up to 0.4% relative error). By default the frequencies are rounded the same way when computing in
    another dtype (float32 on CPU), so positions are encoded as in training. None keeps exact float32 frequencies,
    which is what the reference CPU path does.
    """
    from safetensors.torch import load_file
    from transformers import AutoConfig, AutoModel
    weights, config_dir, _ = base_layout(base_dir)
    if expected_sha256 and _sha256(weights) != expected_sha256:
        raise ModelError("base weights sha256 differs from the release manifest (expected %s)" % expected_sha256)
    config = AutoConfig.from_pretrained(str(config_dir), trust_remote_code=False, local_files_only=True)
    if config.model_type != "modernbert" or getattr(config, "is_encoder_decoder", False):
        raise ModelError("the base encoder must be a ModernBERT model")
    model = AutoModel.from_config(config, dtype=torch.float32, attn_implementation=attn_implementation)
    state = load_file(str(weights))
    encoder_state = {key[len("encoder."):]: value for key, value in state.items() if key.startswith("encoder.")} or state
    missing, unexpected = model.load_state_dict(encoder_state, strict=False)
    unexpected = [key for key in unexpected if not key.startswith(("head.", "decoder.", "cls.", "classifier."))]
    if missing or unexpected:
        raise ModelError("base encoder weights do not match the configuration (missing %d, unexpected %d)" % (len(missing), len(unexpected)))
    del state, encoder_state
    model = model.to(dtype=dtype)
    rotary = getattr(model, "rotary_emb", None)
    if rope_dtype is not None and rotary is not None:
        with torch.no_grad():
            for name, buffer in rotary.named_buffers():
                if name.endswith("inv_freq") and buffer.is_floating_point():
                    buffer.copy_(buffer.to(rope_dtype).to(buffer.dtype))
    return model.to(device)


def _resolve_dtype(dtype: str | torch.dtype, device: torch.device) -> torch.dtype:
    if isinstance(dtype, torch.dtype):
        return dtype
    if dtype == "auto":  # compute_dtype: bfloat16 on CUDA, float32 elsewhere
        return torch.bfloat16 if device.type == "cuda" else torch.float32
    try:
        return {"float32": torch.float32, "bfloat16": torch.bfloat16, "float16": torch.float16}[dtype]
    except KeyError:
        raise ModelError("dtype must be 'auto', 'float32', 'bfloat16' or 'float16'") from None


def _tf32_enabled() -> bool:
    """True when float32 matmuls may run in TF32 ("high"/"medium"); the stored scores used the default "highest"."""
    try:
        return torch.get_float32_matmul_precision() != "highest"
    except Exception:  # pragma: no cover - API differences across torch versions
        return False


def _read_json(path: Path) -> dict:
    with open(path, "rb") as stream:
        return json.loads(stream.read())


# ----------------------------------------------------------------------------- public model

class SystemOneDecisionModel:
    """A typed-decision encoder, adapter and calibrated marker-scoring head."""

    def __init__(self, encoder, tokenizer, head: LayaDecisionHead, calibration: dict, head_config: dict, release: dict,
                 device: torch.device, dtype: torch.dtype, max_seq_length: int, question: dict | None,
                 offset_field: str | None, offset_hidden: bool):
        self.encoder, self.tokenizer, self.head = encoder, tokenizer, head
        self.calibration, self.head_config, self.release = calibration, head_config, release
        self.device, self.dtype, self.max_seq_length = device, dtype, max_seq_length
        self.question, self.offset_field, self.offset_hidden = question, offset_field, offset_hidden

    def __repr__(self) -> str:
        return "SystemOneDecisionModel(%s, device=%s, dtype=%s, max_seq_length=%d, offset_field=%r)" % (
            self.release.get("model", "unnamed"), self.device, str(self.dtype).replace("torch.", ""), self.max_seq_length, self.offset_field)

    @classmethod
    def from_pretrained(cls, path: str | os.PathLike, *, base: str | os.PathLike | None = None, base_revision: str | None = None,
                        revision: str | None = None, question: dict | None = None, device: str | torch.device | None = None,
                        dtype: str | torch.dtype = "auto", verify_base: bool = True, max_seq_length: int | None = None,
                        offset_field: str | None | bool = True, rope_dtype: str | torch.dtype | None = "bfloat16") -> "SystemOneDecisionModel":
        """Load a published model folder (or a training-run `adapter/` directory) plus the Laya base.

        path: model folder or HF repo id with adapter_config.json, adapter_model.safetensors and
            system_one/{config.json, head.safetensors, calibration.json[, release.json]}.
        base: Laya base (local dir in HF or flat layout, or a repo id). Default: release.json's base, else
            convaiinnovations/laya at the pinned revision.
        question: the exact question the model was trained with. Default: release.json's question; required for
            a raw training-run adapter directory.
        dtype: "auto" reproduces the reference scorer (bfloat16 on CUDA, float32 on CPU).
        verify_base: check the base weights' sha256 against the release manifest (or the pinned Laya hash).
        offset_field: True uses the trained setting; None/False disables the offset (ablation only); a string
            overrides it (testing only). Trained models with an offset need state[offset_field] in [0, 1].
        rope_dtype: "bfloat16" (default) rounds the rotary frequencies as in training and benchmarking, also
            on float32 CPU runs; None keeps float32 frequencies (the reference CPU behaviour).
        """
        from peft import PeftModel
        from safetensors.torch import load_file
        from transformers import AutoTokenizer
        model_dir = _local_dir(path, revision)
        s1 = model_dir / "system_one"
        head_config = _read_json(s1 / "config.json")
        calibration = _read_json(s1 / "calibration.json")
        release = _read_json(s1 / "release.json") if (s1 / "release.json").is_file() else {}
        if head_config.get("schema") != HEAD_SCHEMA or calibration.get("schema") != CALIBRATION_SCHEMA:
            raise ModelError("unsupported decision head or calibration schema")
        if head_config.get("head") != "laya" or head_config.get("backbone") != "encoder" or head_config.get("prompt_template") != PROMPT_TEMPLATE:
            raise ModelError("only encoder candidates with the Laya head and the laya-v1 prompt are supported")
        if release and release.get("schema") != RELEASE_SCHEMA:
            raise ModelError("unsupported release manifest schema")
        device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
        dtype = _resolve_dtype(dtype, device)
        if device.type == "cuda" and _tf32_enabled():
            warnings.warn("TF32 matmul is enabled; the float32 head then differs from the stored benchmark scores")

        base_info = release.get("base", {})
        base_source = base if base is not None else base_info.get("repo_id", LAYA_BASE_REPO)
        base_dir = _local_dir(base_source, base_revision or base_info.get("revision", LAYA_BASE_REVISION), LAYA_BASE_FILES)
        expected = base_info.get("sha256", LAYA_BASE_SHA256) if verify_base else None
        encoder = load_laya_encoder(base_dir, dtype, device, expected,
                                    rope_dtype=None if rope_dtype is None else _resolve_dtype(rope_dtype, device))
        _, _, tokenizer_dir = base_layout(base_dir)
        if not (tokenizer_dir / "tokenizer.json").is_file() and (model_dir / "tokenizer.json").is_file():
            tokenizer_dir = model_dir
        tokenizer = AutoTokenizer.from_pretrained(str(tokenizer_dir), trust_remote_code=False, local_files_only=True)
        if tokenizer.pad_token_id is None:
            tokenizer.pad_token = tokenizer.eos_token
        tokenizer.padding_side = "right"

        encoder = PeftModel.from_pretrained(encoder, str(model_dir), is_trainable=False, local_files_only=True, autocast_adapter_dtype=True)
        encoder.eval()
        hidden = int(encoder.get_input_embeddings().weight.shape[1])
        if head_config.get("hidden_size") != hidden:
            raise ModelError("decision head width differs from the base encoder")
        weights = load_file(str(s1 / "head.safetensors"))
        actions = int(weights["act_head.2.weight"].shape[0]) if "act_head.2.weight" in weights else 2
        head = LayaDecisionHead(hidden, int(head_config.get("laya_head_layers") or 2), actions).float()
        expected_keys = set(head.state_dict())
        if set(weights) != expected_keys:
            raise ModelError("head weights do not match the head configuration (missing %d, unexpected %d)"
                             % (len(expected_keys - set(weights)), len(set(weights) - expected_keys)))
        head.load_state_dict({key: value.to(torch.float32) for key, value in weights.items()}, strict=True)
        head.to(device).eval()

        if offset_field is True:
            field = head_config.get("offset_field") or None
        elif offset_field in (None, False):
            field = None
        else:
            field = str(offset_field)
        limit = int(max_seq_length or head_config.get("max_seq_length") or 2048)
        return cls(encoder, tokenizer, head, calibration, head_config, release, device, dtype, limit,
                   question or release.get("question"), field, bool(head_config.get("offset_hidden", False)))

    # --- scoring ---

    def _question(self, question: dict | None) -> dict:
        question = question or self.question
        if not isinstance(question, dict):
            raise ModelError("no question: pass question={'type': 'noul', 'description': ...} or ship system_one/release.json")
        return question

    def temperature(self, question: dict | None = None) -> float:
        question = self._question(question)
        count = len(question_slots(question))
        return float(self.calibration.get("buckets", {}).get(bucket_key(question, count), self.calibration.get("global_temperature", 1.0)))

    def encode(self, state: Any, question: dict | None = None) -> Encoded:
        """Token ids, marker positions and offset for one state (for golden tests and truncation checks)."""
        return encode_decision(self.tokenizer, state, self._question(question), self.max_seq_length, self.offset_field, self.offset_hidden)

    @torch.inference_mode()
    def _logits(self, rows: Sequence[Encoded]) -> list[torch.Tensor]:
        """Raw logits (offset included, temperature not applied)."""
        length = max(len(row.ids) for row in rows)
        pad = self.tokenizer.pad_token_id
        ids = torch.tensor([row.ids + [pad] * (length - len(row.ids)) for row in rows], device=self.device)
        mask = torch.tensor([[1] * len(row.ids) + [0] * (length - len(row.ids)) for row in rows], device=self.device)
        hidden = self.encoder(input_ids=ids, attention_mask=mask).last_hidden_state.float()
        out = []
        for index, row in enumerate(rows):
            z = self.head.score(hidden[index], mask[index], row.positions, row.qtype)
            if row.offset:
                # noul slots are [false, true]: the upstream log-odds move the "true" answer only
                z = z + torch.tensor([0.0, row.offset], device=z.device, dtype=z.dtype)
            out.append(z)
        return out