"""Segment-aware interactions between policy, evidence, and claim vectors."""
import torch
from torch import Tensor, nn


class TriVectorInteractionHead(nn.Module):
    """Classify policy/evidence/claim vectors using seven ordered feature blocks."""

    def __init__(self, hidden_dim: int, num_classes: int = 4, dropout: float = .1):
        super().__init__()
        for name, value in (("hidden_dim", hidden_dim), ("num_classes", num_classes)):
            if isinstance(value, bool) or not isinstance(value, int):
                raise TypeError(f"{name} must be an integer")
            if value <= 0:
                raise ValueError(f"{name} must be positive")
        self.hidden_dim = hidden_dim
        self.classifier = nn.Sequential(
            nn.Dropout(dropout),
            nn.Linear(7 * hidden_dim, 2 * hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(2 * hidden_dim, num_classes),
        )

    def construct_features(self, h_P: Tensor, h_E: Tensor, h_C: Tensor) -> Tensor:
        """Return [P, E, C, |P-E|, P*E, |E-C|, E*C] with shape [B, 7H]."""
        for vector in (h_P, h_E, h_C):
            if not isinstance(vector, Tensor) or not vector.is_floating_point():
                raise TypeError("segment vectors must be floating-point tensors")
            if vector.ndim != 2 or vector.shape[-1] != self.hidden_dim:
                raise ValueError("segment vectors must have shape [B, hidden_dim]")
            if vector.shape != h_P.shape or vector.dtype != h_P.dtype or vector.device != h_P.device:
                raise ValueError("segment vectors must share shape, dtype, and device")
        return torch.cat((h_P, h_E, h_C, (h_P - h_E).abs(), h_P * h_E,
                          (h_E - h_C).abs(), h_E * h_C), dim=-1)

    def forward(self, h_P: Tensor, h_E: Tensor, h_C: Tensor) -> Tensor:
        return self.classifier(self.construct_features(h_P, h_E, h_C))

    def pool_segments(self, hidden_states: Tensor, segments: Tensor,
                      attention_mask: Tensor) -> tuple[Tensor, Tensor, Tensor]:
        """Mean-pool [B,L,H] states into policy (1), evidence (2), and claim (3).

        Segment 0 denotes delimiters/padding. Only tokens with a binary attention
        mask of 1 contribute. Every row must contain an unmasked token in each
        segment; empty segments raise ValueError rather than producing NaNs.
        Float16/bfloat16 sums and division use float32 to avoid overflow and
        count rounding; pooled vectors retain the input dtype and device.
        """
        if not isinstance(hidden_states, Tensor) or not hidden_states.is_floating_point():
            raise TypeError("hidden_states must be a floating-point tensor")
        if hidden_states.ndim != 3 or hidden_states.shape[-1] != self.hidden_dim:
            raise ValueError("hidden_states must have shape [B, L, hidden_dim]")
        for name, tensor in (("segments", segments), ("attention_mask", attention_mask)):
            if not isinstance(tensor, Tensor):
                raise TypeError(f"{name} must be a tensor")
            if tensor.shape != hidden_states.shape[:2]:
                raise ValueError(f"{name} must have shape [B, L]")
            if tensor.device != hidden_states.device:
                raise ValueError(f"{name} must share the hidden_states device")
        if segments.dtype not in (torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64):
            raise TypeError("segments must contain integer IDs")
        if ((segments < 0) | (segments > 3)).any():
            raise ValueError("segment IDs must be 0, 1, 2, or 3")
        if not ((attention_mask == 0) | (attention_mask == 1)).all():
            raise ValueError("attention_mask must be binary")
        accumulation_dtype = (torch.float32 if hidden_states.dtype in (torch.float16, torch.bfloat16)
                              else hidden_states.dtype)
        pooled = []
        for segment_id in (1, 2, 3):
            selected = (segments == segment_id) & attention_mask.bool()
            counts = selected.sum(dim=1, keepdim=True)
            if (counts == 0).any():
                raise ValueError(f"empty segment {segment_id} in at least one batch row")
            masked = hidden_states.masked_fill(~selected.unsqueeze(-1), 0)
            mean = masked.sum(dim=1, dtype=accumulation_dtype) / counts.to(accumulation_dtype)
            pooled.append(mean.to(hidden_states.dtype))
        return tuple(pooled)

    def forward_from_hidden(self, hidden_states: Tensor, segments: Tensor,
                            attention_mask: Tensor) -> Tensor:
        """Pool segment states and return [B, num_classes] classification logits."""
        return self(*self.pool_segments(hidden_states, segments, attention_mask))