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

task_names=(
"turning_on_radio"
"picking_up_trash"
"hiding_Easter_eggs"
"wash_a_baseball_cap"
"hanging_pictures"
"attach_a_camera_to_a_tripod"
"make_microwave_popcorn"
"bringing_water"
"tidying_bedroom"
"putting_shoes_on_rack"
"setting_the_fire"
"cook_hot_dogs"
)

node_per_task=5
num_test_per_episode=2
# ckpt_name="pt12-2w-comet2-refactor-20260213065325"
# ckpt_name="pi05-b1kpt12-cs32"
ckpt_name="pt12-comet-5w-comet4-b1knorm-20260309094349"
gpu_per_node=4
episodes_per_task=10

total_work=$((episodes_per_task * num_test_per_episode))
base=$((total_work / node_per_task))
rem=$((total_work % node_per_task))

submitted=0

for task_name in "${task_names[@]}"; do
  for (( node = 0; node < node_per_task; node++ )); do
    if (( node < rem )); then
      node_trials=$((base + 1))
      node_start=$((node * (base + 1)))
    else
      node_trials=$base
      node_start=$((rem * (base + 1) + (node - rem) * base))
    fi

    job_command="$(cat <<EOF
export TASK_NAME="${task_name}"
export CKPT_NAME="${ckpt_name}"
export NODE_ID="${node}"
export PATH_TO_CKPT=/tmp/\${CKPT_NAME}
export OPENPI_DATA_HOME=/opt/openpi-cache
export EPISODES_PER_TASK=${episodes_per_task}
export EVAL_START=${node_start}
export EVAL_COUNT=${node_trials}
export NUM_GPU=${gpu_per_node}

apt-get install -y libxt6

EVAL_ROOT=/opt/eval

cd \${EVAL_ROOT}/BEHAVIOR-1K

# prepare b1k env
mkdir -p /opt/miniconda3/envs/behavior
hf download shangkuns/b1k-behavior-env behavior-env.tar --local-dir /tmp/hf-cache
tar xf /tmp/hf-cache/behavior-env.tar -C /opt/miniconda3/envs/behavior

source /opt/miniconda3/etc/profile.d/conda.sh
conda activate behavior
python -c "from omnigibson.utils.asset_utils import download_omnigibson_robot_assets; download_omnigibson_robot_assets()"
python -c "from omnigibson.utils.asset_utils import download_behavior_1k_assets; download_behavior_1k_assets(accept_license=True)"
python -c "from omnigibson.utils.asset_utils import download_2025_challenge_task_instances; download_2025_challenge_task_instances()"
conda deactivate

hf download shangkuns/\${CKPT_NAME} --local-dir \${PATH_TO_CKPT}

rm -rf \${EVAL_ROOT}/b1k-eval-utils
hf download shangkuns/b1k-eval-utils --local-dir \${EVAL_ROOT}/b1k-eval-utils
cp \${EVAL_ROOT}/b1k-eval-utils/run_dual_eval.sh /opt/eval/BEHAVIOR-1K/run_dual_eval.sh
cp \${EVAL_ROOT}/b1k-eval-utils/eval_b1k.sh /opt/eval/BEHAVIOR-1K/eval_b1k.sh
cp \${EVAL_ROOT}/b1k-eval-utils/eval_openpi.sh /opt/eval/openpi-comet/eval_openpi.sh

mkdir -p \${PATH_TO_CKPT}/assets/behavior-1k/2025-challenge-demos
cp \${EVAL_ROOT}/b1k-eval-utils/norm_stats.json \${PATH_TO_CKPT}/assets/behavior-1k/2025-challenge-demos/norm_stats.json

bash \${EVAL_ROOT}/b1k-eval-utils/patch_walltime.sh \${EVAL_ROOT}/BEHAVIOR-1K/OmniGibson/omnigibson/learning
bash run_dual_eval.sh

hf upload shangkuns/\${CKPT_NAME} /opt/eval/eval_output \${TASK_NAME} --repo-type dataset
EOF
)"

    lep job create \
      --resource-shape "my.${gpu_per_node}xl40s" \
      --node-group oci-ord-lepton-001 \
      --num-workers 1 \
      --container-image littlespray/b1k-eval-light:v4 \
      --command "${job_command}" \
      --intra-job-communication=true \
      --env NVIDIA_DRIVER_CAPABILITIES=all \
      --env DEBIAN_FRONTEND=noninteractive \
      --secret HUGGING_FACE_HUB_TOKEN=SHANGKUN_HF_NV.shangkuns \
      --ttl-seconds-after-finished 259200 \
      --shared-memory-size 786432 \
      --log-collection true \
      --queue-priority 7 \
      --can-preempt \
      --name "$(printf '%s-n%s' "$task_name" "$node" | tr '[:upper:]_' '[:lower:]-' | tail -c 36)"

    submitted=$((submitted + 1))
  done
done

echo "Submitted ${submitted} jobs."