import torch
import torch.nn as nn
import math
import copy

class PositionalEncoding(nn.Module):
    def __init__(self, d_model: int, dropout: float, max_len: int = 5000):
        super().__init__()
        self.dropout = nn.Dropout(dropout)

        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_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)
        pe = pe.unsqueeze(0)
        self.register_buffer("pe", pe)

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

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

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

class DecoderLayer(nn.Module):
    def __init__(self, d_model: int, num_heads: int, d_ff: int, dropout: float):
        super().__init__()
        self.self_attn = nn.MultiheadAttention(
            d_model, num_heads, dropout=dropout, batch_first=True
        )
        self.src_attn = nn.MultiheadAttention(
            d_model, num_heads, dropout=dropout, batch_first=True
        )

        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.ReLU(),
            nn.Linear(d_ff, d_model),
        )

        self.sublayers = nn.ModuleList(
            [SublayerConnection(d_model, dropout) for _ in range(3)]
        )

    def forward(
        self,
        x: torch.Tensor,
        memory: torch.Tensor,
        src_mask: torch.Tensor,
        tgt_mask: torch.Tensor,
    ) -> torch.Tensor:

        def self_attn_layer(x):
            attn_output, _ = self.self_attn(
                x, x, x, attn_mask=tgt_mask
            )
            return attn_output

        def src_attn_layer(x):
            attn_output, _ = self.src_attn(
                x, memory, memory, key_padding_mask=None
            )
            return attn_output

        x = self.sublayers[0](x, self_attn_layer)
        x = self.sublayers[1](x, src_attn_layer)
        x = self.sublayers[2](x, self.ffn)
        return x

class EmbeddingsAndPositionalEncoding(nn.Module):
    def __init__(self, vocab_size: int, d_model: int, dropout: float):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.pe = PositionalEncoding(d_model, dropout)
        self.d_model = d_model

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

class Decoder(nn.Module):
    def __init__(
        self,
        vocab_size: int,
        N: int,
        d_model: int,
        num_heads: int,
        d_ff: int,
        dropout: float,
    ):
        super().__init__()
        self.embed = EmbeddingsAndPositionalEncoding(vocab_size, d_model, dropout)
        layer = DecoderLayer(d_model, num_heads, d_ff, dropout)
        self.layers = nn.ModuleList([copy.deepcopy(layer) for _ in range(N)])

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

if __name__ == "__main__":
    def subsequent_mask(size):
        mask = torch.triu(torch.ones(size, size), diagonal=1)
        return mask.bool()

    VOCAB_SIZE = 10000
    N_LAYERS = 6
    D_MODEL = 512
    NUM_HEADS = 8
    D_FF = 2048
    BATCH_SIZE = 4
    SRC_SEQ_LEN = 100
    TGT_SEQ_LEN = 50
    DROPOUT = 0.1

    decoder = Decoder(
        vocab_size=VOCAB_SIZE,
        N=N_LAYERS,
        d_model=D_MODEL,
        num_heads=NUM_HEADS,
        d_ff=D_FF,
        dropout=DROPOUT,
    )

    dummy_tgt_ids = torch.randint(
        low=1, high=VOCAB_SIZE, size=(BATCH_SIZE, TGT_SEQ_LEN)
    )
    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, SRC_SEQ_LEN).bool()

    output = decoder(
        tgt_ids=dummy_tgt_ids,
        memory=dummy_memory,
        src_mask=dummy_src_mask,
        tgt_mask=dummy_tgt_mask,
    )

    print(f"入力 (tgt_ids) 形状: {dummy_tgt_ids.shape}")
    print(f"メモリ 形状: {dummy_memory.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
# Decoder.py
# 第2部 第13章