"""Masked-scope scoring helpers for LLaDA-Guard inference."""

from __future__ import annotations

import torch
import torch.nn.functional as F
from tqdm import tqdm

from .data import PROMPT_BLOCK_ID, RESPONSE_BLOCK_ID, build_eval_loader
from .model import MASK_ID

_MASK_SEED_EXAMPLE_STRIDE = 104729


def _build_cached_eval_loader(
    examples,
    tokenizer,
    max_seq_length,
    batch_size,
    gen_only=False,
    skip_cls=False,
    prompt_format: str = "aegis_taxonomy",
):
    """Compatibility wrapper for the deterministic public eval loader."""
    del skip_cls
    return build_eval_loader(
        examples,
        tokenizer,
        max_seq_length,
        batch_size,
        gen_only=gen_only,
        prompt_format=prompt_format,
    )


def _mask_seed(batch, batch_idx: int, base_seed: int, t_idx: int) -> int | torch.Tensor:
    """Stable mask seed keyed by example id when available."""
    example_index = batch.get("example_index")
    if torch.is_tensor(example_index) and bool((example_index >= 0).all().item()):
        return example_index.to(dtype=torch.long) * _MASK_SEED_EXAMPLE_STRIDE + int(base_seed) + int(t_idx)
    return int(base_seed) + int(batch_idx) * 97 + int(t_idx)


def _base_llada_lm(model):
    get_base_model = getattr(model, "get_base_model", None)
    return get_base_model() if callable(get_base_model) else model


def _prefer_native_rms_norm(model) -> None:
    """Patch LLaDA RMSLayerNorm modules to PyTorch's native RMSNorm fast path."""
    if not hasattr(F, "rms_norm"):
        return
    modules = getattr(_base_llada_lm(model), "modules", None)
    if not callable(modules):
        return
    for module in modules():
        if module.__class__.__name__ != "RMSLayerNorm" or getattr(module, "_native_rms_norm", False):
            continue
        normalized_shape = getattr(module, "normalized_shape", None)
        weight = getattr(module, "weight", None)
        eps = getattr(module, "eps", None)
        if normalized_shape is None or weight is None or eps is None:
            continue

        def native_forward(x: torch.Tensor, rms_module=module) -> torch.Tensor:
            y = F.rms_norm(x, rms_module.normalized_shape, rms_module.weight, rms_module.eps)
            bias = getattr(rms_module, "bias", None)
            if bias is not None:
                y = y + bias
            return y

        module.forward = native_forward
        module._native_rms_norm = True


def _prefer_inplace_silu(model) -> None:
    """Enable in-place SiLU on existing activation modules."""
    modules = getattr(_base_llada_lm(model), "modules", None)
    if not callable(modules):
        return
    for module in modules():
        if module.__class__.__name__ == "SiLU":
            module.inplace = True


def apply_eval_inference_fastpaths(model) -> None:
    """Apply idempotent inference-only module fast paths used by eval scoring."""
    _prefer_native_rms_norm(model)
    _prefer_inplace_silu(model)


def _mask_ratio_upper_bound(eps: float, max_t: float) -> float:
    return max(float(eps), min(1.0, float(max_t)))


def _scale_mask_ratio(t_values, scale: float, eps: float, max_t: float = 1.0):
    """Scale a masking ratio while keeping it inside a valid range."""
    upper = _mask_ratio_upper_bound(eps, max_t)
    if abs(float(scale) - 1.0) <= 1e-12 and upper >= 1.0:
        return t_values
    if torch.is_tensor(t_values):
        return torch.clamp(t_values.to(dtype=torch.float32) * float(scale), min=eps, max=upper)
    return max(eps, min(upper, float(t_values) * float(scale)))


def _sample_relative_mask_indices(
    block_ids: torch.Tensor,
    real_lengths: torch.Tensor,
    eps: float = 1e-3,
    target_block_ids: tuple[int, ...] = (1, 2),
    t_override: float | torch.Tensor | None = None,
    max_t: float = 1.0,
    seed: int | torch.Tensor | list[int] | tuple[int, ...] | None = None,
) -> tuple[list[torch.Tensor], torch.Tensor]:
    """Sample exact-k relative mask positions for the selected target region."""
    device = block_ids.device
    batch_size = block_ids.shape[0]
    upper = _mask_ratio_upper_bound(eps, max_t)
    if t_override is None:
        base = torch.rand(batch_size, device=device)
        offset = torch.arange(batch_size, device=device).float() / max(batch_size, 1)
        t_values = (upper - eps) * ((base / max(batch_size, 1) + offset) % 1.0) + eps
    elif torch.is_tensor(t_override):
        t_values = torch.clamp(t_override.to(device=device, dtype=torch.float32), min=eps, max=upper)
    else:
        t_values = torch.full(
            (batch_size,),
            max(eps, min(upper, float(t_override))),
            device=device,
        )

    row_seeds = None
    generator = None
    if torch.is_tensor(seed):
        row_seeds = seed.detach().to(device="cpu", dtype=torch.long).tolist()
    elif isinstance(seed, (list, tuple)):
        row_seeds = [int(value) for value in seed]
    elif seed is not None:
        generator = torch.Generator(device=device)
        generator.manual_seed(int(seed))

    if row_seeds is not None and len(row_seeds) != batch_size:
        raise ValueError(f"Expected {batch_size} row seeds, got {len(row_seeds)}")

    positions = torch.arange(block_ids.shape[1], device=device).unsqueeze(0)
    selections: list[torch.Tensor] = []
    for idx in range(batch_size):
        target_mask = torch.zeros(block_ids.shape[1], dtype=torch.bool, device=device)
        for block_id in target_block_ids:
            target_mask |= block_ids[idx] == block_id
        target_positions = (
            target_mask & (positions[0] < real_lengths[idx])
        ).nonzero(as_tuple=True)[0]
        target_len = target_positions.numel()
        if target_len == 0:
            selections.append(torch.empty(0, dtype=torch.long, device=device))
            continue
        n_mask = max(1, min(target_len, int(round(float(t_values[idx]) * target_len))))
        row_generator = generator
        if row_seeds is not None:
            row_generator = torch.Generator(device=device)
            row_generator.manual_seed(int(row_seeds[idx]))
        perm = torch.randperm(target_len, generator=row_generator, device=device)
        selections.append(perm[:n_mask])
    return selections, t_values


def _relative_to_absolute_mask(
    block_ids: torch.Tensor,
    real_lengths: torch.Tensor,
    relative_indices: list[torch.Tensor],
    target_block_ids: tuple[int, ...] = (1, 2),
) -> torch.Tensor:
    """Map relative target-region positions to absolute sequence positions."""
    device = block_ids.device
    batch_size, seq_len = block_ids.shape
    positions = torch.arange(seq_len, device=device).unsqueeze(0)
    mask = torch.zeros(batch_size, seq_len, dtype=torch.bool, device=device)

    for idx in range(batch_size):
        rel_idx = relative_indices[idx]
        if rel_idx.numel() == 0:
            continue
        target_mask = torch.zeros(seq_len, dtype=torch.bool, device=device)
        for block_id in target_block_ids:
            target_mask |= block_ids[idx] == block_id
        target_positions = (
            target_mask & (positions[0] < real_lengths[idx])
        ).nonzero(as_tuple=True)[0]
        mask[idx, target_positions[rel_idx]] = True

    return mask


def _full_target_mask(
    block_ids: torch.Tensor,
    real_lengths: torch.Tensor,
    target_block_ids: tuple[int, ...],
) -> torch.Tensor:
    """Build a mask covering every token in the selected target region."""
    device = block_ids.device
    positions = torch.arange(block_ids.shape[1], device=device).unsqueeze(0)
    target_mask = torch.zeros_like(block_ids, dtype=torch.bool)
    for block_id in target_block_ids:
        target_mask |= block_ids == block_id
    return target_mask & (positions < real_lengths.unsqueeze(1))


def _pad_to_match_length(
    tensor: torch.Tensor,
    target_len: int,
    pad_id: int,
) -> torch.Tensor:
    """Right-pad a batch tensor to a target length."""
    if tensor.shape[1] >= target_len:
        return tensor
    pad_len = target_len - tensor.shape[1]
    pad = torch.full((tensor.shape[0], pad_len), pad_id, dtype=tensor.dtype, device=tensor.device)
    return torch.cat([tensor, pad], dim=1)


def _pad_mask_to_match_length(mask: torch.Tensor, target_len: int) -> torch.Tensor:
    """Right-pad a boolean mask to a target length."""
    if mask.shape[1] >= target_len:
        return mask
    pad_len = target_len - mask.shape[1]
    pad = torch.zeros((mask.shape[0], pad_len), dtype=mask.dtype, device=mask.device)
    return torch.cat([mask, pad], dim=1)


def _build_attention_mask_from_lengths(
    real_lengths: torch.Tensor,
    target_len: int,
) -> torch.Tensor:
    """Build a right-padded attention mask from per-example real lengths."""
    positions = torch.arange(target_len, device=real_lengths.device).unsqueeze(0)
    return (positions < real_lengths.unsqueeze(1)).to(dtype=torch.long)


def _compute_masked_token_losses_from_logits(
    logits: torch.Tensor,
    input_ids: torch.Tensor,
    mask: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute masked token CE, aggregated per example, from precomputed logits."""
    device = input_ids.device
    batch_size = input_ids.shape[0]

    if not mask.any():
        zeros = logits.float().sum(dim=(1, 2)) * 0.0
        return zeros, zeros

    raw_ce = F.cross_entropy(logits[mask], input_ids[mask], reduction="none").float()
    example_idx = (
        torch.arange(batch_size, device=device).unsqueeze(1).expand_as(mask)[mask]
    )
    per_ex_sum = torch.zeros(batch_size, device=device, dtype=raw_ce.dtype)
    per_ex_sum.scatter_add_(0, example_idx, raw_ce)
    denom = mask.sum(dim=1).float()
    per_ex_loss = per_ex_sum / denom.clamp(min=1)
    return per_ex_loss, -per_ex_loss


def _compute_paired_masked_token_outputs(
    model,
    safe_input_ids: torch.Tensor,
    safe_mask: torch.Tensor,
    safe_real_lengths: torch.Tensor,
    unsafe_input_ids: torch.Tensor,
    unsafe_mask: torch.Tensor,
    unsafe_real_lengths: torch.Tensor,
    pad_id: int,
    mask_id: int = MASK_ID,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    """Run one paired safe/unsafe forward and return padded ids + logits."""
    max_len = max(safe_input_ids.shape[1], unsafe_input_ids.shape[1])
    safe_ids = _pad_to_match_length(safe_input_ids, max_len, pad_id)
    unsafe_ids = _pad_to_match_length(unsafe_input_ids, max_len, pad_id)
    safe_attention_mask = _build_attention_mask_from_lengths(safe_real_lengths, max_len)
    unsafe_attention_mask = _build_attention_mask_from_lengths(unsafe_real_lengths, max_len)
    safe_mask_padded = _pad_mask_to_match_length(safe_mask, max_len)
    unsafe_mask_padded = _pad_mask_to_match_length(unsafe_mask, max_len)

    combined_ids = torch.cat([safe_ids, unsafe_ids], dim=0)
    combined_attention_mask = torch.cat([safe_attention_mask, unsafe_attention_mask], dim=0)
    combined_mask = torch.cat([safe_mask_padded, unsafe_mask_padded], dim=0)
    noisy = combined_ids.clone()
    noisy[combined_mask] = mask_id
    logits = model(input_ids=noisy, attention_mask=combined_attention_mask).logits

    half = safe_input_ids.shape[0]
    return safe_ids, logits[:half], unsafe_ids, logits[half:]


def _target_token_counts(
    block_ids: torch.Tensor,
    real_lengths: torch.Tensor,
    *,
    target_block_ids: tuple[int, ...],
) -> torch.Tensor:
    """Count valid tokens in the scored target region for each row."""
    positions = torch.arange(block_ids.shape[1], device=block_ids.device).unsqueeze(0)
    target_mask = torch.zeros_like(block_ids, dtype=torch.bool)
    for block_id in target_block_ids:
        target_mask |= block_ids == block_id
    valid_target = target_mask & (positions < real_lengths.unsqueeze(1))
    return valid_target.sum(dim=1).clamp(min=1).float()


def _adaptive_optimal_t_values(
    batch,
    *,
    target_block_ids: tuple[int, ...],
    repeats: int,
    sigma_over_c: float,
    min_t: float,
    max_t: float,
) -> tuple[torch.Tensor, ...]:
    """Bias-variance optimal t per row: (sigma^2 / (2 c^2 R n))^(1/3)."""
    repeats = max(1, int(repeats))
    counts = _target_token_counts(
        batch["gen_safe_block_ids"],
        batch["gen_safe_real_lengths"],
        target_block_ids=target_block_ids,
    )
    t = ((float(sigma_over_c) ** 2) / (2.0 * repeats * counts)).pow(1.0 / 3.0)
    t = torch.clamp(t, min=float(min_t), max=float(max_t))
    return tuple(t.clone() for _ in range(repeats))


def _slice_batch_rows(batch: dict[str, torch.Tensor], row_mask: torch.Tensor) -> dict[str, torch.Tensor]:
    """Select a subset of rows from a batch dict."""
    row_mask = row_mask.to(dtype=torch.bool)
    sliced = {}
    for key, value in batch.items():
        if torch.is_tensor(value) and value.dim() > 0 and value.shape[0] == row_mask.shape[0]:
            sliced[key] = value[row_mask]
        else:
            sliced[key] = value
    return sliced


def _score_batch_masked_scope_sequential(
    model,
    batch,
    batch_idx,
    t_values,
    pad_id,
    response_mask_scale=1.0,
    adaptive_t_repeats: int | None = None,
    adaptive_t_sigma_over_c: float = 1.0,
    adaptive_t_min: float = 1e-3,
    adaptive_t_max: float = 0.5,
):
    """Score a batch with prompt/response scoped masks, one t at a time."""
    device = batch["gen_safe_input_ids"].device
    batch_size = batch["gen_safe_input_ids"].shape[0]
    has_response = batch.get("has_response")
    if has_response is None:
        has_response = (batch["gen_safe_block_ids"] == RESPONSE_BLOCK_ID).any(dim=1).float()
    has_response = has_response > 0.5
    has_prompt_only = bool((~has_response).any().item())
    has_with_response = bool(has_response.any().item())

    if has_prompt_only and has_with_response:
        scores = torch.zeros(batch_size, device=device)
        prompt_mask = ~has_response
        response_mask = has_response
        if prompt_mask.any():
            scores[prompt_mask] = _score_batch_masked_scope_sequential(
                model,
                _slice_batch_rows(batch, prompt_mask),
                batch_idx,
                t_values,
                pad_id,
                response_mask_scale=response_mask_scale,
                adaptive_t_repeats=adaptive_t_repeats,
                adaptive_t_sigma_over_c=adaptive_t_sigma_over_c,
                adaptive_t_min=adaptive_t_min,
                adaptive_t_max=adaptive_t_max,
            )
        if response_mask.any():
            scores[response_mask] = _score_batch_masked_scope_sequential(
                model,
                _slice_batch_rows(batch, response_mask),
                batch_idx,
                t_values,
                pad_id,
                response_mask_scale=response_mask_scale,
                adaptive_t_repeats=adaptive_t_repeats,
                adaptive_t_sigma_over_c=adaptive_t_sigma_over_c,
                adaptive_t_min=adaptive_t_min,
                adaptive_t_max=adaptive_t_max,
            )
        return scores.detach()

    active_t_values = t_values
    target_block_ids = (PROMPT_BLOCK_ID,) if has_prompt_only else (RESPONSE_BLOCK_ID,)
    if adaptive_t_repeats is not None:
        active_t_values = _adaptive_optimal_t_values(
            batch,
            target_block_ids=target_block_ids,
            repeats=adaptive_t_repeats,
            sigma_over_c=adaptive_t_sigma_over_c,
            min_t=adaptive_t_min,
            max_t=adaptive_t_max,
        )

    score_accum = torch.zeros(batch_size, device=device)
    for t_idx, t_value in enumerate(active_t_values):
        if has_prompt_only:
            prompt_relative_indices, _ = _sample_relative_mask_indices(
                batch["gen_safe_block_ids"],
                batch["gen_safe_real_lengths"],
                t_override=t_value,
                seed=_mask_seed(batch, batch_idx, 1337, t_idx),
                target_block_ids=(PROMPT_BLOCK_ID,),
            )
            safe_mask = _relative_to_absolute_mask(
                batch["gen_safe_block_ids"],
                batch["gen_safe_real_lengths"],
                prompt_relative_indices,
                target_block_ids=(PROMPT_BLOCK_ID,),
            )
            unsafe_mask = _relative_to_absolute_mask(
                batch["gen_unsafe_block_ids"],
                batch["gen_unsafe_real_lengths"],
                prompt_relative_indices,
                target_block_ids=(PROMPT_BLOCK_ID,),
            )
            scope_safe_mask = safe_mask | _full_target_mask(
                batch["gen_safe_block_ids"],
                batch["gen_safe_real_lengths"],
                target_block_ids=(RESPONSE_BLOCK_ID,),
            )
            scope_unsafe_mask = unsafe_mask | _full_target_mask(
                batch["gen_unsafe_block_ids"],
                batch["gen_unsafe_real_lengths"],
                target_block_ids=(RESPONSE_BLOCK_ID,),
            )
        else:
            response_t_value = _scale_mask_ratio(t_value, response_mask_scale, 1e-3)
            response_relative_indices, _ = _sample_relative_mask_indices(
                batch["gen_safe_block_ids"],
                batch["gen_safe_real_lengths"],
                t_override=response_t_value,
                seed=_mask_seed(batch, batch_idx, 7331, t_idx),
                target_block_ids=(RESPONSE_BLOCK_ID,),
            )
            safe_mask = _relative_to_absolute_mask(
                batch["gen_safe_block_ids"],
                batch["gen_safe_real_lengths"],
                response_relative_indices,
                target_block_ids=(RESPONSE_BLOCK_ID,),
            )
            unsafe_mask = _relative_to_absolute_mask(
                batch["gen_unsafe_block_ids"],
                batch["gen_unsafe_real_lengths"],
                response_relative_indices,
                target_block_ids=(RESPONSE_BLOCK_ID,),
            )
            scope_safe_mask = safe_mask
            scope_unsafe_mask = unsafe_mask

        safe_ids, safe_logits, unsafe_ids, unsafe_logits = _compute_paired_masked_token_outputs(
            model,
            batch["gen_safe_input_ids"],
            scope_safe_mask,
            batch["gen_safe_real_lengths"],
            batch["gen_unsafe_input_ids"],
            scope_unsafe_mask,
            batch["gen_unsafe_real_lengths"],
            pad_id,
        )
        max_len = safe_ids.shape[1]
        safe_mask_padded = _pad_mask_to_match_length(safe_mask, max_len)
        unsafe_mask_padded = _pad_mask_to_match_length(unsafe_mask, max_len)
        _, safe_logp = _compute_masked_token_losses_from_logits(
            safe_logits,
            safe_ids,
            safe_mask_padded,
        )
        _, unsafe_logp = _compute_masked_token_losses_from_logits(
            unsafe_logits,
            unsafe_ids,
            unsafe_mask_padded,
        )
        score_accum += unsafe_logp - safe_logp

    return (score_accum / len(active_t_values)).detach()


def _score_batch_masked_scope_tiled(
    model,
    batch,
    batch_idx,
    t_values,
    pad_id,
    response_mask_scale=1.0,
    adaptive_t_repeats: int | None = None,
    adaptive_t_sigma_over_c: float = 1.0,
    adaptive_t_min: float = 1e-3,
    adaptive_t_max: float = 0.5,
):
    """Score a batch with prompt/response scoped masks, tiled over t values."""
    device = batch["gen_safe_input_ids"].device
    has_response = batch.get("has_response")
    if has_response is None:
        has_response = (batch["gen_safe_block_ids"] == RESPONSE_BLOCK_ID).any(dim=1).float()
    has_response = has_response > 0.5
    has_prompt_only = bool((~has_response).any().item())
    has_with_response = bool(has_response.any().item())
    batch_size = batch["gen_safe_input_ids"].shape[0]

    if has_prompt_only and has_with_response:
        scores = torch.zeros(batch_size, device=device)
        prompt_mask = ~has_response
        response_mask = has_response
        if prompt_mask.any():
            scores[prompt_mask] = _score_batch_masked_scope_tiled(
                model,
                _slice_batch_rows(batch, prompt_mask),
                batch_idx,
                t_values,
                pad_id,
                response_mask_scale=response_mask_scale,
                adaptive_t_repeats=adaptive_t_repeats,
                adaptive_t_sigma_over_c=adaptive_t_sigma_over_c,
                adaptive_t_min=adaptive_t_min,
                adaptive_t_max=adaptive_t_max,
            )
        if response_mask.any():
            scores[response_mask] = _score_batch_masked_scope_tiled(
                model,
                _slice_batch_rows(batch, response_mask),
                batch_idx,
                t_values,
                pad_id,
                response_mask_scale=response_mask_scale,
                adaptive_t_repeats=adaptive_t_repeats,
                adaptive_t_sigma_over_c=adaptive_t_sigma_over_c,
                adaptive_t_min=adaptive_t_min,
                adaptive_t_max=adaptive_t_max,
            )
        return scores.detach()

    target_block_ids = (PROMPT_BLOCK_ID,) if has_prompt_only else (RESPONSE_BLOCK_ID,)
    active_t_values = t_values
    if adaptive_t_repeats is not None:
        active_t_values = _adaptive_optimal_t_values(
            batch,
            target_block_ids=target_block_ids,
            repeats=adaptive_t_repeats,
            sigma_over_c=adaptive_t_sigma_over_c,
            min_t=adaptive_t_min,
            max_t=adaptive_t_max,
        )

    all_safe_masks = []
    all_unsafe_masks = []
    for t_idx, t_value in enumerate(active_t_values):
        if has_prompt_only:
            relative_indices, _ = _sample_relative_mask_indices(
                batch["gen_safe_block_ids"],
                batch["gen_safe_real_lengths"],
                t_override=t_value,
                seed=_mask_seed(batch, batch_idx, 1337, t_idx),
                target_block_ids=(PROMPT_BLOCK_ID,),
            )
        else:
            response_t_value = _scale_mask_ratio(t_value, response_mask_scale, 1e-3)
            relative_indices, _ = _sample_relative_mask_indices(
                batch["gen_safe_block_ids"],
                batch["gen_safe_real_lengths"],
                t_override=response_t_value,
                seed=_mask_seed(batch, batch_idx, 7331, t_idx),
                target_block_ids=(RESPONSE_BLOCK_ID,),
            )

        all_safe_masks.append(
            _relative_to_absolute_mask(
                batch["gen_safe_block_ids"],
                batch["gen_safe_real_lengths"],
                relative_indices,
                target_block_ids=target_block_ids,
            )
        )
        all_unsafe_masks.append(
            _relative_to_absolute_mask(
                batch["gen_unsafe_block_ids"],
                batch["gen_unsafe_real_lengths"],
                relative_indices,
                target_block_ids=target_block_ids,
            )
        )

    num_t_values = len(active_t_values)

    def _repeat_batch_tensor(x: torch.Tensor) -> torch.Tensor:
        return torch.cat([x] * num_t_values, dim=0)

    repeated_safe_ids = _repeat_batch_tensor(batch["gen_safe_input_ids"])
    repeated_unsafe_ids = _repeat_batch_tensor(batch["gen_unsafe_input_ids"])
    repeated_safe_lengths = _repeat_batch_tensor(batch["gen_safe_real_lengths"])
    repeated_unsafe_lengths = _repeat_batch_tensor(batch["gen_unsafe_real_lengths"])
    tiled_safe_mask = torch.cat(all_safe_masks, dim=0)
    tiled_unsafe_mask = torch.cat(all_unsafe_masks, dim=0)

    if has_prompt_only:
        full_response_safe_mask = _full_target_mask(
            batch["gen_safe_block_ids"],
            batch["gen_safe_real_lengths"],
            target_block_ids=(RESPONSE_BLOCK_ID,),
        )
        full_response_unsafe_mask = _full_target_mask(
            batch["gen_unsafe_block_ids"],
            batch["gen_unsafe_real_lengths"],
            target_block_ids=(RESPONSE_BLOCK_ID,),
        )
        scope_safe_mask = tiled_safe_mask | _repeat_batch_tensor(full_response_safe_mask)
        scope_unsafe_mask = tiled_unsafe_mask | _repeat_batch_tensor(full_response_unsafe_mask)
    else:
        scope_safe_mask = tiled_safe_mask
        scope_unsafe_mask = tiled_unsafe_mask

    safe_ids, safe_logits, unsafe_ids, unsafe_logits = _compute_paired_masked_token_outputs(
        model,
        repeated_safe_ids,
        scope_safe_mask,
        repeated_safe_lengths,
        repeated_unsafe_ids,
        scope_unsafe_mask,
        repeated_unsafe_lengths,
        pad_id,
    )
    max_len = safe_ids.shape[1]
    tiled_safe_mask = _pad_mask_to_match_length(tiled_safe_mask, max_len)
    tiled_unsafe_mask = _pad_mask_to_match_length(tiled_unsafe_mask, max_len)
    _, safe_logp = _compute_masked_token_losses_from_logits(
        safe_logits,
        safe_ids,
        tiled_safe_mask,
    )
    _, unsafe_logp = _compute_masked_token_losses_from_logits(
        unsafe_logits,
        unsafe_ids,
        tiled_unsafe_mask,
    )
    return (unsafe_logp - safe_logp).view(num_t_values, batch_size).mean(dim=0).detach()


def _select_eval_labels_and_valid_mask(batch, *, masked_scope: bool):
    """Choose labels/masks to match prompt-only or response-conditioned scoring."""
    if masked_scope:
        has_response = batch.get("has_response")
        prompt_labels = batch.get("prompt_is_unsafe")
        prompt_valid = batch.get("prompt_label_mask")
        response_labels = batch.get("response_is_unsafe")
        response_valid = batch.get("response_label_mask")
        if (
            has_response is not None
            and prompt_labels is not None
            and prompt_valid is not None
            and response_labels is not None
            and response_valid is not None
        ):
            has_response_mask = has_response > 0.5
            labels = torch.where(has_response_mask, response_labels, prompt_labels)
            valid = torch.where(has_response_mask, response_valid > 0.5, prompt_valid > 0.5)
            return labels, valid

    labels = batch["is_unsafe"]
    valid = batch.get("conv_label_mask", torch.ones_like(labels)) > 0.5
    return labels, valid


@torch.no_grad()
def collect_eval_scores(
    model,
    eval_loader,
    device,
    t_values=(0.15, 0.15, 0.15),
    pad_id: int = 0,
    tile_t_values: bool = True,
    masked_scope: bool = True,
    response_mask_scale: float = 1.0,
    return_indices: bool = False,
    adaptive_t_repeats: int | None = None,
    adaptive_t_sigma_over_c: float = 1.0,
    adaptive_t_min: float = 1e-3,
    adaptive_t_max: float = 0.5,
    progress: bool = True,
):
    """Collect masked-scope unsafe-minus-safe scores for an eval loader."""
    if not masked_scope:
        raise ValueError("This release supports masked_scope=True scoring only.")

    model.eval()
    apply_eval_inference_fastpaths(model)
    y_true = []
    scores = []
    example_indices = []

    for batch_idx, batch in enumerate(tqdm(eval_loader, desc="LLaDA eval", disable=not progress)):
        batch = {key: value.to(device, non_blocking=True) for key, value in batch.items()}
        if tile_t_values:
            try:
                batch_scores = _score_batch_masked_scope_tiled(
                    model,
                    batch,
                    batch_idx,
                    t_values,
                    pad_id,
                    response_mask_scale=response_mask_scale,
                    adaptive_t_repeats=adaptive_t_repeats,
                    adaptive_t_sigma_over_c=adaptive_t_sigma_over_c,
                    adaptive_t_min=adaptive_t_min,
                    adaptive_t_max=adaptive_t_max,
                )
            except torch.cuda.OutOfMemoryError:
                torch.cuda.empty_cache()
                batch_scores = _score_batch_masked_scope_sequential(
                    model,
                    batch,
                    batch_idx,
                    t_values,
                    pad_id,
                    response_mask_scale=response_mask_scale,
                    adaptive_t_repeats=adaptive_t_repeats,
                    adaptive_t_sigma_over_c=adaptive_t_sigma_over_c,
                    adaptive_t_min=adaptive_t_min,
                    adaptive_t_max=adaptive_t_max,
                )
        else:
            batch_scores = _score_batch_masked_scope_sequential(
                model,
                batch,
                batch_idx,
                t_values,
                pad_id,
                response_mask_scale=response_mask_scale,
                adaptive_t_repeats=adaptive_t_repeats,
                adaptive_t_sigma_over_c=adaptive_t_sigma_over_c,
                adaptive_t_min=adaptive_t_min,
                adaptive_t_max=adaptive_t_max,
            )

        labels, valid = _select_eval_labels_and_valid_mask(batch, masked_scope=masked_scope)
        if not valid.any():
            continue
        scores.extend(batch_scores[valid].detach().cpu().tolist())
        y_true.extend(labels[valid].detach().cpu().tolist())
        if return_indices:
            example_indices.extend(batch["example_index"][valid].detach().cpu().tolist())

    if return_indices:
        return y_true, scores, example_indices
    return y_true, scores