"""Biaffine attention module for dependency parsing.

Based on Dozat and Manning (2017): "Deep Biaffine Attention for Neural Dependency Parsing"
"""

import torch
import torch.nn as nn
from torch import Tensor


class Biaffine(nn.Module):
    """Biaffine attention layer for scoring pairs of representations.

    Computes: s[i,j] = x[i]^T W y[j] + b

    For dependency parsing:
    - x represents dependent representations
    - y represents head representations
    - s[i,j] is the score of word j being the head of word i
    """

    def __init__(
        self,
        in_features: int,
        out_features: int = 1,
        bias_x: bool = True,
        bias_y: bool = True,
    ):
        """Initialize Biaffine layer.

        Args:
            in_features: Size of input features.
            out_features: Size of output features (1 for arc scoring, n_rels for relation scoring).
            bias_x: Whether to add bias to x (dependent).
            bias_y: Whether to add bias to y (head).
        """
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.bias_x = bias_x
        self.bias_y = bias_y

        # Weight matrix: (out_features, in_features + bias_x, in_features + bias_y)
        self.weight = nn.Parameter(
            torch.zeros(out_features, in_features + int(bias_x), in_features + int(bias_y))
        )
        self.reset_parameters()

    def reset_parameters(self) -> None:
        """Initialize parameters with Xavier uniform."""
        nn.init.xavier_uniform_(self.weight)

    def forward(self, x: Tensor, y: Tensor) -> Tensor:
        """Compute biaffine attention scores.

        Args:
            x: Dependent representations [batch, seq_len, in_features]
            y: Head representations [batch, seq_len, in_features]

        Returns:
            scores: Attention scores
                - If out_features == 1: [batch, seq_len, seq_len]
                - Otherwise: [batch, seq_len, seq_len, out_features]
        """
        # Add bias terms
        if self.bias_x:
            x = torch.cat([x, torch.ones_like(x[..., :1])], dim=-1)
        if self.bias_y:
            y = torch.cat([y, torch.ones_like(y[..., :1])], dim=-1)

        # Efficient biaffine computation using einsum
        # x: [batch, dep_len, in_features+1]
        # y: [batch, head_len, in_features+1]
        # weight: [out_features, in_features+1, in_features+1]
        # Result: [batch, dep_len, head_len, out_features]
        scores = torch.einsum("bxi,oij,byj->bxyo", x, self.weight, y)

        # Squeeze if out_features == 1
        if self.out_features == 1:
            scores = scores.squeeze(-1)

        return scores

    def extra_repr(self) -> str:
        return (
            f"in_features={self.in_features}, "
            f"out_features={self.out_features}, "
            f"bias_x={self.bias_x}, "
            f"bias_y={self.bias_y}"
        )
