"""Dependency parsing head using biaffine attention."""

import torch.nn as nn
from torch import Tensor

from .biaffine import Biaffine


class DependencyHead(nn.Module):
    """Biaffine dependency parsing head for arc and relation prediction."""

    def __init__(
        self,
        hidden_size: int,
        arc_hidden_dim: int = 256,
        rel_hidden_dim: int = 128,
        num_deprels: int = 50,
        dropout: float = 0.1,
    ):
        """Initialize DependencyHead.

        Args:
            hidden_size: Size of input hidden states.
            arc_hidden_dim: Hidden dimension for arc prediction MLPs.
            rel_hidden_dim: Hidden dimension for relation prediction MLPs.
            num_deprels: Number of dependency relation types.
            dropout: Dropout probability.
        """
        super().__init__()
        self.num_deprels = num_deprels

        # Arc prediction MLPs
        self.arc_head_mlp = nn.Sequential(
            nn.Linear(hidden_size, arc_hidden_dim),
            nn.LayerNorm(arc_hidden_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
        )
        self.arc_dep_mlp = nn.Sequential(
            nn.Linear(hidden_size, arc_hidden_dim),
            nn.LayerNorm(arc_hidden_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
        )

        # Relation prediction MLPs
        self.rel_head_mlp = nn.Sequential(
            nn.Linear(hidden_size, rel_hidden_dim),
            nn.LayerNorm(rel_hidden_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
        )
        self.rel_dep_mlp = nn.Sequential(
            nn.Linear(hidden_size, rel_hidden_dim),
            nn.LayerNorm(rel_hidden_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
        )

        # Biaffine scorers
        self.arc_biaffine = Biaffine(arc_hidden_dim, out_features=1)
        self.rel_biaffine = Biaffine(rel_hidden_dim, out_features=num_deprels)

    def forward(self, word_repr: Tensor) -> tuple[Tensor, Tensor]:
        """Compute arc and relation scores.

        Args:
            word_repr: Word-level representations [batch, seq_len, hidden_size]

        Returns:
            arc_logits: Arc scores [batch, seq_len, seq_len]
                        arc_logits[b, i, j] = score of j being head of i
            rel_logits: Relation scores [batch, seq_len, seq_len, num_deprels]
        """
        # Arc scoring
        arc_head = self.arc_head_mlp(word_repr)  # [batch, seq, arc_hidden]
        arc_dep = self.arc_dep_mlp(word_repr)  # [batch, seq, arc_hidden]
        arc_logits = self.arc_biaffine(arc_dep, arc_head)  # [batch, dep, head]

        # Relation scoring
        rel_head = self.rel_head_mlp(word_repr)  # [batch, seq, rel_hidden]
        rel_dep = self.rel_dep_mlp(word_repr)  # [batch, seq, rel_hidden]
        rel_logits = self.rel_biaffine(rel_dep, rel_head)  # [batch, dep, head, n_rels]

        # Embedding: concat head and dep projections (like old COMBO)
        import torch
        embedding = torch.cat([rel_head, rel_dep], dim=-1)  # [batch, seq, 2*rel_hidden]

        return arc_logits, rel_logits, embedding
