#!/bin/bash
export PYTHONPATH=$PYTHONPATH:/storage/ice-shared/ae8803che/hxue/data/world_model
source /storage/ice-shared/ae8803che/hxue/data/wm/bin/activate
cd /storage/ice-shared/ae8803che/hxue/data/world_model

JOBID=4052280
CONFIG=wm/config/fulltraj_dit/lang_table.yaml
CKPT_DIR=checkpoints/lang_table_fulltraj_dit_v1
BATCH_SIZE=64

echo "Starting evaluation sequence..."

# 1. 10 Steps
srun --overlap --jobid $JOBID python wm/eval/eval_checkpoints.py --config $CONFIG --ckpt_dir $CKPT_DIR --inference_steps 10 --label 10steps --batch_size $BATCH_SIZE

# 2. 20 Steps
srun --overlap --jobid $JOBID python wm/eval/eval_checkpoints.py --config $CONFIG --ckpt_dir $CKPT_DIR --inference_steps 20 --label 20steps --batch_size $BATCH_SIZE

# 3. 50 Steps
srun --overlap --jobid $JOBID python wm/eval/eval_checkpoints.py --config $CONFIG --ckpt_dir $CKPT_DIR --inference_steps 50 --label 50steps --batch_size $BATCH_SIZE

# 4. 100 Steps
srun --overlap --jobid $JOBID python wm/eval/eval_checkpoints.py --config $CONFIG --ckpt_dir $CKPT_DIR --inference_steps 100 --label 100steps --batch_size $BATCH_SIZE

# 5. 50 Steps with Noise
srun --overlap --jobid $JOBID python wm/eval/eval_checkpoints.py --config $CONFIG --ckpt_dir $CKPT_DIR --inference_steps 50 --noise_level 0.1 --label 50steps_noise0.1 --batch_size $BATCH_SIZE

# 6. Plot comparisons
python wm/eval/plot_comparisons.py

echo "Evaluation sequence finished."