"""Cut from amos/decoder.py: only the definitions amos.serve reaches (see sources.json)."""
import argparse
import glob
import hashlib
import itertools
import json
import math
import os
import random
import shutil
import string
import time
from pathlib import Path
import torch

ANSWER = "<answer>"

IMAGE_PAD = "<|image_pad|>"

LETTERS = string.ascii_uppercase + string.ascii_lowercase  # the first 52 option codes; runs before amos_codes use these alone

KINDS = ("noul", "choice", "score")  # answered by option letters

HEADED = ("span", "region")  # answered by Heads

SPAN_LENGTH_PENALTY = 1.6

SPAN_CAP = 30

HINTS = {"score": " (levels from lowest to highest)", "span": " (quote the case)", "region": " (point to it in the image)"}

IMAGE_PIXELS = int(os.environ.get("AMOS_IMAGE_PIXELS", 1 << 20))

HEADS = "heads.pt"

REGION_KEEP = 0.5  # a point's box: cells joined to the peak holding this share of its probability (vision.peak_box)

SYSTEM = ("You make decisions about a case. Read the case, then answer every question with the letter of the option "
          "that is most likely right.")

def options(q: dict) -> list[str]:
    """Option lines in answer order; yes/no is yes first, so its target is [P(yes), P(no)]."""
    if q["type"] == "noul":
        c = q.get("criteria") or {}
        return [f"yes: {c['true']}" if c.get("true") else "yes", f"no: {c['false']}" if c.get("false") else "no"]
    return list(q["options"])

def render(state: str, questions: list[dict], images=(), codes=LETTERS) -> list[dict]:
    """Chat messages for a case. Images (the state's after "Case:", image options where they are listed, a pick-any
    option's image after its question) make the user message a list of text and image items; without any it is plain
    text."""
    parts = ["Case:\n", *images, f"{state}\n\nQuestions:\n"]
    for i, q in enumerate(questions, start=1):
        parts += ["\n" if i > 1 else "", f"q{i}. {q['text']}" + HINTS.get(q["type"], "")]
        if "subject" in q:
            parts += [" ", q["subject"]]
        for j, o in enumerate(options(q) if q["type"] in KINDS else ()):
            parts += [f"\n{codes[j]}) ", o]
    merged = []
    for x in parts:
        if isinstance(x, str) and merged and isinstance(merged[-1], str):
            merged[-1] += x
        else:
            merged.append(x)
    user = merged[0] if len(merged) == 1 else [{"type": "text", "text": x} if isinstance(x, str) else {"type": "image", "image": x}
                                               for x in merged if x != ""]
    answers = "\n".join(f"q{i}: {ANSWER}" for i in range(1, len(questions) + 1))
    reply = {"role": "assistant", "content": answers}
    if len(questions) == 1 and questions[0].get("reasoning"):  # a trace goes in the thinking block, as think() reads
        reply["reasoning_content"] = questions[0]["reasoning"]
    return [{"role": "system", "content": SYSTEM}, {"role": "user", "content": user}, reply]

def expand(questions: list[dict]) -> tuple[list[dict], list[int]]:
    """Pick-any questions become one yes/no per option ("subject") in the same prompt, each with its own answer ->
    (questions, index of each one's original question)."""
    flat, owner = [], []
    for k, q in enumerate(questions):
        if q["type"] != "multi":
            flat.append(q)
            owner.append(k)
            continue
        for j, o in enumerate(q["options"]):
            flat.append({"type": "noul", "text": q["text"] + " Does this option apply?", "subject": o,
                         **({"target": q["target"][j]} if "target" in q else {})})
            owner.append(k)
    return flat, owner

PLAIN_TEMPLATE = ("{% for m in messages %}{{ m['role'] }}:\n{% if m['content'] is string %}{{ m['content'] }}{% else %}"
                  "{% for x in m['content'] %}{{ x['text'] if x['type'] == 'text' else '' }}{% endfor %}{% endif %}"
                  "\n\n{% endfor %}")

def option_codes(tokenizer, n: int) -> list[str]:
    """The first n option codes: the 52 letters, then two-letter codes the tokenizer keeps as one token."""
    two = ("".join(p) for p in itertools.product(string.ascii_uppercase, repeat=2))
    codes = list(LETTERS)[:n] + [c for c in two if len(tokenizer.encode(c, add_special_tokens=False)) == 1][:max(0, n - 52)]
    if len(codes) != n:
        raise ValueError(f"the tokenizer has only {len(codes)} one-token option codes, not {n}")
    return codes

class Codec:
    def __init__(self, tokenizer, image_processor=None, n_codes: int = len(LETTERS)):
        self.tok, self.image_processor = tokenizer, image_processor
        if getattr(tokenizer, "chat_template", "") is None:
            tokenizer.chat_template = PLAIN_TEMPLATE
        if ANSWER not in tokenizer.get_added_vocab():
            tokenizer.add_special_tokens({"additional_special_tokens": [ANSWER]})
        self.marker = tokenizer.convert_tokens_to_ids(ANSWER)
        self.codes = option_codes(tokenizer, n_codes)
        ids = [tokenizer.encode(x, add_special_tokens=False) for x in self.codes]
        if any(len(i) != 1 for i in ids):
            raise ValueError("every option letter must be one token")
        self.letters = [i[0] for i in ids]
        self.pad = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id
        # A unified model (Gemma 4) comes with its whole processor, which turns each image placeholder into its own run
        # of tokens; Qwen comes with an image processor alone, and the placeholder is expanded here
        self.unified = image_processor is not None and hasattr(image_processor, "image_token")
        pad = image_processor.image_token if self.unified else IMAGE_PAD
        self.image_pad = tokenizer.convert_tokens_to_ids(pad) if image_processor is not None else None
        self.echo = False  # set by load_heads when the quote head was trained on an echoed case (see encode_case)
        # the tokenizer's added tokens (the answer marker, chat roles, <think>, image placeholders) as user text would
        # tokenize as the tokens themselves and break the prompt; defuse() swaps their "<" for a look-alike
        self.specials = sorted((t for t in tokenizer.get_added_vocab() if t.startswith("<")), key=len, reverse=True)

    def defuse(self, text):
        """User text with any added-token string made ordinary text; same length, so character offsets still hold."""
        if not isinstance(text, str) or "<" not in text:
            return text
        for t in self.specials:
            text = text.replace(t, "‹" + t[1:])
        return text

    def safe(self, state: str, questions: list[dict]) -> tuple[str, list[dict]]:
        """state and questions with every user-written string defused (images and other values pass through)."""
        d = self.defuse
        qs = [{**q, **{k: d(q[k]) for k in ("text", "subject") if k in q},
               **({"options": [d(o) for o in q["options"]]} if "options" in q else {}),
               **({"criteria": {k: d(v) for k, v in q["criteria"].items()}} if isinstance(q.get("criteria"), dict) else {})}
              for q in questions]
        return d(state), qs

    def encode(self, state: str, questions: list[dict]) -> tuple[list[int], list[int]]:
        """-> (token ids, the position whose next-token logits answer each question)."""
        state, questions = self.safe(state, questions)
        text = self.tok.apply_chat_template(render(state, questions, codes=self.codes), tokenize=False,
                                            enable_thinking=False)
        ids = self.tok(text, add_special_tokens=False)["input_ids"]
        at = [i - 1 for i, t in enumerate(ids) if t == self.marker]
        if len(at) != len(questions):
            raise ValueError("one answer marker per question")
        return ids, at

    def reason_span(self, ids: list[int]) -> tuple[int, int]:
        """The tokens inside the thinking block, (first, end): the trace a reasoning record teaches the model to write."""
        start, end = (self.tok.convert_tokens_to_ids(t) for t in ("<think>", "</think>"))
        return ids.index(start) + 1, ids.index(end)

    def encode_case(self, state: str, questions: list[dict], images=()) -> tuple[list[int], list[int], dict]:
        """-> (token ids, answer positions, extra); without images the ids are exactly encode()'s. extra["state"] is
        (first state token, each state token's char offsets in the state). With images extra also holds the
        image processor's pixel_values and image_grid_thw, and "images": (first token, rows, cols) of each image's
        merged token grid; each image placeholder is expanded to its token count, as the model's processor does."""
        state, questions = self.safe(state, questions)
        messages = render(state, questions, images, self.codes)
        user = messages[1]["content"]
        pics = [] if isinstance(user, str) else [x["image"] for x in user if x["type"] == "image"]
        text = self.tok.apply_chat_template(messages, tokenize=False, enable_thinking=False)
        enc = self.tok(text, add_special_tokens=False, return_offsets_mapping=True)
        ids, spans = enc["input_ids"], enc["offset_mapping"]
        # an empty state has no text to find (and Gemma's template trims the newlines that would surround it)
        s0 = text.index(state + "\n\nQuestions:\n") if state else 0
        s1 = s0 + len(state)
        toks = [i for i, (a, b) in enumerate(spans) if a < s1 and b > s0]
        offsets = [(max(spans[i][0], s0) - s0, min(spans[i][1], s1) - s0) for i in toks]
        where, extra = list(range(len(ids))), {}
        if pics and self.unified:
            ids, where, extra = self._unified_images(text, ids, pics)
        elif pics:
            if self.image_processor is None:
                raise ValueError("this model has no vision tower")
            ip = self.image_processor
            v = ip(images=[p.convert("RGB") for p in pics], max_pixels=IMAGE_PIXELS, return_tensors="pt")
            grid, m = v["image_grid_thw"], ip.merge_size
            counts, out, where, starts = iter((grid.prod(-1) // m**2).tolist()), [], [], []
            for t in ids:
                where.append(len(out))
                if t == self.image_pad:
                    starts.append(len(out))
                    out += [t] * next(counts)
                else:
                    out.append(t)
            if len(starts) != len(pics):
                raise ValueError("one image placeholder per image")
            ids = out
            extra = {"pixel_values": v["pixel_values"], "image_grid_thw": grid,
                     "images": [(s, int(g[1]) // m, int(g[2]) // m) for s, g in zip(starts, grid)]}
        extra["state"] = (where[toks[0]] if toks else 0, offsets)
        at = [i - 1 for i, t in enumerate(ids) if t == self.marker]
        if len(at) != len(questions):
            raise ValueError("one answer marker per question")
        if self.echo and toks and any(q["type"] == "span" for q in questions):
            # Echo: the case again after every answer marker, and quotes are read from this copy. Reading left to
            # right, a case token cannot see what follows it ("Total due" after a line item), but every copied token
            # has seen the whole case and each question. Attention only looks back, so no earlier answer changes.
            sep = self.tok("\n\nThe case again:\n", add_special_tokens=False)["input_ids"]
            copy = self.tok(state, add_special_tokens=False, return_offsets_mapping=True)
            extra["state"] = (len(ids) + len(sep), [tuple(o) for o in copy["offset_mapping"]])
            ids = ids + sep + copy["input_ids"]
        return ids, at, extra

    def _unified_images(self, text: str, ids: list[int], pics) -> tuple[list[int], list[int], dict]:
        """A unified model's processor expands each image placeholder into <|image>, its image tokens, <image|>; ->
        (expanded ids, each original token's new position, the model's image inputs). The images have no separate
        grid of merged tokens here, so "images" is empty and the point and box heads do not apply."""
        ip = self.image_processor
        v = ip(text=text, images=[p.convert("RGB") for p in pics], return_tensors="pt", add_special_tokens=False)
        new = v["input_ids"][0].tolist()
        run = {self.image_pad, self.tok.convert_tokens_to_ids(ip.boi_token), self.tok.convert_tokens_to_ids(ip.eoi_token)}
        where, j = [], 0
        for t in ids:
            where.append(j)
            if t == self.image_pad:
                while j < len(new) and new[j] in run:
                    j += 1
            elif j < len(new) and new[j] == t:
                j += 1
            else:
                raise ValueError("the processor changed a text token")
        if j != len(new):
            raise ValueError("the processor added tokens after the text")
        return new, where, {"pixel_values": v["pixel_values"], "image_position_ids": v["image_position_ids"],
                            "mm_token_type_ids": v["mm_token_type_ids"][0], "images": []}

def load_model(path: str, device: str, *, dtype=None, vision=None):
    """-> (model, codec). vision=None loads the image-text model when the checkpoint has a vision tower and an image
    processor config beside it, else the text model alone. Qwen's image processor runs on PIL, never torchvision
    (not installed, and the Auto classes demand it), so training and serving resize alike on every machine."""
    from transformers import (
        AutoConfig,
        AutoModelForCausalLM,
        AutoModelForImageTextToText,
        AutoTokenizer,
        Qwen2VLImageProcessorPil,
    )

    if device == "cuda":  # as in amos.train: cuDNN attention fails to load in bf16 with the pip CUDA wheels
        torch.backends.cuda.enable_cudnn_sdp(False)

    tok = AutoTokenizer.from_pretrained(path)
    config = AutoConfig.from_pretrained(path)
    unified = config.model_type.startswith("gemma4")  # images through its own processor (which needs torchvision)
    if vision is None:
        vision = hasattr(config, "vision_config") and (Path(path) / ("processor_config.json" if unified else
                                                                       "preprocessor_config.json")).exists()
    processor = None
    if vision:
        from transformers import AutoProcessor

        processor = AutoProcessor.from_pretrained(path) if unified else Qwen2VLImageProcessorPil.from_pretrained(path)
    codec = Codec(tok, processor, getattr(config, "amos_codes", len(LETTERS)))
    if dtype is None:
        dtype = torch.float32 if getattr(config, "amos_export_dtype", None) == "float32" else torch.bfloat16
    from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES

    # some image-text models (Muse Glimmer) have no causal-LM class: load the whole model, its vision tower unused
    causal = config.model_type in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
    model = (AutoModelForImageTextToText if vision or not causal else AutoModelForCausalLM).from_pretrained(
        path, config=config, dtype=dtype)
    if vision and not unified:
        patchify_as_matmul(model.model.visual.patch_embed)
    if model.get_input_embeddings().weight.shape[0] < len(tok):
        model.resize_token_embeddings(len(tok), mean_resizing=True)
    return model.to(device), codec

def patchify_as_matmul(patch_embed):
    """Qwen's patch embedding is a conv3d whose kernel is the whole flattened patch, so it is one matmul with the same
    weights; run it as one. cuDNN's conv3d fails to load on some hosts and PyTorch's own takes ~5 s per photo."""
    proj = patch_embed.proj
    patch_embed.forward = lambda x: torch.nn.functional.linear(x.to(proj.weight.dtype), proj.weight.flatten(1), proj.bias)

def noncausal(model) -> bool:
    """Whether the run was trained with bidirectional full attention (train --noncausal), so it must be read that way."""
    return bool(getattr(model.config, "amos_noncausal", False))

def bidirectional_mask(lengths: list[int], width: int, device) -> torch.Tensor:
    """[B, 1, L, L], True where the key is one of its row's real tokens: every query sees the whole case, never pads."""
    keys = torch.arange(width, device=device) < torch.tensor(lengths, device=device)[:, None]
    return keys[:, None, None, :].expand(-1, 1, width, -1)

PREFILL_CHUNK = 16384

def chunked_hidden(decoder, x: torch.Tensor) -> torch.Tensor:
    """Last hidden states [1, L, d] for one causal text row, read PREFILL_CHUNK tokens at a time."""
    from transformers import DynamicCache

    cache, out = DynamicCache(config=decoder.config), []
    for s in range(0, x.shape[1], PREFILL_CHUNK):
        out.append(decoder(input_ids=x[:, s : s + PREFILL_CHUNK], past_key_values=cache, use_cache=True).last_hidden_state)
    del cache
    return torch.cat(out, 1)

def layerwise_hidden(decoder, x: torch.Tensor) -> torch.Tensor:
    """Last hidden states [1, L, d] for one noncausal text row, read a layer at a time.

    A noncausal run's full-attention layers let every token see the whole case, so it cannot be read in chunks the way
    a causal one can. Instead the whole case goes through each layer before the next, PREFILL_CHUNK tokens at a time:
    a linear-attention layer is a forward recurrence, so its chunks carry its conv and recurrent state; a full-attention
    layer first takes keys and values for every token, then each chunk of queries attends to all of them. The same
    computation as one pass, holding one layer's activations for one chunk instead of every layer's for the whole case."""
    import torch.nn.functional as F
    from transformers import DynamicCache
    from transformers.models.qwen3_5.modeling_qwen3_5 import apply_rotary_pos_emb

    n, C = x.shape[1], PREFILL_CHUNK
    h = decoder.embed_tokens(x)
    pos = torch.arange(n, device=x.device).view(1, 1, -1).expand(4, 1, -1)
    cos, sin = decoder.rotary_emb(h, pos[1:])
    text_pos, cache = pos[0], DynamicCache(config=decoder.config)
    for i, layer in enumerate(decoder.layers[: decoder.config.num_hidden_layers]):
        if decoder.config.layer_types[i] == "linear_attention":
            for s in range(0, n, C):  # each chunk is overwritten after it is read; later chunks are still this layer's input
                h[:, s : s + C] = layer(h[:, s : s + C], position_embeddings=(cos[:, s : s + C], sin[:, s : s + C]),
                                        attention_mask=None, position_ids=text_pos[:, s : s + C], past_key_values=cache)
            continue
        at = layer.self_attn
        keys, values = [], []
        for s in range(0, n, C):
            hn = layer.input_layernorm(h[:, s : s + C])
            shape = (*hn.shape[:-1], -1, at.head_dim)
            k = at.k_norm(at.k_proj(hn).view(shape)).transpose(1, 2)
            _, k = apply_rotary_pos_emb(k, k, cos[:, s : s + C], sin[:, s : s + C])
            keys.append(k)
            values.append(at.v_proj(hn).view(shape).transpose(1, 2))
        keys, values = torch.cat(keys, 2), torch.cat(values, 2)
        for s in range(0, n, C):
            res = h[:, s : s + C]
            hn = layer.input_layernorm(res)
            shape = hn.shape[:-1]
            q, gate = torch.chunk(at.q_proj(hn).view(*shape, -1, at.head_dim * 2), 2, dim=-1)
            q = at.q_norm(q).transpose(1, 2)
            q, _ = apply_rotary_pos_emb(q, q, cos[:, s : s + C], sin[:, s : s + C])
            o = F.scaled_dot_product_attention(q, keys, values, scale=at.scaling, enable_gqa=True)
            o = at.o_proj(o.transpose(1, 2).reshape(*shape, -1) * torch.sigmoid(gate.reshape(*shape, -1)))
            mid = res + o
            h[:, s : s + C] = mid + layer.mlp(layer.post_attention_layernorm(mid))
        del keys, values
    del cache
    for s in range(0, n, C):
        h[:, s : s + C] = decoder.norm(h[:, s : s + C])
    return h

def hidden(model, codec, batch: list[tuple]) -> torch.Tensor:
    """Last hidden states [B, L, d] for a batch of (ids, answer positions, questions[, extra]).

    Sequences are right-padded without a mask: the model is causal (its linear-attention layers are recurrences that
    only run forward), so padding after a sequence cannot change anything read from inside it. A noncausal run's
    full-attention layers see the whole sequence, so they get a mask hiding the pads; the recurrences stay causal.
    Entries with images go through the image-text model, which fills image tokens from the vision tower and gives
    them 2-D positions."""
    width = max(len(e[0]) for e in batch)
    x = torch.full((len(batch), width), codec.pad, dtype=torch.long, device=model.device)
    for b, e in enumerate(batch):
        x[b, : len(e[0])] = torch.tensor(e[0], device=model.device)
    vis = [e[3] for e in batch if len(e) > 3 and "pixel_values" in e[3]]
    kw = {} if not noncausal(model) else {"attention_mask": {
        "full_attention": bidirectional_mask([len(e[0]) for e in batch], width, x.device), "linear_attention": None}}
    if not vis and len(batch) == 1 and width > PREFILL_CHUNK:
        return layerwise_hidden(model.get_decoder(), x) if kw else chunked_hidden(model.get_decoder(), x)
    if not vis:
        return model.get_decoder()(input_ids=x, **kw).last_hidden_state
    if codec.unified and kw:
        raise ValueError("noncausal attention is built for Qwen3.5's layer types, not Gemma 4's")
    if codec.unified:  # Gemma 4: image patches go through the same layers, placed by their token types and positions
        types = torch.zeros_like(x)
        for b, e in enumerate(batch):
            if len(e) > 3 and "mm_token_type_ids" in e[3]:
                types[b, : len(e[3]["mm_token_type_ids"])] = e[3]["mm_token_type_ids"].to(x.device)
        return model.model(input_ids=x, mm_token_type_ids=types,
                           pixel_values=torch.cat([v["pixel_values"] for v in vis]).to(model.device, model.dtype),
                           image_position_ids=torch.cat([v["image_position_ids"] for v in vis]).to(model.device)
                           ).last_hidden_state
    grid, types = torch.cat([v["image_grid_thw"] for v in vis]).to(model.device), (x == codec.image_pad).int()
    if kw:  # its rope index would read a dict mask as per-row padding, so the positions come unmasked, as causal
        kw["position_ids"] = model.model.get_rope_index(x, types, image_grid_thw=grid)[0]
    return model.model(input_ids=x, pixel_values=torch.cat([v["pixel_values"] for v in vis]).to(model.device),
                       image_grid_thw=grid, mm_token_type_ids=types, **kw).last_hidden_state

def finish_logits(model, logits: torch.Tensor) -> torch.Tensor:
    """The model's own last step from raw output-layer logits, as its forward would apply it: Granite divides by
    logits_scaling, Gemma caps with final_logit_softcapping. Qwen has neither, so its logits pass unchanged."""
    cfg = getattr(model.config, "text_config", None) or model.config
    if scaling := getattr(cfg, "logits_scaling", None):
        logits = logits / scaling
    if cap := getattr(cfg, "final_logit_softcapping", None):
        logits = cap * torch.tanh(logits / cap)
    return logits

def distributions(model, codec, batch: list[tuple], h: torch.Tensor | None = None) -> list[torch.Tensor]:
    """Log-probabilities over the option letters of each yes/no, choice and score question in a batch of (ids, answer
    positions, questions[, extra]), in order. Only the 52 letter logits are computed, from the tied output
    embeddings."""
    h = hidden(model, codec, batch) if h is None else h
    picks = [(b, p, q) for b, e in enumerate(batch) for p, q in zip(e[1], e[2]) if q["type"] in KINDS]
    if not picks:
        return []
    rows = torch.tensor([b for b, _, _ in picks], device=h.device)
    cols = torch.tensor([p for _, p, _ in picks], device=h.device)
    logits = finish_logits(model, h[rows, cols].float() @ model.get_output_embeddings().weight[codec.letters].float().T)
    return [logits[k, : len(options(q))].log_softmax(-1) for k, (_, _, q) in enumerate(picks)]

class BoxHead(torch.nn.Module):
    """A point's box, predicted rather than grown from cells: reads the question's vector, the image features pooled
    by the point head's own cell distribution, and that distribution's centre and spread, and moves the centre and
    sets the size. Cells alone draw a box around only the few cells the point head is sure of (a car's tail light
    rather than the car); this head is trained on overlap with the whole object."""

    def __init__(self, d: int, width: int = 512):
        super().__init__()
        self.net = torch.nn.Sequential(torch.nn.Linear(2 * d + 4, width), torch.nn.GELU(), torch.nn.Linear(width, 4))
        torch.nn.init.zeros_(self.net[-1].weight)  # starts as a half-image box on the cells' centre
        torch.nn.init.zeros_(self.net[-1].bias)

    def forward(self, hq: torch.Tensor, cells: torch.Tensor, logits: torch.Tensor, rows: int, cols: int) -> torch.Tensor:
        p = logits.detach().float().softmax(-1)  # the point head stays as it is; only the box learns
        ys, xs = torch.meshgrid((torch.arange(rows, device=p.device) + 0.5) / rows,
                                (torch.arange(cols, device=p.device) + 0.5) / cols, indexing="ij")
        xs, ys = xs.flatten(), ys.flatten()
        mx, my = (p * xs).sum(), (p * ys).sum()
        sx, sy = (p * (xs - mx) ** 2).sum().sqrt(), (p * (ys - my) ** 2).sum().sqrt()
        pooled = (p[:, None] * cells.float()).sum(0)
        o = self.net(torch.cat([hq.float(), pooled, torch.stack([mx, my, sx, sy])]).to(self.net[0].weight.dtype)).float()
        cx, cy = (mx + 0.5 * torch.tanh(o[0])).clamp(0, 1), (my + 0.5 * torch.tanh(o[1])).clamp(0, 1)
        w, hgt = torch.sigmoid(o[2]), torch.sigmoid(o[3])
        return torch.stack([cx - w / 2, cy - hgt / 2, cx + w / 2, cy + hgt / 2]).clamp(0, 1)

class Heads(torch.nn.Module):
    """Quote and point heads, as amos.model has them: a question's answer-position vector scored against the state's
    tokens (start and end of a quote) or against the first state image's tokens (one cell each); `keep` is the
    region readout's share (decode_head). With a BoxHead, a region's logits carry four more values: its box."""

    def __init__(self, d: int, kinds=HEADED, keep: float = REGION_KEEP, box: bool = False):
        super().__init__()
        self.kinds, self.keep = list(kinds), keep
        self.span = (SPAN_CAP, SPAN_LENGTH_PENALTY)  # the quote decoder's cap and penalty; serving sets the run's own
        self.span_q, self.span_k = torch.nn.Linear(d, 2 * d), torch.nn.Linear(d, 2 * d)
        self.region_q, self.region_k = torch.nn.Linear(d, d), torch.nn.Linear(d, d)
        for m in (self.span_q, self.span_k, self.region_q, self.region_k):
            torch.nn.init.normal_(m.weight, std=0.02)
            torch.nn.init.zeros_(m.bias)
        self.box = BoxHead(d) if box else None

    def forward(self, h: torch.Tensor, batch: list[tuple]) -> list[torch.Tensor]:
        """Logits of each span ([2, state tokens]) and region ([cells], then the 4 box edges when there is a box head)
        question in the batch, in order."""
        h, d, out = h.to(self.span_q.weight.dtype), h.shape[-1], []
        for b, e in enumerate(batch):
            for p, q in zip(e[1], e[2]):
                if q["type"] == "span":
                    first, offsets = e[3]["state"]
                    k = self.span_k(h[b, first : first + len(offsets)]).view(-1, 2, d)
                    out.append(torch.einsum("cd,lcd->cl", self.span_q(h[b, p]).view(2, d), k) / math.sqrt(d))
                elif q["type"] == "region":
                    if not e[3].get("images"):
                        raise ValueError("a region question needs a state image")
                    start, rows, cols = e[3]["images"][0]
                    cells = h[b, start : start + rows * cols]
                    x = self.region_k(cells) @ self.region_q(h[b, p]) / math.sqrt(d)
                    if self.box is not None:
                        x = torch.cat([x, self.box(h[b, p], cells, x, rows, cols).to(x.dtype)])
                    out.append(x)
        return out

def decode_head(q: dict, x: torch.Tensor, extra: dict, state: str, keep: float = REGION_KEEP,
                span: tuple[int, float] = (SPAN_CAP, SPAN_LENGTH_PENALTY)) -> dict:
    """A span or region result in amos.model.decode's format; a quote is always cut from the state, at most span[0]
    tokens long, each token past the start scaled by exp(-span[1]). A point is the most likely cell and the box of the
    cells around it holding `keep` of its probability, with the box's mass."""
    from types import SimpleNamespace

    from .model import decode_span
    from .vision import peak_box

    x = x.float().cpu()
    if q["type"] == "span":
        x = torch.cat([x.new_full((2, 1), -torch.inf), x], 1)  # decode_span skips a leading [CLS] slot
        a, z, _, p = decode_span(x[0], x[1], SimpleNamespace(offsets=extra["state"][1], state=state), span[1], span[0])
        # Qwen fuses a JSON string's closing quote with the last word ('detail."', 'approves.",'), so a quote cut at
        # token edges can end (or start) on an unmatched quote mark; drop it and what follows it
        while state[a:z].count('"') % 2:
            if state[z - 1] in '",}]':
                z -= 1
            elif state[a] in '"{[':
                a += 1
            else:
                break
        t = state[a:z]
        a, z = a + len(t) - len(t.lstrip()), z - (len(t) - len(t.rstrip()))
        return {"type": "span", "start": a, "end": max(a, z), "text": state[a : max(a, z)], "p": p}
    _, rows, cols = extra["images"][0]
    cells = x[: rows * cols].softmax(-1)
    if len(x) > rows * cols:  # a box head's box; its mass is the cells it covers
        box = [round(v, 4) for v in x[rows * cols :].tolist()]
        r0, r1 = int(box[1] * rows), max(int(box[1] * rows) + 1, math.ceil(box[3] * rows))
        c0, c1 = int(box[0] * cols), max(int(box[0] * cols) + 1, math.ceil(box[2] * cols))
        return {"type": "region", "cell": int(cells.argmax()), "box": box,
                "p": float(cells.view(rows, cols)[r0:r1, c0:c1].sum())}
    box, p = peak_box(cells, rows, cols, keep)
    return {"type": "region", "cell": int(cells.argmax()), "box": box, "p": p}

def load_heads(path, model, codec: Codec | None = None) -> Heads | None:
    """The trained heads saved beside a checkpoint, or None. A quote head trained on echoed cases switches the codec's
    echo on, since it reads nothing else."""
    f = Path(path) / HEADS
    if not f.exists():
        return None
    saved = torch.load(f, map_location="cpu", weights_only=True)
    heads = Heads(model.get_input_embeddings().weight.shape[1], saved["kinds"], saved.get("keep", REGION_KEEP),
                  box=any(k.startswith("box.") for k in saved["state"]))
    heads.load_state_dict(saved["state"])
    if codec is not None:
        codec.echo = bool(saved.get("echo", False))
    return heads.to(model.device).eval()

def tempered(lp: torch.Tensor, t: float | None = None) -> list[float]:
    """Probabilities from log-probabilities, sharpened (t < 1) or softened (t > 1) by a fitted temperature."""
    return (lp / t).softmax(-1).tolist() if t else lp.exp().tolist()

def rotate(q: dict, r: int) -> dict:
    """A choice question with its options rotated r places; anything else as it is (a scale's order is its meaning)."""
    if q["type"] != "choice" or not r:
        return q
    s = r % len(q["options"])
    return {**q, "options": q["options"][s:] + q["options"][:s]}

def read(model, codec, state: str, questions: list[dict], images=(), max_len: int | None = None,
         temperature: dict | None = None):
    """-> (option distributions for letter-answered questions, input tokens), one forward pass."""
    ids, at, extra = codec.encode_case(state, questions, images)
    if max_len and len(ids) > max_len:
        raise ValueError(f"maximum context length: request has {len(ids)} tokens, limit {max_len}")
    entry = [(ids, at, questions, extra)]
    t = temperature or {}
    picked = [q for q in questions if q["type"] in KINDS]
    return [tempered(lp, t.get(q["type"])) for lp, q in
            zip(distributions(model, codec, entry, hidden(model, codec, entry)), picked)], len(ids)

def wide_choice(model, codec, state: str, q: dict, images=(), max_len: int | None = None,
                temperature: dict | None = None):
    """A choice over more options than there are letters, in two passes -> (distribution, most input tokens).

    First the options are split into chunks, each asked as its own question; then the leading options of every chunk
    are asked again together. Finalists share their first-pass mass (each chunk weighted equally) in proportion to the
    second pass; every other option keeps its first-pass share, so the distribution sums to one and its top answer is
    the second pass's."""
    n, per = len(q["options"]), len(codec.codes)
    size = -(-n // -(-n // per))  # balanced chunks of at most `per`
    chunks = [list(range(s, min(s + size, n))) for s in range(0, n, size)]
    first, t1 = read(model, codec, state, [{**q, "options": [q["options"][i] for i in c]} for c in chunks], images, max_len,
                     temperature)
    prior = [0.0] * n
    for c, p in zip(chunks, first):
        for i, x in zip(c, p):
            prior[i] = x / len(chunks)
    keep = per // len(chunks)
    finalists = sorted(i for c, p in zip(chunks, first) for i, _ in sorted(zip(c, p), key=lambda t: -t[1])[:keep])
    (second,), t2 = read(model, codec, state, [{**q, "options": [q["options"][i] for i in finalists]}], images, max_len,
                         temperature)
    mass, out = sum(prior[i] for i in finalists), prior[:]
    for i, x in zip(finalists, second):
        out[i] = mass * x
    return out, max(t1, t2)

@torch.no_grad()
def sureness(r: dict) -> float:
    """How sure a yes/no, choice or score result is: the probability of its answer."""
    return max(r["p_yes"], 1 - r["p_yes"]) if r["type"] == "noul" else max(r["probs"])

@torch.no_grad()
def think_many(model, codec, state: str, qs: list[dict], max_tokens: int, temperature: dict | None = None,
               generate=None):
    """think for several questions about one case at once: their thoughts are generated as one left-padded batch, so a
    case with eight unsure questions costs about one generation instead of eight. generate(prefixes, max_tokens) ->
    [(text, tokens)] writes the thoughts elsewhere instead (a vLLM server: serve's AMOS_THINK_URL); the letters are
    still read here, so the answers are the same either way."""
    if noncausal(model):
        raise ValueError("a noncausal run cannot think: generation reads a causal model's cache")
    state, qs = codec.safe(state, qs)
    prefixes = [codec.tok.apply_chat_template(render(state, [q], codes=codec.codes)[:2], tokenize=False,
                                              add_generation_prompt=True, enable_thinking=True) for q in qs]
    if generate is not None:
        thoughts = generate(prefixes, max_tokens)
    else:
        ids = [codec.tok(p, add_special_tokens=False)["input_ids"] for p in prefixes]
        width, pad = max(map(len, ids)), codec.tok.pad_token_id
        batch = torch.tensor([[pad] * (width - len(x)) + x for x in ids], device=model.device)
        mask = torch.tensor([[0] * (width - len(x)) + [1] * len(x) for x in ids], device=model.device)
        out = model.generate(batch, attention_mask=mask, max_new_tokens=max_tokens, do_sample=False, use_cache=True,
                             pad_token_id=pad)
        thoughts = [(codec.tok.decode(row, skip_special_tokens=False).replace(codec.tok.pad_token or "\0", ""),
                     int((row != pad).sum())) for row in out[:, width:]]
    results = []
    for prefix, q, (text, n) in zip(prefixes, qs, thoughts):
        reasoning = text.split("</think>")[0].replace("<|im_end|>", "").strip()
        full = codec.tok(prefix + reasoning + "\n</think>\n\n" + f"q1: {ANSWER}", add_special_tokens=False)["input_ids"]
        at = [i - 1 for i, t in enumerate(full) if t == codec.marker]
        lp = distributions(model, codec, [(full, at, [q])])[0]
        results.append((tempered(lp, (temperature or {}).get(q["type"])), n))
    return results

def results(model, codec, state: str, questions: list[dict], images=(), heads: Heads | None = None,
            max_len: int | None = None, rotations: int = 1, temperature: dict | None = None,
            think_below: float | None = None, think_tokens: int = 512, generate=None) -> tuple[list[dict], int]:
    """One typed result per question in amos.model.decode's format (noul p_yes; choice/score probs, choice, expected;
    multi independent probs; span start/end/text/p; region cell/box/p), and the number of input tokens. A choice over
    more options than there are letters is answered on its own (wide_choice). rotations and temperature: see
    one_pass. think_below: a yes/no, choice or score answer less sure than this is asked again after the model
    thinks (think), as GLiDE spends compute only where a decision is uncertain; the result then carries "thought"."""
    if think_below and noncausal(model):  # on every request, not only the unsure ones
        raise ValueError("a noncausal run cannot think: generation reads a causal model's cache")
    out, tokens = _results(model, codec, state, questions, images, heads, max_len, rotations, temperature)
    if think_below and not images:
        unsure = [k for k, (q, r) in enumerate(zip(questions, out))
                  if q["type"] in KINDS and len(options(q)) <= len(codec.codes) and sureness(r) < think_below]
        for k, (p, n) in zip(unsure, think_many(model, codec, state, [questions[k] for k in unsure], think_tokens,
                                                temperature, generate) if unsure else []):
            q = questions[k]
            if q["type"] == "noul":
                out[k] = {"type": "noul", "p_yes": p[0], "thought": n}
            else:
                out[k] = {"type": q["type"], "probs": p, "choice": max(range(len(p)), key=p.__getitem__), "thought": n}
                if q["type"] == "score":
                    out[k]["expected"] = sum(i * x for i, x in enumerate(p))
    return out, tokens

def _results(model, codec, state: str, questions: list[dict], images=(), heads: Heads | None = None,
             max_len: int | None = None, rotations: int = 1, temperature: dict | None = None) -> tuple[list[dict], int]:
    wide = {k for k, q in enumerate(questions) if q["type"] == "choice" and len(q["options"]) > len(codec.codes)}
    out, tokens = [None] * len(questions), 0
    for k in sorted(wide):
        p, n = wide_choice(model, codec, state, questions[k], images, max_len, temperature)
        out[k], tokens = {"type": "choice", "probs": p, "choice": max(range(len(p)), key=p.__getitem__)}, max(tokens, n)
    rest = [k for k in range(len(questions)) if k not in wide]
    if rest:
        sub, n = one_pass(model, codec, state, [questions[k] for k in rest], images, heads, max_len, rotations,
                          temperature)
        for k, r in zip(rest, sub):
            out[k] = r
        tokens = max(tokens, n)
    return out, tokens

ROTATION_BATCH_TOKENS = 65536

def one_pass(model, codec, state: str, questions: list[dict], images=(), heads: Heads | None = None,
             max_len: int | None = None, rotations: int = 1, temperature: dict | None = None) -> tuple[list[dict], int]:
    """results() for questions that fit the letters, in one forward pass.

    rotations > 1 (and a choice question present): the case is also read with every choice question's options rotated,
    each rotation one more row of the same batched forward, and choice distributions are averaged in the original order.
    temperature: {question type: t} fitted by `decoder calibrate` (pick-any's per-option yes/no uses noul's)."""
    flat, owner = expand(questions)
    rotations = rotations if any(q["type"] == "choice" for q in flat) else 1
    entries = []
    for r in range(rotations):
        qs = [rotate(q, r) for q in flat]
        ids, at, extra = codec.encode_case(state, qs, images)
        if max_len and len(ids) > max_len:
            raise ValueError(f"maximum context length: request has {len(ids)} tokens, limit {max_len}")
        entries.append((ids, at, qs, extra))
    ids, extra = entries[0][0], entries[0][3]
    if heads is None and any(q["type"] in HEADED for q in flat):
        raise ValueError("this model has no span or region head")
    if sum(len(e[0]) for e in entries) <= ROTATION_BATCH_TOKENS:
        h = hidden(model, codec, entries)
        lps = distributions(model, codec, entries, h)
    else:  # a long case: one rotation at a time, so memory stays that of a single row
        h, lps = None, []
        for e in entries:
            he = hidden(model, codec, [e])
            lps += distributions(model, codec, [e], he)
            h = he if h is None else h
    head = iter(heads(h[:1], entries[:1]) if heads is not None else [])
    picked = [j for j, q in enumerate(flat) if q["type"] in KINDS]
    temps = temperature or {}
    mean = {}
    for r in range(rotations):
        for m, j in enumerate(picked):
            p, n = tempered(lps[r * len(picked) + m], temps.get(flat[j]["type"])), len(options(flat[j]))
            s = r % n if flat[j]["type"] == "choice" else 0
            acc = mean.setdefault(j, [0.0] * n)
            for i in range(n):  # rotated position (i - s) mod n holds original option i
                acc[i] += p[(i - s) % n] / rotations
    lps = iter(mean[j] for j in picked)
    out = [None] * len(questions)
    for q, k in zip(flat, owner):
        t = questions[k]["type"]
        if q["type"] in HEADED:
            out[k] = decode_head(q, next(head), extra, state, heads.keep, heads.span)
            continue
        p = next(lps)
        if t == "multi":
            out[k] = out[k] or {"type": t, "probs": []}
            out[k]["probs"].append(p[0])
        elif t == "noul":
            out[k] = {"type": t, "p_yes": p[0]}
        else:
            out[k] = {"type": t, "probs": p, "choice": max(range(len(p)), key=p.__getitem__)}
            if t == "score":
                out[k]["expected"] = sum(i * x for i, x in enumerate(p))
    return out, len(ids)

CALIBRATION = "calibration.json"

def calibration(run) -> dict | None:
    """The per-type temperatures `decoder calibrate` fitted for a run, or None."""
    f = Path(run) / CALIBRATION
    return json.loads(f.read_text()).get("temperature") if f.exists() else None

def span_decoding(run) -> tuple[int, float]:
    """The quote decoder's (cap, penalty) fitted for a run on held-out quote records, or the defaults."""
    f = Path(run) / CALIBRATION
    s = (json.loads(f.read_text()).get("span") or {}) if f.exists() else {}
    return int(s.get("cap", SPAN_CAP)), float(s.get("penalty", SPAN_LENGTH_PENALTY))

def multi_cut(run) -> float | None:
    """The pick-any cut fitted on our own held-out pick-any questions (an option is picked at p >= cut), or None."""
    f = Path(run) / CALIBRATION
    return json.loads(f.read_text()).get("multi_cut") if f.exists() else None