import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from typing import Optional

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, *args, **kwargs) -> torch.Tensor:
        output = sublayer(x, *args, **kwargs)
        output = self.dropout(output)
        return self.norm(x + output)

class MaskedSelfAttention(nn.Module):
    def __init__(self, d_model: int, num_heads: int, dropout_rate: float):
        super().__init__()
        self.mha = self._build_mha(d_model, num_heads, dropout_rate)

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

                self.attn_dropout = nn.Dropout(dropout_rate)

            def forward(self, query, key, value, mask):
                B, S_q, D = query.shape

                q = self.w_q(query).view(B, S_q, self.num_heads, self.d_k).transpose(1, 2)
                k = self.w_k(key).view(B, -1, self.num_heads, self.d_k).transpose(1, 2)
                v = self.w_v(value).view(B, -1, self.num_heads, self.d_k).transpose(1, 2)

                scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
                scores = scores.masked_fill(mask == 0, -1e9)

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

                x = torch.matmul(attn, v)
                x = x.transpose(1, 2).contiguous().view(B, S_q, self.d_model)
                return self.w_o(x)

        return SimpleMHA(d_model, num_heads, dropout_rate)

    def forward(self, x: torch.Tensor, tgt_mask: torch.Tensor) -> torch.Tensor:
        return self.mha(query=x, key=x, value=x, mask=tgt_mask)

class SimpleCrossAttention(nn.Module):
    def __init__(self, d_model: int, num_heads: int, dropout_rate: float):
        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)

        self.attn_dropout = nn.Dropout(dropout_rate)

    def forward(self, query: torch.Tensor, memory: torch.Tensor,
                mask: Optional[torch.Tensor] = None) -> torch.Tensor:
        B, S_q, D = query.shape

        q = self.w_q(query).view(B, S_q, self.num_heads, self.d_k).transpose(1, 2)
        k = self.w_k(memory).view(B, -1, self.num_heads, self.d_k).transpose(1, 2)
        v = self.w_v(memory).view(B, -1, self.num_heads, self.d_k).transpose(1, 2)

        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
        if mask is not None:
            mask = mask.unsqueeze(2)
            scores = scores.masked_fill(mask == 0, -1e9)

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

        x = torch.matmul(attn, v)
        x = x.transpose(1, 2).contiguous().view(B, S_q, self.d_model)
        return self.w_o(x)

class SimplePositionwiseFeedForward(nn.Module):
    def __init__(self, d_model: int, d_ff: int, dropout_rate: float):
        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 DecoderLayer(nn.Module):
    def __init__(self, d_model: int, num_heads: int, d_ff: int, dropout_rate: float):
        super().__init__()
        self.self_attn = MaskedSelfAttention(d_model, num_heads, dropout_rate)
        self.cross_attn = SimpleCrossAttention(d_model, num_heads, dropout_rate)
        self.ffn = SimplePositionwiseFeedForward(d_model, d_ff, dropout_rate)

        self.sublayer = nn.ModuleList([
            SublayerConnection(d_model, dropout_rate),
            SublayerConnection(d_model, dropout_rate),
            SublayerConnection(d_model, dropout_rate),
        ])

    def forward(self, x: torch.Tensor, memory: torch.Tensor,
                src_mask: torch.Tensor, tgt_mask: torch.Tensor) -> torch.Tensor:
        x = self.sublayer[0](x, self.self_attn, tgt_mask)
        x = self.sublayer[1](x, self.cross_attn, memory, src_mask)
        x = self.sublayer[2](x, self.ffn)
        return x

if __name__ == '__main__':
    def subsequent_mask(size):
        attn_shape = (1, size, size)
        subsequent_mask = torch.triu(torch.ones(attn_shape), diagonal=1).bool()
        return ~subsequent_mask

    D_MODEL = 512
    NUM_HEADS = 8
    D_FF = 2048
    BATCH_SIZE = 4
    SRC_SEQ_LEN = 100
    TGT_SEQ_LEN = 50
    DROPOUT = 0.1

    decoder_layer = DecoderLayer(D_MODEL, NUM_HEADS, D_FF, DROPOUT)

    dummy_x = torch.randn(BATCH_SIZE, TGT_SEQ_LEN, D_MODEL)
    dummy_memory = torch.randn(BATCH_SIZE, SRC_SEQ_LEN, D_MODEL)

    dummy_tgt_mask = subsequent_mask(TGT_SEQ_LEN)
    dummy_src_mask = torch.ones(BATCH_SIZE, 1, SRC_SEQ_LEN).bool()
    dummy_src_mask[:, :, -20:] = 0

    output = decoder_layer(
        x=dummy_x,
        memory=dummy_memory,
        src_mask=dummy_src_mask,
        tgt_mask=dummy_tgt_mask
    )

    print(f"入力 (x) 形状: {dummy_x.shape}")
    print(f"メモリ 形状: {dummy_memory.shape}")
    print(f"ターゲットマスク形状 (tgt_mask): {dummy_tgt_mask.shape}")
    print(f"ソースマスク形状 (src_mask): {dummy_src_mask.shape}")
    print(f"出力形状 (output): {output.shape}")

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

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