#!/bin/bash
# 自适应Prompt预训练 - 一键启动脚本

set -e

export CUDA_VISIBLE_DEVICES=0
export PYTHONPATH="PROJECT_ROOT:$PYTHONPATH"

TRAIN_DIR="PROJECT_ROOT/training/pretrain"
mkdir -p $TRAIN_DIR
cd $TRAIN_DIR

echo "======================================"
echo "自适应Prompt VLM预训练"
echo "策略: 根据annotation长度调整prompt"
echo "======================================"
echo ""

# Step 1: 分析annotations
echo "======================================"
echo "Step 1: 分析Annotation质量"
echo "======================================"

ANALYSIS_FILE="PROJECT_ROOT/data/dataset/pretrain/train/annotation_analysis.json"

if [ ! -f "$ANALYSIS_FILE" ]; then
    echo "运行annotation分析..."
    python analyze_annotations.py
    
    if [ $? -ne 0 ]; then
        echo "❌ 分析失败"
        exit 1
    fi
    echo "✓ 分析完成"
else
    echo "✓ 分析文件已存在: $ANALYSIS_FILE"
fi

echo ""

# Step 2: 准备数据
echo "======================================"
echo "Step 2: 准备自适应Prompt数据"
echo "======================================"

DATA_FILE="PROJECT_ROOT/data/dataset/pretrain/train/pretrain_data_adaptive.pkl"

if [ ! -f "$DATA_FILE" ]; then
    echo "准备训练数据..."
    python prepare_pretrain_data_adaptive.py
    
    if [ $? -ne 0 ]; then
        echo "❌ 数据准备失败"
        exit 1
    fi
    echo "✓ 数据准备完成"
else
    echo "✓ 数据文件已存在: $DATA_FILE"
fi

echo ""

# Step 3: 测试数据加载 (可选)
echo "======================================"
echo "Step 3: 测试数据加载"
echo "======================================"

read -p "是否测试数据加载? (y/n): " -n 1 -r
echo
if [[ $REPLY =~ ^[Yy]$ ]]; then
    python pretrain_dataset_adaptive.py
    echo "✓ 数据加载测试完成"
fi

echo ""

# Step 4: 训练
echo "======================================"
echo "Step 4: 开始训练"
echo "======================================"

MODEL=$1
if [ -z "$MODEL" ]; then
    echo "用法: bash run_adaptive_pretrain.sh [qwen2.5-vl-3b|qwen2.5-vl-7b] [options]"
    echo ""
    echo "示例:"
    echo "  bash run_adaptive_pretrain.sh qwen2.5-vl-3b"
    echo "  bash run_adaptive_pretrain.sh qwen2.5-vl-3b --wandb"
    echo "  bash run_adaptive_pretrain.sh qwen2.5-vl-3b --epochs 10"
    exit 1
fi

OUTPUT_DIR="PROJECT_ROOT/checkpoints/pretrain"
mkdir -p $OUTPUT_DIR

echo "模型: $MODEL"
echo "数据: $DATA_FILE"
echo "输出: $OUTPUT_DIR"
echo ""

# 构建训练命令
CMD="python train_pretrain_adaptive.py --model $MODEL"

# 添加额外参数
shift
while [[ $# -gt 0 ]]; do
    case $1 in
        --wandb)
            CMD="$CMD --wandb"
            echo "启用 WandB"
            shift
            ;;
        --epochs)
            CMD="$CMD --epochs $2"
            echo "Epochs: $2"
            shift 2
            ;;
        --batch_size)
            CMD="$CMD --batch_size $2"
            echo "Batch Size: $2"
            shift 2
            ;;
        --lr)
            CMD="$CMD --lr $2"
            echo "Learning Rate: $2"
            shift 2
            ;;
        *)
            echo "未知参数: $1"
            shift
            ;;
    esac
done

echo ""
echo "======================================"
echo "执行命令:"
echo "$CMD"
echo "======================================"
echo ""

# 运行训练
eval $CMD

echo ""
echo "======================================"
echo "训练完成！"
echo "======================================"
echo "Checkpoints: $OUTPUT_DIR/$MODEL"
echo ""
echo "Prompt策略总结:"
echo "  短标注 (<20字符): 简单prompt (识别对象)"
echo "  详细标注 (>=20字符): 详细prompt (完整描述)"