#!/bin/bash
# WaveGen 训练可视化工具启动脚本

# 脚本所在目录
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
cd "$SCRIPT_DIR"

echo "============================================================"
echo "🎨 WaveGen 训练可视化工具"
echo "============================================================"
echo ""

# 检查依赖
echo "检查依赖..."
python -c "import viser" 2>/dev/null
if [ $? -ne 0 ]; then
    echo "❌ 缺少 viser 包"
    echo "请安装: pip install viser"
    exit 1
fi

python -c "import scipy" 2>/dev/null
if [ $? -ne 0 ]; then
    echo "❌ 缺少 scipy 包"
    echo "请安装: pip install scipy"
    exit 1
fi

echo "✅ 依赖检查通过"
echo ""

# 检查 core_space 目录
CORE_SPACE="${1:-core_space}"

# 如果提供了相对路径，转换为绝对路径
if [[ ! "$CORE_SPACE" = /* ]]; then
    CORE_SPACE="$SCRIPT_DIR/$CORE_SPACE"
fi

# 尝试在多个位置查找 core_space
if [ ! -d "$CORE_SPACE" ]; then
    echo "⚠️  未找到: $CORE_SPACE"
    echo "尝试其他位置..."

    # 尝试 WaveGen_Augustus_v1/core_space
    ALT_PATH="$SCRIPT_DIR/../WaveGen_Augustus_v1/core_space"
    if [ -d "$ALT_PATH" ]; then
        CORE_SPACE="$ALT_PATH"
        echo "✅ 找到: $CORE_SPACE"
    else
        echo "❌ 未找到 core_space 目录"
        echo ""
        echo "请指定 core_space 目录:"
        echo "  ./visualize.sh /path/to/core_space"
        echo ""
        echo "或者确保以下位置之一存在:"
        echo "  - ./core_space"
        echo "  - ../WaveGen_Augustus_v1/core_space"
        exit 1
    fi
fi

echo "📁 使用 core_space: $CORE_SPACE"
echo ""

# 启动可视化工具
echo "🚀 启动可视化工具..."
echo ""

python utils/visualize_training.py --core-space "$CORE_SPACE"