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_and_split(x: torch.Tensor, linear: nn.Linear) -> torch.Tensor:
            return linear(x).view(B, -1, self.num_heads, self.d_k).transpose(1, 2)

        query = project_and_split(query, self.w_q)
        key   = project_and_split(key, self.w_k)
        value = project_and_split(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)

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.w_2(self.dropout(F.relu(self.w_1(x))))

class SublayerConnection(nn.Module):
    def __init__(self, size: int, dropout_rate: float):
        super().__init__()
        self.norm = nn.LayerNorm(size)
        self.dropout = nn.Dropout(dropout_rate)

    def forward(self, x: torch.Tensor, sublayer: nn.Module, **kwargs) -> torch.Tensor:
        return self.norm(x + self.dropout(sublayer(x, **kwargs)))

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.sublayer = nn.ModuleList([
            SublayerConnection(d_model, dropout_rate) for _ in range(2)
        ])

    def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor:
        x = self.sublayer[0](
            x, lambda q: self.self_attn(query=q, key=q, value=q, mask=mask)
        )
        x = self.sublayer[1](x, self.ffn)
        return x

class SinusoidalPositionalEncoding(nn.Module):
    def __init__(self, d_model: int, max_len: int = 5000):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(
            torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
        )
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        self.register_buffer('pe', pe.unsqueeze(0))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x + self.pe[:, :x.size(1)]

class EmbeddingsAndPositionalEncoding(nn.Module):
    def __init__(self, vocab_size: int, d_model: int, dropout_rate: float):
        super().__init__()
        self.token_embedding = nn.Embedding(vocab_size, d_model)
        self.positional_encoding = SinusoidalPositionalEncoding(d_model)
        self.dropout = nn.Dropout(dropout_rate)
        self.d_model = d_model

    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
        x = self.token_embedding(input_ids)
        x = x * math.sqrt(self.d_model)
        x = self.positional_encoding(x)
        return self.dropout(x)

class Encoder(nn.Module):
    def __init__(
        self,
        vocab_size: int,
        N: int,
        d_model: int,
        num_heads: int,
        d_ff: int,
        dropout_rate: float
    ):
        super().__init__()
        self.embed = EmbeddingsAndPositionalEncoding(vocab_size, d_model, dropout_rate)
        self.layers = nn.ModuleList([
            EncoderLayer(d_model, num_heads, d_ff, dropout_rate) for _ in range(N)
        ])

    def forward(self, src_ids: torch.Tensor, src_mask: torch.Tensor) -> torch.Tensor:
        x = self.embed(src_ids)
        for layer in self.layers:
            x = layer(x, src_mask)
        return x

if __name__ == '__main__':
    VOCAB_SIZE = 10000
    N_LAYERS = 6
    D_MODEL = 512 
    NUM_HEADS = 8
    D_FF = 2048 
    BATCH_SIZE = 4
    SEQ_LEN = 100 
    DROPOUT = 0.1

    encoder = Encoder(VOCAB_SIZE, N_LAYERS, D_MODEL, NUM_HEADS, D_FF, DROPOUT)
    dummy_src_ids = torch.randint(low=1, high=VOCAB_SIZE, size=(BATCH_SIZE, SEQ_LEN))
    dummy_src_mask = torch.ones(BATCH_SIZE, 1, SEQ_LEN).to(torch.uint8) 
    dummy_src_mask[:, :, -10:] = 0 
    output = encoder(dummy_src_ids, dummy_src_mask)

    print(f"入力形状: {dummy_src_ids.shape}")
    print(f"マスク形状 (src_mask): {dummy_src_mask.shape}")
    print(f"出力形状: {output.shape}")

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

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