#!/bin/bash
# DreamZero T-Rex Training Script with Wan2.2-TI2V-5B backbone
#
# T-Rex: Dexmate Vega-1 dual-arm + 2x Sharpa Wave dexterous hands.
# State/action are 58-dim (7 arm + 22 hand per side) -> max_action_dim=64.
#
# Usage:
#   bash scripts/train/trex_training_wan22.sh
#
# Prerequisites:
#   - T-Rex dataset converted to LeRobot v2 + GEAR format at TREX_DATA_ROOT
#     (scripts/data/convert_trex_v3_to_v2.py then scripts/data/convert_lerobot_to_gear.py)
#   - Wan2.2-TI2V-5B weights (download from HuggingFace)
#     hf download Wan-AI/Wan2.2-TI2V-5B --local-dir ./checkpoints/Wan2.2-TI2V-5B
#   - Image encoder (CLIP) from Wan2.1 - Wan2.2-TI2V-5B does not include it
#     hf download Wan-AI/Wan2.1-I2V-14B-480P --local-dir ./checkpoints/Wan2.1-I2V-14B-480P
#   - umt5-xxl tokenizer
#     hf download google/umt5-xxl --local-dir ./checkpoints/umt5-xxl

export HYDRA_FULL_ERROR=1

# Repo root: must be a directory that contains groot/ (so experiment.py can be found).
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
SCRIPT_REPO_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)"
if [ -n "$DREAMZERO_ROOT" ] && [ -d "$DREAMZERO_ROOT/groot" ]; then
    : # keep existing and valid
elif [ -d "$SCRIPT_REPO_ROOT/groot" ]; then
    DREAMZERO_ROOT="$SCRIPT_REPO_ROOT"
else
    echo "ERROR: Set DREAMZERO_ROOT to the dreamzero repo root that contains groot/."
    exit 1
fi

# ============ USER CONFIGURATION ============
TREX_DATA_ROOT=${TREX_DATA_ROOT:-"$DREAMZERO_ROOT/data/trex_small"}
OUTPUT_DIR=${OUTPUT_DIR:-"$DREAMZERO_ROOT/checkpoints/dreamzero_trex_wan22_lora"}

# Wan2.2-TI2V-5B checkpoint (contains: diffusion weights, T5, VAE)
WAN22_CKPT_DIR=${WAN22_CKPT_DIR:-"$DREAMZERO_ROOT/checkpoints/Wan2.2-TI2V-5B"}

# Image encoder: Wan2.2-TI2V-5B does NOT include CLIP - use Wan2.1's or standalone
IMAGE_ENCODER_DIR=${IMAGE_ENCODER_DIR:-"$DREAMZERO_ROOT/checkpoints/Wan2.1-I2V-14B-480P"}

TOKENIZER_DIR=${TOKENIZER_DIR:-"$DREAMZERO_ROOT/checkpoints/umt5-xxl"}
# =============================================

# ============ AUTO-DOWNLOAD WEIGHTS ============
if [ ! -d "$WAN22_CKPT_DIR" ] || [ -z "$(ls -A "$WAN22_CKPT_DIR" 2>/dev/null)" ]; then
    echo "Wan2.2-TI2V-5B not found at $WAN22_CKPT_DIR. Downloading from HuggingFace..."
    hf download Wan-AI/Wan2.2-TI2V-5B --local-dir "$WAN22_CKPT_DIR"
fi

if [ ! -d "$TOKENIZER_DIR" ] || [ -z "$(ls -A "$TOKENIZER_DIR" 2>/dev/null)" ]; then
    echo "umt5-xxl tokenizer not found at $TOKENIZER_DIR. Downloading from HuggingFace..."
    hf download google/umt5-xxl --local-dir "$TOKENIZER_DIR"
fi

# Image encoder: download Wan2.1 if not present (only need CLIP from it)
if [ ! -f "$IMAGE_ENCODER_DIR/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth" ]; then
    echo "Image encoder not found. Downloading Wan2.1-I2V-14B-480P (for CLIP only)..."
    hf download Wan-AI/Wan2.1-I2V-14B-480P --local-dir "$IMAGE_ENCODER_DIR"
fi
# ================================================

# Validate dataset exists and is GEAR-converted
if [ ! -d "$TREX_DATA_ROOT" ]; then
    echo "ERROR: T-Rex dataset not found at $TREX_DATA_ROOT"
    exit 1
fi
if [ ! -f "$TREX_DATA_ROOT/meta/embodiment.json" ]; then
    echo "ERROR: $TREX_DATA_ROOT/meta/embodiment.json missing - run convert_lerobot_to_gear.py first"
    exit 1
fi
python "$DREAMZERO_ROOT/scripts/data/check_trex_dataset_ready.py" \
    --dataset-root "$TREX_DATA_ROOT"

EXPERIMENT_PY="$DREAMZERO_ROOT/groot/vla/experiment/experiment.py"
cd "$DREAMZERO_ROOT"

torchrun --nproc_per_node=gpu --standalone "$EXPERIMENT_PY" \
    report_to=wandb \
    data=dreamzero/trex_relative_wan22 \
    wandb_project=dreamzero \
    train_architecture=lora \
    num_frames=33 \
    action_horizon=24 \
    max_action_dim=64 \
    ++action_head_cfg.config.diffusion_model_cfg.action_dim=64 \
    num_views=3 \
    model=dreamzero/vla \
    model/dreamzero/action_head=wan_flow_matching_action_tf_wan22 \
    model/dreamzero/transform=dreamzero_cotrain \
    num_frame_per_block=2 \
    num_action_per_block=24 \
    num_state_per_block=1 \
    seed=42 \
    training_args.learning_rate=1e-5 \
    training_args.deepspeed="groot/vla/configs/deepspeed/zero2.json" \
    save_steps=8000 \
    training_args.warmup_ratio=0.05 \
    output_dir="$OUTPUT_DIR" \
    per_device_train_batch_size=1 \
    max_steps=100000 \
    weight_decay=1e-5 \
    save_total_limit=5 \
    upload_checkpoints=false \
    bf16=true \
    tf32=true \
    eval_bf16=true \
    dataloader_pin_memory=false \
    dataloader_num_workers=1 \
    save_lora_only=true \
    max_chunk_size=4 \
    save_strategy=steps \
    enable_wandb_video_reconstruction=true \
    wandb_video_reconstruction_steps=1000 \
    wandb_video_reconstruction_episode=0 \
    wandb_video_reconstruction_num_chunks=4 \
    trex_data_root="$TREX_DATA_ROOT" \
    dit_version="$WAN22_CKPT_DIR" \
    text_encoder_pretrained_path="$WAN22_CKPT_DIR/models_t5_umt5-xxl-enc-bf16.pth" \
    image_encoder_pretrained_path="$IMAGE_ENCODER_DIR/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth" \
    vae_pretrained_path="$WAN22_CKPT_DIR/Wan2.2_VAE.pth" \
    tokenizer_path="$TOKENIZER_DIR" \
    "$@"