import numpy as np
import matplotlib.pyplot as plt

# 그래프에 한글 라벨(토큰 이름)을 쓸 것이므로 한글 폰트를 지정해둔다
# Windows 기본 한글 폰트인 맑은 고딕 사용, 없는 환경이면 이 줄만 원하는 폰트로 바꾸면 됨
plt.rcParams["font.family"] = "Malgun Gothic"
plt.rcParams["axes.unicode_minus"] = False  # 마이너스 기호가 네모로 깨지는 것 방지

np.random.seed(42)  # 이 노트북 전체에서 재현 가능하도록 시드 고정

def softmax(x, axis=-1):
    # 오버플로우 방지를 위해 최댓값을 빼고 지수화 (앞으로 계속 재사용)
    e_x = np.exp(x - np.max(x, axis=axis, keepdims=True))
    return e_x / np.sum(e_x, axis=axis, keepdims=True)

words = ["I", "am", "a", "student"]

# 첫 번째 성분을 "사람을 가리키는 정도"라고 생각하고 값을 직접 넣어봤다
# I와 student가 겹치는 성분(0.8)을 갖도록 만들어서, 실제로 서로 더 주목하는지 확인해본다
embeddings = np.array([
    [1.0, 0.0, 0.0, 0.0],   # I
    [0.0, 1.0, 0.0, 0.0],   # am
    [0.0, 0.0, 1.0, 0.0],   # a
    [0.8, 0.0, 0.0, 1.0],   # student
])

# 가중치 행렬(Wq/Wk) 없이(=단위행렬로 대체) 순수 유사도만 확인 -> Q=K=V=embeddings
d_k = embeddings.shape[1]
scores = embeddings @ embeddings.T / np.sqrt(d_k)
weights = softmax(scores, axis=-1)

print("단어별 attention weight (자기 자신 제외하고 가장 높은 값에 주목):")
for i, w in enumerate(words):
    row = ", ".join(f"{words[j]}={weights[i, j]:.2f}" for j in range(len(words)))
    print(f"{w:>8} ->{row}")

seq_len = 4   # 토큰 개수 (문장 길이)
d_model = 8   # 임베딩 차원

# 토큰 4개짜리 문장의 임베딩이라고 가정하고 랜덤 생성
X = np.random.randn(seq_len, d_model)
print("입력 X (토큰 임베딩), shape:", X.shape)
print(X)

# Q, K, V로 변환하는 가중치 행렬 (실제 학습되는 파라미터에 해당)
d_k = d_model  # 이 예제에서는 Q/K/V 차원을 입력과 동일하게 둔다
Wq = np.random.randn(d_model, d_k)
Wk = np.random.randn(d_model, d_k)
Wv = np.random.randn(d_model, d_k)

Q = X @ Wq
K = X @ Wk
V = X @ Wv

print("Q shape:", Q.shape)
print("K shape:", K.shape)
print("V shape:", V.shape)

# 토큰끼리의 유사도 점수 (QK^T), sqrt(d_k)로 스케일링
scores = (Q @ K.T) / np.sqrt(d_k)
print("scores (스케일링 전 유사도), shape:", scores.shape)
print(scores)

# softmax로 각 토큰이 다른 토큰에 주목하는 비율(가중치)로 변환 (softmax는 위에서 정의해둔 걸 재사용)
attn_weights = softmax(scores, axis=-1)
print("\nattention weights (각 행의 합은 1):")
print(attn_weights)
print("행별 합 검증:", attn_weights.sum(axis=-1))

# 가중치로 V를 가중합 -> 최종 Self-Attention 출력
output = attn_weights @ V
print("Self-Attention 출력, shape:", output.shape)
print(output)

tokens = [f"토큰{i}" for i in range(seq_len)]

fig, ax = plt.subplots(figsize=(5, 4))
im = ax.imshow(attn_weights, cmap="viridis")

ax.set_xticks(range(seq_len))
ax.set_yticks(range(seq_len))
ax.set_xticklabels(tokens)
ax.set_yticklabels(tokens)
ax.set_xlabel("Key (주목 대상)")
ax.set_ylabel("Query (주목하는 주체)")
ax.set_title("Self-Attention Weights")

# 칸마다 숫자 표시
for i in range(seq_len):
    for j in range(seq_len):
        ax.text(j, i, f"{attn_weights[i, j]:.2f}", ha="center", va="center", color="white")

fig.colorbar(im)
plt.tight_layout()
plt.show()

# 비교 1: Q = K = X (별도 가중치 없이 그대로 사용) -> 대각선이 지배적인 경향
scores_qk_equal = (X @ X.T) / np.sqrt(d_model)
weights_qk_equal = softmax(scores_qk_equal, axis=-1)

# 비교 2: 기존처럼 Wq, Wk가 독립적인 랜덤 행렬인 경우 (attn_weights, 4번 섹션에서 이미 계산해둔 값 재사용)

fig, axes = plt.subplots(1, 2, figsize=(11, 4))

for ax, weights, title in zip(
    axes,
    [weights_qk_equal, attn_weights],
    ["Q = K = X\n(대각선이 지배적)", "Q, K가 독립적인 랜덤 행렬\n(랜덤한 패턴)"],
):
    im = ax.imshow(weights, cmap="viridis", vmin=0, vmax=1)
    ax.set_xticks(range(seq_len))
    ax.set_yticks(range(seq_len))
    ax.set_xticklabels(tokens)
    ax.set_yticklabels(tokens)
    ax.set_title(title)
    for i in range(seq_len):
        for j in range(seq_len):
            ax.text(j, i, f"{weights[i, j]:.2f}", ha="center", va="center", color="white", fontsize=8)

fig.colorbar(im, ax=axes, fraction=0.046, pad=0.04)
plt.show()

print("Q=K=X일 때 대각선 값:        ", np.diag(weights_qk_equal).round(3))
print("Q,K가 독립적인 랜덤일 때 대각선 값:", np.diag(attn_weights).round(3))

import torch
import torch.nn as nn

class SelfAttention(nn.Module):
    def __init__(self, d_model, d_k):
        super().__init__()
        self.d_k = d_k
        # bias=False: 원래 Transformer 논문의 attention projection도 bias를 안 씀
        self.Wq = nn.Linear(d_model, d_k, bias=False)
        self.Wk = nn.Linear(d_model, d_k, bias=False)
        self.Wv = nn.Linear(d_model, d_k, bias=False)

    def forward(self, x, mask=None):
        # x: (batch, seq_len, d_model)
        Q = self.Wq(x)
        K = self.Wk(x)
        V = self.Wv(x)

        scores = Q @ K.transpose(-2, -1) / (self.d_k ** 0.5)

        if mask is not None:
            # mask==0인 위치는 -inf로 채워서 softmax 후 가중치가 0이 되게 함
            scores = scores.masked_fill(mask == 0, float("-inf"))

        attn_weights = torch.softmax(scores, dim=-1)
        output = attn_weights @ V
        return output, attn_weights

torch.manual_seed(42)
x = torch.randn(1, seq_len, d_model)  # batch=1

attn = SelfAttention(d_model=d_model, d_k=d_model)
out, weights = attn(x)

print("출력 shape:", out.shape)          # (1, seq_len, d_k)
print("가중치 shape:", weights.shape)     # (1, seq_len, seq_len)

class ManualMultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        assert d_model % num_heads == 0, "d_model은 num_heads로 나누어떨어져야 head를 균등하게 쪼갤 수 있음"
        self.num_heads = num_heads
        self.d_head = d_model // num_heads

        self.Wq = nn.Linear(d_model, d_model, bias=False)
        self.Wk = nn.Linear(d_model, d_model, bias=False)
        self.Wv = nn.Linear(d_model, d_model, bias=False)
        self.Wo = nn.Linear(d_model, d_model, bias=False)  # concat 후 최종 projection

    def split_heads(self, x):
        # (batch, seq_len, d_model) -> (batch, num_heads, seq_len, d_head)
        batch, seq_len, d_model = x.shape
        x = x.view(batch, seq_len, self.num_heads, self.d_head)
        return x.transpose(1, 2)

    def forward(self, x):
        batch, seq_len, d_model = x.shape

        Q = self.split_heads(self.Wq(x))
        K = self.split_heads(self.Wk(x))
        V = self.split_heads(self.Wv(x))

        # head마다 독립적으로 Scaled Dot-Product Attention
        scores = Q @ K.transpose(-2, -1) / (self.d_head ** 0.5)
        weights = torch.softmax(scores, dim=-1)
        head_outputs = weights @ V  # (batch, num_heads, seq_len, d_head)

        # head들을 다시 이어붙임(concat) -> (batch, seq_len, d_model)
        concat = head_outputs.transpose(1, 2).contiguous().view(batch, seq_len, d_model)
        output = self.Wo(concat)
        return output, weights

torch.manual_seed(42)
num_heads = 2  # d_model=8을 2개 head로 나누면 head당 차원은 4

manual_mha = ManualMultiHeadAttention(d_model=d_model, num_heads=num_heads)
manual_out, manual_weights = manual_mha(x)

print("최종 출력 shape:", manual_out.shape)        # (1, seq_len, d_model) - head로 쪼개기 전과 동일
print("head별 가중치 shape:", manual_weights.shape)  # (1, num_heads, seq_len, seq_len) - head마다 따로 있음
print("\nhead 0의 attention weights:")
print(manual_weights[0, 0].detach().numpy().round(2))
print("\nhead 1의 attention weights (head마다 다른 패턴에 주목함을 확인):")
print(manual_weights[0, 1].detach().numpy().round(2))

mha = nn.MultiheadAttention(embed_dim=d_model, num_heads=num_heads, batch_first=True)

# Self-Attention이므로 query, key, value에 전부 같은 입력을 넣는다
mha_out, mha_weights = mha(x, x, x)

print("Multi-Head 출력 shape:", mha_out.shape)                 # (1, seq_len, d_model) - 직접 구현한 것과 동일
print("Multi-Head 평균 가중치 shape:", mha_weights.shape)       # (1, seq_len, seq_len), 기본값은 head 평균

# 하삼각(lower-triangular) 마스크: 1이면 볼 수 있음, 0이면 못 봄
causal_mask = torch.tril(torch.ones(seq_len, seq_len)).bool()
print("causal mask:")
print(causal_mask.int())

out_masked, weights_masked = attn(x, mask=causal_mask)

print("\n마스킹 전 가중치 (앞서 계산해둔 weights 재사용):")
print(weights.squeeze(0).detach().numpy().round(2))
print("\n마스킹 후 가중치 (오른쪽 위 삼각형이 전부 0이 됨 = 미래 토큰을 못 봄):")
print(weights_masked.squeeze(0).detach().numpy().round(2))

sentence = ["나는", "오늘", "랭체인을", "공부한다"]
vocab = {w: i for i, w in enumerate(sentence)}

torch.manual_seed(0)
embedding = nn.Embedding(num_embeddings=len(vocab), embedding_dim=d_model)

token_ids = torch.tensor([[vocab[w] for w in sentence]])  # (1, seq_len)
sent_x = embedding(token_ids)  # (1, seq_len, d_model)

_, sent_weights = attn(sent_x)
sent_weights = sent_weights.squeeze(0).detach().numpy()

fig, ax = plt.subplots(figsize=(5, 4))
im = ax.imshow(sent_weights, cmap="viridis")
ax.set_xticks(range(len(sentence)))
ax.set_yticks(range(len(sentence)))
ax.set_xticklabels(sentence)
ax.set_yticklabels(sentence)
ax.set_title("문장에 대한 Self-Attention (학습 전, 랜덤 가중치)")
for i in range(len(sentence)):
    for j in range(len(sentence)):
        ax.text(j, i, f"{sent_weights[i, j]:.2f}", ha="center", va="center", color="white")
fig.colorbar(im)
plt.tight_layout()
plt.show()

import torch.optim as optim

sentence9 = ["나는", "오늘", "랭체인을", "공부한다"]
vocab9 = {w: i for i, w in enumerate(sentence9)}
vocab_size9 = len(vocab9)

torch.manual_seed(0)
embedding9 = nn.Embedding(vocab_size9, d_model)
attn9 = SelfAttention(d_model=d_model, d_k=d_model)
output_head9 = nn.Linear(d_model, vocab_size9)  # attention 출력을 단어별 점수(logit)로 변환

token_ids9 = torch.tensor([[vocab9[w] for w in sentence9]])  # (1, seq_len)
causal_mask9 = torch.tril(torch.ones(len(sentence9), len(sentence9))).bool()
targets9 = token_ids9[:, 1:]  # 각 위치의 "다음 단어" 정답, 마지막 토큰은 다음이 없으므로 제외

# 비교를 위해 학습 전 attention weight를 미리 찍어둔다
with torch.no_grad():
    x9 = embedding9(token_ids9)
    _, weights_before = attn9(x9, mask=causal_mask9)

optimizer = optim.Adam(
    list(embedding9.parameters()) + list(attn9.parameters()) + list(output_head9.parameters()),
    lr=0.05,
)
loss_fn = nn.CrossEntropyLoss()

num_epochs = 300
loss_history = []  # Part 2(18번)에서 학습 레시피를 비교할 때 재사용
for epoch in range(num_epochs):
    x9 = embedding9(token_ids9)
    out9, weights_after = attn9(x9, mask=causal_mask9)
    logits = output_head9(out9)  # (1, seq_len, vocab_size)

    # 마지막 위치는 정답(다음 단어)이 없으므로 제외하고 손실 계산
    loss = loss_fn(logits[:, :-1].reshape(-1, vocab_size9), targets9.reshape(-1))

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    loss_history.append(loss.item())

    if epoch % 50 == 0:
        print(f"epoch{epoch:3d}: loss={loss.item():.4f}")

print("최종 loss:", loss.item())

# 학습 전/후 attention weight 비교
w_before = weights_before.squeeze(0).detach().numpy()
w_after = weights_after.squeeze(0).detach().numpy()

fig, axes = plt.subplots(1, 2, figsize=(11, 4))
for ax, w, title in zip(
    axes,
    [w_before, w_after],
    ["학습 전 (랜덤 초기화)", f"학습 후 ({num_epochs}epoch, loss={loss.item():.3f})"],
):
    im = ax.imshow(w, cmap="viridis", vmin=0, vmax=1)
    ax.set_xticks(range(len(sentence9)))
    ax.set_yticks(range(len(sentence9)))
    ax.set_xticklabels(sentence9)
    ax.set_yticklabels(sentence9)
    ax.set_title(title)
    for i in range(len(sentence9)):
        for j in range(len(sentence9)):
            ax.text(j, i, f"{w[i, j]:.2f}", ha="center", va="center", color="white", fontsize=8)
fig.colorbar(im, ax=axes, fraction=0.046, pad=0.04)
plt.show()

# 실제로 문장을 외웠는지, 다음 단어 예측으로 확인해보자
with torch.no_grad():
    x9 = embedding9(token_ids9)
    out9, _ = attn9(x9, mask=causal_mask9)
    logits = output_head9(out9)
    preds = logits.argmax(dim=-1).squeeze(0).tolist()

print("입력:          ", sentence9)
print("다음 단어 예측:  ", [sentence9[p] for p in preds])
print("(앞 3개 예측이 순서대로 '오늘', '랭체인을', '공부한다'와 같으면 문장을 정확히 외운 것)")

def positional_encoding(seq_len, d_model):
    # (seq_len, d_model) 크기의 위치 인코딩 행렬을 만든다
    position = np.arange(seq_len)[:, np.newaxis]         # (seq_len, 1)
    i = np.arange(d_model)[np.newaxis, :]                 # (1, d_model)
    angle_rates = 1 / np.power(10000, (2 * (i // 2)) / d_model)
    angle_rads = position * angle_rates

    pe = np.zeros((seq_len, d_model))
    pe[:, 0::2] = np.sin(angle_rads[:, 0::2])   # 짝수 인덱스 -> sin
    pe[:, 1::2] = np.cos(angle_rads[:, 1::2])   # 홀수 인덱스 -> cos
    return pe

pe = positional_encoding(seq_len=seq_len, d_model=d_model)
print("Positional Encoding shape:", pe.shape)
print(pe.round(3))

fig, ax = plt.subplots(figsize=(6, 4))
im = ax.imshow(pe, cmap="RdBu", aspect="auto")
ax.set_xlabel("임베딩 차원 (i)")
ax.set_ylabel("위치 (pos)")
ax.set_title("Positional Encoding 값")
fig.colorbar(im)
plt.tight_layout()
plt.show()

def self_attention_np(X, Wq, Wk, Wv):
    Q, K, V = X @ Wq, X @ Wk, X @ Wv
    scores = (Q @ K.T) / np.sqrt(Wq.shape[1])
    weights = softmax(scores, axis=-1)
    return weights @ V

# 토큰 순서를 완전히 뒤집은 입력을 만든다 (같은 토큰들, 순서만 반대)
X_reordered = X[::-1].copy()

# Positional Encoding 없이 -> 같은 토큰(원래 0번, 뒤집으면 마지막 자리)의 출력이 똑같은지 확인
out_original = self_attention_np(X, Wq, Wk, Wv)
out_reordered = self_attention_np(X_reordered, Wq, Wk, Wv)
print("[PE 없음] 원래 0번 토큰 출력 == 뒤집었을 때 그 토큰의 출력?",
      np.allclose(out_original[0], out_reordered[-1]))

# Positional Encoding을 더한 뒤에는 같은 토큰이라도 위치가 다르면 결과가 달라지는지 확인
out_original_pe = self_attention_np(X + pe, Wq, Wk, Wv)
out_reordered_pe = self_attention_np(X_reordered + pe, Wq, Wk, Wv)
print("[PE 적용] 원래 0번 토큰 출력 == 뒤집었을 때 그 토큰의 출력?",
      np.allclose(out_original_pe[0], out_reordered_pe[-1]))
print("(PE 덕분에 '같은 토큰, 다른 위치'가 이제 다른 결과를 낸다는 뜻)")

class PositionwiseFeedForward(nn.Module):
    def __init__(self, d_model, d_ff):
        super().__init__()
        self.linear1 = nn.Linear(d_model, d_ff)
        self.linear2 = nn.Linear(d_ff, d_model)

    def forward(self, x):
        return self.linear2(torch.relu(self.linear1(x)))

# 논문은 d_ff=2048(d_model=512의 4배)을 쓰는데, 우리 예제도 그 비율(4배)만 맞춘다
d_ff = d_model * 4
ffn = PositionwiseFeedForward(d_model=d_model, d_ff=d_ff)

test_out = ffn(x)  # x: 6번 섹션에서 만든 (1, seq_len, d_model) 텐서 재사용
print("FFN 입력 shape:", x.shape)
print("FFN 출력 shape:", test_out.shape)  # 입력과 동일한 shape 유지

# PyTorch에 이미 구현되어 있어서 그대로 가져다 쓰면 된다
layer_norm = nn.LayerNorm(d_model)

def add_and_norm(x, sublayer_output, norm):
    # "Add & Norm" 패턴을 함수로 만들어두면 앞으로 계속 재사용하기 편하다
    return norm(x + sublayer_output)

# 확인: Self-Attention을 통과시키고 Add & Norm 적용
attn_out, _ = attn(x)  # 6번 섹션에서 만든 SelfAttention 재사용
normed = add_and_norm(x, attn_out, layer_norm)

print("정규화 후 전체 평균(0에 가까워야 함):", normed.mean().item())
print("정규화 후 전체 표준편차(1에 가까워야 함):", normed.std().item())

class EncoderLayer(nn.Module):
    def __init__(self, d_model, num_heads, d_ff):
        super().__init__()
        self.mha = ManualMultiHeadAttention(d_model, num_heads)
        self.ffn = PositionwiseFeedForward(d_model, d_ff)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)

    def forward(self, x):
        # 1) Multi-Head Self-Attention + Add & Norm
        attn_out, _ = self.mha(x)
        x = self.norm1(x + attn_out)

        # 2) Position-wise FFN + Add & Norm
        ffn_out = self.ffn(x)
        x = self.norm2(x + ffn_out)
        return x

class Encoder(nn.Module):
    def __init__(self, d_model, num_heads, d_ff, num_layers):
        super().__init__()
        self.layers = nn.ModuleList([
            EncoderLayer(d_model, num_heads, d_ff) for _ in range(num_layers)
        ])

    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        return x

torch.manual_seed(42)
# 논문은 N=6이지만, 실습이니 2층으로 축소
encoder = Encoder(d_model=d_model, num_heads=2, d_ff=d_ff, num_layers=2)
encoder_out = encoder(x)

print("Encoder 입력 shape:", x.shape)
print("Encoder 출력 shape:", encoder_out.shape)  # 입력과 완전히 동일한 shape

class CrossAttention(nn.Module):
    def __init__(self, d_model, d_k):
        super().__init__()
        self.d_k = d_k
        self.Wq = nn.Linear(d_model, d_k, bias=False)
        self.Wk = nn.Linear(d_model, d_k, bias=False)
        self.Wv = nn.Linear(d_model, d_k, bias=False)

    def forward(self, decoder_x, encoder_out):
        # Query는 디코더에서, Key/Value는 인코더 출력에서 가져온다
        Q = self.Wq(decoder_x)
        K = self.Wk(encoder_out)
        V = self.Wv(encoder_out)
        scores = Q @ K.transpose(-2, -1) / (self.d_k ** 0.5)
        weights = torch.softmax(scores, dim=-1)
        output = weights @ V
        return output, weights

class DecoderLayer(nn.Module):
    def __init__(self, d_model, num_heads, d_ff):
        super().__init__()
        self.self_attn = SelfAttention(d_model, d_model)     # 1) Masked Self-Attention
        self.cross_attn = CrossAttention(d_model, d_model)   # 2) Encoder-Decoder Attention
        self.ffn = PositionwiseFeedForward(d_model, d_ff)    # 3) FFN
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)

    def forward(self, decoder_x, encoder_out, causal_mask):
        self_out, _ = self.self_attn(decoder_x, mask=causal_mask)
        decoder_x = self.norm1(decoder_x + self_out)

        cross_out, cross_weights = self.cross_attn(decoder_x, encoder_out)
        decoder_x = self.norm2(decoder_x + cross_out)

        ffn_out = self.ffn(decoder_x)
        decoder_x = self.norm3(decoder_x + ffn_out)
        return decoder_x, cross_weights

torch.manual_seed(42)
decoder_layer = DecoderLayer(d_model=d_model, num_heads=2, d_ff=d_ff)
decoder_input = torch.randn(1, seq_len, d_model)  # 디코더 쪽 입력(예: 번역 중인 문장)
causal_mask_dec = torch.tril(torch.ones(seq_len, seq_len)).bool()

decoder_out, cross_weights = decoder_layer(decoder_input, encoder_out, causal_mask_dec)
print("Decoder 출력 shape:", decoder_out.shape)
print("Cross-Attention 가중치 shape:", cross_weights.shape)  # (batch, decoder_seq, encoder_seq)

class TiedEmbedding(nn.Module):
    def __init__(self, vocab_size, d_model):
        super().__init__()
        self.d_model = d_model
        self.embedding = nn.Embedding(vocab_size, d_model)
        # 출력층을 따로 만들지 않고, 임베딩 테이블(가중치)을 그대로 전치해서 재사용한다

    def encode(self, token_ids):
        return self.embedding(token_ids) * (self.d_model ** 0.5)  # 논문의 √d_model 스케일링

    def decode_to_logits(self, x):
        # (batch, seq, d_model) @ (d_model, vocab_size) -> (batch, seq, vocab_size)
        # embedding.weight shape: (vocab_size, d_model) 이므로 전치해서 곱한다 (가중치 공유)
        return x @ self.embedding.weight.T

tied = TiedEmbedding(vocab_size=vocab_size9, d_model=d_model)  # 10번 섹션의 4단어 문장 재사용

scaled_emb = tied.encode(token_ids9)
logits = tied.decode_to_logits(scaled_emb)

print("공유 가중치로 만든 logits shape:", logits.shape)  # (1, seq_len, vocab_size)
print("임베딩 벡터 크기(스케일링 전) 평균:", tied.embedding(token_ids9).norm(dim=-1).mean().item())
print("임베딩 벡터 크기(스케일링 후) 평균:", scaled_emb.norm(dim=-1).mean().item())
print("스케일 비율(√d_model이어야 함):", d_model ** 0.5)

def warmup_lr(step_num, d_model, warmup_steps=10):
    step_num = max(step_num, 1)  # 0으로 나누기 방지
    return d_model ** -0.5 * min(step_num ** -0.5, step_num * warmup_steps ** -1.5)

steps = np.arange(1, 200)
lrs = [warmup_lr(s, d_model=d_model, warmup_steps=20) for s in steps]

fig, ax = plt.subplots(figsize=(7, 4))
ax.plot(steps, lrs, color="#F5A623", linewidth=2)
ax.axvline(20, color="gray", linestyle="--", linewidth=1, label="warmup_steps=20")
ax.set_xlabel("step")
ax.set_ylabel("learning rate")
ax.set_title("Warmup 학습률 스케줄")
ax.legend()
plt.tight_layout()
plt.show()

class SelfAttentionWithDropout(nn.Module):
    # 10번 섹션 SelfAttention에 dropout만 추가한 버전
    def __init__(self, d_model, d_k, dropout=0.1):
        super().__init__()
        self.d_k = d_k
        self.Wq = nn.Linear(d_model, d_k, bias=False)
        self.Wk = nn.Linear(d_model, d_k, bias=False)
        self.Wv = nn.Linear(d_model, d_k, bias=False)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        Q, K, V = self.Wq(x), self.Wk(x), self.Wv(x)
        scores = Q @ K.transpose(-2, -1) / (self.d_k ** 0.5)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float("-inf"))
        weights = torch.softmax(scores, dim=-1)
        weights = self.dropout(weights)  # attention weight에도 dropout 적용 (논문 방식)
        output = weights @ V
        return output, weights

torch.manual_seed(0)
embedding10 = nn.Embedding(vocab_size9, d_model)
attn10 = SelfAttentionWithDropout(d_model=d_model, d_k=d_model, dropout=0.1)
output_head10 = nn.Linear(d_model, vocab_size9)

optimizer10 = optim.Adam(
    list(embedding10.parameters()) + list(attn10.parameters()) + list(output_head10.parameters()),
    lr=1.0,  # 실제 학습률은 매 스텝 스케줄러가 다시 정하므로 여기선 기준값 1.0으로 둔다
)
loss_fn_smooth = nn.CrossEntropyLoss(label_smoothing=0.1)  # Label Smoothing

warmup_steps = 20
loss_history10 = []
for step in range(1, 301):
    lr = warmup_lr(step, d_model=d_model, warmup_steps=warmup_steps)
    for g in optimizer10.param_groups:
        g["lr"] = lr

    x10 = embedding10(token_ids9)
    out10, _ = attn10(x10, mask=causal_mask9)
    logits10 = output_head10(out10)
    loss = loss_fn_smooth(logits10[:, :-1].reshape(-1, vocab_size9), targets9.reshape(-1))

    optimizer10.zero_grad()
    loss.backward()
    optimizer10.step()
    loss_history10.append(loss.item())

    if step % 50 == 0:
        print(f"step{step:3d}: lr={lr:.5f}, loss={loss.item():.4f}")

print("최종 loss:", loss_history10[-1])
print("(Label Smoothing 때문에 10번 섹션과 달리 loss가 정확히 0까지는 안 내려감 — 의도된 동작)")

fig, ax = plt.subplots(figsize=(7, 4))
ax.plot(loss_history, label="10번 섹션 (기본 Adam, lr 고정)", color="#6B7280")
ax.plot(loss_history10, label="18번 섹션 (Warmup + Dropout + Label Smoothing)", color="#F5A623")
ax.set_xlabel("epoch / step")
ax.set_ylabel("loss")
ax.set_title("학습 레시피 비교")
ax.legend()
plt.tight_layout()
plt.show()

class TransformerEncoder(nn.Module):
    def __init__(self, vocab_size, d_model, num_heads, d_ff, num_layers, max_len=50):
        super().__init__()
        self.d_model = d_model
        self.embedding = nn.Embedding(vocab_size, d_model)
        pe_table = positional_encoding(max_len, d_model)
        self.register_buffer("pe", torch.tensor(pe_table, dtype=torch.float32))
        self.layers = nn.ModuleList([
            EncoderLayer(d_model, num_heads, d_ff) for _ in range(num_layers)
        ])

    def forward(self, token_ids):
        seq_len = token_ids.shape[1]
        x = self.embedding(token_ids) * (self.d_model ** 0.5)    # 16번: √d_model 스케일링
        x = x + self.pe[:seq_len]                                 # 11번: Positional Encoding
        for layer in self.layers:                                # 14번: EncoderLayer × N
            x = layer(x)
        return x

torch.manual_seed(42)
full_encoder = TransformerEncoder(
    vocab_size=vocab_size9, d_model=d_model, num_heads=2, d_ff=d_ff, num_layers=2,
)

final_output = full_encoder(token_ids9)
print("최종 Encoder 출력 shape:", final_output.shape)
print(final_output[0].detach().numpy().round(3))

print("\n이 노트북에서 조립한 조각들:")
print("- Positional Encoding (11번)")
print("- Position-wise FFN (12번)")
print("- Residual + LayerNorm (13번)")
print("- Multi-Head Self-Attention (7번)")
print("=> 이 전부를 합친 EncoderLayer x N이 바로 지금 이 TransformerEncoder 입니다.")

np.random.seed(42)

seq_len = 8
d_model = 8

X = np.random.randn(seq_len, d_model)
d_k = d_model
Wq = np.random.randn(d_model, d_k)
Wk = np.random.randn(d_model, d_k)
Wv = np.random.randn(d_model, d_k)

Q, K, V = X @ Wq, X @ Wk, X @ Wv
scores = (Q @ K.T) / np.sqrt(d_k)
attn_weights = softmax(scores, axis=-1)
output = attn_weights @ V

print("attn_weights shape:", attn_weights.shape)  # (4,4) -> (8,8)로 커진 것 확인

scores_unscaled = Q @ K.T                      # 스케일링 없음
scores_scaled = Q @ K.T / np.sqrt(d_k)          # 스케일링 있음

weights_unscaled = softmax(scores_unscaled, axis=-1)
weights_scaled = softmax(scores_scaled, axis=-1)

print("스케일링 없음 - 행별 최댓값:", weights_unscaled.max(axis=-1).round(3))
print("스케일링 있음 - 행별 최댓값:", weights_scaled.max(axis=-1).round(3))

num_heads_list = [1, 4, 8]

for heads in num_heads_list:
    torch.manual_seed(42)
    manual_mha = ManualMultiHeadAttention(d_model=d_model, num_heads=heads)
    manual_out, manual_weights = manual_mha(x)
    print("[직접 구현] 가중치 shape:", manual_weights.shape)

    mha = nn.MultiheadAttention(embed_dim=d_model, num_heads=heads, batch_first=True)
    mha_out, mha_weights = mha(x, x, x)
    print("[내장 구현] 가중치 shape (기본: head 평균):", mha_weights.shape)

row_before = weights.squeeze(0)[3].detach().numpy()
row_after = weights_masked.squeeze(0)[3].detach().numpy()

print("토큰3 마스킹 전:", row_before.round(3))
print("토큰3 마스킹 후:", row_after.round(3))
print("완전히 동일한가?", np.allclose(row_before, row_after))

class CrossSelfAttention(SelfAttention):
    # 부모와 같은 파라미터 이름/순서(x, mask)를 유지하고, kv_x는 맨 뒤에 추가
    def forward(self, x, mask=None, kv_x=None):
        if kv_x is None:
            kv_x = x  # 안 주면 원래 SelfAttention과 동일하게 동작 (하위 호환)
        Q = self.Wq(x)
        K = self.Wk(kv_x)
        V = self.Wv(kv_x)
        scores = Q @ K.transpose(-2, -1) / (self.d_k ** 0.5)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float("-inf"))
        weights = torch.softmax(scores, dim=-1)
        return weights @ V, weights

query_words = ["랭체인", "뭐야"]                              # 예: 질문
doc_words = ["랭체인", "은", "LLM", "쉽게", "다루는", "도구"]    # 예: 검색된 문서

query_embedding = nn.Embedding(len(query_words), d_model)
doc_embedding = nn.Embedding(len(doc_words), d_model)

query_x = query_embedding(torch.arange(len(query_words)).unsqueeze(0))
kv_x = doc_embedding(torch.arange(len(doc_words)).unsqueeze(0))

cross_attn = CrossSelfAttention(d_model=d_model, d_k=d_model)
out, weights = cross_attn(query_x, kv_x=kv_x)

print("출력 shape:", out.shape)        # (1, 2, 8)
print("가중치 shape:", weights.shape)  # (1, 2, 6)

sentence9_reversed = sentence9[::-1]
reversed_ids = torch.tensor([[vocab9[w] for w in sentence9_reversed]])

with torch.no_grad():
    x_original = embedding9(token_ids9)
    x_reversed = embedding9(reversed_ids)

    out_original_noPE, _ = attn9(x_original)
    out_reversed_noPE, _ = attn9(x_reversed)
    same_noPE = torch.allclose(out_original_noPE[0, 0], out_reversed_noPE[0, -1], atol=1e-5)
    print("[학습된 모델, PE 없음] 같은 토큰(나는)의 출력이 자리만 바뀌어도 동일한가?", same_noPE)

    pe = torch.tensor(positional_encoding(len(sentence9), d_model), dtype=torch.float32)
    out_original_PE, _ = attn9(x_original + pe)
    out_reversed_PE, _ = attn9(x_reversed + pe)
    same_PE = torch.allclose(out_original_PE[0, 0], out_reversed_PE[0, -1], atol=1e-5)
    print("[학습된 모델, PE 적용] 이제는 동일한가?", same_PE)

for n_layers in [2, 4, 6]:
    torch.manual_seed(42)
    encoder_test = Encoder(d_model=d_model, num_heads=2, d_ff=d_ff, num_layers=n_layers)
    out = encoder_test(x)
    num_params = sum(p.numel() for p in encoder_test.parameters())
    print(f"num_layers={n_layers}: 출력 shape={tuple(out.shape)}, 파라미터 수={num_params:,}")

sentence_A = ["나는", "오늘", "랭체인을", "공부한다"]
sentence_B = ["오늘", "날씨가", "정말", "좋다"]

encoder_out_A = encoder(src_embedding(ids_A))
encoder_out_B = encoder(src_embedding(ids_B))

out_with_A, cross_w_A = decoder_layer(decoder_input, encoder_out_A, causal_mask_dec)
out_with_B, cross_w_B = decoder_layer(decoder_input, encoder_out_B, causal_mask_dec)

print("decoder_out 완전히 동일한가?", torch.allclose(out_with_A, out_with_B))
print("두 출력의 차이(절댓값 평균):", (out_with_A - out_with_B).abs().mean().item())