#!/usr/bin/env bash
set -euo pipefail

ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
DATASET_SOURCE="${DATASET_SOURCE:-/mnt/local_nvme/zoubin/cz/self_forcing_predictor_v4_1000_seed0}"
TRAIN_SOURCE="${TRAIN_SOURCE:-/mnt/local_nvme/zoubin/cz/self_forcing_predictor_v4_training_1000_seed0}"
DATASET_DESTINATION="${DATASET_DESTINATION:-$ROOT_DIR/data/self_forcing_predictor_v4_1000_seed0}"
TRAIN_DESTINATION="${TRAIN_DESTINATION:-$ROOT_DIR/checkpoints/predictor_v4_1000_seed0}"

for source in "$DATASET_SOURCE" "$TRAIN_SOURCE"; do
  if [[ ! -d "$source" ]]; then
    echo "错误：找不到 NVMe 源目录：$source" >&2
    exit 1
  fi
done
if [[ ! -f "$DATASET_SOURCE/train_manifest.jsonl" ]]; then
  echo "错误：数据构建尚未完成，缺少 train_manifest.jsonl。" >&2
  exit 1
fi
if [[ ! -f "$TRAIN_SOURCE/predictor_step_02000.safetensors" ]]; then
  echo "错误：训练尚未完成，缺少 predictor_step_02000.safetensors。" >&2
  exit 1
fi

mkdir -p "$DATASET_DESTINATION" "$TRAIN_DESTINATION"
echo "同步数据：$DATASET_SOURCE -> $DATASET_DESTINATION"
rsync -a --partial --human-readable --info=stats2 \
  "$DATASET_SOURCE/" "$DATASET_DESTINATION/"
echo "同步训练产物：$TRAIN_SOURCE -> $TRAIN_DESTINATION"
rsync -a --partial --human-readable --info=stats2 \
  "$TRAIN_SOURCE/" "$TRAIN_DESTINATION/"

dataset_changes="$(
  rsync -a --dry-run --itemize-changes \
    "$DATASET_SOURCE/" "$DATASET_DESTINATION/" | wc -l
)"
training_changes="$(
  rsync -a --dry-run --itemize-changes \
    "$TRAIN_SOURCE/" "$TRAIN_DESTINATION/" | wc -l
)"
if (( dataset_changes != 0 || training_changes != 0 )); then
  echo "错误：同步后仍存在差异：dataset=$dataset_changes training=$training_changes" >&2
  exit 1
fi
echo "NVMe 产物已完整同步回 Self-Forcing 项目。"