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.Module] = 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 scores.dim() == 4 and mask.dim() == 3:
            mask = mask.unsqueeze(1)
        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 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: torch.Tensor, layer: nn.Linear) -> torch.Tensor:
            return layer(x).view(B, -1, self.num_heads, self.d_k).transpose(1, 2)

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

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

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

class PositionwiseFeedForward(nn.Module):
    def __init__(self, d_model: int, d_ff: int, dropout_rate: float = 0.1):
        super().__init__()
        self.w_1 = nn.Linear(d_model, d_ff)
        self.w_2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout_rate)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.dropout(self.w_2(F.relu(self.w_1(x))))

class EncoderLayer(nn.Module):
    def __init__(self, d_model: int, num_heads: int, d_ff: int, dropout_rate: float):
        super().__init__()
        self.self_attn = MultiHeadAttention(d_model, num_heads, dropout_rate)
        self.ffn = PositionwiseFeedForward(d_model, d_ff, dropout_rate)

        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)

        self.dropout1 = nn.Dropout(dropout_rate)
        self.dropout2 = nn.Dropout(dropout_rate)

    def forward(
        self,
        x: torch.Tensor,
        mask: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        attn_out = self.self_attn(x, x, x, mask)
        x = self.norm1(x + self.dropout1(attn_out))

        ff_out = self.ffn(x)
        x = self.norm2(x + self.dropout2(ff_out))

        return x

if __name__ == '__main__':
    D_MODEL = 512
    NUM_HEADS = 8
    D_FF = 2048
    BATCH_SIZE = 4
    SEQ_LEN = 60
    DROPOUT = 0.1

    encoder_layer = EncoderLayer(
        d_model=D_MODEL,
        num_heads=NUM_HEADS,
        d_ff=D_FF,
        dropout_rate=DROPOUT
    )

    dummy_input = torch.randn(BATCH_SIZE, SEQ_LEN, D_MODEL)
    dummy_mask = torch.ones(BATCH_SIZE, 1, SEQ_LEN).to(torch.uint8)
    dummy_mask[:, :, -10:] = 0

    output = encoder_layer(x=dummy_input, mask=dummy_mask)

    print(f"入力 (x) 形状: {dummy_input.shape}")
    print(f"マスク形状 (mask): {dummy_mask.shape}")
    print(f"出力形状 (output): {output.shape}")

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

# Attention Is All You Need
# 3.1 Encoder and Decoder Stacks
# EncoderLayer.py
# 第2部 第8章