#!/usr/bin/env bash
# 验证 HRM 和 TRM 都能在共享数据集上启动训练循环
# 不跑全量训练，每边跑 90s 检查崩没崩
set -euo pipefail
REPO_ROOT="$(cd "$(dirname "$0")/.." && pwd)"
source "$(conda info --base)/etc/profile.d/conda.sh"
conda activate rrm

if [[ ! -d "$REPO_ROOT/data/sudoku-extreme-1k-aug-1000" ]]; then
  echo "datasets not built; run scripts/build_datasets.sh sudoku first"; exit 1
fi

echo "=== HRM smoke ==="
cd "$REPO_ROOT/hrm"
( timeout 90 env WANDB_MODE=offline OMP_NUM_THREADS=4 CUDA_VISIBLE_DEVICES=0 \
    python pretrain.py \
      data_path="$REPO_ROOT/data/sudoku-extreme-1k-aug-1000" \
      epochs=1 eval_interval=1 global_batch_size=64 \
      lr=7e-5 puzzle_emb_lr=7e-5 weight_decay=1.0 puzzle_emb_weight_decay=1.0 \
      > /tmp/hrm_smoke.log 2>&1 ) || true
grep -iE "error|traceback" /tmp/hrm_smoke.log && { echo "HRM smoke FAILED"; tail -30 /tmp/hrm_smoke.log; exit 1; } || echo "HRM smoke OK"

echo "=== TRM smoke ==="
cd "$REPO_ROOT/trm"
( timeout 120 env WANDB_MODE=offline OMP_NUM_THREADS=4 CUDA_VISIBLE_DEVICES=0 \
    python pretrain.py \
      arch=trm \
      data_paths="[$REPO_ROOT/data/sudoku-extreme-1k-aug-1000]" \
      evaluators="[]" \
      epochs=1 eval_interval=1 global_batch_size=64 \
      lr=1e-4 puzzle_emb_lr=1e-4 weight_decay=1.0 puzzle_emb_weight_decay=1.0 \
      arch.mlp_t=True arch.pos_encodings=none \
      arch.L_layers=2 arch.H_cycles=3 arch.L_cycles=6 \
      +run_name=smoke_test ema=True \
      > /tmp/trm_smoke.log 2>&1 ) || true
grep -iE "error|traceback" /tmp/trm_smoke.log && { echo "TRM smoke FAILED"; tail -30 /tmp/trm_smoke.log; exit 1; } || echo "TRM smoke OK"

echo "==> all smoke tests passed"