"""Attention mechanism implementations for sequence models."""

from __future__ import annotations

from typing import Any

import numpy as np


def SelfAttention(hidden_dim: int, **kwargs: Any) -> Any:
    """Single-head self-attention layer: Q, K, V projections and scaled dot-product. Returns nn.Module. Requires torch."""
    try:
        import torch
        from torch import nn
    except ImportError as e:
        raise ImportError("SelfAttention requires PyTorch. Install with: pip install torch") from e

    class _SelfAttention(nn.Module):
        def __init__(self, dim: int):
            super().__init__()
            self.dim = dim
            self.q = nn.Linear(dim, dim, **kwargs)
            self.k = nn.Linear(dim, dim, **kwargs)
            self.v = nn.Linear(dim, dim, **kwargs)

        def forward(self, x: Any) -> tuple[Any, Any]:
            # x: (batch, seq_len, dim)
            q, k, v = self.q(x), self.k(x), self.v(x)
            scores = torch.matmul(q, k.transpose(-2, -1)) / (self.dim ** 0.5)
            attn = torch.softmax(scores, dim=-1)
            out = torch.matmul(attn, v)
            return out, attn

    return _SelfAttention(hidden_dim)


def SentimentLSTMWithAttention(
    vocab_size: int,
    embedding_dim: int = 128,
    hidden_dim: int = 256,
    num_layers: int = 2,
    num_classes: int = 3,
    dropout: float = 0.3,
    padding_idx: int = 0,
) -> Any:
    """LSTM followed by self-attention and linear layer for classification. Returns nn.Module. Requires torch."""
    try:
        import torch
        from torch import nn
    except ImportError as e:
        raise ImportError("SentimentLSTMWithAttention requires PyTorch. Install with: pip install torch") from e

    class _LSTMAttn(nn.Module):
        def __init__(self):
            super().__init__()
            self.embed = nn.Embedding(vocab_size, embedding_dim, padding_idx=padding_idx)
            self.lstm = nn.LSTM(
                embedding_dim,
                hidden_dim,
                num_layers=num_layers,
                batch_first=True,
                dropout=dropout if num_layers > 1 else 0,
            )
            self.attention = SelfAttention(hidden_dim)
            self.fc = nn.Linear(hidden_dim, num_classes)
            self.dropout = nn.Dropout(dropout)

        def forward(self, x: Any) -> Any:
            emb = self.embed(x)
            lstm_out, _ = self.lstm(emb)
            attn_out, _ = self.attention(lstm_out)
            last = attn_out[:, -1, :]
            return self.fc(self.dropout(last))

    return _LSTMAttn()
