#!/bin/bash
# submit_qm7x_enns.bsub ,  ENN baselines on QM7-X polarizability components.
#
# Array 1..18 = 3 methods × 2 targets × 3 seeds.
#   method ∈ {mace, schnet, e3nn}
#   target ∈ {alpha_E, alpha_T2}
#   seed   ∈ {42, 43, 44}
#
# index mapping (1-based):
#   method_idx = (idx - 1) // 6     -> 0..2
#   target_idx = ((idx - 1) % 6) // 3 -> 0..1
#   seed_idx   = (idx - 1) % 3      -> 0..2

#BSUB -J starg_qm7x[1-18]
#BSUB -o logs/qm7x_%I_%J.out
#BSUB -e logs/qm7x_%I_%J.err
#BSUB -q normal
#BSUB -n 16
#BSUB -gpu "num=1:mode=exclusive_process"
#BSUB -W 24:00
#BSUB -M 128GB

set -uo pipefail
export TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD=1
mkdir -p logs results

METHODS=(mace schnet e3nn)
TARGETS=(alpha_E alpha_T2)
SEEDS=(42 43 44)

METHOD="${METHODS[$(( (LSB_JOBINDEX - 1) / 6 ))]}"
TARGET="${TARGETS[$(( ((LSB_JOBINDEX - 1) % 6) / 3 ))]}"
SEED="${SEEDS[$(( (LSB_JOBINDEX - 1) % 3 ))]}"

WORKDIR=$HOME/starg/python/large_scale
QM7X_DIR=${QM7X_DIR:-/u/$USER/data/qm7x}

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

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

case "$METHOD" in
  mace)
    python3 train_baseline_mace.py \
        --dataset  qm7x \
        --target   "$TARGET" \
        --qm7x_dir "$QM7X_DIR" \
        --seed     "$SEED" \
        --out_dir  results/ \
        --device   cuda
    ;;
  schnet)
    python3 train_baseline_pyg_schnet.py \
        --dataset  qm7x \
        --target   "$TARGET" \
        --qm7x_dir "$QM7X_DIR" \
        --seed     "$SEED" \
        --out_dir  results/ \
        --device   cuda
    ;;
  e3nn)
    python3 train_baseline_e3nn.py \
        --dataset  qm7x \
        --target   "$TARGET" \
        --qm7x_dir "$QM7X_DIR" \
        --seed     "$SEED" \
        --out_dir  results/ \
        --device   cuda
    ;;
  *)
    echo "unknown method: $METHOD"
    exit 1
    ;;
esac

echo "[$(date)] done"