"""Utilities for projecting intermediate hidden states into vocabulary logits."""

from __future__ import annotations

from collections.abc import Iterable

from torch import Tensor

from ..models.adapters import ModelAdapter


def apply_final_norm(adapter: ModelAdapter, hidden_states: Tensor) -> Tensor:
    """Apply the model's final norm if one exists."""

    final_norm = adapter.get_final_norm()
    if final_norm is None:
        return hidden_states
    return final_norm(hidden_states)


def project_hidden_to_logits(
    adapter: ModelAdapter,
    hidden_states: Tensor,
    *,
    apply_norm: bool = True,
) -> Tensor:
    """Project hidden states into vocabulary logits using the adapter's LM head."""

    if apply_norm:
        hidden_states = apply_final_norm(adapter, hidden_states)
    lm_head = adapter.get_lm_head()
    return lm_head(hidden_states)


def project_layer_logits(
    adapter: ModelAdapter,
    hidden_states_by_layer: dict[int, Tensor],
    *,
    source_site: str = "block_output",
) -> dict[int, Tensor]:
    """Project a set of per-layer hidden states into per-layer logits."""

    apply_norm = source_site != "final_norm"
    return {
        layer_idx: project_hidden_to_logits(adapter, hidden_states, apply_norm=apply_norm)
        for layer_idx, hidden_states in hidden_states_by_layer.items()
    }


def select_layers(
    hidden_states_by_layer: dict[int, Tensor],
    layers: Iterable[int] | None = None,
) -> dict[int, Tensor]:
    """Return a filtered copy of per-layer hidden states."""

    if layers is None:
        return dict(hidden_states_by_layer)
    selected = set(layers)
    return {
        layer_idx: hidden_states
        for layer_idx, hidden_states in hidden_states_by_layer.items()
        if layer_idx in selected
    }

