import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from typing import Optional, Tuple

def create_look_ahead_mask(size: int) -> torch.Tensor:
    return torch.triu(torch.ones(size, size, dtype=torch.bool), diagonal=1).unsqueeze(0)

def scaled_dot_product_attention(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    mask: Optional[torch.Tensor] = None,
    dropout: Optional[nn.Dropout] = None
) -> Tuple[torch.Tensor, torch.Tensor]:
    d_k = query.size(-1)
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)

    if mask is not None:
        if mask.dim() == 3:
            mask = mask.unsqueeze(1)
        scores = scores.masked_fill(mask, float('-inf'))

    p_attn = F.softmax(scores, dim=-1)

    if dropout is not None:
        p_attn = dropout(p_attn)

    output = torch.matmul(p_attn, value)
    return output, p_attn

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model: int, num_heads: int, dropout_rate: float = 0.1):
        super().__init__()
        assert d_model % num_heads == 0

        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads

        self.w_q = nn.Linear(d_model, d_model)
        self.w_k = nn.Linear(d_model, d_model)
        self.w_v = nn.Linear(d_model, d_model)
        self.w_o = nn.Linear(d_model, d_model)

        self.attn_dropout = nn.Dropout(dropout_rate)

    def forward(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        mask: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        B = query.size(0)

        def project(x, layer):
            return layer(x).view(B, -1, self.num_heads, self.d_k).transpose(1, 2)

        query = project(query, self.w_q)
        key   = project(key,   self.w_k)
        value = project(value, self.w_v)

        x, _ = scaled_dot_product_attention(
            query, key, value, mask, dropout=self.attn_dropout
        )

        x = x.transpose(1, 2).contiguous().view(B, -1, self.d_model)
        return self.w_o(x)

if __name__ == '__main__':
    D_MODEL = 512
    NUM_HEADS = 8
    BATCH_SIZE = 4
    SEQ_LEN = 10
    DROPOUT = 0.1

    mha = MultiHeadAttention(d_model=D_MODEL, num_heads=NUM_HEADS, dropout_rate=DROPOUT)
    dummy_input = torch.randn(BATCH_SIZE, SEQ_LEN, D_MODEL)
    
    look_ahead_mask = create_look_ahead_mask(SEQ_LEN)
    decoder_self_mask = look_ahead_mask
    
    output = mha(query=dummy_input, key=dummy_input, value=dummy_input, mask=decoder_self_mask)
    
    print(f"入力形状 (Q=K=V): {dummy_input.shape}")
    print(f"先読みマスク形状: {decoder_self_mask.shape}")
    print(f"マスク行列（最初の4x4）:\n{decoder_self_mask[0, :4, :4].int()}")
    print(f"出力形状: {output.shape}")

    expected_shape = (BATCH_SIZE, SEQ_LEN, D_MODEL)
    assert output.shape == expected_shape

# Attention Is All You Need
# 3.2.3 Attention in the Decoder
# MaskedDecoderSelfAttention.py
# 第2部 第10章