#!/bin/bash
# ════════════════════════════════════════════════════════
# NeuroFlow — 阿里天池 JupyterLab 一键部署脚本
# 使用前必读：
#   1. 在天池控制台完成登录，并把登录保持时间建议改为 24 小时并重新登录
#   2. 避免在 JupyterLab 工具运行中登录态失效导致 kernel 中断
#   3. GPU 实例按时计费，不使用时及时停止
# 用法: 在天池 JupyterLab 终端运行:
#   bash scripts/deploy_tianchi.sh
# ════════════════════════════════════════════════════════

set -e

echo "╔══════════════════════════════════════════════════╗"
echo "║     NeuroFlow 天池 JupyterLab 自动部署脚本        ║"
echo "╚══════════════════════════════════════════════════╝"
echo ""
echo "⚠️  使用前请先在阿里云控制台完成登录，"
echo "   建议把登录保持时间改为 24 小时并重新登录。"
echo "   若 kernel 后续中断，通常就是登录态过期导致。"
echo ""

# ── 0. 检查 GPU ──
echo "🔍 [1/6] 检查 GPU 环境..."
nvidia-smi || {
    echo "❌ 未检测到 GPU，请确认当前实例已挂载 GPU。"
    exit 1
}
nvidia-smi --query-gpu=name,memory.total --format=csv,noheader

# ── 1. 安装依赖 ──
echo ""
echo "📦 [2/6] 安装编译依赖..."
apt-get update -qq
apt-get install -y -qq cmake build-essential git python3 python3-pip libomp-dev 2>/dev/null

# CUDA 提示
if command -v nvcc &>/dev/null; then
    echo "   CUDA: $(nvcc --version | grep 'release' | awk '{print $6}' | tr -d ',')"
else
    echo "   ⚠️ nvcc 未找到，若天池镜像未预装，请使用官方 CUDA 镜像/环境。"
fi

# ── 2. 获取代码 ──
echo ""
echo "📥 [3/6] 获取 NeuroFlow 源码..."
REPO_URL="https://github.com/chenzhiwenhphp12-afk/neuroflow-model.git"
if [ -d "neuroflow-model" ]; then
    echo "   检测到已有目录 neuroflow-model，尝试拉取更新..."
    cd neuroflow-model || true
    git pull || true
else
    git clone "$REPO_URL" neuroflow-model || {
        echo "   ⚠️ GitHub 克隆失败，请检查网络；若不可达，请手动上传 neuroflow-model 仓库。"
        exit 1
    }
    cd neuroflow-model
fi

echo "   当前目录: $(pwd)"

# ── 3. 准备训练数据 ──
echo ""
echo "📝 [4/6] 准备训练数据..."
if [ ! -f "data/train.txt" ]; then
    mkdir -p data
    python3 - <<'PY'
import os
if not os.path.exists('data/train.txt'):
    samples = []
    for i in range(2000):
        samples.append(f'这是第{i}条训练数据，用于NeuroFlow模型的预训练和微调任务。')
    with open('data/train.txt', 'w', encoding='utf-8') as f:
        f.write('\n'.join(samples))
    print(f'   已生成 {len(samples)} 条训练样本 -> data/train.txt')
else:
    print('   训练数据已存在: data/train.txt')
PY
else
    echo "   训练数据已存在: data/train.txt"
fi

# ── 4. 编译 ──
echo ""
echo "🔧 [5/6] 编译 NeuroFlow (CUDA 模式)..."
HAS_CUDA=false
if [ -f /usr/local/cuda/include/cuda.h ] || [ -d /usr/local/cuda ]; then
    HAS_CUDA=true
fi

if [ "$HAS_CUDA" = true ]; then
    echo "   使用 CUDA 后端编译..."
    cmake -B build_cuda \
        -DNEUROFLOW_USE_CUDA=ON \
        -DNEUROFLOW_USE_AVX2=ON \
        -DNEUROFLOW_USE_BLAS=OFF \
        -DCMAKE_BUILD_TYPE=Release
    cmake --build build_cuda -j"$(nproc)"
    BUILD_DIR="build_cuda"
else
    echo "   ⚠️ 未检测到 CUDA，使用 CPU OpenMP 编译..."
    cmake -B build_cpu \
        -DNEUROFLOW_USE_CUDA=OFF \
        -DNEUROFLOW_USE_AVX2=ON \
        -DNEUROFLOW_USE_BLAS=OFF \
        -DCMAKE_BUILD_TYPE=Release
    cmake --build build_cpu -j"$(nproc)"
    BUILD_DIR="build_cpu"
fi

echo "✅ 编译完成！二进制文件在 ./$BUILD_DIR/"

# ── 5. 运行训练 ──
echo ""
echo "🚀 [6/6] 开始训练验证..."
echo "════════════════════════════════════════════════════"
echo "  配置: configs/config_distill.json"
echo "  数据: data/train.txt"
echo "  GPU:  $(nvidia-smi --query-gpu=name --format=csv,noheader)"
echo "  显存: $(nvidia-smi --query-gpu=memory.total --format=csv,noheader)"
echo "════════════════════════════════════════════════════"

./$BUILD_DIR/neuroflow_train_v2 \
    --config configs/config_distill.json \
    --data data/train.txt \
    --output output_tianchi \
    --epochs 5 \
    --batch-size 64 \
    --lr 0.001 \
    --use-cuda \
    --adam \
    --log-interval 10 || {
        echo ""
        echo "❌ 训练启动失败。"
        echo "   建议先检查：登录态是否过期、kernel 是否断开、CUDA 是否可用。"
        exit 1
    }

echo ""
echo "🎉 训练完成！"
echo "   模型保存在: output_tianchi/"
echo ""
echo "   🔜 下一步"
echo "   1) 全量训练示例："
echo "      ./$BUILD_DIR/neuroflow_train_v2 \\"
echo "          --config configs/config.json \\"
echo "          --data data/train.txt \\"
echo "          --output output_full \\"
echo "          --epochs 100 \\"
echo "          --batch-size 64 \\"
echo "          --lr 0.0001 \\"
echo "          --use-cuda --adam"
echo ""
echo "   2) DLC 任务提交："
echo "      如需我继续帮你准备天池 DLC 提交配置，回复我即可。"