#!/bin/bash
# ═══════════════════════════════════════════════════════════════
# NeuroFlow — 阿里天池 DSW CUDA 一键部署/训练脚本
# 适用: DSW GPU 实例, A10 / V100 / 其他 sm_80+ 显卡
# 用法: bash deploy_dsw.sh
# ═══════════════════════════════════════════════════════════════

set -euo pipefail

REPO_DIR="neuroflow-C++"
BUILD_DIR="build_cuda"
GIT_REPO="https://github.com/chenzhiwenhphp12-afk/neuroflow-model.git"

echo "╔══════════════════════════════════════════════════╗"
echo "║   NeuroFlow DSW 部署 + 蒸馏训练脚本              ║"
echo "╚══════════════════════════════════════════════════╝"

# ── 0. 检查 GPU ──
echo ""
echo "🔍 [0/7] 检查 GPU / CUDA..."
if ! command -v nvidia-smi &>/dev/null; then
    echo "❌ nvidia-smi 未找到，请确认已开启 GPU 实例"
    exit 1
fi
nvidia-smi
echo ""

if ! command -v nvcc &>/dev/null; then
    echo "⚠️  nvcc 未找到。DSW 镜像里 CUDA 通常装在 /usr/local/cuda"
    export CUDA_HOME=/usr/local/cuda
    export PATH="$CUDA_HOME/bin:$PATH"
    if ! command -v nvcc &>/dev/null; then
        echo "❌ 仍未找到 nvcc，请确认 CUDA Toolkit 已安装"
        exit 1
    fi
fi
nvcc --version | grep "release"
echo "✅ CUDA 就绪"
echo ""

# ── 1. 安装系统依赖 ──
echo "📦 [1/7] 安装编译依赖..."
apt-get update -qq
apt-get install -y -qq cmake build-essential python3 python3-pip git 2>/dev/null || true
echo "✅ 依赖安装完成"
echo ""

# ── 2. 获取代码 ──
echo "📥 [2/7] 获取 NeuroFlow 源码..."
if [ -d "$REPO_DIR" ]; then
    echo "   检测到已有目录，执行 git pull..."
    cd "$REPO_DIR"
    git pull || true
    cd ..
else
    echo "   正在克隆: $GIT_REPO"
    git clone "$GIT_REPO" "$REPO_DIR"
fi
cd "$REPO_DIR"
echo "✅ 代码就绪: $(pwd)"
echo ""

# ── 3. 数据格式转换 ──
echo "📝 [3/7] 准备训练数据..."
DEEPSEEK_JSONL="${DEEPSEEK_JSONL:-}"
DATA_TXT="data/distill_train.txt"

if [ ! -f "$DATA_TXT" ]; then
    mkdir -p data

    if [ -n "$DEEPSEEK_JSONL" ] && [ -f "$DEEPSEEK_JSONL" ]; then
        echo "   使用 DeepSeek 蒸馏数据: $DEEPSEEK_JSONL"
        python3 scripts/preprocess_distill.py "$DEEPSEEK_JSONL" "$DATA_TXT" 240000
    else
        echo "   ⚠️  未找到蒸馏数据，生成 5000 条测试样本..."
        python3 -c "
samples = []
for i in range(5000):
    samples.append(f'这是第{i}条训练数据，用于NeuroFlow模型测试。')
with open('$DATA_TXT', 'w', encoding='utf-8') as f:
    f.write('\n'.join(samples))
print(f'已生成 {len(samples)} 条样本 -> $DATA_TXT')
"
    fi
else
    echo "   训练数据已存在: $DATA_TXT"
fi
ls -lh "$DATA_TXT"
echo ""

# ── 4. 修复 CMake CUDA 编译 ──
echo "🔧 [4/7] 修复 CMake CUDA 编译配置..."
# 本项目要求所有 src/*.cpp 都以 CUDA 语言编译，否则训练回退 CPU
# 将非 CUDA 源文件强制设为 CUDA 语言
CMAKE_FILE="CMakeLists.txt"
if [ -f "$CMAKE_FILE" ]; then
    # 在 set_source_files_properties(src/cuda_context.cpp ...) 后插入全量 CUDA 映射
    if ! grep -q "set_source_files_properties(src/tensor_ops.cpp" "$CMAKE_FILE"; then
        echo "   为所有 src/*.cpp 添加 CUDA 编译属性..."
        python3 -c "
import re
p = '$CMAKE_FILE'
with open(p, 'r', encoding='utf-8') as f:
    c = f.read()
anchor = 'set_source_files_properties(src/cuda_context.cpp PROPERTIES LANGUAGE CUDA)'
if anchor in c and 'NEUROFLOW_ALL_CUDA_SOURCES' not in c:
    srcs = '''src/train_v2.cpp src/infer_v2.cpp src/tensor.cpp src/model.cpp src/weight_io.cpp
src/tokenizer.cpp src/sampling.cpp src/causal_lm.cpp src/tensor_ops.cpp src/generative_model.cpp
src/rope.cpp src/swiglu.cpp src/rms_norm.cpp src/adamw.cpp src/scheduler.cpp src/train_lm.cpp
src/grad_scaler.cpp src/cuda_context.cpp'''.split()
    block = '\n'.join([f'set_source_files_properties({s} PROPERTIES LANGUAGE CUDA)' for s in srcs])
    c = c.replace(anchor, anchor + '\n' + block + '\n', 1)
    with open(p, 'w', encoding='utf-8') as f:
        f.write(c)
    print('   CMake 已更新')
else:
    print('   跳过 (已配置或缺少锚点)')
"
    fi
else
    echo "❌ 未找到 $CMAKE_FILE"
    exit 1
fi
echo ""

# ── 5. 编译 CUDA 版本 ──
echo "🔨 [5/7] 编译 NeuroFlow (CUDA 模式)..."
rm -rf "$BUILD_DIR"
mkdir -p "$BUILD_DIR"
cd "$BUILD_DIR"

cmake .. \
    -DNEUROFLOW_USE_CUDA=ON \
    -DNEUROFLOW_USE_BLAS=OFF \
    -DNEUROFLOW_USE_AVX2=OFF \
    -DCMAKE_BUILD_TYPE=Release

cmake --build . -j"$(nproc)"
echo "✅ 编译完成"
echo ""

# ── 6. 小规模验证 ──
echo "🧪 [6/7] 小规模验证 (5 epochs)..."
cd "$REPO_DIR"
./"$BUILD_DIR"/neuroflow_train_v2 \
    --config configs/config_distill.json \
    --data "$DATA_TXT" \
    --output output_dsw_verify \
    --epochs 5 \
    --batch-size 16 \
    --lr 0.001 \
    --use-cuda \
    --adam \
    --log-interval 10 \
    --save-interval 100 || {
        echo "⚠️  验证失败，查看上方错误"
        exit 1
    }
echo "✅ 验证完成"
echo ""

# ── 7. 正式训练 (蒸馏) ──
echo "🚀 [7/7] 开始正式蒸馏训练..."
echo "════════════════════════════════════════════════════"
echo "  配置: configs/config.json (128K vocab)"
echo "  数据: $DATA_TXT"
echo "  输出: output_dsw_distill"
echo "════════════════════════════════════════════════════"

./"$BUILD_DIR"/neuroflow_train_v2 \
    --config configs/config.json \
    --data "$DATA_TXT" \
    --output output_dsw_distill \
    --epochs 10 \
    --batch-size 64 \
    --lr 0.0003 \
    --grad-accum 4 \
    --use-cuda \
    --adam \
    --log-interval 10 \
    --save-interval 2000 \
    --replay-buffer 10000 \
    --replay-ratio 0.25

echo ""
echo "╔══════════════════════════════════════════════════╗"
echo "║           训练完成!                              ║"
echo "╚══════════════════════════════════════════════════╝"
echo ""
echo "📁 模型目录: output_dsw_distill/"
ls -lh output_dsw_distill/ 2>/dev/null || true
echo ""
echo "🔜 续训命令:"
echo "   ./$BUILD_DIR/neuroflow_train_v2 \\"
echo "     --config configs/config.json \\"
echo "     --data $DATA_TXT \\"
echo "     --output output_dsw_distill \\"
echo "     --resume output_dsw_distill/model_final.nfv1 \\"
echo "     --epochs 20 \\"
echo "     --batch-size 64 \\"
echo "     --lr 0.0003 \\"
echo "     --grad-accum 4 \\"
echo "     --use-cuda --adam"