#!/bin/bash
#
# InstanceV 训练启动脚本
#
# 使用说明：
# 1. 前台运行: bash run_instancev_training.sh
# 2. 后台运行: bash run_instancev_training.sh --background
# 3. 断点续跑: bash run_instancev_training.sh --resume
#    指定检查点: bash run_instancev_training.sh --resume-from /path/to/ckpt.safetensors
#    关闭断点续跑: bash run_instancev_training.sh --no-resume
#

set -e

# ==================== 运行参数解析 ====================

RUN_BACKGROUND=false
RUN_INTERNAL=false
RESUME_MODE="auto"   # auto | off | path
RESUME_FROM=""

while [[ $# -gt 0 ]]; do
    case "$1" in
        --background|-bg)
            RUN_BACKGROUND=true
            shift
            ;;
        --run-internal)
            RUN_INTERNAL=true
            shift
            ;;
        --resume)
            RESUME_MODE="auto"
            shift
            ;;
        --no-resume)
            RESUME_MODE="off"
            shift
            ;;
        --resume-from)
            if [[ -z "${2:-}" ]]; then
                echo "错误: --resume-from 需要指定 checkpoint 路径"
                exit 1
            fi
            RESUME_MODE="path"
            RESUME_FROM="$2"
            shift 2
            ;;
        --resume-from=*)
            RESUME_MODE="path"
            RESUME_FROM="${1#*=}"
            shift
            ;;
        *)
            echo "未知参数: $1"
            exit 1
            ;;
    esac
done

# ==================== 配置区域 ====================

# 项目根目录
PROJECT_ROOT="/data/rczhang/PencilFolder/DiffSynth-Studio"
cd "$PROJECT_ROOT"

# 数据路径
DATA_BASE_PATH="/data/rczhang/PencilFolder/data"
IGROUND_BASE="${DATA_BASE_PATH}/iGround"
IGROUND_JSONL="${IGROUND_BASE}/iGround_train_set_processed.jsonl"
IGROUND_CLIPS_DIR="${IGROUND_BASE}/Clips/train"
IGROUND_MASK_ROOT="${IGROUND_BASE}/InstanceMasks/train"
METADATA_PATH="${IGROUND_BASE}/instancev_iground_train.jsonl"
FORCE_REBUILD_METADATA=true      # 数据更新时强制重建 metadata
MIN_INSTANCES=1
MAX_INSTANCES=5

# 模型路径
MODEL_ID_WITH_ORIGIN_PATHS="Wan-AI/Wan2.1-T2V-1.3B:diffusion_pytorch_model*.safetensors,Wan-AI/Wan2.1-T2V-1.3B:models_t5_umt5-xxl-enc-bf16.pth,Wan-AI/Wan2.1-T2V-1.3B:Wan2.1_VAE.pth"

# 时间戳
TIMESTAMP=$(date +"%Y%m%d_%H%M%S")

# 输出路径（每次训练生成新目录）
OUTPUT_PATH="${PROJECT_ROOT}/models/train/instancev_iground_${TIMESTAMP}"
LOG_DIR="${OUTPUT_PATH}/logs"
mkdir -p "$LOG_DIR"
LOG_FILE="${LOG_DIR}/train_${TIMESTAMP}.log"

# 训练参数
NUM_FRAMES=81                    # 视频帧数 (4n+1 格式)
HEIGHT=480                       # 视频高度
WIDTH=832                        # 视频宽度
DATASET_REPEAT=1                 # 数据集重复次数

GRADIENT_ACCUMULATION_STEPS=4    # 梯度累积步数
LEARNING_RATE=1e-4               # 学习率
NUM_EPOCHS=5                     # 训练 Epoch 数
SAVE_STEPS=500                   # 保存间隔

# InstanceV 特有参数
SAUG_DROP_PROB=0.1               # SAUG dropout 概率（论文推荐 0.1）
SAUG_SCALE=0.0                   # 训练时不使用 unconditional guidance

# 显存使用相关（保持默认，不做显存优化）
USE_GRADIENT_CHECKPOINTING=false
MIXED_PRECISION="no"

# Wandb 配置
WANDB_PROJECT="instancev-training"
WANDB_RUN_NAME="instancev_${TIMESTAMP}"

# 其他参数
TRAINABLE_MODELS="dit"
TASK="sft"

# GPU 配置
export CUDA_VISIBLE_DEVICES=0

# 训练环境变量
export CUDA_LAUNCH_BLOCKING=0

# 加速器（优先使用 diffsyn 环境的 accelerate）
ACCELERATE_BIN="/home/rczhang/miniconda3/envs/diffsyn/bin/accelerate"
if [[ ! -x "${ACCELERATE_BIN}" ]]; then
    ACCELERATE_BIN="accelerate"
fi

# Python（优先使用 diffsyn 环境）
PYTHON_BIN="/home/rczhang/miniconda3/envs/diffsyn/bin/python"
if [[ ! -x "${PYTHON_BIN}" ]]; then
    PYTHON_BIN="python"
fi

# ==================== Wandb 登录 ====================
# 如果需要登录，取消下面注释并填入你的 API key
# wandb login YOUR_WANDB_API_KEY

# ==================== 函数定义 ====================

print_config() {
    echo "=============================================="
    echo "        InstanceV Training Configuration      "
    echo "=============================================="
    echo ""
    echo "[数据配置]"
    echo "  - 数据路径: ${DATA_BASE_PATH}"
    echo "  - Metadata: ${METADATA_PATH}"
    echo "  - 重新生成 Metadata: ${FORCE_REBUILD_METADATA}"
    echo ""
    echo "[模型配置]"
    echo "  - 模型: Wan2.1-T2V-1.3B"
    echo "  - 可训练模块: STAPE, IMCA, mv, norm_imca"
    echo "  - 输出路径: ${OUTPUT_PATH}"
    echo ""
    echo "[训练参数]"
    echo "  - 分辨率: ${WIDTH}x${HEIGHT}"
    echo "  - 帧数: ${NUM_FRAMES}"
    echo "  - 学习率: ${LEARNING_RATE}"
    echo "  - Epochs: ${NUM_EPOCHS}"
    echo "  - 梯度累积: ${GRADIENT_ACCUMULATION_STEPS}"
    echo "  - 梯度检查点: ${USE_GRADIENT_CHECKPOINTING}"
    echo "  - 混合精度: ${MIXED_PRECISION}"
    echo "  - 保存间隔: ${SAVE_STEPS} steps"
    echo ""
    echo "[InstanceV 参数]"
    echo "  - SAUG Dropout: ${SAUG_DROP_PROB}"
    echo "  - SAUG Scale: ${SAUG_SCALE}"
    echo ""
    echo "[Wandb]"
    echo "  - Project: ${WANDB_PROJECT}"
    echo "  - Run Name: ${WANDB_RUN_NAME}"
    echo ""
    echo "[断点续跑]"
    echo "  - 模式: ${RESUME_MODE}"
    if [[ "${RESUME_MODE}" == "path" ]]; then
        echo "  - 指定检查点: ${RESUME_FROM}"
    fi
    echo ""
    echo "[GPU]"
    echo "  - CUDA_VISIBLE_DEVICES: ${CUDA_VISIBLE_DEVICES}"
    echo ""
    echo "=============================================="
}

run_training() {
    echo ""
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] Step 1: 检查训练数据..."
    
    if [[ "${FORCE_REBUILD_METADATA}" == "true" ]] || [ ! -f "$METADATA_PATH" ]; then
        echo "[$(date '+%Y-%m-%d %H:%M:%S')] 运行数据预处理，生成 metadata..."
        "${PYTHON_BIN}" examples/wanvideo/model_training/prepare_instancev_iground.py \
            --iground_jsonl "${IGROUND_JSONL}" \
            --clips_dir "${IGROUND_CLIPS_DIR}" \
            --mask_root_dir "${IGROUND_MASK_ROOT}" \
            --output_metadata "$METADATA_PATH" \
            --dataset_base_path "${DATA_BASE_PATH}" \
            --min_instances "${MIN_INSTANCES}" \
            --max_instances "${MAX_INSTANCES}"
    else
        SAMPLE_COUNT=$(wc -l < "$METADATA_PATH")
        echo "[$(date '+%Y-%m-%d %H:%M:%S')] 找到训练 metadata: ${METADATA_PATH}"
        echo "[$(date '+%Y-%m-%d %H:%M:%S')] 样本数量: ${SAMPLE_COUNT}"
    fi
    
    echo ""
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] Step 2: 启动 InstanceV 训练..."
    echo ""
    
    GC_FLAG=""
    if [[ "${USE_GRADIENT_CHECKPOINTING}" == "true" ]]; then
        GC_FLAG="--use_gradient_checkpointing"
    fi
    
    RESUME_ARGS=()
    RESUME_PATH=""
    if [[ "${RESUME_MODE}" != "off" ]]; then
        if [[ "${RESUME_MODE}" == "path" ]]; then
            if [[ ! -f "${RESUME_FROM}" ]]; then
                echo "[$(date '+%Y-%m-%d %H:%M:%S')] 指定的 checkpoint 不存在: ${RESUME_FROM}"
                exit 1
            fi
            RESUME_PATH="${RESUME_FROM}"
        else
            RESUME_PATH=$(ls -1t "${OUTPUT_PATH}"/step-*.safetensors "${OUTPUT_PATH}"/epoch-*.safetensors 2>/dev/null | head -n 1 || true)
        fi
        
        if [[ -n "${RESUME_PATH}" ]]; then
            echo "[$(date '+%Y-%m-%d %H:%M:%S')] 断点续跑: ${RESUME_PATH}"
            RESUME_ARGS+=(--resume_from_checkpoint "${RESUME_PATH}")
        else
            echo "[$(date '+%Y-%m-%d %H:%M:%S')] 未找到可用 checkpoint，开始新训练"
        fi
    fi

    "${ACCELERATE_BIN}" launch \
        --num_processes 1 \
        --num_machines 1 \
        --mixed_precision="${MIXED_PRECISION}" \
        examples/wanvideo/model_training/train_instancev.py \
        --dataset_base_path "${DATA_BASE_PATH}" \
        --dataset_metadata_path "${METADATA_PATH}" \
        --data_file_keys "video" \
        --height ${HEIGHT} \
        --width ${WIDTH} \
        --num_frames ${NUM_FRAMES} \
        --dataset_repeat ${DATASET_REPEAT} \
        --model_id_with_origin_paths "${MODEL_ID_WITH_ORIGIN_PATHS}" \
        --learning_rate ${LEARNING_RATE} \
        --num_epochs ${NUM_EPOCHS} \
        --gradient_accumulation_steps ${GRADIENT_ACCUMULATION_STEPS} \
        --save_steps ${SAVE_STEPS} \
        --output_path "${OUTPUT_PATH}" \
        --remove_prefix_in_ckpt "pipe.dit." \
        --trainable_models "${TRAINABLE_MODELS}" \
        --task "${TASK}" \
        --saug_drop_prob ${SAUG_DROP_PROB} \
        --saug_scale ${SAUG_SCALE} \
        ${GC_FLAG} \
        --use_wandb \
        --wandb_project "${WANDB_PROJECT}" \
        --wandb_run_name "${WANDB_RUN_NAME}" \
        "${RESUME_ARGS[@]}"
    
    echo ""
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] 训练完成!"
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] Checkpoints 保存至: ${OUTPUT_PATH}"
}

# ==================== 主程序 ====================

# 检查是否是内部调用（后台运行模式的实际执行）
if [[ "${RUN_INTERNAL}" == "true" ]]; then
    # 内部调用，直接运行训练
    print_config
    run_training
    exit 0
fi

# 正常启动
print_config

# 检查是否后台运行
if [[ "${RUN_BACKGROUND}" == "true" ]]; then
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] 后台运行模式，日志输出到: ${LOG_FILE}"
    echo ""
    
    # 使用 nohup 调用脚本自身，传递 --run-internal 参数
    INTERNAL_ARGS=(--run-internal)
    if [[ "${RESUME_MODE}" == "off" ]]; then
        INTERNAL_ARGS+=(--no-resume)
    elif [[ "${RESUME_MODE}" == "path" ]]; then
        INTERNAL_ARGS+=(--resume-from "${RESUME_FROM}")
    fi
    
    nohup bash "$0" "${INTERNAL_ARGS[@]}" > "${LOG_FILE}" 2>&1 &
    PID=$!
    
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] 训练进程已启动，PID: ${PID}"
    echo ""
    echo "查看日志: tail -f ${LOG_FILE}"
    echo "停止训练: kill ${PID}"
    echo ""
    
    # 保存 PID 到文件方便后续管理
    echo "${PID}" > "${LOG_DIR}/train_${TIMESTAMP}.pid"
else
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] 前台运行模式"
    echo ""
    run_training
fi