#!/bin/bash
# DreamZero 单节点快速启动脚本 (5 步验证)
#
# 用法:
#   bash scripts/train/quick_start.sh libero          # 单卡验证 LIBERO
#   bash scripts/train/quick_start.sh manifeel 4      # 4 卡 ManiFeel
#   bash scripts/train/quick_start.sh robotwin 2      # 2 卡 RoboTwin
#
# 前置条件: 权重文件已下载到 checkpoints/ 目录

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)"

BENCHMARK="${1:-libero}"
NUM_GPUS="${2:-1}"
OUTPUT_DIR="${OUTPUT_DIR:-$REPO_ROOT/output/quickstart_${BENCHMARK}}"
PER_DEVICE_BS="${PER_DEVICE_BS:-1}"
MAX_STEPS="${MAX_STEPS:-5}"

echo "============================================"
echo " DreamZero Quick Start"
echo "============================================"
echo " Benchmark:   $BENCHMARK"
echo " GPUs:        $NUM_GPUS"
echo " Max steps:   $MAX_STEPS"
echo " Output:      $OUTPUT_DIR"
echo "============================================"

# ============ 验证权重文件 ============
CKPT_DIR="${CHECKPOINT_DIR:-$REPO_ROOT/checkpoints}"
WAN22_DIR="$CKPT_DIR/Wan2.2-TI2V-5B"
TOKENIZER_DIR="$CKPT_DIR/umt5-xxl"
CLIP_DIR="$CKPT_DIR/clip-encoder"

if [ ! -d "$WAN22_DIR" ]; then
    echo "ERROR: 未找到 Wan2.2 权重: $WAN22_DIR"
    echo "请先运行: bash scripts/data/download_checkpoints.sh"
    exit 1
fi
if [ ! -d "$TOKENIZER_DIR" ]; then
    echo "ERROR: 未找到 tokenizer: $TOKENIZER_DIR"
    echo "请先运行: bash scripts/data/download_checkpoints.sh"
    exit 1
fi
if [ ! -f "$CLIP_DIR/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth" ]; then
    echo "ERROR: 未找到 CLIP encoder: $CLIP_DIR"
    echo "请先运行: bash scripts/data/download_checkpoints.sh"
    exit 1
fi

# ============ 数据集选择 ============
case "$BENCHMARK" in
    libero)
        DATA_CFG="dreamzero/libero"
        NUM_FRAMES=12
        ACTION_HORIZON=12
        NUM_VIEWS=1
        MAX_STATE_DIM=44
        MAX_ACTION_DIM=32
        NUM_FRAME_PER_BLOCK=2
        NUM_ACTION_PER_BLOCK=12
        LR=1e-5
        ;;
    manifeel)
        DATA_CFG="dreamzero/manifeel"
        NUM_FRAMES=12
        ACTION_HORIZON=12
        NUM_VIEWS=3
        MAX_STATE_DIM=44
        MAX_ACTION_DIM=32
        NUM_FRAME_PER_BLOCK=2
        NUM_ACTION_PER_BLOCK=12
        LR=1e-5
        ;;
    robotwin)
        DATA_CFG="dreamzero/robotwin"
        NUM_FRAMES=12
        ACTION_HORIZON=12
        NUM_VIEWS=1
        MAX_STATE_DIM=44
        MAX_ACTION_DIM=32
        NUM_FRAME_PER_BLOCK=2
        NUM_ACTION_PER_BLOCK=12
        LR=1e-5
        ;;
    *)
        echo "ERROR: 未知 benchmark: $BENCHMARK (可选: libero, manifeel, robotwin)"
        exit 1
        ;;
esac

# ============ 自动选择 DeepSpeed 配置 ============
if [ "$NUM_GPUS" -le 2 ]; then
    DEEPSPEED_CFG="groot/vla/configs/deepspeed/zero2.json"
elif [ "$NUM_GPUS" -le 8 ]; then
    DEEPSPEED_CFG="groot/vla/configs/deepspeed/zero2_offload.json"
else
    DEEPSPEED_CFG="groot/vla/configs/deepspeed/zero3_multinode.json"
fi

cd "$REPO_ROOT"

torchrun --standalone --nproc_per_node "$NUM_GPUS" \
    groot/vla/experiment/experiment.py \
    report_to=none \
    data="$DATA_CFG" \
    train_architecture=full \
    model=dreamzero/vla \
    model/dreamzero/action_head=wan_flow_matching_action_tf_wan22 \
    model/dreamzero/transform=dreamzero_cotrain \
    num_frames="$NUM_FRAMES" \
    action_horizon="$ACTION_HORIZON" \
    num_views="$NUM_VIEWS" \
    max_state_dim="$MAX_STATE_DIM" \
    max_action_dim="$MAX_ACTION_DIM" \
    num_frame_per_block="$NUM_FRAME_PER_BLOCK" \
    num_action_per_block="$NUM_ACTION_PER_BLOCK" \
    num_state_per_block=1 \
    max_chunk_size=4 \
    frame_seqlen=50 \
    image_resolution_width=320 \
    image_resolution_height=160 \
    per_device_train_batch_size="$PER_DEVICE_BS" \
    max_steps="$MAX_STEPS" \
    save_strategy=no \
    optim=adamw_bnb_8bit \
    training_args.learning_rate="$LR" \
    training_args.deepspeed="$DEEPSPEED_CFG" \
    training_args.bf16=true \
    training_args.tf32=true \
    training_args.eval_bf16=true \
    output_dir="$OUTPUT_DIR" \
    dit_version="$WAN22_DIR" \
    text_encoder_pretrained_path="$WAN22_DIR/models_t5_umt5-xxl-enc-bf16.pth" \
    image_encoder_pretrained_path="$CLIP_DIR/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth" \
    vae_pretrained_path="$WAN22_DIR/Wan2.2_VAE.pth" \
    tokenizer_path="$TOKENIZER_DIR"

echo ""
echo "============================================"
echo " Quick start 完成! (benchmark: $BENCHMARK)"
echo "============================================"