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

def scaled_dot_product_attention(query, key, value, mask=None):
    d_k = query.size(-1)
    scores = torch.matmul(query, key.transpose(-2, -1))
    scores = scores / math.sqrt(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)
    
    return output, p_attn

if __name__ == '__main__':
    D_MODEL = 512
    NUM_HEADS = 8
    D_K = D_MODEL // NUM_HEADS
    BATCH_SIZE = 4
    Q_LEN = 10
    K_LEN = 20

    dummy_q = torch.randn(BATCH_SIZE, NUM_HEADS, Q_LEN, D_K)
    dummy_k = torch.randn(BATCH_SIZE, NUM_HEADS, K_LEN, D_K)
    dummy_v = torch.randn(BATCH_SIZE, NUM_HEADS, K_LEN, D_K)

    dummy_mask = torch.ones(BATCH_SIZE, 1, K_LEN).to(torch.uint8)
    dummy_mask[:, :, -5:] = 0
    dummy_mask = dummy_mask.unsqueeze(2)

    output, attn_weights = scaled_dot_product_attention(
        query=dummy_q, key=dummy_k, value=dummy_v, mask=dummy_mask
    )

    print(f"Q, K, V 形状: ({BATCH_SIZE}, {NUM_HEADS}, {Q_LEN}/{K_LEN}, {D_K})")
    print(f"マスク形状 (input): {dummy_mask.shape}")
    print(f"出力 (output) 形状: {output.shape}")
    print(f"アテンションウェイト形状: {attn_weights.shape}")

    expected_output_shape = (BATCH_SIZE, NUM_HEADS, Q_LEN, D_K)
    assert output.shape == expected_output_shape

    first_q_weights = attn_weights[0, 0, 0, :]
    print(f"\n最初のクエリに対するウェイトの例 (Key長={K_LEN}):")
    print(f"  マスクされていない部分のウェイト合計: {first_q_weights[:-5].sum().item():.4f}")
    print(f"  マスクされた部分のウェイト合計: {first_q_weights[-5:].sum().item():.10f} ")

# Attention Is All You Need
# 3.2.1 Scaled Dot-Product Attention
# ScaledDotProductAttention.py
# 第2部 第5章