import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_seq_len=5000):
        super().__init__()
        self.d_model = d_model
        pe = torch.zeros(max_seq_len, d_model)
        position = torch.arange(max_seq_len).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2) * (-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):
        return x + self.pe[:, :x.size(1)]

class StandardSelfAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        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)

    def forward(self, x, mask=None):
        B, S, D = x.size()
        
        query = self.w_q(x).view(B, S, self.num_heads, self.d_k).transpose(1, 2)
        key   = self.w_k(x).view(B, S, self.num_heads, self.d_k).transpose(1, 2)
        value = self.w_v(x).view(B, S, self.num_heads, self.d_k).transpose(1, 2)
        
        scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(self.d_k)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)
        
        p_attn = F.softmax(scores, dim=-1)
        output = torch.matmul(p_attn, value)
        
        output = output.transpose(1, 2).contiguous().view(B, S, D)
        return self.w_o(output)

if __name__ == '__main__':
    D_MODEL = 512
    NUM_HEADS = 8
    MAX_SEQ_LEN = 60
    SEQ_LEN = 10
    BATCH_SIZE = 4
    DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    abs_pe = PositionalEncoding(d_model=D_MODEL, max_seq_len=MAX_SEQ_LEN)
    s_attn = StandardSelfAttention(d_model=D_MODEL, num_heads=NUM_HEADS)
    
    dummy_input_embedding = torch.randn(BATCH_SIZE, SEQ_LEN, D_MODEL)
    input_with_pe = abs_pe(dummy_input_embedding * math.sqrt(D_MODEL))
    output = s_attn(x=input_with_pe)
    
    print(f"入力埋め込み形状: {dummy_input_embedding.shape}")
    print(f"位置エンコーディング(PE)形状: {abs_pe.pe.shape}")
    print(f"PE加算後の入力形状: {input_with_pe.shape}")
    print(f"出力形状: {output.shape}")

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

# Attention Is All You Need
# 3.5 Positional Encoding
# AbsolutePositionalEncoding.py
# 第2部 第2章