#!/bin/bash
# submit_starg_neural.bsub ,  Neural ★_G on full QM9
#
# Job array index 1..18 = (target × seed). Identical layout to
# submit_starg_ridge.bsub.

#BSUB -J starg_neural[1-18]
#BSUB -o logs/starg_neural_%I_%J.out
#BSUB -e logs/starg_neural_%I_%J.err
#BSUB -q normal
#BSUB -n 4
#BSUB -gpu "num=1:mode=exclusive_process"
#BSUB -W 12:00
#BSUB -M 32GB

set -uo pipefail
mkdir -p logs results

TARGETS=(gap alpha mu zpve mu_vector alpha_tensor)
SEEDS=(0 1 2)
N_SEEDS=${#SEEDS[@]}
TARGET="${TARGETS[$(( (LSB_JOBINDEX - 1) / N_SEEDS ))]}"
SEED="${SEEDS[$(( (LSB_JOBINDEX - 1) % N_SEEDS ))]}"

case "$TARGET" in
  mu_vector|alpha_tensor) GROUP=octahedral; GROUP_PARAM=24 ;;
  *)                      GROUP=cyclic;     GROUP_PARAM=12 ;;
esac

WORKDIR=$HOME/starg/python/large_scale
QM9_DIR=${QM9_DIR:-/u/$USER/data/qm9/dsgdb9nsd}

cd "$WORKDIR"
export PYTHONPATH=".:${PYTHONPATH:-}"

echo "[$(date)] host=$(hostname) array=$LSB_JOBINDEX target=$TARGET seed=$SEED group=$GROUP"
nvidia-smi --query-gpu=name,memory.total --format=csv,noheader || true

python3 train_starg.py \
    --method     neural \
    --target     "$TARGET" \
    --group      "$GROUP" \
    --group_param "$GROUP_PARAM" \
    --qm9_dir    "$QM9_DIR" \
    --seed       "$SEED" \
    --out_dir    results/ \
    --device     cuda

echo "[$(date)] done"