#!/bin/bash
# DreamZero RoboTwin Full Fine-Tuning Script (Wan2.2-TI2V-5B, custom RobotWinDataset)
export HYDRA_FULL_ERROR=1

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

# Activate DreamZero virtual environment
VENV_DIR="/root/autodl-tmp/venvs/dreamzero"
if [ -f "$VENV_DIR/bin/activate" ]; then
    source "$VENV_DIR/bin/activate"
fi

# ============ CONFIGURATION ============
ROBOTWIN_DATA_DIR=${ROBOTWIN_DATA_DIR:-"/root/autodl-tmp/data/robotwin_gear"}
OUTPUT_DIR=${OUTPUT_DIR:-"$REPO_ROOT/checkpoints/dreamzero_robotwin_full"}

if [ -z "${NUM_GPUS:-}" ]; then
  NUM_GPUS=$(nvidia-smi -L 2>/dev/null | wc -l)
fi
PER_DEVICE_BS=${PER_DEVICE_BS:-1}

CKPT_DIR="${CHECKPOINT_DIR:-/root/autodl-tmp/checkpoints}"
WAN22_DIR="$CKPT_DIR/Wan2.2-TI2V-5B"
TOKENIZER_DIR="$CKPT_DIR/umt5-xxl"
CLIP_DIR="$CKPT_DIR/DreamZero-DROID"
# ========================================

# Auto-download weights if missing
source /etc/network_turbo 2>/dev/null
if [ ! -d "$WAN22_DIR" ] || [ -z "$(ls -A "$WAN22_DIR" 2>/dev/null)" ]; then
    echo "Downloading Wan2.2-TI2V-5B..."
    huggingface-cli download Wan-AI/Wan2.2-TI2V-5B --local-dir "$WAN22_DIR"
fi
if [ ! -d "$TOKENIZER_DIR" ] || [ -z "$(ls -A "$TOKENIZER_DIR" 2>/dev/null)" ]; then
    echo "Downloading umt5-xxl..."
    huggingface-cli download google/umt5-xxl --local-dir "$TOKENIZER_DIR"
fi
if [ ! -d "$ROBOTWIN_DATA_DIR" ]; then
    echo "ERROR: RoboTwin data not found at $ROBOTWIN_DATA_DIR"; exit 1
fi

cd "$REPO_ROOT"

# ZeRO config
if [ "$NUM_GPUS" -le 2 ]; then
    DEEPSPEED_CFG=${DEEPSPEED_CFG:-"groot/vla/configs/deepspeed/zero2_offload.json"}
else
    DEEPSPEED_CFG=${DEEPSPEED_CFG:-"groot/vla/configs/deepspeed/zero2.json"}
fi

torchrun --standalone --nproc_per_node "$NUM_GPUS" \
    groot/vla/experiment/experiment.py \
    report_to="${REPORT_TO:-none}" \
    data=dreamzero/robotwin \
    wandb_project="${WANDB_PROJECT:-dreamzero-robotwin-sft}" \
    train_architecture=full \
    model=dreamzero/vla \
    model/dreamzero/action_head=wan_flow_matching_action_tf_wan22 \
    model/dreamzero/transform=dreamzero_cotrain \
    num_frames=12 action_horizon=12 state_horizon=1 num_views=1 \
    num_frame_per_block=2 num_action_per_block=12 \
    num_state_per_block=1 max_chunk_size=4 frame_seqlen=50 \
    image_resolution_width=320 image_resolution_height=160 \
    max_state_dim=44 max_action_dim=32 \
    seed=42 \
    training_args.learning_rate="${LR:-1e-5}" \
    training_args.deepspeed="$DEEPSPEED_CFG" \
    save_steps="${SAVE_STEPS:-2000}" \
    training_args.warmup_ratio=0.05 \
    output_dir="$OUTPUT_DIR" \
    per_device_train_batch_size="$PER_DEVICE_BS" \
    max_steps="${MAX_STEPS:-200000}" \
    weight_decay=1e-5 save_total_limit=2 \
    upload_checkpoints=false bf16=true tf32=true eval_bf16=true \
    dataloader_pin_memory=true dataloader_num_workers=4 \
    save_lora_only=false save_strategy=steps \
    robotwin_dataset_dir="$ROBOTWIN_DATA_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 " RoboTwin training finished! Output: $OUTPUT_DIR"
echo "============================================"