// --- DOM Elements ---
const battleBtn = document.getElementById('battle-btn');
const enemyMessage = document.getElementById('enemy-message');
const enemyImage = document.getElementById('enemy-image');
const battleLog = document.getElementById('battle-log');
const playerModelLayers = document.getElementById('player-model-layers');
const layerInventory = document.getElementById('layer-inventory');
const messageLogArea = document.getElementById('message-log-area');
const animationModal = document.getElementById('animation-modal');
const vizArea = document.getElementById('visualization-area');
const predictionResult = document.getElementById('prediction-result');
const closeModalBtn = document.getElementById('close-modal-btn');
const playerHpBar = document.getElementById('player-hp-bar');
const enemyHpBar = document.getElementById('enemy-hp-bar');
const itemSelectionModal = document.getElementById('item-selection-modal');
const itemChoices = document.getElementById('item-choices');
const restartBtn = document.getElementById('restart-btn');
const layerDescription = document.getElementById('layer-description');
const itemInfoArea = document.getElementById('item-info-area');
const previewAnimationArea = document.getElementById('preview-animation-area');
const choiceInfoArea = document.getElementById('choice-info-area');
const choiceDescription = document.getElementById('choice-description');
const choicePreviewArea = document.getElementById('choice-preview-area');
const titleScreen = document.getElementById('title-screen');
const startGameBtn = document.getElementById('start-game-btn');
const gameContainer = document.getElementById('game-container');

// --- Game State & Config ---
let playerLayers = []; // 現在モデルに配置されているレイヤー
let playerHP = 100;
let enemyHP = 100;
let currentStage = 1;
let isBattleInProgress = false;
let draggedItem = null;
let dragOverIndex = null; // 並び替え先のインデックス
let wasDroppedSuccessfully = false; // ★★★ このフラグを追加
let currentEnemy = { image_b64: null, label: null }; // ★★★ クライアント側で敵の状態を保持
let ENEMY_MAX_HP = 100;
let selectedLayerForMove = null; // タップで移動対象として選択されたレイヤーのインデックス

const PLAYER_MAX_HP = 100;

const allAvailableLayers = [
    // --- Normal Items ---
    {
        id: 0, name: '畳み込み層 (3x3, 4フィルタ)', type: 'Conv2d', icon: 'fa-th-large', params: { out_channels: 4, kernel_size: 3 },
        rarity: 'normal',
        desc: '画像から特徴(エッジ等)を抽出するCNNの心臓部。<br><br><b>使い方:</b> 画像の「パーツ」を見つける専門家です。モデルの<span class="text-yellow-300">最初の方</span>に置いて、画像から形の特徴を捉えさせましょう。'
    },
    {
        id: 2, name: 'ReLU活性化関数', type: 'ReLU', icon: 'fa-chart-line', params: {},
        rarity: 'normal',
        desc: '負の値を0に変換し、モデルに非線形性を与え表現力を高めます。<br><br><b>使い方:</b> モデルが複雑な判断をするための「スイッチ」です。<span class="text-yellow-300">畳み込み層や全結合層の直後</span>に挟むのが定石です。'
    },
    {
        id: 3, name: '最大プーリング (2x2)', type: 'MaxPool2d', icon: 'fa-compress-arrows-alt', params: { kernel_size: 2 },
        rarity: 'normal',
        desc: '情報を圧縮し、位置ズレに強いモデルを作ります。<br><br><b>使い方:</b> 画像の「要約」を行い、重要な部分だけを残します。<span class="text-yellow-300">畳み込み層(とReLU)の後</span>に入れると、より頑健なモデルになります。'
    },
    {
        id: 4, name: '平坦化層', type: 'Flatten', icon: 'fa-stream', params: {},
        rarity: 'normal',
        desc: '2次元の画像データを1次元に変換し、全結合層に渡せるようにします。<br><br><b>使い方:</b> 画像処理パートから最終判断パートへの「橋渡し」役です。モデルの<span class="text-yellow-300">中盤に必ず1つ</span>だけ配置してください。'
    },
    {
        id: 5, name: '全結合層 (16ノード)', type: 'Linear', icon: 'fa-braille', params: { out_features: 16 },
        rarity: 'normal',
        desc: '全ての特徴を結合し、最終的な分類を行います。<br><br><b>使い方:</b> これまでの情報を元に「最終判断」を下す賢者です。モデルの<span class="text-yellow-300">最後の方</span>、平坦化層の後に置きます。'
    },
    {
        id: 8, name: '平均プーリング (2x2)', type: 'AvgPool2d', icon: 'fa-wave-square', params: { kernel_size: 2 },
        rarity: 'normal',
        desc: '範囲内の特徴を「平均化」して情報を圧縮します。<br><br><b>使い方:</b> 最大プーリングと似ていますが、より滑らかに情報を要約します。<span class="text-yellow-300">畳み込み層(とReLU)の後</span>に使い、最大プーリングと使い比べてみましょう。'
    },

    // --- Gold Rare Items ---
    {
        id: 1, name: '畳み込み層 (5x5, 8フィルタ)', type: 'Conv2d', icon: 'fa-border-all', params: { out_channels: 8, kernel_size: 5 },
        rarity: 'gold',
        desc: 'より広い範囲の特徴を抽出する強力な畳み込み層。<br><br><b>使い方:</b> 3x3より広い範囲を見るため、より大局的な特徴を捉えます。これもモデルの<span class="text-yellow-300">最初の方</span>に置きます。'
    },
    {
        id: 6, name: '全結合層 (64ノード)', type: 'Linear', icon: 'fa-sitemap', params: { out_features: 64 },
        rarity: 'gold',
        desc: 'より多くのパラメータを持つ、より強力な全結合層。<br><br><b>使い方:</b> 16ノードより賢い賢者ですが、学習に少し時間がかかります。これも<span class="text-yellow-300">最後の方</span>に置きます。'
    },
    {
        id: 7, name: 'ドロップアウト (p=0.5)', type: 'Dropout', icon: 'fa-random', params: { p: 0.5 },
        rarity: 'gold',
        desc: '学習中にノードをランダムに無効化し、過学習を防ぎます。<br><br><b>使い方:</b> モデルが特定の情報に頼りすぎるのを防ぐ「保険」です。未知の敵に強くなります。<span class="text-yellow-300">全結合層の間</span>に挟むのが効果的です。'
    },
    {
        id: 9, name: '残差ブロック', type: 'ResidualBlock', icon: 'fa-project-diagram', params: {},
        rarity: 'gold',
        desc: '入力情報を「近道」させ、深いモデルの学習を安定させます。<br><br><b>使い方:</b> 情報を失わずに変換を加える特殊なブロックです。モデルが深くなりすぎた時に<span class="text-yellow-300">全結合層の代わり</span>に入れると、性能が改善することがあります。'
    }
];

// ★★★ インベントリを所持数管理に変更
let playerInventory = {}; // { layerId: count, ... }
const ANIMATION_SPEED = 0.7;

// --- Helper & UI Functions ---
const sleep = (ms) => new Promise(resolve => setTimeout(resolve, ms * ANIMATION_SPEED));

async function animateBattleLog(text, clear = true) {
    if (clear) battleLog.textContent = '';
    for (let i = 0; i < text.length; i++) {
        battleLog.textContent += text[i];
        await sleep(30);
    }
}

function logMessage(message, type = 'info') {
    const colors = { info: 'text-gray-400', success: 'text-green-400', error: 'text-red-400', action: 'text-yellow-400' };
    const p = document.createElement('p');
    // メッセージが追加されるときに少し遅延させることで、アニメーションが目に見えるようにする
    p.style.animationDelay = `${messageLogArea.childElementCount * 50}ms`;
    p.innerHTML = `> ${message}`;
    p.className = colors[type];
    messageLogArea.appendChild(p);
    messageLogArea.scrollTop = messageLogArea.scrollHeight;
}

function initializeUI() {
    logMessage('Machine Learning RPGへようこそ！');
    logMessage('インベントリのアイテムをドラッグしてモデルを構築しましょう。');
    startGame();
}

function drawLayerConnections() {
    const canvas = document.getElementById('layer-connections'); // Get canvas element
    if (!canvas) return;
    const ctx = canvas.getContext('2d');
    const parent = playerModelLayers;

    // Canvasのサイズを親要素に合わせる
    canvas.width = parent.clientWidth;
    canvas.height = parent.scrollHeight;

    ctx.clearRect(0, 0, canvas.width, canvas.height);

    const layers = parent.querySelectorAll('.player-layer');
    if (layers.length < 2) return;

    ctx.strokeStyle = 'rgba(0, 242, 255, 0.4)';
    ctx.lineWidth = 2;
    ctx.shadowColor = 'rgba(0, 242, 255, 0.8)';
    ctx.shadowBlur = 10;

    for (let i = 0; i < layers.length - 1; i++) {
        const startEl = layers[i];
        const endEl = layers[i + 1];

        const startRect = startEl.getBoundingClientRect();
        const endRect = endEl.getBoundingClientRect();
        const parentRect = parent.getBoundingClientRect();

        const startX = startRect.left + startRect.width / 2 - parentRect.left;
        const startY = startRect.bottom - parentRect.top + parent.scrollTop;
        const endX = endRect.left + endRect.width / 2 - parentRect.left;
        const endY = endRect.top - parentRect.top + parent.scrollTop;

        ctx.beginPath();
        ctx.moveTo(startX, startY);
        ctx.bezierCurveTo(startX, startY + 30, endX, endY - 30, endX, endY);
        ctx.stroke();
    }
}


function removeLayer(index) {
    const removedLayer = playerLayers.splice(index, 1)[0];
    returnInventoryItem(removedLayer);
    logMessage(`Layer Removed: <span class="text-indigo-300">${removedLayer.name}</span>`, 'action');
    updatePlayerModelUI();
    updatePlayerInventoryUI();
}

function updatePlayerModelUI() {
    const modelArea = playerModelLayers;
    // Canvas以外の要素をクリア
    Array.from(modelArea.children).forEach(child => {
        if (child.tagName !== 'CANVAS') {
            child.remove();
        }
    });

    const createDropZone = (index) => {
        const zone = document.createElement('div');
        zone.className = 'player-layer-drop-zone h-4'; // スマホで見えるように少し高さを付ける
        zone.dataset.index = index;
        zone.addEventListener('click', () => handleLayerMove(index));
        modelArea.appendChild(zone);
    };


    if (playerLayers.length === 0) {
        const p = document.createElement('p');
        p.className = 'text-gray-400 text-center p-4';
        p.textContent = 'インベントリからレイヤー（層）をここにドラッグ＆ドロップ or タップしてください。';
        modelArea.appendChild(p);
    } else {
        createDropZone(0); // 最初のレイヤーの上にドロップゾーンを作成
        playerLayers.forEach((layer, index) => {
            const div = document.createElement('div');
            div.className = `player-layer ${layer.rarity === 'gold' ? 'gold-rare' : ''}`;
            div.draggable = true;
            div.dataset.index = index;

            // ★★★ 選択状態のスタイルを適用 ★★★
            if (selectedLayerForMove === index) {
                div.classList.add('border-cyan-400', 'scale-105');
            }

            div.innerHTML = `<span><i class="fas ${layer.icon} mr-2"></i>${layer.name}</span><button class="text-red-400 hover:text-red-200 text-lg">&times;</button>`;
            div.querySelector('button').onclick = (e) => {
                e.stopPropagation(); // 親要素のクリックイベントを発火させない
                removeLayer(index);
            };

            // ドラッグ＆ドロップイベント (PC用)
            div.addEventListener('dragstart', (e) => handleDragStart(e, index, layer));
            div.addEventListener('dragend', handleDragEnd);
            div.addEventListener('dragenter', (e) => handleDragEnterModelItem(e, index));
            div.addEventListener('dragleave', (e) => e.currentTarget.classList.remove('is-dragged-over'));

            // ★★★ タップイベント (スマホ用) ★★★
            div.addEventListener('click', () => handleLayerSelectForMove(index));

            modelArea.appendChild(div);
            createDropZone(index + 1); // 各レイヤーの下にドロップゾーンを作成
        });
    }
    battleBtn.disabled = playerLayers.length === 0;
    requestAnimationFrame(drawLayerConnections);
}

function updatePlayerInventoryUI() {
    layerInventory.innerHTML = '';
    Object.keys(playerInventory).forEach(layerId => {
        const count = playerInventory[layerId];
        if (count > 0) {
            const layer = allAvailableLayers.find(l => l.id == layerId);
            const item = document.createElement('div');
            item.className = `layer-item min-h-16 relative ${layer.rarity === 'gold' ? 'gold-rare' : ''}`;
            item.draggable = true;
            item.dataset.id = layer.id;

            // ★★★ HTML構造を新しいレイアウト用に変更
            item.innerHTML = `
                <i class="fas ${layer.icon} layer-icon"></i>
                <p>${layer.name}</p>
                <div class="ml-auto bg-indigo-500 text-white text-xs font-bold rounded-full h-6 w-6 flex items-center justify-center">${count}</div>
            `;

            // イベントリスナー
            item.addEventListener('click', () => {
                // PCでの説明表示機能は維持
                showDescription(layer, document.getElementById('layer-description'));
                runPreviewAnimation(layer, previewAnimationArea);
                // スマホ用の追加機能
                addLayerFromInventory(layer);
            });
            item.addEventListener('dragstart', (e) => handleInventoryDragStart(e, layer));
            item.addEventListener('dragend', handleDragEnd);

            // ゴールドレアの場合、キラキラエフェクトを追加
            if (layer.rarity === 'gold') {
                const sparkleContainer = document.createElement('div');
                sparkleContainer.className = 'sparkle-container';
                for (let i = 0; i < 3; i++) {
                    const sparkle = document.createElement('div');
                    sparkle.className = 'sparkle';
                    sparkle.style.top = `${Math.random() * 100}%`;
                    sparkle.style.left = `${Math.random() * 100}%`;
                    sparkle.style.animationDelay = `${Math.random() * 1.5}s`;
                    sparkleContainer.appendChild(sparkle);
                }
                item.appendChild(sparkleContainer);
            }

            layerInventory.appendChild(item);
        }
    });
}

function showDescription(layer, element) {
    element.innerHTML = `<p class="font-bold text-cyan-300">${layer.name}</p><p>${layer.desc}</p>`;
}

function updateHpBars() {
    const playerHpPercent = Math.max(0, (playerHP / PLAYER_MAX_HP) * 100);
    playerHpBar.style.width = `${playerHpPercent}%`;
    // ★★★ 表示を「現在HP / 最大HP」に変更
    playerHpBar.textContent = `${Math.max(0, playerHP)} / ${PLAYER_MAX_HP}`;

    const enemyHpPercent = Math.max(0, (enemyHP / ENEMY_MAX_HP) * 100);
    enemyHpBar.style.width = `${enemyHpPercent}%`;
    // ★★★ 表示を「現在HP / 最大HP」に変更
    enemyHpBar.textContent = `${Math.max(0, enemyHP)} / ${ENEMY_MAX_HP}`;

    // Flash effect on change
    if (playerHpBar.dataset.lastHp && playerHpBar.dataset.lastHp != playerHP) {
        playerHpBar.parentElement.classList.add('flash-damage');
        setTimeout(() => playerHpBar.parentElement.classList.remove('flash-damage'), 300);
    }
    if (enemyHpBar.dataset.lastHp && enemyHpBar.dataset.lastHp != enemyHP) {
        enemyHpBar.parentElement.classList.add('flash-damage');
        setTimeout(() => enemyHpBar.parentElement.classList.remove('flash-damage'), 300);
    }
    playerHpBar.dataset.lastHp = playerHP;
    enemyHpBar.dataset.lastHp = enemyHP;
}

async function fetchNewEnemy() {
    const response = await fetch('/api/get_enemy');
    const enemyData = await response.json();

    // ★★★ グローバル変数に保存
    currentEnemy = {
        image_b64: enemyData.image_b64,
        label: enemyData.label
    };

    enemyMessage.textContent = '野生のMNISTモンスターが現れた！';
    enemyImage.src = currentEnemy.image_b64;
    enemyImage.classList.remove('hidden');
    await animateBattleLog('', true);
}

// --- D&D Functions ---

// ドラッグ開始時の処理
function handleDragStart(e, index, layer) {
    draggedItem = { type: 'model', index: index, layer: layer, element: e.target };
    setTimeout(() => e.target.classList.add('dragging'), 0);
}

function handleInventoryDragStart(e, layer) {
    // ★★★ 新しい一意なインスタンスを作成
    const layerInstance = {
        ...layer,
        instanceId: `inst_${Date.now()}_${Math.random()}`
    };
    draggedItem = { type: 'inventory', layer: layerInstance, element: e.target };
    wasDroppedSuccessfully = false; // ★★★ ドラッグ開始時にフラグをリセット
    setTimeout(() => e.target.classList.add('dragging'), 0);
}

// ドラッグ終了時の処理
function handleDragEnd(e) {
    // ★★★ バグ修正：キャンセル時のクリーンアップ処理 ★★★
    if (draggedItem && draggedItem.type === 'inventory' && !wasDroppedSuccessfully) {
        // インベントリからのドラッグが、モデルエリアにドロップされずに終了した場合
        const tempIndex = playerLayers.findIndex(l => l.instanceId === draggedItem.layer.instanceId);
        if (tempIndex > -1) {
            // モデルデータから仮追加されたアイテムを削除
            playerLayers.splice(tempIndex, 1);
            logMessage('モデルへの追加をキャンセルしました。', 'info');
            // UIを再描画して見た目を元に戻す
            updatePlayerModelUI();
        }
    }

    // --- 既存のクリーンアップ処理 ---
    // is-dragged-over クラスをすべて削除
    document.querySelectorAll('.is-dragged-over').forEach(el => el.classList.remove('is-dragged-over'));
    // dragging クラスを削除
    if (draggedItem && draggedItem.element) {
        draggedItem.element.classList.remove('dragging');
    }
    draggedItem = null;
}

// ドロップを許可するエリアの処理
function allowDrop(ev) {
    ev.preventDefault();
    const modelArea = document.getElementById('player-model-layers');
    if (modelArea.contains(ev.target)) {
        modelArea.classList.add('drag-over');
    }
}

// モデルエリアへのドロップ処理 (メインロジック)
function dropOnModelArea(ev) {
    ev.preventDefault();
    const modelArea = document.getElementById('player-model-layers');
    modelArea.classList.remove('drag-over');

    if (!draggedItem) return;

    // ★★★ ドロップが成功した（試みられた）ことを記録
    if (draggedItem.type === 'inventory') {
        wasDroppedSuccessfully = true;
    }

    if (draggedItem.type === 'inventory') {
        const originalLayer = allAvailableLayers.find(l => l.id === draggedItem.layer.id);
        if (!useInventoryItem(originalLayer)) {
            logMessage(`インベントリに ${draggedItem.layer.name} がありません`, 'error');
            const tempIndex = playerLayers.findIndex(l => l.instanceId === draggedItem.layer.instanceId);
            if (tempIndex > -1) playerLayers.splice(tempIndex, 1);
        } else {
            logMessage(`レイヤー追加: ${draggedItem.layer.name}`, 'action');
        }
    } else {
        logMessage('モデルの順序を変更しました。', 'info');
    }

    updatePlayerModelUI();
    updatePlayerInventoryUI();
    handleDragEnd({ target: ev.target });
}

// モデルエリア内のアイテムにドラッグが入ったときの処理
function handleDragEnterModelItem(e, targetIndex) {
    e.preventDefault();
    const targetItem = e.currentTarget;

    if (!draggedItem || (draggedItem.type === 'model' && draggedItem.layer.instanceId === playerLayers[targetIndex].instanceId)) {
        return;
    }

    document.querySelectorAll('.is-dragged-over').forEach(el => el.classList.remove('is-dragged-over'));
    targetItem.classList.add('is-dragged-over');

    let currentIndex = -1;
    // ★★★ instanceId を使ってドラッグ中のアイテムを検索
    if (draggedItem.layer.instanceId) {
        currentIndex = playerLayers.findIndex(l => l.instanceId === draggedItem.layer.instanceId);
    }

    let movedLayer;
    if (currentIndex > -1) {
        [movedLayer] = playerLayers.splice(currentIndex, 1);
    } else {
        movedLayer = draggedItem.layer;
    }

    playerLayers.splice(targetIndex, 0, movedLayer);

    if (draggedItem.type === 'model') {
        draggedItem.index = targetIndex;
    }

    updatePlayerModelUI();
}

// ★★★ モデルエリアの「何もない部分」にドラッグが入ったときの処理を追加
function handleDragEnterModelArea(e) {
    if (e.target.id !== 'player-model-layers') {
        return;
    }
    document.querySelectorAll('.is-dragged-over').forEach(el => el.classList.remove('is-dragged-over'));

    if (!draggedItem) return;

    // ★★★ instanceId を使ってドラッグ中のアイテムを検索
    let currentIndex = -1;
    if (draggedItem.layer.instanceId) {
        currentIndex = playerLayers.findIndex(l => l.instanceId === draggedItem.layer.instanceId);
    }

    if (currentIndex > -1) {
        if (currentIndex === playerLayers.length - 1) return; // 既に末尾なら何もしない
        const [movedLayer] = playerLayers.splice(currentIndex, 1);
        playerLayers.push(movedLayer);
        if (draggedItem.type === 'model') {
            draggedItem.index = playerLayers.length - 1;
        }
    } else {
        playerLayers.push(draggedItem.layer);
    }

    updatePlayerModelUI();
}

// ★★★ モデルエリアからドラッグが出たときのクリーンアップ処理を追加
function handleDragLeaveModelArea(e) {
    const modelArea = document.getElementById('player-model-layers');
    // ★★★ エリア外に出た判定のみ残し、データ操作ロジックは削除
    if (!modelArea.contains(e.relatedTarget)) {
        modelArea.classList.remove('drag-over');
    }
}

// --- Game Flow (バグ修正) ---
function startGame() {
    playerHP = PLAYER_MAX_HP;
    currentStage = 1;
    playerInventory = {};
    playerLayers = []; // ★★★ ここでモデルの状態もリセットする
    isBattleInProgress = false; // バトル状態をリセット

    // 初期アイテム
    addInventoryItem(allAvailableLayers.find(l => l.id === 0));
    addInventoryItem(allAvailableLayers.find(l => l.id === 5));
    logMessage(`ゲーム開始！初期インベントリを獲得しました。`, 'success');

    // ... (UIリセット処理は変更なし) ...
    battleBtn.classList.remove('hidden');
    restartBtn.classList.add('hidden');
    battleBtn.disabled = true; // モデルが空なので最初は無効
    battleBtn.innerHTML = '<i class="fas fa-fist-raised"></i> バトル開始！';
    battleLog.textContent = '';

    startStage();
}

// --- Game Flow (タイトル画面対応) ---

// ★★★ ゲーム初期化とゲーム開始を分離
function initializeGame() {
    logMessage('Machine Learning RPGへようこそ！');
    logMessage('インベントリのアイテムをドラッグしてモデルを構築しましょう。');

    playerHP = PLAYER_MAX_HP;
    currentStage = 1;
    playerInventory = {};
    playerLayers = [];
    isBattleInProgress = false;

    addInventoryItem(allAvailableLayers.find(l => l.id === 0));
    addInventoryItem(allAvailableLayers.find(l => l.id === 5));
    logMessage(`ゲーム開始！初期インベントリを獲得しました。`, 'success');

    battleBtn.classList.remove('hidden');
    restartBtn.classList.add('hidden');
    restartBtn.textContent = 'タイトルへ戻る'; // テキストを統一
    battleBtn.disabled = true;
    battleBtn.innerHTML = '<i class="fas fa-fist-raised"></i> バトル開始！';
    battleLog.textContent = '';
    messageLogArea.innerHTML = ''; // メッセージログもクリア

    startStage();
}

function handleGameOver() {
    logMessage('ゲームオーバー...', 'error');
    animateBattleLog('敗北...'); // アニメーション付きに変更
    battleBtn.classList.add('hidden');
    restartBtn.classList.remove('hidden');
    // ★★★ ボタンのテキストを明示的に変更
    restartBtn.innerHTML = '<i class="fas fa-undo"></i> タイトルへ戻る';
    isBattleInProgress = false;
}

function startStage() {
    logMessage(`--- Stage ${currentStage} Start ---`, 'action');
    playerHP = PLAYER_MAX_HP;

    // ★★★ ステージに応じて敵のHPを計算・設定
    ENEMY_MAX_HP = 100 * currentStage;
    enemyHP = ENEMY_MAX_HP;

    // ★★★ 構築済みモデルをインベントリに戻す
    if (playerLayers.length > 0) {
        // playerLayersの各アイテムをインベントリに戻す
        playerLayers.forEach(layer => returnInventoryItem(layer));
        playerLayers = []; // モデルを空にする
        logMessage('Your previous model has been returned to inventory.', 'info');
    }

    updateHpBars();
    updatePlayerModelUI();
    fetchNewEnemy();
    updatePlayerInventoryUI();

    // ★★★ 最終ステージクリア後はアイテム選択画面を出さない
    if (currentStage <= 5) {
        showItemSelection();
    }
}

function showItemSelection() {
    itemSelectionModal.classList.remove('hidden');
    itemSelectionModal.classList.add('flex');
    itemChoices.innerHTML = '';

    const confirmBtn = document.getElementById('confirm-choice-btn');
    const choiceFooter = document.getElementById('choice-selection-footer');
    const selectedChoiceInput = document.getElementById('selected-choice-id');
    confirmBtn.disabled = true;
    choiceFooter.classList.add('opacity-0');
    selectedChoiceInput.value = '';

    // ★★★ ステージに応じたゴールドレア出現確率を計算
    // ステージ1: 10%, ステージ2: 20%, ステージ3: 30%, ステージ4: 40%, ステージ5: 50%
    const goldChance = Math.min(0.1 * currentStage, 0.5);

    // プレイヤーがまだ持っていないレイヤーをフィルタリング
    const unownedLayers = allAvailableLayers.filter(layer => !playerInventory[layer.id] || playerInventory[layer.id] === 0);
    const normalChoices = unownedLayers.filter(l => l.rarity === 'normal');
    const goldChoices = unownedLayers.filter(l => l.rarity === 'gold');

    const finalChoices = [];
    const numChoices = 3;

    for (let i = 0; i < numChoices; i++) {
        let chosenLayer = null;
        // 確率判定でゴールドレアを引くか、通常枠しか残っていない場合
        if (Math.random() < goldChance && goldChoices.length > 0) {
            const index = Math.floor(Math.random() * goldChoices.length);
            chosenLayer = goldChoices.splice(index, 1)[0];
        }
        // 通常枠を引くか、ゴールド枠が空の場合
        else if (normalChoices.length > 0) {
            const index = Math.floor(Math.random() * normalChoices.length);
            chosenLayer = normalChoices.splice(index, 1)[0];
        }
        // それでも選択肢がなければ、残っている方から引く
        else if (goldChoices.length > 0) {
            const index = Math.floor(Math.random() * goldChoices.length);
            chosenLayer = goldChoices.splice(index, 1)[0];
        }

        if (chosenLayer) {
            finalChoices.push(chosenLayer);
        }
    }

    if (finalChoices.length === 0) {
        logMessage('全てのレイヤーを収集しました！', 'success');
        itemSelectionModal.classList.add('hidden');
        itemSelectionModal.classList.remove('flex');
        return;
    }

    // 選択肢のUIを生成
    finalChoices.forEach(layer => {
        const item = document.createElement('div');
        item.className = `layer-item w-48 h-48 flex flex-col justify-center ${layer.rarity === 'gold' ? 'gold-rare' : ''}`;
        item.innerHTML = `<i class="fas ${layer.icon} layer-icon text-5xl"></i><p class="text-lg mt-2">${layer.name}</p>`;
        item.dataset.id = layer.id; // ★★★ IDをデータ属性として保持

        // ★★★ ゴールドレアの場合、キラキラエフェクトを追加
        if (layer.rarity === 'gold') {
            const sparkleContainer = document.createElement('div');
            sparkleContainer.className = 'sparkle-container';
            for (let i = 0; i < 5; i++) { // モーダルでは少し多めに
                const sparkle = document.createElement('div');
                sparkle.className = 'sparkle';
                sparkle.style.top = `${Math.random() * 100}%`;
                sparkle.style.left = `${Math.random() * 100}%`;
                sparkle.style.animationDelay = `${Math.random() * 1.5}s`;
                sparkleContainer.appendChild(sparkle);
            }
            item.appendChild(sparkleContainer);
        }

                // ★★★ クリック/タップ時の動作を変更 ★★★
        item.addEventListener('click', () => {
            // 他のアイテムの選択状態を解除
            document.querySelectorAll('#item-choices .layer-item').forEach(el => {
                el.classList.remove('border-cyan-400', 'scale-105');
            });
            // このアイテムを選択状態にする
            item.classList.add('border-cyan-400', 'scale-105');

            // 説明とプレビューを表示
            showDescription(layer, choiceDescription);
            runPreviewAnimation(layer, choicePreviewArea);

            // 決定ボタンを有効化し、選択したIDを保持
            selectedChoiceInput.value = layer.id;
            confirmBtn.disabled = false;
            choiceFooter.classList.remove('opacity-0');
        });

        itemChoices.appendChild(item);
    });
}

function selectItem(selectedLayer) {
    addInventoryItem(selectedLayer);
    logMessage(`You got a new layer: <span class="text-indigo-300">${selectedLayer.name}</span>!`, 'success');
    updatePlayerInventoryUI();
    itemSelectionModal.classList.add('hidden');
    itemSelectionModal.classList.remove('flex');

    // ★★★ アイテム選択後、モデルをインベントリに戻す
    if (playerLayers.length > 0) {
        playerLayers.forEach(layer => returnInventoryItem(layer));
        playerLayers = [];
        logMessage('Model reset. Please rebuild your model.', 'info');
        updatePlayerModelUI();
        updatePlayerInventoryUI();
    }
}

// --- PREVIEW ANIMATION ENGINE ---
function generateDummyData(shape = [4, 14, 14]) {
    // [channels, height, width]
    return Array.from({ length: shape[0] }, () =>
        Array.from({ length: shape[1] }, () =>
            Array.from({ length: shape[2] }, () => Math.random() * 2 - 1)
        )
    );
}

// --- PREVIEW ANIMATION ENGINE ---
async function runPreviewAnimation(layer, previewArea) {
    const currentAnimationId = Date.now();
    previewArea.dataset.animationId = currentAnimationId;
    previewArea.innerHTML = '';
    const vizArea = previewArea;
    const checkInterrupted = () => previewArea.dataset.animationId != currentAnimationId;

    // ★★★ プレビュー要素に適用する基本スタイル
    const previewElementStyle = (el) => {
        el.style.position = 'absolute';
        el.style.left = '50%';
        el.style.top = '50%';
        el.style.transform = 'translate(-50%, -50%)';
    };

    let fromEl, toEl, inputData;
    switch (layer.type) {
        case 'Conv2d':
        case 'MaxPool2d':
            fromEl = document.createElement('div');
            previewElementStyle(fromEl);
            previewArea.appendChild(fromEl);

            inputData = (layer.type === 'Conv2d') ? generateDummyData([1, 28, 28]) : generateDummyData([4, 28, 28]);
            await animateGrid(fromEl, inputData, { isInput: true, duration: 500 }, vizArea);
            if (checkInterrupted()) return;

            toEl = document.createElement('div');
            previewElementStyle(toEl);
            previewArea.appendChild(toEl);

            // fromElをすぐに消さず、toElが生成されるのを待つ
            fromEl.style.transition = 'opacity 0.5s ease-out 0.5s'; // 少し遅れてフェードアウト
            fromEl.style.opacity = '0';

            if (layer.type === 'Conv2d') {
                await animateConv(fromEl, toEl, generateDummyData([layer.params.out_channels, 28, 28]), { duration: 1200, isPreview: true }, vizArea);
            } else {
                await animatePool(fromEl, toEl, generateDummyData([4, 14, 14]), { duration: 1200, isPreview: true }, vizArea);
            }
            break;

        case 'AvgPool2d':
            fromEl = document.createElement('div');
            previewElementStyle(fromEl);
            previewArea.appendChild(fromEl);

            inputData = (layer.type === 'Conv2d') ? generateDummyData([1, 28, 28]) : generateDummyData([4, 28, 28]);
            await animateGrid(fromEl, inputData, { isInput: true, duration: 500 }, vizArea);
            if (checkInterrupted()) return;

            toEl = document.createElement('div');
            previewElementStyle(toEl);
            previewArea.appendChild(toEl);

            fromEl.style.transition = 'opacity 0.5s ease-out 0.5s';
            fromEl.style.opacity = '0';

            if (layer.type === 'Conv2d') {
                await animateConv(fromEl, toEl, generateDummyData([layer.params.out_channels, 28, 28]), { duration: 1200, isPreview: true }, vizArea);
            } else {
                const isAverage = layer.type === 'AvgPool2d';
                await animatePool(fromEl, toEl, generateDummyData([4, 14, 14]), { duration: 1200, isPreview: true, isAverage }, vizArea);
            }
            break;

        case 'ReLU':
        case 'Dropout':
            fromEl = document.createElement('div');
            fromEl.style.position = 'absolute';
            fromEl.style.left = '50%';
            fromEl.style.top = '50%';
            previewArea.appendChild(fromEl);
            const dummyDataWithNegatives = generateDummyData([4, 10, 10]);
            await animateGrid(fromEl, dummyDataWithNegatives, { duration: 500 }, vizArea); // ★ 時間を延長
            if (checkInterrupted()) return;

            if (layer.type === 'ReLU') {
                await animateReLU(fromEl, { duration: 800 }); // ★ 時間を延長
            } else {
                await animateDropout(fromEl, { duration: 800 }); // ★ 時間を延長
            }
            break;

        case 'Flatten':
            fromEl = document.createElement('div');
            previewArea.appendChild(fromEl);
            await animateGrid(fromEl, generateDummyData([4, 14, 14]), { duration: 600 }, vizArea);
            if (checkInterrupted()) return;

            toEl = document.createElement('div');
            previewArea.appendChild(toEl);
            await animateFlatten(fromEl, toEl, 4 * 14 * 14, { duration: 1000 }, vizArea); // ★ 時間を延長
            if (checkInterrupted()) return;

            fromEl.style.opacity = 0;
            break;

        case 'Linear':
            fromEl = document.createElement('div'); // Flattened bar
            previewElementStyle(fromEl); // 中央に配置
            fromEl.style.left = '25%'; // 左側に配置
            previewArea.appendChild(fromEl);
            await animateFlatten(null, fromEl, 100, { duration: 500 }, vizArea);
            if (checkInterrupted()) return;

            toEl = document.createElement('div'); // Nodes
            previewElementStyle(toEl); // 中央に配置
            toEl.style.left = '75%'; // 右側に配置
            previewArea.appendChild(toEl);
            await animateLinear(fromEl, toEl, Array(layer.params.out_features).fill(0), null, null, { duration: 1000 }, vizArea);
            if (checkInterrupted()) return;

            fromEl.style.opacity = 0;
            break;

        case 'ResidualBlock':
            fromEl = document.createElement('div');
            previewElementStyle(fromEl);
            fromEl.style.left = '25%'; // 左側に配置
            previewArea.appendChild(fromEl);
            await animateLinear(null, fromEl, Array(32).fill(0), null, null, { duration: 500, nodeSize: 8 }, vizArea);
            if (checkInterrupted()) return;

            toEl = document.createElement('div');
            previewElementStyle(toEl);
            toEl.style.left = '75%'; // 右側に配置
            previewArea.appendChild(toEl);
            // プレビューなので skipFromEl と fromEl は同じものを渡す
            await animateResidual(fromEl, fromEl, toEl, Array(32).fill(0), { duration: 1200, nodeSize: 8 }, vizArea);
            if (checkInterrupted()) return;

            fromEl.style.opacity = 0;
            break;
    }
}

// ★★★ インベントリからのタップでレイヤーを追加する関数 ★★★
function addLayerFromInventory(layer) {
    // 新しい一意なインスタンスを作成して追加
    const layerInstance = {
        ...layer,
        instanceId: `inst_${Date.now()}_${Math.random()}`
    };
    if (useInventoryItem(layerInstance)) {
        playerLayers.push(layerInstance);
        logMessage(`レイヤー追加: ${layerInstance.name}`, 'action');
        updatePlayerModelUI();
        updatePlayerInventoryUI();
    } else {
        logMessage(`インベントリに ${layer.name} がありません`, 'error');
    }
}

// ★★★ モデル内のレイヤーを移動のために「選択」する関数 ★★★
function handleLayerSelectForMove(index) {
    if (selectedLayerForMove === index) {
        // 同じレイヤーを再度タップしたら選択解除
        selectedLayerForMove = null;
        logMessage('レイヤーの移動をキャンセルしました。', 'info');
    } else {
        selectedLayerForMove = index;
        logMessage(`レイヤー '${playerLayers[index].name}' を選択しました。移動先の青いエリアをタップしてください。`, 'action');
    }
    updatePlayerModelUI(); // UIを再描画して選択状態を反映
}

// ★★★ 選択したレイヤーをドロップゾーンに「移動」する関数 ★★★
function handleLayerMove(targetIndex) {
    if (selectedLayerForMove === null) return; // 何も選択されていなければ何もしない

    // 移動するレイヤーを取得
    const [movedLayer] = playerLayers.splice(selectedLayerForMove, 1);

    // 削除によってインデックスがずれるのを補正
    const adjustedTargetIndex = selectedLayerForMove < targetIndex ? targetIndex - 1 : targetIndex;

    // 新しい場所にレイヤーを挿入
    playerLayers.splice(adjustedTargetIndex, 0, movedLayer);

    logMessage('レイヤーを移動しました。', 'success');
    selectedLayerForMove = null; // 移動が終わったら選択状態を解除
    updatePlayerModelUI();
}

// --- Main Game Logic ---
function addInventoryItem(layer) {
    playerInventory[layer.id] = (playerInventory[layer.id] || 0) + 1;
}

function useInventoryItem(layer) {
    if (playerInventory[layer.id] && playerInventory[layer.id] > 0) {
        playerInventory[layer.id]--;
        return true;
    }
    return false;
}

function returnInventoryItem(layer) {
    playerInventory[layer.id]++;
}

// ★★★ モデル構築ロジックを所持数システムに対応
function addLayer(layer) {
    if (useInventoryItem(layer)) {
        playerLayers.push(layer);
        logMessage(`Layer Added: <span class="text-indigo-300">${layer.name}</span>`, 'action');
        updatePlayerModelUI();
        updatePlayerInventoryUI();
    } else {
        logMessage(`You don't have any more <span class="text-indigo-300">${layer.name}</span>`, 'error');
    }
}


function validateArchitecture(layers) {
    if (layers.length === 0) {
        return { isValid: false, message: 'モデルが空です。' };
    }

    let isFlattened = false; // テンソルが平坦化されたかどうかを追跡

    for (let i = 0; i < layers.length; i++) {
        const currentLayerType = layers[i].type;

        if (isFlattened) {
            // 平坦化された後に入れることができない層
            if (['Conv2d', 'MaxPool2d', 'AvgPool2d', 'Flatten'].includes(currentLayerType)) {
                return {
                    isValid: false,
                    message: `無効な順序: ${layers[i - 1].name} の後には ${currentLayerType} を配置できません。一度平坦化すると元に戻せません。`
                };
            }
        } else {
            // 平坦化される前にしか入れられない層
            if (['Conv2d', 'MaxPool2d', 'AvgPool2d'].includes(currentLayerType)) {
                // OK
            } else if (currentLayerType === 'Flatten') {
                isFlattened = true;
            } else if (['Linear', 'Dropout', 'ResidualBlock'].includes(currentLayerType)) { // ★★★ ResidualBlockを追加
                isFlattened = true;
            }
        }
    }
    return { isValid: true, message: '有効なアーキテクチャです。' };
}

// --- Main Battle Logic ---
async function handleBattle() {
    const validationResult = validateArchitecture(playerLayers);
    if (!validationResult.isValid) {
        logMessage(`エラー: ${validationResult.message}`, 'error');
        await animateBattleLog(`モデル構成エラー！`);
        await sleep(2000);
        await animateBattleLog('', true);
        return;
    }

    if (isBattleInProgress) return;
    isBattleInProgress = true;
    battleBtn.disabled = true;

    // // --- フェーズ1: 訓練 ---
    // battleBtn.innerHTML = '<i class="fas fa-spinner fa-spin"></i> モデルを訓練中...';
    // await animateBattleLog('戦闘準備... モデルを訓練中...');
    // logMessage('モデルの訓練を開始しました...', 'info');

    // // EelからFetch APIに変更
    // const trainResponse = await fetch('/api/train_player_model', {
    //     method: 'POST',
    //     headers: {
    //         'Content-Type': 'application/json',
    //     },
    //     body: JSON.stringify(playerLayers),
    // });
    // const trainResult = await trainResponse.json();

    // if (!trainResult.success) {
    //     await animateBattleLog(`エラー: ${trainResult.message}`);
    //     logMessage(`訓練エラー: ${trainResult.message}`, 'error');
    //     isBattleInProgress = false;
    //     updatePlayerModelUI();
    //     battleBtn.innerHTML = '<i class="fas fa-fist-raised"></i> バトル開始！';
    //     return;
    // }
    // logMessage(trainResult.message, 'success');
    // await sleep(500);

    // --- フェーズ2: 戦闘ループ ---
    while (playerHP > 0 && enemyHP > 0) {
        battleBtn.innerHTML = '<i class="fas fa-spinner fa-spin"></i> 攻撃中...';
        await animateBattleLog('新たな敵をスキャン... 推論実行...');

        const inferenceResponse = await fetch('/api/run_inference', {
            method: 'POST',
            headers: { 'Content-Type': 'application/json' },
            body: JSON.stringify({
                layer_configs: playerLayers,
                enemy_image_b64: currentEnemy.image_b64,
                enemy_label: currentEnemy.label,
            }),
        });
        const result = await inferenceResponse.json();

        if (result.error) {
            // ★★★ エラーメッセージを正しく表示するように修正 ★★★
            // 「訓練エラー: undefined」の問題はここで発生していた
            await animateBattleLog(`エラー: ${result.error}`);
            logMessage(`推論エラー: ${result.error}`, 'error');
            break; // エラーが発生したらループを抜ける
        }

        enemyImage.src = result.image_b64;
        logMessage('推論完了！攻撃を可視化します...', 'success');

        animationModal.classList.remove('hidden');
        await runDynamicAnimation(result);

        const damage = Math.max(1, Math.round(result.confidence * 100));

        if (result.is_correct) {
            enemyHP -= damage;
            await animateBattleLog(`攻撃成功！ ${damage} のダメージ！`);
            logMessage(`正解！ ${damage} のダメージを与えた。(自信度: ${(result.confidence * 100).toFixed(1)}%) 敵HP: ${Math.max(0, enemyHP)}`, 'success');
        } else {
            playerHP -= damage;
            await animateBattleLog(`攻撃失敗！ ${damage} のダメージ！`);
            logMessage(`不正解！ ${damage} のダメージを受けた。(自信度: ${(result.confidence * 100).toFixed(1)}%) プレイヤーHP: ${Math.max(0, playerHP)}`, 'error');
        }
        updateHpBars();

        if (enemyHP > 0 && playerHP > 0) {
            await sleep(1500);
            await fetchNewEnemy(); // 次の敵を取得
        }
    }

    // --- フェーズ3: バトル終了後処理 ---
    if (enemyHP <= 0) {
        // ★★★ ステージ5をクリアしたか判定
        if (currentStage >= 5) {
            logMessage(`CONGRATULATIONS! 全てのステージをクリアしました！`, 'success');
            await animateBattleLog('GAME CLEAR!');
            enemyMessage.textContent = '全てのMNISTモンスターを倒した！';
            battleBtn.classList.add('hidden');
            restartBtn.classList.remove('hidden');
        } else {
            // 通常のステージクリア処理
            logMessage(`勝利！ 敵を倒した。`, 'success');
            await animateBattleLog('勝利！');
            await sleep(2000);
            currentStage++;
            startStage();
        }
    } else if (playerHP <= 0) {
        handleGameOver();
    } else {
        await animateBattleLog('バトル中断');
    }

    isBattleInProgress = false;
    updatePlayerModelUI();
    battleBtn.innerHTML = '<i class="fas fa-fist-raised"></i> バトル開始！';
}

function returnToTitle() {
    gameContainer.style.opacity = 0;
    setTimeout(() => {
        gameContainer.classList.add('hidden');
        titleScreen.classList.remove('hidden');
        titleScreen.style.opacity = 1;
    }, 500);
}

// --- Event Listeners & Initial Load ---

// ★★★ スタートボタンのイベントリスナー
startGameBtn.addEventListener('click', () => {
    titleScreen.style.opacity = 0;
    setTimeout(() => {
        titleScreen.classList.add('hidden');
        gameContainer.classList.remove('hidden');
        gameContainer.style.display = 'flex'; // flexを再適用
        gameContainer.style.opacity = 1;
        initializeGame();
    }, 1000); // フェードアウトを待つ
});

// ★★★ リスタートボタンのイベントリスナーをタイトルへ戻るように変更
restartBtn.addEventListener('click', returnToTitle);

// --- DYNAMIC ANIMATION ENGINE (変更なし) ---
// ... (以下、アニメーション関連の長いコードは変更がないため省略します)
const valueToColor = (value, maxAbs) => {
    if (maxAbs === 0) return 'rgb(128, 128, 128)';
    const intensity = Math.min(Math.abs(value) / maxAbs, 1);
    if (value > 0) {
        const r = 128 + 127 * intensity; const g = 128 - 128 * intensity; const b = 128 + 127 * intensity;
        return `rgb(${r}, ${g}, ${b})`;
    } else {
        const r = 128 - 128 * intensity; const g = 128 - 128 * intensity; const b = 128 + 127 * intensity;
        return `rgb(${r}, ${g}, ${b})`;
    }
};

async function runDynamicAnimation(data) {
    vizArea.innerHTML = '<div id="labels-container" class="absolute inset-x-0 top-0 h-24 pointer-events-none flex items-center"></div><div id="prediction-result" class="absolute text-2xl font-bold right-8 top-8 opacity-0 transition-opacity duration-500"></div>';
    document.getElementById('prediction-result').textContent = '';
    document.getElementById('prediction-result').classList.remove('opacity-100');

    const { architecture, outputs, weights } = data;
    const vizWidth = vizArea.clientWidth;
    const vizHeight = vizArea.clientHeight;
    const stageCount = architecture.length + 1; // +1 for output probabilities

    let currentElement = null;
    // ★★★ 視覚的な要素をスタックで管理
    const vizElementStack = [];

    for (let i = 0; i < architecture.length; i++) {
        const layerInfo = architecture[i];
        const xPercent = (i + 1) / stageCount;

        if (layerInfo.type === 'ReLU' || layerInfo.type === 'Dropout') {
            await showStageLabel(layerInfo.type, xPercent, null);
            // スタックのトップにある要素に対してアニメーションを適用
            const targetElement = vizElementStack[vizElementStack.length - 1];
            if (layerInfo.type === 'ReLU') {
                await animateReLU(targetElement);
            } else {
                await animateDropout(targetElement);
            }
            continue;
            // ★★★ 修正: ReLU/DropoutではcurrentElementのopacityを変えないように変更
        }

        const nextElement = document.createElement('div');
        nextElement.className = 'layer-viz absolute';
        nextElement.style.left = `${xPercent * 100}%`;
        nextElement.style.top = '50%';
        nextElement.style.transform = 'translate(-50%, -50%)';
        vizArea.appendChild(nextElement);

        if (currentElement) {
            currentElement.style.transition = 'opacity 0.3s';
            currentElement.style.opacity = '0.3';
        }

        await showStageLabel(layerInfo.type, xPercent, layerInfo.shape);

        switch (layerInfo.type) {
            case 'Input':
                currentData = outputs.input[0];
                await animateGrid(nextElement, currentData, { isInput: true }, vizArea);
                break;
            case 'Conv2d':
                currentData = outputs[layerInfo.name][0];
                await animateConv(currentElement, nextElement, currentData, {}, vizArea);
                break;
            case 'MaxPool2d':
            // ★★★ AvgPool2d を追加
            case 'AvgPool2d':
                currentData = outputs[layerInfo.name][0];
                const isAverage = layerInfo.type === 'AvgPool2d';
                await animatePool(currentElement, nextElement, currentData, { isAverage }, vizArea);
                break;
            case 'Flatten':
                await animateFlatten(currentElement, nextElement, layerInfo.shape[0], {}, vizArea);
                break;
            case 'Linear':
                // ★★★ 修正: スタックから最新の視覚要素を取得
                const sourceElement = vizElementStack[vizElementStack.length - 1];

                currentData = outputs[layerInfo.name][0];
                const linearWeights = weights[layerInfo.name + '_w'];
                const linearBiases = weights[layerInfo.name + '_b'];
                await animateLinear(sourceElement, nextElement, currentData, linearWeights, linearBiases, {}, vizArea);
                break;
            // ★★★ ResidualBlock を追加
            case 'ResidualBlock':
                const skipFromElement = vizElementStack.length > 1 ? vizElementStack[vizElementStack.length - 2] : currentElement;
                currentData = outputs[layerInfo.name][0];
                await animateResidual(skipFromElement, currentElement, nextElement, currentData, {}, vizArea);
                break;

            default: // Input, Conv2d, MaxPool2d
                currentData = outputs[layerInfo.name] ? outputs[layerInfo.name][0] : outputs.input[0];
                await animateGrid(nextElement, currentData, { isInput: layerInfo.type === 'Input' }, vizArea);
        }

        currentElement = nextElement;
        // ★★★ 視覚的に意味のある要素だけをスタックに積む
        if (layerInfo.type !== 'Flatten') {
            vizElementStack.push(currentElement);
        }
    }

    // --- Final Output/Softmax Animation ---
    const finalLayerInfo = architecture[architecture.length - 1];
    const finalLayerName = finalLayerInfo.name;

    // ★★★ デバッグ用ログと安全なアクセス
    console.log(`Accessing final output with key: ${finalLayerName}`);
    const logits = outputs[finalLayerName] ? outputs[finalLayerName][0] : [];

    const finalSourceElement = vizElementStack[vizElementStack.length - 1];
    await animateSoftmax(finalSourceElement, data.prediction, data.label, logits, {}, vizArea);

    await sleep(1500);
    animationModal.classList.add('hidden');
}

// --- Layer-specific Animation Functions ---

const showStageLabel = async (text, xPercent, shape) => {
    const labelsContainer = document.getElementById('labels-container');
    labelsContainer.querySelectorAll('.stage-label').forEach(l => {
        l.style.opacity = '0'; l.style.transform = 'translateY(20px)';
    });
    await sleep(200);
    labelsContainer.innerHTML = '';

    const label = document.createElement('div');
    label.className = 'stage-label';
    let labelText = text;
    if (shape) {
        labelText += `<br><span class="text-sm text-gray-400">(${shape.join(' &times; ')})</span>`;
    }
    label.innerHTML = labelText;
    label.style.left = `${xPercent * 100}%`;
    labelsContainer.appendChild(label);

    await sleep(50);
    label.style.opacity = '1';
    label.style.transform = 'translateX(-50%) translateY(-100%)';
};

// Simplified grid creation for any 2D/3D data
async function animateGrid(container, data3d, options = {}, vizArea = window.vizArea) {
    container.style.transform = 'translate(-50%, -50%) perspective(1000px) rotateY(-90deg)';
    container.style.opacity = '0';
    const duration = options.duration || 500;
    container.style.transition = `transform ${duration / 1000}s ease-out, opacity ${duration / 1000}s`;

    const vizHeight = vizArea.clientHeight;
    const displayChannels = Math.min(data3d.length, 6);
    const displaySize = Math.min(data3d[0].length, 14);

    const availableHeight = vizHeight * (options.isInput ? 0.4 : 0.6);
    const margin = 10;
    const singleGridMaxHeight = (availableHeight - (displayChannels - 1) * margin) / displayChannels;
    const gridSize = Math.max(20, singleGridMaxHeight); // ★★★ 正しい変数名はこちら
    const cellSize = gridSize / displaySize;

    const totalHeight = displayChannels * gridSize + (displayChannels - 1) * margin;

    // ★★★ エラー修正箇所: `totalGridSize` を `gridSize` に修正
    container.style.width = `${gridSize}px`;
    container.style.height = `${totalHeight}px`;
    container.style.transform = `translate(-50%, -${totalHeight / 2}px) perspective(1000px) rotateY(-90deg)`;

    const maxAbs = Math.max(...data3d.flat().flat().map(Math.abs));

    for (let i = 0; i < displayChannels; i++) {
        const featureMap = document.createElement('div');
        featureMap.className = 'absolute';
        featureMap.style.width = `${gridSize}px`;
        featureMap.style.height = `${gridSize}px`;
        featureMap.style.top = `${i * (gridSize + margin)}px`;

        for (let r = 0; r < displaySize; r++) {
            for (let c = 0; c < displaySize; c++) {
                // Sample from original data if larger
                const origR = Math.floor(r * data3d[i].length / displaySize);
                const origC = Math.floor(c * data3d[i][0].length / displaySize);
                const val = data3d[i][origR][origC];

                const cell = document.createElement('div');
                cell.className = 'grid-cell absolute';
                cell.style.width = `${cellSize}px`; cell.style.height = `${cellSize}px`;
                cell.style.left = `${c * cellSize}px`; cell.style.top = `${r * cellSize}px`;
                cell.style.backgroundColor = valueToColor(val, maxAbs);
                cell.dataset.value = val;
                featureMap.appendChild(cell);
            }
        }
        container.appendChild(featureMap);
    }

    await sleep(50);
    container.style.transform = `translate(-50%, -${totalHeight / 2}px) perspective(1000px) rotateY(0deg)`;
    container.style.opacity = '1';
    await sleep(duration);
}

async function animateConv(fromEl, toEl, toData, options = {}, vizArea = window.vizArea) {
    const duration = options.duration || 500;

    // Kernel animation on 'from' element
    if (options.isPreview) {
        // fromElは背景として表示され続けるので、カーネルアニメーションは不要
    } else {
        const firstMap = fromEl.querySelector(':scope > div');
        if (firstMap) {
            const kernelHighlight = document.createElement('div');
            kernelHighlight.style.position = 'absolute';
            kernelHighlight.style.border = '2px solid #a5b4fc';
            kernelHighlight.style.transition = 'all 0.1s linear';
            const mapSize = firstMap.clientWidth;
            const kernelDisplaySize = mapSize / 7; // e.g., 14px -> 2px kernel
            kernelHighlight.style.width = `${kernelDisplaySize}px`;
            kernelHighlight.style.height = `${kernelDisplaySize}px`;
            firstMap.appendChild(kernelHighlight);

            for (let i = 0; i <= 6; i++) {
                kernelHighlight.style.top = `${i * kernelDisplaySize}px`;
                kernelHighlight.style.left = `${(i % 3) * mapSize / 3}px`;
                await sleep(duration * 0.05); // ★ 全体時間に対する割合でsleep
            }
            kernelHighlight.remove();
        }
    }

    await sleep(options.duration * 0.4 || 200);
    await animateGrid(toEl, toData, { duration: duration * 0.6 }, vizArea);
}

async function animatePool(fromEl, toEl, toData, options = {}, vizArea = window.vizArea) {
    const duration = options.duration || 500;

    if (!options.isPreview) {
        const maps = fromEl.querySelectorAll(':scope > div');
        maps.forEach((map) => {
            map.style.transition = `all ${duration * 0.4 / 1000}s ease-in-out`;
            // isAverageの場合、ぼかしエフェクトを追加
            if (options.isAverage) {
                map.style.filter = 'blur(2px)';
            }
            map.style.transform = `scale(0.5)`;
            map.style.opacity = '0';
        });
        await sleep(duration * 0.4);
    }

    await animateGrid(toEl, toData, { duration: duration * 0.8 }, vizArea);
}

async function animateDropout(element, options = {}) {
    if (!element) return;
    const cells = element.querySelectorAll('.grid-cell');
    const originalColors = new Map();
    const cellsToDrop = Array.from(cells).sort(() => 0.5 - Math.random()).slice(0, cells.length / 2);

    cellsToDrop.forEach(cell => {
        originalColors.set(cell, cell.style.backgroundColor); // 元の色を保存
        cell.style.transition = 'all 0.2s ease-in-out';
        cell.style.backgroundColor = 'rgb(40, 40, 40)'; // 暗い色（非アクティブ）に
        cell.style.transform = 'scale(0.8)';
    });

    await sleep(options.duration || 400);

    cellsToDrop.forEach(cell => {
        // 元の色に戻す
        cell.style.backgroundColor = originalColors.get(cell);
        cell.style.transform = 'scale(1)';
    });
    await sleep(options.duration ? options.duration * 0.5 : 200);
}

async function animateFlatten(fromEl, toEl, toShape, options = {}, vizArea = window.vizArea) {
    // ★★★ エラー修正箇所 ★★★
    // fromElがnullの場合（Linearプレビューなど）は、ソース要素のアニメーションをスキップ
    if (fromEl) {
        fromEl.style.transition = 'transform 0.4s ease-in-out, opacity 0.4s';
        fromEl.style.transform += ' scale(0)';
        fromEl.style.opacity = '0';
    }

    // Display a simplified, representative bar
    const vizHeight = vizArea.clientHeight;
    // ... (以降のロジックは変更なし)
    const displayNodes = Math.min(toShape, 128);
    const nodeHeight = Math.min(2.5, (vizHeight * 0.8) / displayNodes);
    const totalHeight = displayNodes * nodeHeight;
    toEl.style.height = `${totalHeight}px`;
    toEl.style.width = '20px';
    toEl.style.position = 'absolute'; // ★ 念のため追加
    toEl.style.left = '50%';         // ★ 念のため追加
    toEl.style.top = '50%';          // ★ 念のため追加
    toEl.style.transform = `translate(-50%, -${totalHeight / 2}px)`; // ★ 念のため追加

    const maxAbs = 1; // Dummy value for color
    for (let i = 0; i < displayNodes; i++) {
        const cell = document.createElement('div');
        cell.className = 'grid-cell absolute';
        cell.style.width = '100%';
        cell.style.height = `${nodeHeight}px`;
        cell.style.top = `${i * nodeHeight}px`;
        cell.style.backgroundColor = valueToColor(Math.random() * 2 - 1, maxAbs);
        cell.style.transform = 'scale(0)';
        cell.style.transition = 'transform 0.3s';
        toEl.appendChild(cell);

        // プレビュー時のアニメーションが速すぎないように調整
        const delay = (options.duration || 300) / displayNodes;
        await sleep(delay > 5 ? 5 : delay); // 5msより長くは待たない

        cell.style.transform = 'scale(1)';
    }
    await sleep(options.duration * 0.5 || 150);
}

async function animateLinear(fromEl, toEl, toData, weights, biases, options = {}, vizArea = window.vizArea) {
    const vizRect = vizArea.getBoundingClientRect();

    const displayNodes = Math.min(toData.length, 64);
    const vizHeight = vizArea.clientHeight;
    // ★★★ nodeSizeをオプションで受け取れるように
    const nodeSize = options.nodeSize || Math.min(15, (vizHeight * 0.7) / displayNodes);
    const spacing = (vizHeight * 0.7) / displayNodes;
    const totalHeight = displayNodes * spacing;

    toEl.style.height = `${totalHeight}px`;
    // ★★★ transformのY座標計算を修正
    toEl.style.transform = `translate(-50%, -${totalHeight / 2}px)`;

    const toCells = [];
    const maxAbs = Math.max(...toData.map(Math.abs));

    for (let i = 0; i < displayNodes; i++) {
        const cell = document.createElement('div');
        cell.className = 'grid-cell';
        cell.style.width = `${nodeSize}px`; cell.style.height = `${nodeSize}px`;
        cell.style.borderRadius = '50%';
        cell.style.position = 'absolute';
        cell.style.top = `${i * spacing}px`;
        cell.style.backgroundColor = 'rgb(128,128,128)';
        toEl.appendChild(cell);
        toCells.push(cell);
    }

    // fromElがnullの場合（プレビューの初回など）はパーティクルを飛ばさない
    if (!fromEl) {
        await sleep(options.duration || 400);
    } else {
        const maxParticles = 50;
        for (let i = 0; i < maxParticles; i++) {
            // ★★★ 座標計算の基準をvizAreaの左上隅(0,0)に統一
            const fromRect = fromEl.getBoundingClientRect();
            const toRect = toEl.getBoundingClientRect();

            const particle = document.createElement('div');
            particle.className = 'particle';
            vizArea.appendChild(particle);

            const startX = fromRect.right - vizRect.left;
            const startY = (fromRect.top - vizRect.top) + Math.random() * fromRect.height;
            const endX = toRect.left - vizRect.left;
            const endY = (toRect.top - vizRect.top) + Math.random() * toRect.height;

            particle.animate([
                { transform: `translate(${startX}px, ${startY}px) scale(0.5)`, opacity: 1 },
                { transform: `translate(${endX}px, ${endY}px) scale(1)`, opacity: 0 }
            ], {
                duration: 300 + Math.random() * 200,
                easing: 'ease-in-out',
                delay: Math.random() * 200,
            }).onfinish = () => particle.remove();
        }
        await sleep(options.duration || 400);
    }

    toCells.forEach((cell, i) => {
        const dataIdx = Math.floor(i * toData.length / displayNodes);
        cell.style.transition = 'background-color 0.5s';
        cell.style.backgroundColor = valueToColor(toData[dataIdx], maxAbs);
    });
    await sleep(options.duration ? options.duration * 0.5 : 500);
}

async function animateReLU(element, options = {}) {
    if (!element) return;
    const cells = element.querySelectorAll('.grid-cell');

    const animations = [];
    for (const cell of cells) {
        const value = parseFloat(cell.dataset.value || 0);
        if (value < 0) {
            const animationPromise = new Promise(async (resolve) => {
                await sleep(Math.random() * (options.duration || 400) * 0.5); // ランダムな遅延
                cell.style.transition = 'background-color 0.3s';
                cell.style.backgroundColor = 'rgb(128, 128, 128)';
                // ★★★ データを更新して、ReLUが適用されたことを記録
                cell.dataset.value = 0;
                resolve();
            });
            animations.push(animationPromise);
        }
    }
    await Promise.all(animations); // 全てのアニメーションが終わるのを待つ
    await sleep((options.duration || 400) * 0.5);
}

async function animateResidual(skipFromEl, fromEl, toEl, toData, options = {}, vizArea = window.vizArea) {
    const duration = options.duration || 1000;

    // メインパスのアニメーション (Linearと同様)
    animateLinear(fromEl, toEl, toData, null, null, options, vizArea);

    // スキップコネクションのアニメーション
    await sleep(duration * 0.1); // 少し遅れて開始

    const svg = document.createElementNS('http://www.w3.org/2000/svg', 'svg');
    svg.style.position = 'absolute';
    svg.style.top = '0';
    svg.style.left = '0';
    svg.style.width = '100%';
    svg.style.height = '100%';
    svg.style.pointerEvents = 'none';
    svg.style.zIndex = '10';
    vizArea.appendChild(svg);

    const path = document.createElementNS('http://www.w3.org/2000/svg', 'path');
    path.setAttribute('fill', 'none');
    path.setAttribute('stroke', 'url(#skip-gradient)');
    path.setAttribute('stroke-width', '3');

    const vizRect = vizArea.getBoundingClientRect();
    const startRect = skipFromEl.getBoundingClientRect();
    const endRect = toEl.getBoundingClientRect();

    const startX = startRect.right - vizRect.left;
    const startY = startRect.top + startRect.height / 2 - vizRect.top;
    const endX = endRect.left - vizRect.left;
    const endY = endRect.top + endRect.height / 2 - vizRect.top;

    const ctrlYOffset = -80; // 上に膨らむカーブ
    const d = `M ${startX},${startY} C ${startX + 50},${startY + ctrlYOffset} ${endX - 50},${endY + ctrlYOffset} ${endX},${endY}`;
    path.setAttribute('d', d);

    const defs = document.createElementNS('http://www.w3.org/2000/svg', 'defs');
    const gradient = document.createElementNS('http://www.w3.org/2000/svg', 'linearGradient');
    gradient.id = 'skip-gradient';
    gradient.innerHTML = `
        <stop offset="0%" stop-color="#a5b4fc" stop-opacity="0" />
        <stop offset="50%" stop-color="#a5b4fc" stop-opacity="1" />
        <stop offset="100%" stop-color="#a5b4fc" stop-opacity="0" />
    `;
    defs.appendChild(gradient);
    svg.appendChild(defs);
    svg.appendChild(path);

    const length = path.getTotalLength();
    path.style.strokeDasharray = length;
    path.style.strokeDashoffset = length;

    path.animate([
        { strokeDashoffset: length },
        { strokeDashoffset: 0 }
    ], {
        duration: duration * 0.8,
        easing: 'cubic-bezier(0.4, 0, 0.2, 1)'
    }).onfinish = () => {
        setTimeout(() => svg.remove(), 200);
    };

    await sleep(duration);
}

const softmax = (logits) => {
    const maxLogit = Math.max(...logits);
    const exps = logits.map(logit => Math.exp(logit - maxLogit));
    const sumExps = exps.reduce((a, b) => a + b, 0);
    return exps.map(exp => exp / sumExps);
};

async function animateSoftmax(fromEl, prediction, label, logits, options = {}, vizArea = window.vizArea) {
    const probabilities = softmax(logits);
    const toEl = fromEl; // Reuse the last linear layer element
    toEl.innerHTML = '';

    const vizHeight = vizArea.clientHeight;
    const numOutputNodes = 10;
    const nodeHeight = Math.min(40, (vizHeight * 0.8) / numOutputNodes * 0.8);
    const spacing = (vizHeight * 0.8) / numOutputNodes;
    const totalHeight = numOutputNodes * spacing;
    toEl.style.transform = `translate(-50%, -${totalHeight / 2}px)`;

    for (let i = 0; i < numOutputNodes; i++) {
        const prob = probabilities[i];
        const wrapper = document.createElement('div');
        wrapper.className = 'flex items-center relative transition-all duration-300';
        wrapper.style.height = `${spacing}px`;
        wrapper.style.width = `200px`;

        const labelDiv = document.createElement('div');
        labelDiv.className = 'mr-4 font-bold';
        labelDiv.textContent = i;
        labelDiv.style.fontSize = `${nodeHeight * 0.6}px`;

        const barContainer = document.createElement('div');
        barContainer.className = 'flex-grow h-full bg-white/10 rounded overflow-hidden border border-indigo-400/50';

        const bar = document.createElement('div');
        bar.style.width = '0%';
        bar.style.height = '100%';
        bar.style.backgroundColor = '#a5b4fc';
        bar.style.transition = 'width 0.8s ease-out';

        const probText = document.createElement('div');
        probText.className = 'absolute right-0 top-1/2 -translate-y-1/2 font-mono';
        probText.textContent = `${(prob * 100).toFixed(1)}%`;
        probText.style.fontSize = `${nodeHeight * 0.4}px`;

        barContainer.appendChild(bar);
        wrapper.appendChild(labelDiv);
        wrapper.appendChild(barContainer);
        wrapper.appendChild(probText);
        toEl.appendChild(wrapper);

        setTimeout(() => { bar.style.width = `${prob * 100}%`; }, 100);
    }

    await sleep(1000);

    // Highlight prediction
    const predWrapper = toEl.childNodes[prediction];
    predWrapper.style.transform = 'scale(1.1)';
    predWrapper.style.backgroundColor = 'rgba(253, 224, 71, 0.2)';
    predWrapper.style.borderRadius = '8px';

    const resultText = document.getElementById('prediction-result');
    resultText.innerHTML = `Prediction: <span class="text-yellow-300 text-3xl">${prediction}</span> (True: ${label})`;
    resultText.classList.add('opacity-100');
}


// --- Event Listeners & Initial Load ---
battleBtn.addEventListener('click', handleBattle);
restartBtn.addEventListener('click', startGame);
closeModalBtn.addEventListener('click', () => {
    animationModal.classList.add('hidden');
});

// D&Dイベントリスナーのセットアップ
const modelArea = document.getElementById('player-model-layers');
modelArea.addEventListener('dragover', allowDrop);
modelArea.addEventListener('drop', dropOnModelArea);
// ★★★ 修正: dragleave イベントリスナーをより堅牢なものに変更
modelArea.addEventListener('dragleave', handleDragLeaveModelArea);
modelArea.addEventListener('dragenter', handleDragEnterModelArea);

document.addEventListener('drop', (e) => {
    // モデルエリア外へのドロップ処理 (このロジックは重要なので残す)
    if (draggedItem) {
        const modelArea = document.getElementById('player-model-layers');
        if (!modelArea.contains(e.target)) {
            if (draggedItem.type === 'model') {
                // モデル外へのドロップはキャンセルとみなし、UIを更新して元の位置に戻す
                // ★★★ バグ修正: 元の状態に戻すには、playerLayersを再構築する必要がある
                // ただし、この操作は複雑なので、単純にログだけ出すか、何もしないのが安全
                logMessage('モデルの並び替えをキャンセルしました。', 'info');
                // UIを再描画すればOK
                updatePlayerModelUI();
            } else if (draggedItem.type === 'inventory') {
                // ★★★ instanceId を使って検索
                const tempIndex = playerLayers.findIndex(l => l.instanceId === draggedItem.layer.instanceId);
                if (tempIndex > -1) playerLayers.splice(tempIndex, 1);
                updatePlayerModelUI();
            }
        }
    }
});

document.addEventListener('DOMContentLoaded', () => {
    gameContainer.style.opacity = 0;
    gameContainer.style.transition = 'opacity 0.5s ease-in-out';
});

document.getElementById('confirm-choice-btn').addEventListener('click', () => {
    const selectedId = document.getElementById('selected-choice-id').value;
    if (selectedId) {
        const selectedLayer = allAvailableLayers.find(l => l.id == selectedId);
        if (selectedLayer) {
            selectItem(selectedLayer);
        }
    }
});