"""Output heads: answer selection, relation, support, contradiction, confidence.

The selection head is the substantive addition.  Previously ``reason`` returned
*every* path the verifier accepted, so the answer set was "everything reachable"
and exact-match against a specific gold set was unattainable by construction.
Selection is the model's job: :class:`PathSelectionHead` emits a per-path logit
and a learned decision threshold, so the model commits to a subset and is
scored on that commitment.
"""

from __future__ import annotations

from dataclasses import dataclass

import torch
from torch import Tensor, nn


@dataclass(frozen=True)
class SelectionOutput:
    """Per-path selection logits and the resulting boolean decision."""

    logits: Tensor
    keep: Tensor
    threshold: Tensor


class PathSelectionHead(nn.Module):
    """Score candidate paths and decide which ones support the answer.

    The threshold is **query-dependent**: a global bias plus an offset predicted
    from the query and the pooled candidate evidence.  A single global scalar
    cannot express both "this query has one answer" and "this one has seven", and
    the measured symptom was a threshold that drifted without becoming
    discriminative -- selection precision flat near 0.20 while path recall fell
    from 0.761 to 0.368, and a wider beam scoring *worse* than greedy.
    """

    def __init__(self, hidden_dim: int, dropout: float = 0.1) -> None:
        super().__init__()
        self.score = nn.Sequential(
            nn.Linear(hidden_dim * 3, hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.GELU(),
            nn.Linear(hidden_dim // 2, 1),
        )
        self.threshold = nn.Parameter(torch.zeros(1))
        self.threshold_offset = nn.Sequential(
            nn.Linear(hidden_dim * 2, hidden_dim // 2),
            nn.GELU(),
            nn.Linear(hidden_dim // 2, 1),
        )

    def forward(
        self, reasoning_states: Tensor, path_states: Tensor, query_state: Tensor
    ) -> SelectionOutput:
        if query_state.ndim == 1:
            query_state = query_state.unsqueeze(0)
        if query_state.shape[0] == 1 and reasoning_states.shape[0] != 1:
            query_state = query_state.expand(reasoning_states.shape[0], -1)
        features = torch.cat((reasoning_states, path_states, query_state), dim=-1)
        logits = self.score(features).squeeze(-1)
        # One threshold per query, from the query plus the pooled evidence.
        context = torch.cat(
            (query_state[:1], path_states.mean(0, keepdim=True)), dim=-1
        )
        threshold = self.threshold + self.threshold_offset(context).reshape(())
        return SelectionOutput(logits, logits > threshold, threshold)


class NextHopHead(nn.Module):
    """Score each candidate transition for beam expansion (PLAN section 14)."""

    def __init__(self, hidden_dim: int, dropout: float = 0.1) -> None:
        super().__init__()
        self.score = nn.Sequential(
            nn.Linear(hidden_dim * 4, hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, 1),
        )

    def forward(
        self,
        reasoning_states: Tensor,
        source_states: Tensor,
        relation_states: Tensor,
        target_states: Tensor,
    ) -> Tensor:
        features = torch.cat(
            (reasoning_states, source_states, relation_states, target_states), dim=-1
        )
        return self.score(features).squeeze(-1)


class EntityAnswerHead(nn.Module):
    """Rank candidate answer entities against the reasoning state."""

    def __init__(self, hidden_dim: int) -> None:
        super().__init__()
        self.project = nn.Linear(hidden_dim, hidden_dim)
        self.scale = nn.Parameter(torch.ones(1))

    def forward(self, reasoning_state: Tensor, entity_states: Tensor) -> Tensor:
        if reasoning_state.ndim == 1:
            reasoning_state = reasoning_state.unsqueeze(0)
        projected = self.project(reasoning_state)
        return self.scale * (entity_states @ projected.transpose(-2, -1)).squeeze(-1)


class RelationAnswerHead(nn.Module):
    """Predict the relation that answers a relation-mode query."""

    def __init__(self, hidden_dim: int, num_relations: int) -> None:
        super().__init__()
        self.score = nn.Sequential(
            nn.Linear(hidden_dim * 2, hidden_dim),
            nn.GELU(),
            nn.Linear(hidden_dim, num_relations),
        )

    def forward(self, reasoning_state: Tensor, relation_context: Tensor) -> Tensor:
        return self.score(torch.cat((reasoning_state, relation_context), dim=-1))


class VerificationStack(nn.Module):
    """Shared neural verification trunk consumed by the support-side heads.

    PLAN section 20 describes neural verification as consuming the query, the PAT
    path representation, the graph representation, the candidate answer and the
    constraint state.  That is a real inference problem, not a linear readout, so
    it gets a residual MLP trunk of configurable depth and width whose output is
    shared by the support, contradiction, count and confidence heads.

    Sharing one trunk is also what makes the heads bucket reach its section 6
    budget honestly: the capacity does inference work rather than padding a
    per-head linear layer to hit a number.
    """

    def __init__(
        self,
        hidden_dim: int,
        depth: int = 2,
        width_multiplier: int = 4,
        dropout: float = 0.1,
        num_signals: int = 0,
    ) -> None:
        super().__init__()
        if depth <= 0 or width_multiplier <= 0:
            raise ValueError("verification depth and width_multiplier must be positive")
        if num_signals < 0:
            raise ValueError("num_signals must not be negative")
        inner = hidden_dim * width_multiplier
        self.num_signals = num_signals
        self.input_projection = nn.Linear(hidden_dim * 3, hidden_dim)
        # Scalar evidence summaries, projected and added rather than concatenated so
        # the trunk width is unchanged.  These exist because the pooled evidence mean
        # this trunk otherwise receives is built from *different distributions* at
        # training and inference -- gold positives plus sampled negatives during
        # training, the beam's own candidates at inference -- and the support head
        # learned a rule that does not transfer: it abstained on 25% of answerable
        # examples and only 12.5% of unanswerable ones, worse than chance in the
        # direction that matters.  A signal computed identically in both regimes fixes
        # the mismatch rather than reweighting around it.
        self.signal_projection = (
            nn.Linear(num_signals, hidden_dim) if num_signals else None
        )
        self.blocks = nn.ModuleList(
            [
                nn.Sequential(
                    nn.LayerNorm(hidden_dim),
                    nn.Linear(hidden_dim, inner),
                    nn.GELU(),
                    nn.Dropout(dropout),
                    nn.Linear(inner, hidden_dim),
                )
                for _ in range(depth)
            ]
        )
        self.final_norm = nn.LayerNorm(hidden_dim)

    def forward(
        self,
        reasoning_state: Tensor,
        evidence_state: Tensor,
        query_state: Tensor,
        signals: Tensor | None = None,
    ) -> Tensor:
        states = self.input_projection(
            torch.cat((reasoning_state, evidence_state, query_state), dim=-1)
        )
        if self.signal_projection is not None:
            if signals is None:
                raise ValueError(
                    f"this verification stack expects {self.num_signals} evidence "
                    "signals but none were supplied"
                )
            if signals.ndim == 1:
                signals = signals.unsqueeze(0)
            if signals.shape[-1] != self.num_signals:
                raise ValueError(
                    f"expected {self.num_signals} evidence signals, "
                    f"got {signals.shape[-1]}"
                )
            states = states + self.signal_projection(signals.to(states.dtype))
        for block in self.blocks:
            states = states + block(states)
        return self.final_norm(states)


class SupportHead(nn.Module):
    """Estimate ``P(answer supported | query, evidence, graph)``."""

    def __init__(self, hidden_dim: int) -> None:
        super().__init__()
        self.score = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.GELU(),
            nn.Linear(hidden_dim // 2, 1),
        )

    def forward(self, verified_state: Tensor) -> Tensor:
        return self.score(verified_state).squeeze(-1)


class ContradictionHead(nn.Module):
    """Four-way status classifier over the PLAN section 19 output states."""

    def __init__(self, hidden_dim: int) -> None:
        super().__init__()
        self.score = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.GELU(),
            nn.Linear(hidden_dim // 2, 4),
        )

    def forward(self, verified_state: Tensor) -> Tensor:
        return self.score(verified_state)


class CountHead(nn.Module):
    """Predict an aggregate count for count/aggregate answer modes."""

    def __init__(self, hidden_dim: int, max_count: int = 32) -> None:
        super().__init__()
        self.max_count = max_count
        self.score = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.GELU(),
            nn.Linear(hidden_dim // 2, max_count + 1),
        )

    def forward(self, verified_state: Tensor) -> Tensor:
        return self.score(verified_state)


class ConfidenceHead(nn.Module):
    """Emit a calibrated-in-training confidence logit (PLAN section 47)."""

    def __init__(self, hidden_dim: int) -> None:
        super().__init__()
        self.score = nn.Sequential(
            nn.Linear(hidden_dim + 2, hidden_dim // 2),
            nn.GELU(),
            nn.Linear(hidden_dim // 2, 1),
        )

    def forward(
        self,
        verified_state: Tensor,
        support_logit: Tensor,
        selected_fraction: Tensor,
    ) -> Tensor:
        features = torch.cat(
            (
                verified_state.reshape(1, -1),
                support_logit.reshape(1, -1),
                selected_fraction.reshape(1, -1),
            ),
            dim=-1,
        )
        return self.score(features).squeeze(-1)