import torch
import torch.nn as nn
import matplotlib.pyplot as plt
import japanize_matplotlib
from ipywidgets import interact, IntSlider

def visualize_scaling_interactive(seed):
    # 1. 乱数のシードを設定（再現性の確保）
    torch.manual_seed(seed)
    
    # 2. 前提条件の設定
    d_k = 64  # 次元の数
    scaling_factor = torch.sqrt(torch.tensor(d_k, dtype=torch.float32))
    
    # 3. 生のインプット情報（スコア）の生成
    # 10個の単語に対する相性スコア。シードによって値が変化します。
    scores_raw = torch.randn(1, 10) * 10 
    
    # 4. スケーリング処理
    scores_scaled = scores_raw / scaling_factor

    # 5. 注目度（確率）への変換
    softmax = nn.Softmax(dim=-1)
    attn_raw = softmax(scores_raw)
    attn_scaled = softmax(scores_scaled)

    # --- 数値情報の出力 ---
    print(f"現在のシード値: {seed}")
    print(f"スケーリング係数 (√d_k): {scaling_factor.item():.2f}")
    print("-" * 50)
    
    # 表形式のように見せるための数値表示
    print("【1. 生のスコア】（Softmax前）")
    raw_list = [f"{x:.2f}" for x in scores_raw[0].numpy()]
    print(f"値: {' | '.join(raw_list)}")
    
    print("\n【2. スケーリング後のスコア】（Softmax前）")
    scaled_list = [f"{x:.2f}" for x in scores_scaled[0].numpy()]
    print(f"値: {' | '.join(scaled_list)}")
    
    print("-" * 50)
    print("【最終結果：スケーリングなしの注目度】")
    raw_attn_list = [f"{x*100:.1f}%" for x in attn_raw[0].detach().numpy()]
    print(f"確率: {' | '.join(raw_attn_list)}")
    
    print("\n【最終結果：スケーリングありの注目度】")
    scaled_attn_list = [f"{x*100:.1f}%" for x in attn_scaled[0].detach().numpy()]
    print(f"確率: {' | '.join(scaled_attn_list)}")

    # --- グラフの描画 ---
    fig, axes = plt.subplots(1, 2, figsize=(14, 5))
    
    axes[0].bar(range(10), attn_raw[0].detach().numpy(), color='salmon', alpha=0.8)
    axes[0].set_title(f"スケーリングなし（シード {seed}）\n特定の単語が独占しやすい")
    axes[0].set_ylabel("注目度（0.0 - 1.0）")
    axes[0].set_ylim(0, 1)
    axes[0].grid(axis='y', linestyle='--', alpha=0.6)

    axes[1].bar(range(10), attn_scaled[0].detach().numpy(), color='skyblue', alpha=0.8)
    axes[1].set_title(f"スケーリングあり（シード {seed}）\n全体をバランスよく見渡せる")
    axes[1].set_ylabel("注目度（0.0 - 1.0）")
    axes[1].set_ylim(0, 1)
    axes[1].grid(axis='y', linestyle='--', alpha=0.6)

    plt.tight_layout()
    plt.show()

# ウィジェットの設定：0から100までシードを選べる
interact(visualize_scaling_interactive, 
         seed=IntSlider(min=0, max=100, step=1, value=42, description='シード設定:'))
# 第1部
# 第5章：スポットライト ― Attentionの正体と全体図
# 【体験ドリル】Softmaxとスケーリングの効果