"""Hidden-state-only completion head for explicit-step pointer recurrence."""

import torch
from torch import Tensor, nn


class CompletionHead(nn.Module):
    """One stop logit per answer-position recurrent state [B,H].

    No numeric loop index, parsed requested depth, decoded symbol, or target is
    accepted by this interface. The prompt can encode its requested Steps.
    """

    def __init__(self, width: int, intermediate: int = 128) -> None:
        super().__init__()
        if width < 1 or intermediate < 1:
            raise ValueError("Completion-head widths must be positive")
        self.width = width
        self.intermediate = intermediate
        self.network = nn.Sequential(nn.LayerNorm(width), nn.Linear(width, intermediate),
                                     nn.GELU(), nn.Linear(intermediate, 1))

    def forward(self, answer_state: Tensor) -> Tensor:
        if answer_state.ndim != 2 or answer_state.shape[-1] != self.width:
            raise ValueError("Completion head requires answer states [batch, hidden]")
        return self.network(answer_state).squeeze(-1)