#!/bin/bash
# submit_mlp.bsub ,  Standard / Invariant / Augmented MLP baselines on QM9
#
# Job array index 1..54 sweeps (mode × target × seed).
#   modes:   standard, invariant, augmented   (3)
#   targets: gap, alpha, mu, zpve, mu_vector, alpha_tensor   (6)
#   seeds:   0, 1, 2   (3)
#   total:   3 × 6 × 3 = 54

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

set -uo pipefail
mkdir -p logs results

MODES=(standard invariant augmented)
TARGETS=(gap alpha mu zpve mu_vector alpha_tensor)
SEEDS=(0 1 2)
N_T=${#TARGETS[@]}
N_S=${#SEEDS[@]}

IDX0=$((LSB_JOBINDEX - 1))
MODE="${MODES[$(( IDX0 / (N_T * N_S) ))]}"
T_IDX=$(( (IDX0 / N_S) % N_T ))
S_IDX=$(( IDX0 % N_S ))
TARGET="${TARGETS[$T_IDX]}"
SEED="${SEEDS[$S_IDX]}"

case "$TARGET" in
  mu_vector|alpha_tensor) GROUP=octahedral ;;
  *)                      GROUP=cyclic ;;
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 mode=$MODE target=$TARGET seed=$SEED"
nvidia-smi --query-gpu=name,memory.total --format=csv,noheader || true

python3 train_baseline_mlp.py \
    --mode    "$MODE" \
    --target  "$TARGET" \
    --group   "$GROUP" \
    --qm9_dir "$QM9_DIR" \
    --seed    "$SEED" \
    --out_dir results/ \
    --device  cuda

echo "[$(date)] done"