"""Multimodal prototype: typed questions about IMAGES (optionally with a text state).

    from jevlike.mm import MultimodalSystemOne
    m = MultimodalSystemOne.load("checkpoints/jevlike-mm-proto")
    m.predict(None, {"kind": Choice("What type of document is this?", ["invoice", "letter", "resume"]),
                     "signed": Noul("The page is handwritten")}, image="scan.png")

Design (research prototype, see reports/mm_proto.md):

* **Vision**: a frozen SigLIP (default google/siglip-base-patch16-224). Its 14x14 patch grid is
  average-pooled to 7x7 and SigLIP's pooled (attention-pooling head) embedding is prepended, giving
  ``n_img_tokens`` = 50 vectors. A trainable 2-layer MLP projector maps them into ModernBERT's
  token-embedding space (plus a learned per-slot position vector).
* **Input**: image tokens are spliced into the sequence right after [CLS] through
  ``inputs_embeds``, framed by learned [IMG_START] / [IMG_END] vectors::

      [CLS] [IMG_START] <50 image tokens> [IMG_END] <state text> [SEP] <type tag> <instr> [SEP] [MASK] key: text ...

  The question part, the option markers, the decision head, the marker scorer and the escalate
  head are the text model's (jevlike/model.py), unchanged.
* **Similarity prior** (``sim_prior``): SigLIP's own image-text score for every option (option text
  for choice/score, the statement for noul, scored at the ``true`` marker) is turned into a few
  features (raw logit, sigmoid, z-score / softmax / gap-to-max within the question, question type)
  and added, through a zero-initialized MLP, to that option's [MASK] embedding. This hands the
  language side SigLIP's zero-shot knowledge of label names it has never been trained on; the
  encoder still reads the instructions and options and decides. Training drops the prior for a
  fraction of questions (``sim_dropout``), so ``use_sim=False`` is a meaningful ablation.
* **Text-only inputs** go through exactly the text path (no image tokens, no prior), so the
  checkpoint directory is also a valid text ``SystemOne`` checkpoint: ``model.safetensors`` holds
  the DecisionModel, ``mm.safetensors`` + ``mm_config.json`` the multimodal parts; SigLIP weights
  come from ``mm_config.vision`` (HF id / path) or ``vision.safetensors`` if present.

Limitations: 224x224 input (small document text is unreadable), one image per question, long
states are truncated (no chunking) when an image is present.
"""
from __future__ import annotations

import json
import math
import os
import re
import time
from dataclasses import asdict, dataclass, field
from typing import Callable, Mapping, Sequence

import numpy as np
import torch
import torch.nn.functional as F
from safetensors.torch import load_file, save_file
from torch import nn

from jevlike.calibrate import temperature
from jevlike.model import DecisionModel
from jevlike.predict import SystemOne
from jevlike.serialize import TYPE_IDS, Encoded, Serializer, TokenBudgetBatchSampler, collate
from jevlike.types import QTYPES, PublicQuestion, Question, Record

MM_CONFIG, MM_WEIGHTS, VISION_WEIGHTS = "mm_config.json", "mm.safetensors", "vision.safetensors"
N_SIM_FEATS = 6 + len(QTYPES)
IMG_OFFSET = 1  # image block starts right after [CLS]

Questions = Mapping[str, "PublicQuestion | Question"]
ImageLike = "PIL.Image.Image | str | os.PathLike | np.ndarray | None"


@dataclass
class MMConfig:
    vision: str = "google/siglip-base-patch16-224"
    vision_config: dict = field(default_factory=dict)  # SiglipConfig.to_dict(); makes the architecture self-contained
    image_size: int = 224
    pool: int = 2              # average-pool the patch grid by this factor (14x14 -> 7x7)
    use_pooled: bool = True    # prepend SigLIP's pooled (MAP-head) embedding as one more token
    proj_hidden: int = 2048
    sim_prior: bool = True
    sim_dropout: float = 0.3   # training: fraction of image questions whose prior is withheld
    sim_template: str = "{}"
    text_max_len: int = 64     # SigLIP text length (SigLIP expects padding to 64)

    @property
    def grid(self) -> int:
        vc = self.vision_config.get("vision_config", {})
        return (vc.get("image_size", self.image_size) // vc.get("patch_size", 16)) // self.pool

    @property
    def n_img_tokens(self) -> int:
        return self.grid ** 2 + int(self.use_pooled)

    @property
    def n_slots(self) -> int:  # [IMG_START] + tokens + [IMG_END]
        return self.n_img_tokens + 2

    @classmethod
    def load(cls, path: str) -> "MMConfig":
        with open(os.path.join(path, MM_CONFIG)) as f:
            d = json.load(f)
        return cls(**{k: v for k, v in d.items() if k in cls.__dataclass_fields__})


# ---------------------------------------------------------------- images & option prompts

def load_pixels(image, size: int = 224) -> torch.Tensor:
    """PIL image / path / HxWx3 uint8 array -> [3, size, size] float tensor, SigLIP preprocessing
    (squash-resize bicubic, scale to [-1, 1])."""
    from PIL import Image
    if isinstance(image, (str, os.PathLike)):
        image = Image.open(image)
    elif isinstance(image, np.ndarray):
        image = Image.fromarray(image)
    image = image.convert("RGB")
    if image.size != (size, size):
        image = image.resize((size, size), Image.BICUBIC)
    x = torch.from_numpy(np.asarray(image, dtype=np.float32).copy()).permute(2, 0, 1) / 255.0
    return (x - 0.5) / 0.5


_OPAQUE = re.compile(r"^([A-Za-z]|\d+|(opt_|option_|class_|label_|c|cat_|choice_|id_)\d+|[A-Z]{2,4})$")


def option_prompts(q: Question) -> list[str | None]:
    """The text SigLIP scores for each marker: option text (readable key + text) for choice/score,
    the statement at the ``true`` marker for noul (``false`` gets none)."""
    if q.type == "noul":
        return [None, q.instructions.strip()]
    out = []
    for o in q.options:
        key, text = o.key.strip(), o.text.strip()
        if q.type == "score" or _OPAQUE.match(key):
            out.append(text or key)
        else:
            k = key.replace("_", " ")
            out.append(f"{k}: {text}" if text and text.lower() != k.lower() else k)
    return out


class MMSerializer:
    """Text serializer + image slots: image items get ``n_slots`` placeholder positions after [CLS]
    (their text is serialized with ``max_len - n_slots``); markers are shifted accordingly."""

    def __init__(self, tokenizer, max_len: int, head_max_len: int, n_slots: int):
        self.text = Serializer(tokenizer, max_len, head_max_len)
        img_len = max_len - n_slots
        self.img = Serializer(tokenizer, img_len, min(head_max_len, img_len - 64))
        self.n_slots, self.pad = n_slots, self.text.pad

    def encode_many(self, items: Sequence[tuple[str, Question, bool]]) -> list[Encoded]:
        out: list[Encoded | None] = [None] * len(items)
        for has_img, ser in ((False, self.text), (True, self.img)):
            idx = [i for i, it in enumerate(items) if bool(it[2]) == has_img]
            if not idx:
                continue
            for i, e in zip(idx, ser.encode_many([(items[i][0] or "", items[i][1]) for i in idx])):
                if has_img:
                    e = Encoded([e.input_ids[0]] + [self.pad] * self.n_slots + e.input_ids[1:],
                                [m + self.n_slots for m in e.markers], e.qtype, e.trunc)
                out[i] = e
        return out  # type: ignore[return-value]


def mm_collate(encs: Sequence[Encoded], labels, pad_id: int, image_of: Sequence[int], pixels: Sequence[torch.Tensor],
               prompts: Sequence[Sequence[str | None]] | None, use_sim: Sequence[bool] | None = None) -> dict:
    """``serialize.collate`` + image fields.

    image_of[i]: index into ``pixels`` for row i, or -1 (text-only row). prompts[i]: per-marker SigLIP
    prompt (or None); use_sim[i]: False withholds the prior for row i. Adds ``pixel_values`` [N, 3, S, S],
    ``image_rows`` / ``image_index`` [R], ``sim_strings`` (list, unique) and ``sim_marker_text`` [M]
    (index into sim_strings, -1 = none)."""
    b = collate(encs, labels, pad_id)
    used = sorted({j for j in image_of if j >= 0})
    remap = {j: k for k, j in enumerate(used)}
    rows = [i for i, j in enumerate(image_of) if j >= 0]
    b["image_rows"] = torch.tensor(rows, dtype=torch.long)
    b["image_index"] = torch.tensor([remap[image_of[i]] for i in rows], dtype=torch.long)
    b["pixel_values"] = torch.stack([pixels[j] for j in used]) if used else torch.zeros(0, 3, 1, 1)
    strings: dict[str, int] = {}
    mt: list[int] = []
    for i, e in enumerate(encs):
        ps = prompts[i] if prompts is not None and image_of[i] >= 0 and (use_sim is None or use_sim[i]) else None
        for r in range(len(e.markers)):
            s = ps[r] if ps is not None else None
            mt.append(strings.setdefault(s, len(strings)) if s else -1)
    b["sim_strings"] = list(strings)
    b["sim_marker_text"] = torch.tensor(mt, dtype=torch.long)
    return b


# ---------------------------------------------------------------- model

class MultimodalDecisionModel(nn.Module):
    def __init__(self, dm: DecisionModel, mcfg: MMConfig, siglip: nn.Module | None = None,
                 text_tokenize: Callable[[list[str]], torch.Tensor] | None = None):
        super().__init__()
        from transformers import SiglipConfig, SiglipModel
        if siglip is None:  # architecture only (weights loaded by the caller)
            siglip = SiglipModel(SiglipConfig.from_dict(mcfg.vision_config))
        self.dm, self.mcfg = dm, mcfg
        mcfg.vision_config = siglip.config.to_dict()
        self.siglip = siglip.eval().requires_grad_(False)
        self.text_tokenize = text_tokenize
        Dv, D = siglip.config.vision_config.hidden_size, dm.backbone.config.hidden_size
        n = mcfg.n_img_tokens
        self.proj = nn.Sequential(nn.LayerNorm(Dv), nn.Linear(Dv, mcfg.proj_hidden), nn.GELU(),
                                  nn.Linear(mcfg.proj_hidden, D))
        self.img_pos = nn.Parameter(torch.randn(n, D) * 0.02)
        self.img_start = nn.Parameter(torch.randn(D) * 0.02)
        self.img_end = nn.Parameter(torch.randn(D) * 0.02)
        self.sim_mlp = nn.Sequential(nn.Linear(N_SIM_FEATS, 256), nn.GELU(), nn.Linear(256, D))
        nn.init.zeros_(self.sim_mlp[-1].weight)  # starts as "no prior"
        nn.init.zeros_(self.sim_mlp[-1].bias)
        self._text_cache: dict[str, torch.Tensor] = {}

    # the frozen SigLIP never trains (and stays in eval mode)
    def train(self, mode: bool = True):
        super().train(mode)
        self.siglip.eval()
        return self

    def mm_parameters(self) -> list[tuple[str, nn.Parameter]]:
        return [(n, p) for n, p in self.named_parameters() if not n.startswith(("dm.", "siglip."))]

    @property
    def device(self) -> torch.device:
        return self.img_pos.device

    # -------------------------------------------------------- SigLIP side
    @torch.no_grad()
    def image_features(self, pixel_values: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        """-> (tokens [N, n_img_tokens, Dv], normalized image embedding [N, Dv])."""
        vm = self.siglip.vision_model
        dtype = next(vm.parameters()).dtype
        out = vm(pixel_values=pixel_values.to(self.device, dtype))
        h = out.last_hidden_state
        N, P, Dv = h.shape
        g = int(math.isqrt(P))
        grid = F.avg_pool2d(h.transpose(1, 2).reshape(N, Dv, g, g).float(), self.mcfg.pool)
        toks = grid.flatten(2).transpose(1, 2)
        pooled = out.pooler_output.float()
        if self.mcfg.use_pooled:
            toks = torch.cat([pooled[:, None], toks], 1)
        return toks, F.normalize(pooled, dim=-1)

    @torch.no_grad()
    def text_embeds(self, strings: Sequence[str]) -> torch.Tensor:
        """Normalized SigLIP text embeddings [U, Dv] (cached per string)."""
        if not strings:
            return torch.zeros(0, self.siglip.config.text_config.hidden_size, device=self.device)
        todo = [s for s in dict.fromkeys(strings) if s not in self._text_cache]
        for i in range(0, len(todo), 256):
            chunk = todo[i:i + 256]
            ids = self.text_tokenize([self.mcfg.sim_template.format(s) for s in chunk]).to(self.device)
            emb = F.normalize(self.siglip.text_model(input_ids=ids).pooler_output.float(), dim=-1)
            for s, e in zip(chunk, emb):
                self._text_cache[s] = e
        if len(self._text_cache) > 200_000:
            self._text_cache.clear()
        return torch.stack([self._text_cache[s] for s in strings])

    def sim_features(self, b: dict, img_emb: torch.Tensor, B: int, K: int) -> tuple[torch.Tensor, torch.Tensor]:
        """Per-marker prior features. Returns (marker positions into marker_idx, features [m, N_SIM_FEATS])."""
        mt = b["sim_marker_text"]
        sel = (mt >= 0).nonzero().flatten()
        if len(sel) == 0:
            return sel, torch.zeros(0, N_SIM_FEATS, device=self.device)
        txt = self.text_embeds(b["sim_strings"])
        row_img = torch.full((B,), -1, dtype=torch.long, device=self.device)
        row_img[b["image_rows"]] = b["image_index"]
        seg, rank = b["marker_seg"][sel], b["marker_rank"][sel]
        cos = (img_emb[row_img[seg]] * txt[mt[sel]]).sum(-1)
        sl = self.siglip
        logit = cos * sl.logit_scale.float().exp() + sl.logit_bias.float()
        # within-question statistics over the markers that have a prompt
        M = torch.full((B, K), float("nan"), device=self.device).index_put((seg, rank), logit)
        valid = ~M.isnan()
        cnt = valid.sum(-1).clamp(min=1)
        Mz = M.nan_to_num(0.0)
        mean = Mz.sum(-1) / cnt
        std = (((Mz - mean[:, None]) ** 2) * valid).sum(-1).div(cnt).sqrt().clamp(min=1e-3)
        mx = M.masked_fill(~valid, -1e9).max(-1).values
        soft = M.masked_fill(~valid, -1e9).softmax(-1)
        qt = b["qtype"][seg]
        multi = (cnt[seg] > 1).float()
        feats = torch.stack([
            logit / 10, torch.sigmoid(logit),
            multi * (logit - mean[seg]) / std[seg], multi * soft[seg, rank], multi * (logit - mx[seg]) / 10,
            torch.ones_like(logit)], -1)
        feats = torch.cat([feats, F.one_hot(qt, len(QTYPES)).float()], -1)
        return sel, feats

    # -------------------------------------------------------- forward
    def embed(self, b: dict) -> torch.Tensor:
        ids = b["input_ids"]
        emb = self.dm.backbone.get_input_embeddings()(ids)
        B, T, D = emb.shape
        rows = b.get("image_rows")
        if rows is None or len(rows) == 0:
            return emb
        toks, img_emb = self.image_features(b["pixel_values"])
        v = (self.proj(toks) + self.img_pos).to(emb.dtype)
        n = self.mcfg.n_img_tokens
        emb = emb.clone()
        emb[rows, IMG_OFFSET] = self.img_start.to(emb.dtype)
        emb[rows, IMG_OFFSET + 1:IMG_OFFSET + 1 + n] = v[b["image_index"]]
        emb[rows, IMG_OFFSET + 1 + n] = self.img_end.to(emb.dtype)
        if self.mcfg.sim_prior and "sim_marker_text" in b:
            K = int(b["n_options"].max())
            sel, feats = self.sim_features(b, img_emb, B, K)
            if len(sel):
                add = self.sim_mlp(feats).to(emb.dtype)
                emb = emb.reshape(B * T, D).index_add(0, b["marker_idx"][sel], add).reshape(B, T, D)
        return emb

    def forward(self, **b) -> dict[str, torch.Tensor]:
        rows = b.get("image_rows")
        if rows is None or len(rows) == 0:  # pure text batch: exactly the text model
            return self.dm(**b)
        return self.dm(**b, inputs_embeds=self.embed(b))

    # -------------------------------------------------------- io
    def save_pretrained(self, path: str, tokenizer=None, save_vision: bool = False) -> None:
        """model.safetensors / config.json = the DecisionModel (a valid text checkpoint);
        mm.safetensors / mm_config.json = projector, image slots, prior MLP."""
        self.dm.save_pretrained(path, tokenizer)
        with open(os.path.join(path, MM_CONFIG), "w") as f:
            json.dump(asdict(self.mcfg), f, indent=1, default=str)
        save_file({n: p.detach().contiguous().cpu() for n, p in self.mm_parameters()}, os.path.join(path, MM_WEIGHTS))
        if save_vision:
            save_file({k: v.detach().contiguous().cpu() for k, v in self.siglip.state_dict().items()},
                      os.path.join(path, VISION_WEIGHTS))

    @classmethod
    def from_pretrained(cls, path: str, device: str | torch.device = "cpu", vision: str | None = None,
                        text_tokenize=None) -> "MultimodalDecisionModel":
        device = torch.device(device)
        dm = DecisionModel.from_pretrained(path, device)
        mcfg = MMConfig.load(path)
        siglip, tt = load_siglip(path if os.path.exists(os.path.join(path, VISION_WEIGHTS)) else (vision or mcfg.vision),
                                 mcfg, device)
        model = cls(dm, mcfg, siglip, text_tokenize or tt)
        sd = load_file(os.path.join(path, MM_WEIGHTS))
        missing = model.load_state_dict(sd, strict=False)
        bad = [k for k in missing.missing_keys if not k.startswith(("dm.", "siglip."))]
        assert not bad and not missing.unexpected_keys, (bad, missing.unexpected_keys)
        return model.to(device)

    @classmethod
    def from_text_checkpoint(cls, path: str, mcfg: MMConfig | None = None, device="cpu") -> "MultimodalDecisionModel":
        """New multimodal model around a trained text checkpoint (fresh projector / prior MLP)."""
        device = torch.device(device)
        mcfg = mcfg or MMConfig()
        siglip, tt = load_siglip(mcfg.vision, mcfg, device)
        return cls(DecisionModel.from_pretrained(path, device), mcfg, siglip, tt).to(device)


def load_siglip(src: str, mcfg: MMConfig, device: torch.device):
    """SigLIP model (frozen; bf16 on CUDA) + a text tokenize function. ``src`` is an HF id / dir, or a
    checkpoint dir holding vision.safetensors (+ mm_config.json with the architecture)."""
    from transformers import AutoTokenizer, SiglipConfig, SiglipModel
    if os.path.exists(os.path.join(src, VISION_WEIGHTS)):
        model = SiglipModel(SiglipConfig.from_dict(mcfg.vision_config))
        model.load_state_dict(load_file(os.path.join(src, VISION_WEIGHTS)))
        tok_src = mcfg.vision
    else:
        model = SiglipModel.from_pretrained(src)
        tok_src = src
    if device.type == "cuda":
        model = model.to(torch.bfloat16)
    model = model.to(device).eval().requires_grad_(False)
    try:
        tok = AutoTokenizer.from_pretrained(tok_src)
    except Exception:  # noqa: BLE001 - tiny test configs have no tokenizer; caller passes text_tokenize
        tok = None

    def tokenize(strings: list[str]) -> torch.Tensor:
        return tok(strings, padding="max_length", truncation=True, max_length=mcfg.text_max_len,
                   return_tensors="pt")["input_ids"]
    return model, (tokenize if tok is not None else None)


# ---------------------------------------------------------------- data (training / eval)

def read_mm_jsonl(path: str, root: str = ".") -> list[tuple[Record, str | None, dict]]:
    """Record JSONL with an optional extra "image" (path, relative to ``root``) -> (record, image path, raw dict)."""
    out = []
    with open(path, encoding="utf-8") as f:
        for line in f:
            if line.strip():
                d = json.loads(line)
                img = d.get("image")
                out.append((Record.from_json(line), os.path.join(root, img) if img else None, d))
    return out


class MMDataset(torch.utils.data.Dataset):
    """Pre-encoded (record, image) items; images are decoded in ``collate`` (DataLoader workers)."""

    def __init__(self, items: Sequence[tuple[Record, str | None]], ser: MMSerializer, image_size: int = 224,
                 sim_dropout: float = 0.0, seed: int = 0):
        self.records = [r for r, _ in items]
        self.images = [img for _, img in items]
        self.encs = ser.encode_many([(r.state, r.question, img is not None) for r, img in items])
        self.prompts = [option_prompts(r.question) for r in self.records]
        self.labels = [r.label for r in self.records]
        self.lengths = np.array([len(e) for e in self.encs])
        self.pad_id, self.size, self.sim_dropout = ser.pad, image_size, sim_dropout
        self.rng = np.random.default_rng(seed)

    def __len__(self) -> int:
        return len(self.encs)

    def __getitem__(self, i: int) -> int:
        return i

    def collate(self, idx: Sequence[int]) -> dict:
        paths = list(dict.fromkeys(self.images[i] for i in idx if self.images[i]))
        pix = [load_pixels(p, self.size) for p in paths]
        where = {p: k for k, p in enumerate(paths)}
        image_of = [where[self.images[i]] if self.images[i] else -1 for i in idx]
        use_sim = (self.rng.random(len(idx)) >= self.sim_dropout).tolist() if self.sim_dropout else None
        b = mm_collate([self.encs[i] for i in idx], [self.labels[i] for i in idx], self.pad_id, image_of, pix,
                       [self.prompts[i] for i in idx], use_sim)
        b["index"] = torch.tensor(list(idx))
        return b


def to_device(b: dict, device) -> dict:
    out = {k: (v.to(device, non_blocking=True) if isinstance(v, torch.Tensor) else v) for k, v in b.items()}
    out["k_max"] = int(b["n_options"].max())
    return out


@torch.no_grad()
def collect(model: MultimodalDecisionModel, ds: MMDataset, device, max_tokens: int = 8192) -> dict:
    """Raw per-question outputs over a dataset (the train.collect contract, so calibrate.fit and
    train.summarize work on it)."""
    model.eval()
    sampler = TokenBudgetBatchSampler(ds.lengths, max_tokens=max_tokens, max_batch=128, shuffle=False)
    logits: list = [None] * len(ds)
    esc = torch.zeros(len(ds))
    for idx in sampler:
        b = to_device(ds.collate(idx), device)
        with torch.autocast(device.type, dtype=torch.bfloat16, enabled=device.type == "cuda"):
            out = model(**b)
        lg, e = out["logits"].float().cpu(), out["escalate_logit"].float().cpu()
        for j, i in enumerate(idx):
            logits[i] = lg[j, :len(ds.encs[i].markers)]
        esc[idx] = e
    return {"logits": logits, "escalate_logit": esc, "labels": torch.tensor(ds.labels),
            "qtype": torch.tensor([e.qtype for e in ds.encs]),
            "n_options": torch.tensor([len(e.markers) for e in ds.encs]),
            "source": [r.source for r in ds.records]}


# ---------------------------------------------------------------- public API

class MultimodalSystemOne(SystemOne):
    """``SystemOne`` plus images. ``predict(state_or_None, questions, image=...)`` returns the same
    answer dicts as ``SystemOne.predict``. Calls without any image take the text path unchanged
    (including ``long_state="chunk"``)."""

    def __init__(self, model: MultimodalDecisionModel, tokenizer, calibration: dict | None, device: torch.device,
                 **kw):
        super().__init__(model.dm, tokenizer, calibration, device, **kw)
        self.mm = model.eval()
        cfg = model.dm.cfg
        self.mser = MMSerializer(tokenizer, cfg.max_len, cfg.head_max_len, model.mcfg.n_slots)

    @classmethod
    def load(cls, path: str, device: str | torch.device | None = None, vision: str | None = None,
             **kw) -> "MultimodalSystemOne":
        from transformers import AutoTokenizer
        device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
        cal_path = os.path.join(path, "calibration.json")
        cal = json.load(open(cal_path)) if os.path.exists(cal_path) else None
        model = MultimodalDecisionModel.from_pretrained(path, device, vision)
        return cls(model, AutoTokenizer.from_pretrained(path), cal, device, **kw)

    def predict(self, state: str | None, questions: Questions, image: ImageLike = None,  # type: ignore[override]
                calibrated: bool = True, use_sim: bool = True, **kw) -> dict:
        return self.predict_batch([(state, questions, image)], calibrated=calibrated, use_sim=use_sim, **kw)[0]

    @torch.inference_mode()
    def predict_batch(self, items: Sequence[tuple], calibrated: bool = True, use_sim: bool = True,  # type: ignore[override]
                      **kw) -> list[dict]:
        """items: (state, questions) or (state, questions, image); state may be None/"" with an image."""
        items = [tuple(it) + (None,) * (3 - len(it)) for it in items]
        if all(img is None for _, _, img in items):
            return super().predict_batch([(s or "", qs) for s, qs, _ in items], calibrated=calibrated, **kw)
        t0 = time.perf_counter()
        pixels, image_of_item = [], []
        for _, _, img in items:
            if img is None:
                image_of_item.append(-1)
            else:
                image_of_item.append(len(pixels))
                pixels.append(img if isinstance(img, torch.Tensor) else load_pixels(img, self.mm.mcfg.image_size))
        flat: list[tuple[int, str, Question]] = []
        for i, (_, qs, _) in enumerate(items):
            for name, q in qs.items():
                q = q if isinstance(q, Question) else Question.from_public(q)
                try:
                    q.validate()
                except AssertionError as e:
                    raise ValueError(f"question {name!r}: {e}") from None
                flat.append((i, name, q))
        encs = self.mser.encode_many([(items[i][0] or "", q, image_of_item[i] >= 0) for i, _, q in flat])
        prompts = [option_prompts(q) for _, _, q in flat]
        logits: list[torch.Tensor] = [None] * len(encs)  # type: ignore[list-item]
        esc = torch.zeros(len(encs))
        for idx in self.batches([len(e) for e in encs]) if encs else []:
            b = mm_collate([encs[j] for j in idx], None, self.ser.pad, [image_of_item[flat[j][0]] for j in idx],
                           pixels, [prompts[j] for j in idx], [use_sim] * len(idx))
            b = to_device(b, self.device)
            with torch.autocast(self.device.type, dtype=torch.bfloat16, enabled=self.device.type == "cuda"):
                out = self.mm(**b)
            lg, e = out["logits"].float().cpu(), out["escalate_logit"].float().cpu()
            for r, j in enumerate(idx):
                logits[j] = lg[r, :len(encs[j].markers)]
            esc[idx] = e
        cal = self.cal if calibrated else None
        if cal:
            esc = cal["escalate"]["a"] * esc + cal["escalate"]["b"]
        p_wrong = 1 - torch.sigmoid(esc)
        results = [{"answers": {}, "model": self.name, "latency_ms": 0.0} for _ in items]
        for j, (i, name, q) in enumerate(flat):
            p = (logits[j] / temperature(cal, q.type, q.n_options)).softmax(-1)
            results[i]["answers"][name] = self._answer(q, p, float(p_wrong[j]))
        ms = (time.perf_counter() - t0) * 1000
        for r in results:
            r["latency_ms"] = ms
        return results