# 準備（ライブラリと共通関数の定義）
import torch
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import japanize_matplotlib
import math

# 共通で使用する生成関数
def get_positional_encoding(max_seq_len, d_model):
    pe = torch.zeros(max_seq_len, d_model)
    position = torch.arange(0, max_seq_len, dtype=torch.float).unsqueeze(1)
    div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
    pe[:, 0::2] = torch.sin(position * div_term)
    pe[:, 1::2] = torch.cos(position * div_term)
    return pe

# 基本パラメータ
MAX_LEN = 50
D_MODEL = 128
POS_TO_SHOW = 15
pe = get_positional_encoding(MAX_LEN, D_MODEL)


# 全体像の可視化（位置エンコーディング行列）
plt.figure(figsize=(12, 6))
sns.heatmap(pe.numpy(), cmap='RdBu', center=0)
plt.title(f"1. 位置エンコーディング行列 (全体像) [最大長:{MAX_LEN}, 次元数:{D_MODEL}]")
plt.xlabel("次元インデックス (i)")
plt.ylabel("トークンの位置 (pos)")
plt.show()


# 数理的機能の検証（相対距離の抽出）
plt.figure(figsize=(10, 8))
correlation = torch.matmul(pe, pe.transpose(0, 1))
sns.heatmap(correlation.numpy(), cmap='viridis')
plt.title("2. 相対距離の抽出 (PEベクトル同士のドット積による幾何学的相関)")
plt.xlabel("位置 B")
plt.ylabel("位置 A")
plt.show()


# 具体的構造の確認（ベクトル断面図）
plt.figure(figsize=(12, 6))
full_vector = pe[POS_TO_SHOW].numpy()

# 偶数次元 (sin): 青の丸 / 奇数次元 (cos): 緑の三角 / ベースライン(basefmt)は非表示
plt.stem(range(0, D_MODEL, 2), full_vector[0::2],
         linefmt='b-', markerfmt='bo', basefmt=' ', label='偶数次元 (sin)')
plt.stem(range(1, D_MODEL, 2), full_vector[1::2],
         linefmt='g-', markerfmt='g^', basefmt=' ', label='奇数次元 (cos)')

plt.title(f"3. 位置 {POS_TO_SHOW} におけるベクトル断面図 (偶数:● / 奇数:▲)")
plt.xlabel("次元インデックス")
plt.ylabel("エンコーディング値")
plt.legend()
plt.show()