#!/bin/bash
#SBATCH --job-name=dovla_lattice_eval
#SBATCH --account=def-yalda_cpu
#SBATCH --partition=cpubase_bycore_b1
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=4
#SBATCH --mem=12G
#SBATCH --time=00:30:00
#SBATCH --array=0-5%6
#SBATCH --output=outputs/hpc/logs/%x_%A_%a.out
#SBATCH --error=outputs/hpc/logs/%x_%A_%a.err

set -euo pipefail

PROJECT_DIR="${PROJECT_DIR:-$SLURM_SUBMIT_DIR}"
DATASET="${DATASET:?Set DATASET to the matching CIL directory}"
RUN_ROOT="${RUN_ROOT:?Set RUN_ROOT to the matching training root}"
MODE="${MODE:-full}"
PYTHON="${PYTHON:-$PROJECT_DIR/.venv/bin/python}"
EVAL_DEVICE="${EVAL_DEVICE:-cpu}"
TASK_INDEX="${SLURM_ARRAY_TASK_ID:-0}"
USE_VISUAL_CONTAINER="${USE_VISUAL_CONTAINER:-0}"

if [[ "$MODE" == "full" ]]; then
  SEED="$((TASK_INDEX / 2))"
  if (( TASK_INDEX % 2 == 0 )); then
    OBJECTIVE="lattice_field"
  else
    OBJECTIVE="legacy"
  fi
  RUN_DIR="$RUN_ROOT/$OBJECTIVE/seed_$SEED"
elif [[ "$MODE" == "scaling" ]]; then
  K="${K:?Set K for scaling evaluation}"
  SEED="$TASK_INDEX"
  RUN_DIR="$RUN_ROOT/k_$K/seed_$SEED"
elif [[ "$MODE" == "visual" ]]; then
  VISUAL_SEEDS="${VISUAL_SEEDS:-3}"
  if (( TASK_INDEX >= VISUAL_SEEDS )); then
    echo "skip visual eval index $TASK_INDEX: VISUAL_SEEDS=$VISUAL_SEEDS"
    exit 0
  fi
  SEED="$TASK_INDEX"
  RUN_DIR="$RUN_ROOT/lattice_field/seed_$SEED"
  USE_VISUAL_CONTAINER=1
elif [[ "$MODE" == "field_only" ]]; then
  SEED="$TASK_INDEX"
  OBJECTIVE="${OBJECTIVE:-lattice_field}"
  RUN_DIR="$RUN_ROOT/$OBJECTIVE/seed_$SEED"
elif [[ "$MODE" == "baseline" ]]; then
  BASELINE="${BASELINE:?Set BASELINE for baseline evaluation}"
  SEED="$TASK_INDEX"
  RUN_DIR="$RUN_ROOT/$BASELINE/seed_$SEED"
else
  echo "MODE must be full, scaling, visual, or baseline" >&2
  exit 2
fi

EVAL_EXTRA_ARGS=()
if [[ "$MODE" == "scaling" ]]; then
  EVAL_EXTRA_ARGS+=(--training-k "$K" --all-groups)
fi
if [[ "${ALL_GROUPS:-0}" == "1" ]]; then
  EVAL_EXTRA_ARGS+=(--all-groups)
fi

cd "$PROJECT_DIR"
export OMP_NUM_THREADS=1
export OPENBLAS_NUM_THREADS=1
export MKL_NUM_THREADS=1
export DOVLA_TORCH_THREADS=1

if (( USE_VISUAL_CONTAINER )); then
  SCRATCH_ROOT="/scratch/$USER/dovla"
  SIF="${SIF:-$SCRATCH_ROOT/containers/pytorch_2.7.1_cuda12.8.sif}"
  CONTAINER_PYTHON="${CONTAINER_PYTHON:-$SCRATCH_ROOT/envs/maniskill/bin/python}"
  module load StdEnv/2023 apptainer/1.4.5
  APPTAINER_GPU_ARGS=()
  if [[ "$EVAL_DEVICE" == cuda* ]]; then
    APPTAINER_GPU_ARGS+=(--nv)
  fi
  PYTHON_COMMAND=(
    apptainer exec
    "${APPTAINER_GPU_ARGS[@]}"
    --env "OMP_NUM_THREADS=1,OPENBLAS_NUM_THREADS=1,MKL_NUM_THREADS=1,DOVLA_TORCH_THREADS=1"
    -B "$PROJECT_DIR:$PROJECT_DIR"
    -B "/scratch/$USER:/scratch/$USER"
    "$SIF"
    "$CONTAINER_PYTHON"
  )
else
  PYTHON_COMMAND=("$PYTHON")
fi

test -f "$RUN_DIR/best.pt"
"${PYTHON_COMMAND[@]}" scripts/eval_lattice_checkpoint.py \
  --checkpoint "$RUN_DIR/best.pt" \
  --dataset "$DATASET" \
  --out "$RUN_DIR/lattice_eval.json" \
  --device "$EVAL_DEVICE" \
  "${EVAL_EXTRA_ARGS[@]}"