"""Typed path readout: answer from node types and structure, never relation names.

The unseen-ontology holdout renames every relation to an opaque label, so a readout
that matches the question to relation names has nothing to match.  Node types keep
their names, and an oracle that follows only node types reaches 0.99 exact match on
that split, so the information is in the input.  This readout enumerates the paths
from the start entities and scores each path by its node types alone.

Two complementary scorers, blended with a weight chosen on internal holdouts:

* :class:`TypeSequenceScorer` reads the whole question and predicts the hop count and
  the type at each hop.  Robust to new wording of known question types, but it learns
  whole type sequences, so it fails on new *combinations* (0.01 on held-out intents).
* :class:`TypeBagScorer` reads the question token by token: each token names a type
  or nothing, and a path is scored by how likely its multiset of types is under those
  token probabilities, plus the answer type.  Order is left to the graph.  Because a
  word such as "supplier" names the same type in every intent, this composes (0.65 on
  held-out intents), at some cost on new wording.

Types are scored by name embedding, so type names never seen in training work as
long as the name carries meaning.  A small classifier over the readout's own scores
decides whether to answer or abstain, and confidence is isotonic-calibrated per
predicted hop count.
"""

from __future__ import annotations

import math
import re
from collections import Counter
from collections.abc import Callable, Sequence
from dataclasses import dataclass, replace
from pathlib import Path
from typing import Any

import torch
from torch import Tensor, nn
from torch.nn import functional

from graph.graph import GraphStore
from graph.training.calibration import IsotonicCalibrator
from graph.types import (
    EntityAnswer,
    GraphPath,
    GraphQuery,
    ReasoningResult,
    ReasoningStatus,
)

MAX_HOPS = 4
#: Largest per-type count the bag scorer models along one path.
MAX_COUNT = 4
DEFAULT_ENCODER = "sentence-transformers/all-mpnet-base-v2"
#: Features read by the abstain head: v1 has 5, v2 adds novelty and agreement, v3 adds
#: the gate's P(sequence scorer is right).
FEATURES_V1 = 5
FEATURES_V2 = 7
FEATURES_V3 = 8
NUM_FEATURES = FEATURES_V1
#: Index of the template-novelty feature.  It predicts errors on held-out *intents* but
#: not on new wording of known intents, so the abstain head must not read it: with it,
#: the v3 head abstained on 71% of answerable unseen-ontology questions.
NOVELTY_FEATURE = 5
#: Inputs to the gate; see :func:`gate_features`.  Gates trained with the first 5 only
#: still load: :meth:`TypedReadout.gate_probability` slices to the gate's input size.
GATE_FEATURES = 7
#: Hard cap on enumerated paths per query, so a dense graph cannot stall inference.
MAX_PATHS = 200_000
_ENTITY = re.compile(r"\be\d+_\d+\b")

Encoder = Callable[[Sequence[str]], Tensor]
TokenEncoder = Callable[[Sequence[str]], list[Tensor]]


class TypeSequenceScorer(nn.Module):
    """Question embedding -> P(hop count) and P(node type at hop i | hop count)."""

    def __init__(self, dim: int, max_hops: int = MAX_HOPS) -> None:
        super().__init__()
        self.max_hops = max_hops
        self.hop = nn.Sequential(nn.Linear(dim, 256), nn.GELU(), nn.Linear(256, max_hops))
        self.slot = nn.Parameter(torch.randn(max_hops, max_hops, 64) * 0.1)
        self.mix = nn.Sequential(nn.Linear(dim + 64, 512), nn.GELU(), nn.Linear(512, dim))
        self.scale = nn.Parameter(torch.tensor(10.0))

    def hop_logp(self, question: Tensor) -> Tensor:
        return functional.log_softmax(self.hop(question), -1)

    def type_logits(
        self, question: Tensor, hops: int, position: int, type_embeddings: Tensor
    ) -> Tensor:
        """Logits over ``type_embeddings`` rows for the node at ``position``."""
        slot = self.slot[hops - 1, position].expand(question.shape[0], -1)
        query = functional.normalize(self.mix(torch.cat([question, slot], -1)), dim=-1)
        return self.scale * (query @ type_embeddings.T)


class TypeBagScorer(nn.Module):
    """Token-level type mentions -> likelihood of a path's type multiset."""

    def __init__(self, dim: int, max_count: int = MAX_COUNT) -> None:
        super().__init__()
        self.max_count = max_count
        self.tok = nn.Sequential(nn.Linear(dim, dim), nn.GELU(), nn.Linear(dim, dim))
        self.none = nn.Parameter(torch.tensor(2.0))
        self.scale = nn.Parameter(torch.tensor(10.0))
        self.ans = nn.Sequential(nn.Linear(dim, 512), nn.GELU(), nn.Linear(512, dim))
        self.ascale = nn.Parameter(torch.tensor(10.0))

    def token_probs(self, tokens: Tensor, type_embeddings: Tensor) -> Tensor:
        """Per token, P(token names type k); the remainder is P(names no type)."""
        z = functional.normalize(self.tok(tokens), dim=-1)
        logits = self.scale * (z @ type_embeddings.T)
        logits = torch.cat([logits, self.none.expand(logits.shape[0], 1)], -1)
        return torch.softmax(logits, -1)[:, :-1]

    def count_logp(self, tokens: Tensor, type_embeddings: Tensor) -> Tensor:
        """[types, max_count + 1]: log P(type k is named c times), Poisson-binomial."""
        p = self.token_probs(tokens, type_embeddings).clamp(1e-6, 1 - 1e-6)
        k = p.shape[1]
        dist = torch.zeros(k, self.max_count + 1, device=p.device)
        dist[:, 0] = 1.0
        for row in p:
            prob = row[:, None]
            shifted = torch.cat([torch.zeros(k, 1, device=p.device), dist[:, :-1]], 1)
            dist = dist * (1 - prob) + shifted * prob
        return torch.log(dist.clamp_min(1e-12))

    def answer_logp(self, question: Tensor, type_embeddings: Tensor) -> Tensor:
        z = functional.normalize(self.ans(question), dim=-1)
        return functional.log_softmax(self.ascale * (z @ type_embeddings.T), -1)


class GateHead(nn.Module):
    """P(the sequence scorer's answer is right rather than the bag scorer's)."""

    def __init__(self, features: int = GATE_FEATURES) -> None:
        super().__init__()
        self.net = nn.Sequential(nn.Linear(features, 16), nn.GELU(), nn.Linear(16, 1))

    def forward(self, features: Tensor) -> Tensor:
        return self.net(features).squeeze(-1)


class AbstainHead(nn.Module):
    """P(answering is correct) from the readout's own scores."""

    def __init__(self, features: int = NUM_FEATURES) -> None:
        super().__init__()
        self.net = nn.Sequential(nn.Linear(features, 32), nn.GELU(), nn.Linear(32, 1))

    def forward(self, features: Tensor) -> Tensor:
        return self.net(features).squeeze(-1)


@dataclass(frozen=True)
class ReadoutScore:
    """Everything the readout computes for one query before deciding."""

    answers: frozenset[str]
    features: tuple[float, ...]
    best_paths: dict[str, GraphPath]
    endpoint_scores: dict[str, float]
    #: argmax of the hop-count head; known at inference, used to pick a calibrator.
    predicted_hops: int = 0


class TypedReadout:
    """Neuro-symbolic readout over a :class:`GraphStore`; see the module docstring."""

    def __init__(
        self,
        scorer: TypeSequenceScorer,
        abstain: AbstainHead,
        encode: Encoder,
        *,
        feature_mean: Tensor,
        feature_std: Tensor,
        threshold: float,
        calibration: tuple[tuple[float, ...], tuple[float, ...]] | None = None,
        hop_calibration: dict[int, tuple[tuple[float, ...], tuple[float, ...]]] | None = None,
        bag: TypeBagScorer | None = None,
        encode_tokens: TokenEncoder | None = None,
        blend_weight: float = 1.0,
        template_bank: Tensor | None = None,
        gate: GateHead | None = None,
        gate_mean: Tensor | None = None,
        gate_std: Tensor | None = None,
        feature_mask: Sequence[int] | None = None,
        decision_source: str = "chosen",
        model_version: str = "typed-readout-v1",
        device: str | torch.device = "cpu",
    ) -> None:
        if bag is not None and encode_tokens is None:
            raise ValueError("a bag scorer needs a token encoder")
        self.device = torch.device(device)
        self.scorer = scorer.to(self.device).eval()
        self.bag = bag.to(self.device).eval() if bag is not None else None
        self.abstain = abstain.to(self.device).eval()
        self.encode = encode
        self.encode_tokens = encode_tokens
        self.blend_weight = float(blend_weight)
        self.template_bank = (
            functional.normalize(template_bank.to(self.device).float(), dim=-1)
            if template_bank is not None
            else None
        )
        self.gate = gate.to(self.device).eval() if gate is not None else None
        self.gate_mean = (gate_mean if gate_mean is not None else torch.zeros(GATE_FEATURES)).to(self.device)
        self.gate_std = (
            (gate_std if gate_std is not None else torch.ones(GATE_FEATURES)).to(self.device).clamp_min(1e-6)
        )
        #: Indices of score features the abstain head reads (None: all of them).
        self.feature_mask = tuple(feature_mask) if feature_mask is not None else None
        #: Which scores the answer/abstain decision reads: "chosen" (the scorer that
        #: produced the answer) or "sequence" (always the sequence scorer, whose scores
        #: the decision was calibrated on and which transfer under wording shift).
        if decision_source not in ("chosen", "sequence"):
            raise ValueError(f"unknown decision_source {decision_source!r}")
        self.decision_source = decision_source
        self.feature_mean = feature_mean.to(self.device)
        self.feature_std = feature_std.to(self.device).clamp_min(1e-6)
        self.threshold = float(threshold)
        self.calibration = calibration
        self.hop_calibration = dict(hop_calibration or {})
        self.model_version = model_version
        self._type_cache: dict[str, Tensor] = {}

    @property
    def feature_count(self) -> int:
        if self.gate is not None:
            return FEATURES_V3
        return FEATURES_V2 if self.template_bank is not None else FEATURES_V1

    # -- construction -------------------------------------------------------

    @classmethod
    def from_checkpoint(
        cls,
        path: str | Path,
        *,
        device: str | torch.device = "cpu",
        encode: Encoder | None = None,
        encode_tokens: TokenEncoder | None = None,
    ) -> TypedReadout:
        bundle = torch.load(path, map_location="cpu", weights_only=True)
        encoder_name = bundle.get("encoder", DEFAULT_ENCODER)
        if encode is None:
            encode, default_tokens = sentence_encoders(encoder_name, device)
            encode_tokens = encode_tokens or default_tokens
        scorer = TypeSequenceScorer(bundle["dim"], bundle.get("max_hops", MAX_HOPS))
        scorer.load_state_dict(bundle["scorer"])
        bag = None
        if bundle.get("bag") is not None:
            bag = TypeBagScorer(bundle["dim"], bundle.get("max_count", MAX_COUNT))
            bag.load_state_dict(bundle["bag"])
        gate = None
        if bundle.get("gate") is not None:
            gate = GateHead(int(bundle["gate"]["net.0.weight"].shape[1]))
            gate.load_state_dict(bundle["gate"])
        # The head's input size is read from its weights, not ``num_features``: an
        # internal readout saved before its gate existed recorded the wrong count.
        abstain = AbstainHead(int(bundle["abstain"]["net.0.weight"].shape[1]))
        abstain.load_state_dict(bundle["abstain"])
        calibration = bundle.get("calibration")
        hop_calibration = {
            int(hops): (tuple(pair[0]), tuple(pair[1]))
            for hops, pair in (bundle.get("hop_calibration") or {}).items()
        }
        return cls(
            scorer,
            abstain,
            encode,
            feature_mean=bundle["feature_mean"],
            feature_std=bundle["feature_std"],
            threshold=bundle["threshold"],
            calibration=tuple(map(tuple, calibration)) if calibration else None,
            hop_calibration=hop_calibration,
            bag=bag,
            encode_tokens=encode_tokens,
            blend_weight=bundle.get("blend_weight", 1.0),
            template_bank=bundle.get("template_bank"),
            gate=gate,
            gate_mean=bundle.get("gate_mean"),
            gate_std=bundle.get("gate_std"),
            feature_mask=bundle.get("feature_mask"),
            decision_source=bundle.get("decision_source", "chosen"),
            model_version=bundle.get("model_version", "typed-readout-v1"),
            device=device,
        )

    def state(self, *, dim: int, encoder: str) -> dict[str, Any]:
        """Serialisable bundle for :meth:`from_checkpoint`."""
        return {
            "scorer": self.scorer.state_dict(),
            "bag": self.bag.state_dict() if self.bag is not None else None,
            "abstain": self.abstain.state_dict(),
            "feature_mean": self.feature_mean.cpu(),
            "feature_std": self.feature_std.cpu(),
            "threshold": self.threshold,
            "calibration": self.calibration,
            "hop_calibration": self.hop_calibration,
            "blend_weight": self.blend_weight,
            "template_bank": self.template_bank.half().cpu()
            if self.template_bank is not None
            else None,
            "gate": self.gate.state_dict() if self.gate is not None else None,
            "gate_mean": self.gate_mean.cpu() if self.gate is not None else None,
            "gate_std": self.gate_std.cpu() if self.gate is not None else None,
            "feature_mask": list(self.feature_mask) if self.feature_mask is not None else None,
            "decision_source": self.decision_source,
            "dim": dim,
            "max_hops": self.scorer.max_hops,
            "max_count": self.bag.max_count if self.bag is not None else MAX_COUNT,
            "num_features": int(self.abstain.net[0].in_features),
            "encoder": encoder,
            "model_version": self.model_version,
        }

    # -- scoring --------------------------------------------------------------

    def _type_embeddings(self, names: Sequence[str]) -> Tensor:
        missing = [name for name in names if name not in self._type_cache]
        if missing:
            vectors = self.encode([name.replace("_", " ") for name in missing])
            for name, vector in zip(missing, vectors, strict=True):
                self._type_cache[name] = vector.to(self.device).float()
        return torch.stack([self._type_cache[name] for name in names])

    def encode_query(self, query: GraphQuery) -> tuple[Tensor, Tensor | None, Tensor | None]:
        """Sentence vector, token vectors (bag) and masked-template vector (novelty)."""
        text = question_text(query)
        texts = [text, template_text(text)] if self.template_bank is not None else [text]
        vectors = self.encode(texts)
        tokens = self.encode_tokens([text])[0] if self.bag is not None else None
        masked = vectors[1] if self.template_bank is not None else None
        return vectors[0], tokens, masked

    @torch.no_grad()
    def dual(
        self,
        graph: GraphStore,
        query: GraphQuery,
        question: Tensor | None = None,
        tokens: Tensor | None = None,
        masked: Tensor | None = None,
    ) -> dict[str, Any]:
        """Per-endpoint best scores under each scorer, and the shared signals."""
        if question is None:
            question, tokens, masked = self.encode_query(query)
        question = question.to(self.device).float().reshape(1, -1)
        type_names = sorted({name for name in graph.nodes.values() if name})
        type_index = {name: i for i, name in enumerate(type_names)}
        type_embeddings = self._type_embeddings(type_names)
        hop_logp = self.scorer.hop_logp(question)[0]
        hop_table = hop_logp.tolist()
        use_bag = self.bag is not None and tokens is not None
        if use_bag:
            tokens = functional.normalize(tokens.to(self.device).float(), dim=-1)
            count_logp = self.bag.count_logp(tokens, type_embeddings).cpu()
            answer_logp = self.bag.answer_logp(question, type_embeddings)[0].cpu()
            # Every type contributes log P(count = 0) unless the path contains it, so
            # sum that once and adjust only the types on the path: O(hops) per path.
            zero_total = float(count_logp[:, 0].sum())
            count_table = count_logp.tolist()
        starts = set(query.start_entity_ids) or set(query.start_entities)
        per_hop: dict[int, list[Tensor]] = {}
        tables: dict[str, tuple[dict[str, float], dict[str, GraphPath]]] = {
            "sequence": ({}, {}), "bag": ({}, {}), "blend": ({}, {}),
        }

        def keep(name: str, end: str, value: float, path: GraphPath) -> None:
            best, paths = tables[name]
            if value > best.get(end, -math.inf):
                best[end] = value
                paths[end] = path

        max_hops = min(query.max_hops, self.scorer.max_hops)
        for count, path in enumerate(graph.iter_paths(sorted(starts), max_hops)):
            if count >= MAX_PATHS:
                break
            hops = len(path.steps)
            if not hops or hops < query.min_hops:
                continue
            end = path.steps[-1].target_id
            if end in starts:
                continue
            types = [graph.nodes[step.target_id] for step in path.steps]
            if hops not in per_hop:
                per_hop[hops] = [
                    functional.log_softmax(
                        self.scorer.type_logits(question, hops, i, type_embeddings)[0], -1
                    ).tolist()
                    for i in range(hops)
                ]
            sequence = hop_table[hops - 1] + sum(
                per_hop[hops][i][type_index[name]] for i, name in enumerate(types)
            )
            keep("sequence", end, sequence, path)
            if use_bag:
                bag_value = float(answer_logp[type_index[types[-1]]]) + zero_total
                for name, count in Counter(types).items():
                    row = count_table[type_index[name]]
                    bag_value += row[min(count, self.bag.max_count)] - row[0]
                keep("bag", end, bag_value, path)
                keep("blend", end, self.blend_weight * sequence + bag_value, path)
        novelty = 0.0
        if self.template_bank is not None and masked is not None:
            vector = functional.normalize(masked.to(self.device).float(), dim=-1)
            novelty = float((self.template_bank @ vector).max())
        return {"tables": tables, "hop_logp": hop_logp, "novelty": novelty, "use_bag": use_bag}

    @torch.no_grad()
    def gate_probability(self, dual: dict[str, Any]) -> float | None:
        """P(sequence scorer is right), or None without a gate or bag scorer."""
        if self.gate is None or not dual["use_bag"]:
            return None
        values = gate_features(
            dual["tables"]["sequence"][0], dual["tables"]["bag"][0], dual["hop_logp"], dual["novelty"]
        )
        if values is None:
            return None
        width = int(self.gate.net[0].in_features)
        features = torch.tensor(values[:width], device=self.device).reshape(1, -1)
        mean, std = self.gate_mean[:width], self.gate_std[:width]
        return float(torch.sigmoid(self.gate((features - mean) / std))[0])

    @torch.no_grad()
    def score(
        self,
        graph: GraphStore,
        query: GraphQuery,
        question: Tensor | None = None,
        tokens: Tensor | None = None,
        masked: Tensor | None = None,
    ) -> ReadoutScore:
        """Score every endpoint reachable from the query's start entities."""
        dual = self.dual(graph, query, question, tokens, masked)
        tables = dual["tables"]
        gate_p = self.gate_probability(dual)
        if gate_p is not None:
            chosen = "sequence" if gate_p >= 0.5 else "bag"
        elif dual["use_bag"]:
            chosen = "blend"
        else:
            chosen = "sequence"
        best, best_paths = tables[chosen]
        extra: tuple[float, ...] = ()
        if self.template_bank is not None:
            agree = (
                float(_top(tables["sequence"][0]) == _top(tables["bag"][0]))
                if dual["use_bag"]
                else 1.0
            )
            extra = (dual["novelty"], agree)
            if self.gate is not None:
                extra = (*extra, gate_p if gate_p is not None else 1.0)
        result = _finalise(best, best_paths, dual["hop_logp"], extra)
        if self.decision_source == "sequence" and chosen != "sequence":
            sequence_view = _finalise(*tables["sequence"], dual["hop_logp"], extra)
            result = replace(result, features=sequence_view.features)
        return result

    @torch.no_grad()
    def probability(self, score: ReadoutScore) -> float:
        """P(answering is correct), from the abstain head."""
        values = score.features
        if self.feature_mask is not None:
            values = tuple(values[i] for i in self.feature_mask)
        features = torch.tensor(values, device=self.device).reshape(1, -1)
        logit = self.abstain((features - self.feature_mean) / self.feature_std)
        return float(torch.sigmoid(logit)[0])

    def calibrate(self, confidence: float, hops: int | None = None) -> float:
        """Isotonic calibration, per predicted hop count when fitted for it."""
        table = self.hop_calibration.get(hops) if hops is not None else None
        table = table or self.calibration
        if not table:
            return confidence
        calibrator = IsotonicCalibrator(*table)
        return float(calibrator.transform(torch.tensor([confidence]))[0])

    # -- decision -------------------------------------------------------------

    def reason(self, graph: GraphStore, query: GraphQuery) -> ReasoningResult:
        score = self.score(graph, query)
        probability = self.probability(score)
        answer = bool(score.answers) and probability >= self.threshold
        # Confidence describes the decision made: for an abstention, the chance
        # that abstaining was right is approximated by 1 - P(answer correct).
        raw = probability if answer else 1.0 - probability
        confidence = self.calibrate(raw, score.predicted_hops)
        answers = (
            tuple(
                EntityAnswer(entity, confidence)
                for entity in sorted(score.answers, key=lambda e: -score.endpoint_scores[e])
            )
            if answer
            else ()
        )
        status = ReasoningStatus.SUPPORTED if answer else ReasoningStatus.INSUFFICIENT_EVIDENCE
        boolean_answer = None
        aggregate_value = None
        if query.answer_mode == "count":
            aggregate_value = len(answers)
        elif query.answer_mode == "boolean":
            targets = set(query.target_entity_ids)
            boolean_answer = bool(targets & {a.entity_id for a in answers}) if targets else answer
        return ReasoningResult(
            status=status,
            answers=answers,
            supporting_paths=tuple(score.best_paths[a.entity_id] for a in answers)
            if query.return_paths
            else (),
            conflicting_paths=(),
            confidence=confidence,
            calibrated=bool(self.calibration or self.hop_calibration),
            model_version=self.model_version,
            graph_version=graph.graph_version,
            reasoning_steps=max((len(p.steps) for p in score.best_paths.values()), default=0),
            metadata={"readout": "typed", "answer_probability": probability},
            boolean_answer=boolean_answer,
            aggregate_value=aggregate_value,
        )

    __call__ = reason


def gate_features(
    sequence: dict[str, float], bag: dict[str, float], hop_logp: Tensor, novelty: float
) -> tuple[float, ...] | None:
    """Gate inputs: how much each scorer dislikes the other's answer, hop certainty,
    template novelty, whether the two agree, and each scorer's best score.

    The best scores separate the two shifts: on new wording of a known intent the
    sequence scorer still finds a path that fits its predicted types (high best score);
    on a new combination its memorised type sequence rarely exists from the start
    entity (low best score), which is when the bag scorer should be trusted."""
    if not sequence or not bag:
        return None
    seq_top, bag_top = _top(sequence), _top(bag)
    seq_max, bag_max = max(sequence.values()), max(bag.values())
    bag_on_seq = max(bag[e] for e in seq_top) - bag_max
    seq_on_bag = max(sequence[e] for e in bag_top) - seq_max
    return (
        max(bag_on_seq, -50.0),
        max(seq_on_bag, -50.0),
        float(hop_logp.max().exp()),
        novelty,
        float(seq_top == bag_top),
        max(seq_max, -100.0),
        max(bag_max, -100.0),
    )


def _top(scores: dict[str, float]) -> frozenset[str]:
    if not scores:
        return frozenset()
    top = max(scores.values())
    return frozenset(e for e, s in scores.items() if s >= top - 1e-6)


def _finalise(
    best: dict[str, float],
    best_paths: dict[str, GraphPath],
    hop_logp: Tensor,
    extra: tuple[float, ...] = (),
) -> ReadoutScore:
    hops = int(hop_logp.argmax()) + 1
    if not best:
        return ReadoutScore(
            answers=frozenset(),
            features=(-30.0, 0.0, 0.0, float(hop_logp.max()), 0.0, *extra),
            best_paths={},
            endpoint_scores={},
            predicted_hops=hops,
        )
    ranked = sorted(best.values(), reverse=True)
    top = ranked[0]
    answers = frozenset(e for e, s in best.items() if s >= top - 1e-6)
    runner_up = next((s for s in ranked if s < top - 1e-6), top - 30.0)
    features = (
        top,
        top - runner_up,
        float(len(answers)),
        float(hop_logp.max()),
        math.log1p(len(best)),
        *extra,
    )
    return ReadoutScore(answers, features, best_paths, dict(best), hops)


def question_text(query: GraphQuery) -> str:
    return query.text or query.canonical_text() or ""


def template_text(text: str) -> str:
    """Question with entity identifiers masked, for template-novelty lookups."""
    return _ENTITY.sub("X", text).lower().strip()


def sentence_encoders(
    name: str, device: str | torch.device = "cpu"
) -> tuple[Encoder, TokenEncoder]:
    """Frozen sentence-transformer: pooled unit vectors and per-token vectors."""
    from sentence_transformers import SentenceTransformer

    model = SentenceTransformer(name, device=str(device))

    def encode(texts: Sequence[str]) -> Tensor:
        with torch.no_grad():
            return (
                model.encode(
                    list(texts),
                    normalize_embeddings=True,
                    convert_to_tensor=True,
                    batch_size=256,
                )
                .float()
                .clone()
            )

    def encode_tokens(texts: Sequence[str]) -> list[Tensor]:
        with torch.no_grad():
            rows = model.encode(list(texts), output_value="token_embeddings", batch_size=256)
        return [functional.normalize(row.float().clone(), dim=-1) for row in rows]

    return encode, encode_tokens


def sentence_encoder(name: str, device: str | torch.device = "cpu") -> Encoder:
    """Pooled encoder only (the v1 interface)."""
    return sentence_encoders(name, device)[0]