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

def scaled_dot_product_attention(Q, K, V, mask=None, dropout=None):
    d_k = Q.size(-1)
    scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)

    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)

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

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

    output = torch.matmul(attn, V)
    return output, 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.h = 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, key, value, mask=None):
        batch_size = query.size(0)

        def transform(x, linear):
            return (
                linear(x)
                .view(batch_size, -1, self.h, self.d_k)
                .transpose(1, 2)
            )

        Q = transform(query, self.W_Q)
        K = transform(key,   self.W_K)
        V = transform(value, self.W_V)

        x, attn = scaled_dot_product_attention(
            Q, K, V, mask=mask, dropout=self.attn_dropout
        )

        x = (
            x.transpose(1, 2)
             .contiguous()
             .view(batch_size, -1, self.h * self.d_k)
        )

        return self.W_O(x), attn

if __name__ == "__main__":
    D_MODEL = 512
    NUM_HEADS = 8
    BATCH_SIZE = 4
    SEQ_LEN_Q = 60
    SEQ_LEN_K = 60

    query = torch.randn(BATCH_SIZE, SEQ_LEN_Q, D_MODEL)
    key   = torch.randn(BATCH_SIZE, SEQ_LEN_K, D_MODEL)
    value = torch.randn(BATCH_SIZE, SEQ_LEN_K, D_MODEL)

    mask = torch.ones(BATCH_SIZE, 1, SEQ_LEN_Q, SEQ_LEN_K)

    print("入力 query 形状:", query.shape)
    print("入力 key    形状:", key.shape)
    print("入力 value  形状:", value.shape)
    print("mask 形状:", mask.shape)

    mha = MultiHeadAttention(
        d_model=D_MODEL,
        num_heads=NUM_HEADS,
        dropout_rate=0.1
    )

    output, attention_weights = mha(query, key, value, mask)

    print("出力形状:", output.shape)
    print("Attention重み形状:", attention_weights.shape)

    assert output.shape == (BATCH_SIZE, SEQ_LEN_Q, D_MODEL)
    assert attention_weights.shape == (BATCH_SIZE, NUM_HEADS, SEQ_LEN_Q, SEQ_LEN_K)

# Attention Is All You Need
# 3.2.2 Multi-Head Attention (core logic, simplified)
# MultiHeadAttention.py
# 第2部 第6章