#!/bin/bash
# ============================================================================
# ViTViz — SLURM array job: 30 tasks (6 models × 5 seeds)
#
# Config is selected via VITVIZ_CONFIG env var:
#
#   Run 1 (classical attacks, rollout, ~21h wall time on A100):
#     VITVIZ_CONFIG=configs/experiments/exp_run1_classical_rollout.yaml \
#     sbatch --array=0-29%2 scripts/apuana_sweep_array.sh
#
#   Run 2 (TGR+SAGA, rollout, ~52h wall time on A100):
#     VITVIZ_CONFIG=configs/experiments/exp_run2_vit_attacks_rollout.yaml \
#     sbatch --array=0-29%2 scripts/apuana_sweep_array.sh
#
#   Run 3 (all attacks, GMAR-L2, ~50h wall time on A100):
#     VITVIZ_CONFIG=configs/experiments/exp_run3_all_gmar.yaml \
#     sbatch --array=0-29%2 scripts/apuana_sweep_array.sh
#
# %2 = real GPU limit (QoS simple: 2 GPUs max per user simultaneously)
# A100 nodes forced via --nodelist; jobs queue if busy (acceptable).
# ============================================================================

#SBATCH --job-name=vitviz_sweep
#SBATCH --partition=long-simple
#SBATCH --qos=simple
#SBATCH --gres=gpu:1
#SBATCH --nodelist=cluster-node[6-10]
#SBATCH --cpus-per-task=8
#SBATCH --mem=32G
#SBATCH --time=1-06:00:00          # 30h — safe for A100; adjust if needed
#SBATCH --output=logs/sweep_%A_%a.out
#SBATCH --error=logs/sweep_%A_%a.err

# ── Config ───────────────────────────────────────────────────────────────────
CONFIG="${VITVIZ_CONFIG:-configs/experiments/exp_run3_all_gmar.yaml}"

# ── Model × seed mapping ─────────────────────────────────────────────────────
MODELS=(
    "ViT-S/32 · ImageNet-1k"
    "ViT-S/16 · ImageNet-1k"
    "ViT-B/32 · ImageNet-1k"
    "ViT-B/16 · ImageNet-1k"
    "ViT-L/32 · ImageNet-1k"
    "ViT-L/16 · ImageNet-1k"
)
SEEDS=(42 123 456 789 2026)

MODEL_IDX=$(( SLURM_ARRAY_TASK_ID / 5 ))
SEED_IDX=$(( SLURM_ARRAY_TASK_ID % 5 ))

MODEL="${MODELS[$MODEL_IDX]}"
SEED="${SEEDS[$SEED_IDX]}"

echo "======================================================"
echo "Array job   : $SLURM_ARRAY_JOB_ID[$SLURM_ARRAY_TASK_ID]"
echo "Config      : $CONFIG"
echo "Model       : $MODEL"
echo "Seed        : $SEED"
echo "Node        : $(hostname)"
echo "GPU         : $(nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null | head -1)"
echo "======================================================"

# ── Environment ──────────────────────────────────────────────────────────────
cd "$HOME/ViTViz" || { echo "ERROR: ~/ViTViz not found"; exit 1; }

source "$HOME/vitviz_venv/bin/activate" \
    || source ".venv/bin/activate" \
    || { echo "ERROR: venv not found"; exit 1; }

mkdir -p logs

# ── Run ──────────────────────────────────────────────────────────────────────
python experiments/run_attack_sweep.py \
    --config "$CONFIG" \
    --model-name "$MODEL" \
    --seed "$SEED" \
    --resume

EXIT_CODE=$?
echo "======================================================"
echo "Finished with exit code $EXIT_CODE"
echo "======================================================"
exit $EXIT_CODE