"""Pure-tensor helpers for generation-time SAE feature extraction.

No model or transformers imports — all functions take tensors and return
tensors so they're testable without GPU.
"""
from __future__ import annotations

import torch


def build_position_masks(
    prompt_real_lens: torch.Tensor,
    gen_lens: torch.Tensor,
    n_image_patches: int,
    seq_len: int,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute prompt-vs-gen position masks for a left-padded teacher-forced batch.

    LLaVA replaces the single ``<image>`` token with ``n_image_patches`` patch
    embeddings. After image expansion, row ``i`` has a real region of length
    ``n_image_patches + prompt_real_lens[i] - 1 + gen_lens[i]`` flush against
    position ``seq_len - 1`` (left padding puts pads at the start).

    Returns:
        prompt_mask: (B, seq_len) bool, True over the image-patch and prompt-text
            positions of each row's real region.
        gen_mask: (B, seq_len) bool, True over the last ``gen_lens[i]`` positions
            of each row.
    """
    B = prompt_real_lens.shape[0]
    assert gen_lens.shape == (B,)
    device = prompt_real_lens.device
    pos = torch.arange(seq_len, device=device).unsqueeze(0).expand(B, -1)  # (B, S)

    # Real region per row, flush right.
    real_lens = n_image_patches + prompt_real_lens - 1 + gen_lens         # (B,)
    assert (real_lens <= seq_len).all(), \
        f"real region overflows seq_len: max real_len={int(real_lens.max())}, seq_len={seq_len}"
    real_start = (seq_len - real_lens).unsqueeze(1)                       # (B, 1)
    gen_start = (seq_len - gen_lens).unsqueeze(1)                         # (B, 1)
    end = torch.full_like(real_start, seq_len)                            # (B, 1)

    prompt_mask = (pos >= real_start) & (pos < gen_start)
    gen_mask = (pos >= gen_start) & (pos < end)
    return prompt_mask, gen_mask


def left_pad_collate(
    rows: list[torch.Tensor],
    pad_id: int,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Stack 1-D LongTensors with left padding.

    Args:
        rows: list of 1-D tensors with possibly different lengths.
        pad_id: token id to use for padding.

    Returns:
        ids: (B, L_max) LongTensor, left-padded with ``pad_id``.
        attn: (B, L_max) LongTensor, 1 over real tokens, 0 over pads.
    """
    L_max = max(r.shape[0] for r in rows)
    B = len(rows)
    ids = torch.full((B, L_max), pad_id, dtype=rows[0].dtype, device=rows[0].device)
    attn = torch.zeros((B, L_max), dtype=torch.long, device=rows[0].device)
    for i, r in enumerate(rows):
        L = r.shape[0]
        ids[i, L_max - L :] = r
        attn[i, L_max - L :] = 1
    return ids, attn


def gather_gen_positions(
    capture_hidden: dict[int, torch.Tensor],
    gen_mask: torch.Tensor,
) -> tuple[dict[int, torch.Tensor], torch.Tensor]:
    """Gather generated-token hidden states into a packed (B, T_max, D) tensor.

    Unlike ``extract_gen_sae_features`` (which SAE-encodes then max-pools to a
    single vector per row), this keeps the full token sequence so a downstream
    sequence model (e.g. an attention-pooling probe) can read every position.

    The generated region of row ``i`` occupies the last ``gen_mask[i].sum()``
    positions of the left-padded sequence. We left-anchor each row's gen tokens
    at index 0 of the output and mark validity.

    Args:
        capture_hidden: l -> (B, S, D) per-layer hidden states.
        gen_mask: (B, S) bool, True over generated positions.

    Returns:
        feats_seq: l -> (B, T_max, D) packed gen-position hidden states; rows
            shorter than T_max are zero-padded at the tail.
        valid_mask: (B, T_max) bool, True at real gen positions. Rows whose
            ``gen_mask`` is all-False yield an all-False ``valid_mask`` row —
            the caller must filter those out before feeding an attention module
            (an all-masked row produces NaNs).
    """
    gen_counts = gen_mask.sum(dim=1)                                        # (B,)
    B = gen_counts.shape[0]
    T_max = int(gen_counts.max().item()) if B > 0 else 0
    any_h = next(iter(capture_hidden.values()))
    device = any_h.device
    valid = (
        torch.arange(T_max, device=device).unsqueeze(0) < gen_counts.unsqueeze(1)
        if T_max > 0 else torch.zeros((B, 0), dtype=torch.bool, device=device)
    )                                                                       # (B, T_max)
    out: dict[int, torch.Tensor] = {}
    for l, h in capture_hidden.items():
        S, D = h.shape[1], h.shape[2]
        if T_max == 0:
            out[l] = h.new_zeros((B, 0, D))
            continue
        idx = torch.arange(T_max, device=h.device).unsqueeze(0).expand(B, -1)  # (B, T_max)
        gen_starts = (S - gen_counts).unsqueeze(1)                             # (B, 1)
        pos = (gen_starts + idx).clamp(max=S - 1)                              # (B, T_max)
        pos_exp = pos.unsqueeze(-1).expand(-1, -1, D)                          # (B, T_max, D)
        h_gen = torch.gather(h, dim=1, index=pos_exp)                          # (B, T_max, D)
        out[l] = h_gen.masked_fill(~valid.unsqueeze(-1), 0.0)
    return out, valid


def build_scope_mask(
    scope: str,
    prompt_mask: torch.Tensor,   # (B, S) vision-patch + prompt-text region (real, pre-gen)
    gen_mask: torch.Tensor,      # (B, S) generated-token region
    vision_mask: torch.Tensor,   # (B, S) image-patch positions (input_ids == image_token_id)
) -> torch.Tensor:
    """Select which token positions the probe reads, per ``probe_token_scope``.

    Scopes:
      gen        — generated tokens only (default; surgical, output-aligned)
      vision     — image-patch tokens only
      prompt     — prompt TEXT tokens only (prompt region minus vision patches)
      prompt_gen — prompt text + generated tokens (excludes vision patches)
      all        — vision patches + prompt text + generated tokens
    """
    prompt_text = prompt_mask & ~vision_mask
    if scope == "gen":
        return gen_mask
    if scope == "vision":
        return vision_mask
    if scope == "prompt":
        return prompt_text
    if scope == "prompt_gen":
        return prompt_text | gen_mask
    if scope == "all":
        return prompt_mask | gen_mask
    raise ValueError(f"Unknown probe_token_scope: {scope!r}")


def gather_masked_positions(
    capture_hidden: dict[int, torch.Tensor],
    mask: torch.Tensor,
) -> tuple[dict[int, torch.Tensor], torch.Tensor]:
    """Pack the True positions of an ARBITRARY ``mask`` into (B, T_max, D), in order.

    Generalizes ``gather_gen_positions`` (which assumes a tail-flush contiguous
    region) to any selection — e.g. vision-patch tokens, which form a contiguous
    block in the MIDDLE of the sequence, or a union of disjoint spans. Sequence
    order of the selected tokens is preserved.

    Args:
        capture_hidden: l -> (B, S, D) per-layer hidden states.
        mask: (B, S) bool, True at positions to keep.

    Returns:
        feats_seq: l -> (B, T_max, D) packed selected hidden states (left-aligned,
            zero-padded tail).
        valid_mask: (B, T_max) bool, True at real selected positions. Rows whose
            ``mask`` is all-False yield an all-False row — filter before feeding an
            attention module.
    """
    counts = mask.sum(dim=1)                                                # (B,)
    B = counts.shape[0]
    T_max = int(counts.max().item()) if B > 0 else 0
    any_h = next(iter(capture_hidden.values()))
    S = any_h.shape[1]
    device = any_h.device
    # Positions of True entries, in ascending order, padded with sentinel S.
    pos = torch.arange(S, device=device).unsqueeze(0).expand(B, -1)         # (B, S)
    pos_masked = pos.masked_fill(~mask, S)
    sorted_pos, _ = pos_masked.sort(dim=1)                                  # (B, S)
    valid = (
        torch.arange(T_max, device=device).unsqueeze(0) < counts.unsqueeze(1)
        if T_max > 0 else torch.zeros((B, 0), dtype=torch.bool, device=device)
    )                                                                       # (B, T_max)
    out: dict[int, torch.Tensor] = {}
    for l, h in capture_hidden.items():
        D = h.shape[2]
        if T_max == 0:
            out[l] = h.new_zeros((B, 0, D))
            continue
        idx = sorted_pos[:, :T_max].clamp(max=S - 1)                        # (B, T_max)
        h_sel = torch.gather(h, dim=1, index=idx.unsqueeze(-1).expand(-1, -1, D))
        out[l] = h_sel.masked_fill(~valid.unsqueeze(-1), 0.0)
    return out, valid


def masked_max_pool(h: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
    """Per-row max-pool over True positions of ``mask``.

    Args:
        h: (B, S, D) hidden states.
        mask: (B, S) bool, True for positions to include in the pool.

    Returns:
        (B, D) max-pooled features. Rows where ``mask`` is all False return
        -inf (caller is responsible for ensuring at least one True per row,
        or for masking out empty rows downstream).
    """
    assert h.dim() == 3 and mask.dim() == 2
    assert h.shape[:2] == mask.shape
    # float("-inf") not torch.finfo(...).min so torch.isinf detects empty-mask rows.
    masked = h.masked_fill(~mask.unsqueeze(-1), float("-inf"))
    return masked.max(dim=1).values


def masked_mean_pool(h: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
    """Per-row mean-pool over True positions of ``mask``.

    Non-redistributable alternative to ``masked_max_pool``: every active
    position contributes equally, so a downstream suppression objective
    cannot be satisfied by lowering one peak and raising others.

    Args:
        h: (B, S, D) hidden states.
        mask: (B, S) bool, True for positions to include in the pool.

    Returns:
        (B, D) mean-pooled features. Rows where ``mask`` is all False return
        zeros (caller is responsible for masking those rows out downstream).
    """
    assert h.dim() == 3 and mask.dim() == 2
    assert h.shape[:2] == mask.shape
    m = mask.unsqueeze(-1).to(h.dtype)              # (B, S, 1)
    summed = (h * m).sum(dim=1)                     # (B, D)
    count = m.sum(dim=1).clamp(min=1.0)             # (B, 1) — div-by-zero guard
    return summed / count