from __future__ import annotations

from dataclasses import dataclass

import torch
import torch.nn.functional as F
from torch import nn


@dataclass(frozen=True)
class AttentionOutput:
    values: torch.Tensor
    pattern: torch.Tensor | None = None


class CausalSelfAttention(nn.Module):
    def __init__(self, d_model: int, n_heads: int, bias: bool = False) -> None:
        super().__init__()
        if d_model % n_heads != 0:
            raise ValueError("d_model must be divisible by n_heads.")

        self.d_model = d_model
        self.n_heads = n_heads
        self.d_head = d_model // n_heads

        self.W_Q = nn.Linear(d_model, d_model, bias=bias)
        self.W_K = nn.Linear(d_model, d_model, bias=bias)
        self.W_V = nn.Linear(d_model, d_model, bias=bias)
        self.W_O = nn.Linear(d_model, d_model, bias=bias)
        self.register_buffer("_causal_mask", torch.empty(0, 0, dtype=torch.bool), persistent=False)

    def forward(self, x: torch.Tensor, return_pattern: bool = False) -> AttentionOutput:
        batch, seq_len, _ = x.shape
        q = self._split_heads(self.W_Q(x))
        k = self._split_heads(self.W_K(x))
        v = self._split_heads(self.W_V(x))

        scores = torch.matmul(q, k.transpose(-1, -2)) / self.d_head
        if self._causal_mask.shape[0] < seq_len or self._causal_mask.device != x.device:
            self._causal_mask = torch.triu(
                torch.ones(seq_len, seq_len, dtype=torch.bool, device=x.device),
                diagonal=1,
            )
        causal_mask = self._causal_mask[:seq_len, :seq_len]
        scores = scores.masked_fill(causal_mask, float("-inf"))
        pattern = F.softmax(scores, dim=-1)

        attended = torch.matmul(pattern, v) / 3.0
        values = self.W_O(self._merge_heads(attended, batch, seq_len))
        return AttentionOutput(values=values, pattern=pattern if return_pattern else None)

    def _split_heads(self, x: torch.Tensor) -> torch.Tensor:
        batch, seq_len, _ = x.shape
        return x.view(batch, seq_len, self.n_heads, self.d_head).transpose(1, 2)

    def _merge_heads(self, x: torch.Tensor, batch: int, seq_len: int) -> torch.Tensor:
        return x.transpose(1, 2).contiguous().view(batch, seq_len, self.d_model)