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

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:
        scores = scores.masked_fill(mask == 0, -1e9)

    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 EncoderDecoderAttention(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.d_v = 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,
        memory: torch.Tensor,
        mask: Optional[torch.Tensor] = None
    ) -> torch.Tensor:

        B, tgt_len, _ = query.shape
        _, src_len, _ = memory.shape

        q = self.w_q(query).view(B, tgt_len, self.num_heads, self.d_k).transpose(1, 2)
        k = self.w_k(memory).view(B, src_len, self.num_heads, self.d_k).transpose(1, 2)
        v = self.w_v(memory).view(B, src_len, self.num_heads, self.d_v).transpose(1, 2)

        if mask is not None:
            mask = mask.unsqueeze(1).expand(B, self.num_heads, tgt_len, src_len)

        x, _ = scaled_dot_product_attention(
            q, k, v, mask=mask, dropout=self.attn_dropout
        )

        x = x.transpose(1, 2).contiguous().view(B, tgt_len, self.d_model)

        return self.w_o(x)

if __name__ == '__main__':
    D_MODEL = 512
    NUM_HEADS = 8
    BATCH_SIZE = 4
    SRC_SEQ_LEN = 100
    TGT_SEQ_LEN = 50
    DROPOUT = 0.1
    PAD_IDX = 0

    encoder_decoder_attn = EncoderDecoderAttention(
        d_model=D_MODEL,
        num_heads=NUM_HEADS,
        dropout_rate=DROPOUT
    )

    dummy_query = torch.randn(BATCH_SIZE, TGT_SEQ_LEN, D_MODEL)
    dummy_memory = torch.randn(BATCH_SIZE, SRC_SEQ_LEN, D_MODEL)

    dummy_mask = torch.ones(BATCH_SIZE, 1, SRC_SEQ_LEN)
    dummy_mask[:, :, -20:] = 0

    output = encoder_decoder_attn(
        query=dummy_query,
        memory=dummy_memory,
        mask=dummy_mask
    )

    print(f"クエリ形状 (デコーダ入力): {dummy_query.shape}")
    print(f"メモリ形状 (エンコーダ出力): {dummy_memory.shape}")
    print(f"マスク形状 (input to forward): {dummy_mask.shape}")
    print(f"出力形状 (output): {output.shape}")

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

# Attention Is All You Need
# 3.2.3 Attention in the Decoder – Encoder-Decoder Attention
# EncoderDecoderAttention.py
# 第2部 第11章