#!/bin/bash
# submit_pyg_schnet.bsub ,  SchNet via torch_geometric.nn.models.SchNet
# (bypasses schnetpack's sqlite layer that hits KeyError on the CCC env).
#
# Array 1..12 = 4 scalar targets (gap, alpha, mu, zpve) × 3 seeds.

#BSUB -J starg_pyg_schnet[1-12]
#BSUB -o logs/pyg_schnet_%I_%J.out
#BSUB -e logs/pyg_schnet_%I_%J.err
#BSUB -q normal
#BSUB -n 8
#BSUB -gpu "num=1:mode=exclusive_process"
#BSUB -W 12:00
#BSUB -M 64GB

set -uo pipefail
mkdir -p logs results

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

WORKDIR=$HOME/starg/python/large_scale
PYG_ROOT=${PYG_ROOT:-/u/$USER/data/qm9_pyg}

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

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

python3 train_baseline_pyg_schnet.py \
    --target   "$TARGET" \
    --pyg_root "$PYG_ROOT" \
    --seed     "$SEED" \
    --out_dir  results/ \
    --device   cuda

echo "[$(date)] done"