import torch
import torch.nn as nn

# ── 재현성 고정 / Fix reproducibility ─────────────────────────────
torch.manual_seed(42)

# ── 하이퍼파라미터 / Hyperparameters ──────────────────────────────
B, T, C = 2, 4, 6   # batch_size=2, seq_len=4, embed_dim=6

# ── 패딩 포함 입력 생성 / Create padded input ─────────────────────
# shape: (B, T, C) = (2, 4, 6)
# - 문장1: 실제 토큰 3개 + <PAD> 1개
# - 문장2: 실제 토큰 2개 + <PAD> 2개
x = torch.tensor([
    # sentence 1
    [[1., 2., 3., 4., 5., 6.],    # token 1
     [7., 8., 9., 1., 2., 3.],    # token 2
     [4., 5., 6., 7., 8., 9.],    # token 3
     [0., 0., 0., 0., 0., 0.]],   # <PAD>
    # sentence 2
    [[2., 4., 6., 8., 1., 3.],    # token 1
     [5., 7., 9., 2., 4., 6.],    # token 2
     [0., 0., 0., 0., 0., 0.],    # <PAD>
     [0., 0., 0., 0., 0., 0.]],   # <PAD>
])  # shape: (2, 4, 6)

# ════════════════════════════════════════════════════════════════
# 1. Batch Normalization
#    - nn.BatchNorm1d 입력 형태: (N, C) 또는 (N, C, L)
#    - NLP 시퀀스에 적용하려면 (B, T, C) → (B, C, T) permute 필요
#    - train() 모드: 현재 미니배치 통계(μ_B, σ_B²) 사용
# ════════════════════════════════════════════════════════════════
bn = nn.BatchNorm1d(num_features=C, affine=False)  # γ,β 없이 순수 정규화만 확인
bn.train()

x_for_bn = x.permute(0, 2, 1)          # (B, C, T) = (2, 6, 4)
with torch.no_grad():
    out_bn = bn(x_for_bn)               # shape: (2, 6, 4)
out_bn = out_bn.permute(0, 2, 1)        # 다시 (B, T, C) = (2, 4, 6)

# ════════════════════════════════════════════════════════════════
# 2. Layer Normalization
#    - normalized_shape=C: 마지막 차원(feature 축) 방향으로 정규화
#    - 입력 shape 그대로 유지: (B, T, C) → (B, T, C)
#    - 각 토큰 독립적으로 처리 → 패딩 토큰의 간섭 없음
# ════════════════════════════════════════════════════════════════
ln = nn.LayerNorm(normalized_shape=C, elementwise_affine=False)

with torch.no_grad():
    out_ln = ln(x)                      # shape: (2, 4, 6)

# ════════════════════════════════════════════════════════════════
# 3. 결과 출력 및 비교 / Output & Comparison
# ════════════════════════════════════════════════════════════════
print("=" * 60)
print("[입력값 / Input]")
print(f"  문장1 token1 : {x[0, 0].tolist()}")
print(f"  문장1 <PAD>  : {x[0, 3].tolist()}")

print("\n[Batch Normalization 출력]")
print("  ※ 열(Column) 방향 통계 → PAD(0) 값이 통계에 섞임")
print(f"  문장1 token1 : {out_bn[0, 0].numpy().round(4).tolist()}")
print(f"  문장1 <PAD>  : {out_bn[0, 3].numpy().round(4).tolist()}")
# PAD 토큰 출력이 0이 아님 → BN 통계가 PAD에 의해 왜곡되었음을 방증

print("\n[Layer Normalization 출력]")
print("  ※ 행(Row) 방향 통계 → PAD(0) 토큰은 스스로만 정규화됨")
print(f"  문장1 token1 : {out_ln[0, 0].numpy().round(4).tolist()}")
print(f"  문장1 <PAD>  : {out_ln[0, 3].numpy().round(4).tolist()}")
# PAD 토큰 출력 = 0 (모든 feature가 동일값 → σ=0 → NaN 방지용 ε로 처리 후 0)

# ════════════════════════════════════════════════════════════════
# 4. 수동 검증: LN 수식을 직접 계산하여 nn.LayerNorm과 대조
#    LN(x) = (x - μ) / sqrt(σ² + ε)
# ════════════════════════════════════════════════════════════════
print("\n[수동 검증 / Manual Verification  ─  문장1 token1]")
tok = x[0, 0]                                      # [1,2,3,4,5,6]
mu_manual  = tok.mean()
var_manual = tok.var(unbiased=False)               # LN은 biased variance 사용
ln_manual  = (tok - mu_manual) / (var_manual + 1e-5).sqrt()

print(f"  수동 계산    : {ln_manual.numpy().round(4).tolist()}")
print(f"  nn.LayerNorm : {out_ln[0, 0].numpy().round(4).tolist()}")
# 두 결과가 일치함을 확인