#!/bin/bash
# =============================================================================
# Cosmos-Predict2.5 Action-Conditioned Evaluation on GR1 Robot Data
# =============================================================================
#
# 一键评估：指定模型路径（DCP 目录或 _ema_bf16.pt），自动完成转换 + 推理。
#
# 使用方式 (在 GPU 机器上):
#   bash scripts/eval_gr1_robot.sh /path/to/checkpoint
#   bash scripts/eval_gr1_robot.sh /path/to/model_ema_bf16.pt
#   bash scripts/eval_gr1_robot.sh /path/to/iter_000002000         # DCP 目录
#   bash scripts/eval_gr1_robot.sh /path/to/iter_000002000 --debug
#   bash scripts/eval_gr1_robot.sh /path/to/model_ema_bf16.pt --num-episodes 50
#
# 选项:
#   --num-episodes N        评估 episode 数量 (default: 100)
#   --output-dir PATH       输出目录
#   --guidance N            CFG guidance 值 (default: 0)
#   --num-gpus N            使用 GPU 数量 (default: 1)
#   --experiment NAME       Hydra 实验名 (default: 自动推断)
#   --debug                 调试模式 (5 episodes, 1 GPU)
#   -h, --help              显示帮助
# =============================================================================

set -euo pipefail

# ---------------------------------------------------------------------------
# 默认配置
# ---------------------------------------------------------------------------
NUM_EPISODES=100
GUIDANCE=0
OUTPUT_ROOT="/inspire/ssd/project/security-defense-and-attack/25015/world_model/output"
DATASET_PATH="/inspire/ssd/project/security-defense-and-attack/25015/world_model/data/PhysicalAI-Robotics-GR00T-Teleop-GR1/GR1_robot"
COSMOS_ROOT="/inspire/ssd/project/security-defense-and-attack/25015/cosmos-predict2.5"
NUM_GPUS=1
DEBUG=0
EXPERIMENT_NAME="cosmos_predict2p5_2B_action_conditioned_gr00t_gr1_customized_13frame_full_16nodes_release_oss"
CHECKPOINT_INPUT=""

# ---------------------------------------------------------------------------
# 解析参数 (第一个位置参数为 checkpoint 路径)
# ---------------------------------------------------------------------------
if [ $# -eq 0 ]; then
    echo "用法: bash scripts/eval_gr1_robot.sh <checkpoint_path> [选项]"
    echo ""
    echo "checkpoint_path 可以是:"
    echo "  - DCP 目录 (如 .../checkpoints/iter_000002000)"
    echo "  - _ema_bf16.pt 文件 (如 .../iter_000002000/model_ema_bf16.pt)"
    echo ""
    sed -n '3,18p' "$0"
    exit 1
fi

CHECKPOINT_INPUT="$1"
shift

while [[ $# -gt 0 ]]; do
    case $1 in
        --num-episodes)    NUM_EPISODES="$2"; shift 2 ;;
        --output-dir)      OUTPUT_ROOT="$2"; shift 2 ;;
        --guidance)        GUIDANCE="$2"; shift 2 ;;
        --num-gpus)        NUM_GPUS="$2"; shift 2 ;;
        --experiment)      EXPERIMENT_NAME="$2"; shift 2 ;;
        --debug)           DEBUG=1; shift ;;
        -h|--help)
            sed -n '3,18p' "$0"
            exit 0
            ;;
        *) echo "Unknown option: $1"; exit 1 ;;
    esac
done

# 调试模式覆盖
if [ "$DEBUG" -eq 1 ]; then
    NUM_EPISODES=5
    NUM_GPUS=1
    echo "=== DEBUG MODE ==="
fi

# ---------------------------------------------------------------------------
# Step 1: 确定 _ema_bf16.pt 路径 (自动转换 DCP → PT)
# ---------------------------------------------------------------------------
if [ ! -e "$CHECKPOINT_INPUT" ]; then
    echo "ERROR: 路径不存在: ${CHECKPOINT_INPUT}"
    exit 1
fi

if [ -f "$CHECKPOINT_INPUT" ]; then
    # 已经是 .pt 文件，直接使用
    CHECKPOINT_PATH="$CHECKPOINT_INPUT"
    echo "使用已有 PT checkpoint: ${CHECKPOINT_PATH}"
else
    # 是目录，需要找到或转换 _ema_bf16.pt
    CKPT_DIR="$CHECKPOINT_INPUT"

    # 检查是否已有 _ema_bf16.pt
    if [ -f "${CKPT_DIR}/model_ema_bf16.pt" ]; then
        CHECKPOINT_PATH="${CKPT_DIR}/model_ema_bf16.pt"
        echo "找到已有 PT checkpoint: ${CHECKPOINT_PATH}"
    elif [ -d "${CKPT_DIR}/model" ]; then
        # DCP 格式，需要转换
        echo "DCP checkpoint，开始转换..."
        echo "  输入: ${CKPT_DIR}/model/"
        echo "  输出: ${CKPT_DIR}/model_ema_bf16.pt"
        echo ""

        cd "${COSMOS_ROOT}"
        python scripts/convert_distcp_to_pt.py \
            "${CKPT_DIR}/model" \
            "${CKPT_DIR}"

        if [ -f "${CKPT_DIR}/model_ema_bf16.pt" ]; then
            CHECKPOINT_PATH="${CKPT_DIR}/model_ema_bf16.pt"
            echo ""
            echo "转换成功: ${CHECKPOINT_PATH}"
        else
            echo "ERROR: DCP → PT 转换失败"
            exit 1
        fi
    else
        echo "ERROR: 目录中既没有 model_ema_bf16.pt 也没有 model/ 子目录"
        echo "  内容: $(ls "$CKPT_DIR")"
        exit 1
    fi
fi

# ---------------------------------------------------------------------------
# Step 2: 准备评估数据
# ---------------------------------------------------------------------------
EVAL_DATA_DIR="${OUTPUT_ROOT}/eval/gr1_eval_data"

if [ ! -d "${EVAL_DATA_DIR}" ] || [ -z "$(ls -A "${EVAL_DATA_DIR}" 2>/dev/null)" ]; then
    echo ""
    echo "准备评估数据..."
    cd "${COSMOS_ROOT}"
    python scripts/prepare_gr1_eval_data.py \
        --dataset-path "${DATASET_PATH}" \
        --output-dir "${EVAL_DATA_DIR}" \
        --num-episodes "${NUM_EPISODES}"
else
    echo "评估数据已存在: ${EVAL_DATA_DIR}"
fi

# ---------------------------------------------------------------------------
# Step 3: 设置环境变量
# ---------------------------------------------------------------------------
export IMAGINAIRE_OUTPUT_ROOT="${OUTPUT_ROOT}"
export HF_HOME="${HF_HOME:-${OUTPUT_ROOT}/hf_cache}"

# 设置 HF 离线模式，防止尝试下载
export HF_HUB_OFFLINE=1
export TRANSFORMERS_OFFLINE=1

# ---------------------------------------------------------------------------
# Step 4: 运行推理
# ---------------------------------------------------------------------------
cd "${COSMOS_ROOT}"

PREDICT_DIR="${OUTPUT_ROOT}/eval/gr1_predicted"
mkdir -p "${PREDICT_DIR}"

# 日志文件
LOG_DIR="${OUTPUT_ROOT}/logs"
mkdir -p "${LOG_DIR}"
LOG_FILE="${LOG_DIR}/eval_$(date +%Y%m%d_%H%M%S).log"

echo ""
echo "==================================================================="
echo "Cosmos-Predict2.5  Action-Conditioned Evaluation  (GR1 Robot)"
echo "==================================================================="
echo "  Checkpoint:    ${CHECKPOINT_PATH}"
echo "  Experiment:    ${EXPERIMENT_NAME}"
echo "  评估数据:      ${EVAL_DATA_DIR}"
echo "  预测输出:      ${PREDICT_DIR}"
echo "  GPU 数量:      ${NUM_GPUS}"
echo "  Episode 数:    ${NUM_EPISODES}"
echo "  Guidance:      ${GUIDANCE}"
echo "  日志:          ${LOG_FILE}"
echo "==================================================================="

# 统计输入文件数
NUM_MP4=$(ls "${EVAL_DATA_DIR}"/*.mp4 2>/dev/null | wc -l)
echo "  输入视频数:    ${NUM_MP4}"
echo ""

# 设置 CUDA 可见设备
if [ "$NUM_GPUS" -eq 1 ]; then
    CUDA_DEVICES="0"
else
    CUDA_DEVICES=$(seq -s, 0 $((NUM_GPUS - 1)))
fi

# ---------------------------------------------------------------------------
# 本地模型路径 (与训练脚本一致，避免从 HF/S3 下载)
# ---------------------------------------------------------------------------
VAE_PATH="/inspire/ssd/project/security-defense-and-attack/25015/world_model/Wan2.1-VAE/Wan2.1_VAE.pth"
REASON1_PATH="/inspire/ssd/project/security-defense-and-attack/25015/world_model/Cosmos-Reason1-7B"
QWEN_PATH="/inspire/ssd/project/security-defense-and-attack/25015/models/Qwen2.5-VL-7B-Instruct"

# 构建推理命令的本地路径参数
LOCAL_PATH_ARGS=""
if [ -f "$VAE_PATH" ]; then
    LOCAL_PATH_ARGS="${LOCAL_PATH_ARGS} --vae_path ${VAE_PATH}"
    echo "  使用本地 VAE: ${VAE_PATH}"
else
    echo "  WARNING: 本地 VAE 不存在: ${VAE_PATH}"
fi
if [ -d "$REASON1_PATH" ]; then
    LOCAL_PATH_ARGS="${LOCAL_PATH_ARGS} --text_encoder_path ${REASON1_PATH}"
    echo "  使用本地 Cosmos-Reason1-7B: ${REASON1_PATH}"
else
    echo "  WARNING: 本地 Cosmos-Reason1-7B 不存在: ${REASON1_PATH}"
fi
if [ -d "$QWEN_PATH" ]; then
    LOCAL_PATH_ARGS="${LOCAL_PATH_ARGS} --qwen_path ${QWEN_PATH}"
    echo "  使用本地 Qwen2.5-VL-7B-Instruct: ${QWEN_PATH}"
else
    echo "  WARNING: 本地 Qwen2.5-VL-7B-Instruct 不存在: ${QWEN_PATH}"
fi

echo ""
echo "启动推理..."

CUDA_VISIBLE_DEVICES=${CUDA_DEVICES} PYTHONPATH=. python \
    cosmos_predict2/_src/predict2/action/inference/inference_gr00t.py \
    --experiment="${EXPERIMENT_NAME}" \
    --ckpt_path="${CHECKPOINT_PATH}" \
    --input_video_root="${EVAL_DATA_DIR}" \
    --save_root="${PREDICT_DIR}" \
    --resolution 480,832 \
    --guidance ${GUIDANCE} \
    --chunk_size 12 \
    --start 0 \
    --end ${NUM_EPISODES} \
    --fps_downsample_ratio 2 \
    --save_fps 10 \
    ${LOCAL_PATH_ARGS} \
    2>&1 | tee "${LOG_FILE}"

echo ""
echo "==================================================================="
echo "评估完成!"
echo "  预测视频: ${PREDICT_DIR}"
echo "  日志:     ${LOG_FILE}"
echo "==================================================================="