#!/bin/bash
#SBATCH --job-name=dovla_hparam_sweep
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=8
#SBATCH --gres=gpu:1
#SBATCH --mem=64000M
#SBATCH --time=48:00:00
#SBATCH --output=logs/phase_a4_hparam_%A_%a.out
#SBATCH --error=logs/phase_a4_hparam_%A_%a.err
#SBATCH --array=0-8

set -euo pipefail

# Phase A4: Hyperparameter sweep
# Grid: 3 LR x 3 hidden_dim = 9 configs

PROJECT_DIR="${PROJECT_DIR:-$PWD}"
cd "$PROJECT_DIR"

source .venv/bin/activate

DATASET="/scratch/$USER/dovla/experiments/maniskill_presuccess_six_task_collection"
OUT_ROOT="/scratch/$USER/dovla/experiments/phase_a4_hparam_sweep"

# Hyperparameter grid
LRS=(0.0001 0.0003 0.001)
HIDDEN_DIMS=(256 512 1024)

# Map array index to config
IDX=$SLURM_ARRAY_TASK_ID
LR_IDX=$((IDX / 3))
HD_IDX=$((IDX % 3))

LR="${LRS[$LR_IDX]}"
HIDDEN_DIM="${HIDDEN_DIMS[$HD_IDX]}"

OUT_DIR="$OUT_ROOT/lr${LR}_h${HIDDEN_DIM}"
mkdir -p "$OUT_DIR" logs

echo "=== Phase A4: Hyperparameter Sweep ==="
echo "Config $IDX: LR=$LR, Hidden=$HIDDEN_DIM"
echo ""

python scripts/train_dovla.py \
    --dataset "$DATASET" \
    --out "$OUT_DIR" \
    --objective lattice_field \
    --hidden-dim "$HIDDEN_DIM" \
    --epochs 50 \
    --batch-groups 16 \
    --lr "$LR" \
    --device auto \
    --seed 0

echo ""
# Quick eval
python scripts/eval_lattice_checkpoint.py \
    --checkpoint "$OUT_DIR/best.pt" \
    --dataset /scratch/$USER/dovla/experiments/maniskill_presuccess_six_task_collection \
    --out "$OUT_DIR/lattice_eval.json" \
    --mode field_only \
    --all-groups

echo "✅ Phase A4 config $IDX complete"