#!/bin/bash
#SBATCH --job-name=ba-refit-seed
#SBATCH --partition=dgx-b200
#SBATCH --gpus=1
#SBATCH --cpus-per-task=10
#SBATCH --mem=200G
#SBATCH --time=24:00:00
#SBATCH --output=%x_%A_%a.out
#SBATCH --array=0-4          # 5 seeds → indices 0..4

HOME_LOC=~/
SCRIPT_LOC=$HOME_LOC/PeptiVerse/training_classifiers
ALT_EMB_LOC=$HOME_LOC/PeptiVerse/training_data_clean

# ── Configure per submission ──────────────────────────────────────────
BINDER_MODEL='wt'    # chemberta / peptideclm / wt
MODE='pooled'               # pooled / unpooled

# wt-wt
DATA_PATH="${SCRIPT_LOC}/binding_affinity/pair_wt_wt_${MODE}"
BASE_OUT_DIR="${SCRIPT_LOC}/binding_affinity/wt_wt_${MODE}"

# wt-smiles (chemberta or peptideclm)
#DATA_PATH="${ALT_EMB_LOC}/binding_affinity/${BINDER_MODEL}/pair_wt_smiles_${MODE}"
#BASE_OUT_DIR="${SCRIPT_LOC}/binding_affinity/${BINDER_MODEL}_smiles_${MODE}"
# ────────────────────────────────────────────────────────────────────────────

SEEDS=(1986 42 0 123 12345)
SEED=${SEEDS[$SLURM_ARRAY_TASK_ID]}

LOG_LOC=$SCRIPT_LOC/src_bash/logs
mkdir -p $LOG_LOC
DATE=$(date +%m_%d)

cd $SCRIPT_LOC

echo "Running: binder=${BINDER_MODEL}  mode=${MODE}  seed=${SEED}"
echo "  data  : ${DATA_PATH}"
echo "  out   : ${BASE_OUT_DIR}"

START_TIME=$(date +%s%N)

python -u refit_binding_affinity_seed.py \
    --dataset_path "${DATA_PATH}" \
    --base_out_dir "${BASE_OUT_DIR}" \
    --mode         "${MODE}" \
    --seed         "${SEED}" \
    > "${LOG_LOC}/${DATE}_ba_refit_${BINDER_MODEL}_${MODE}_seed${SEED}.log" 2>&1

END_TIME=$(date +%s%N)
ELAPSED_S=$(( (END_TIME - START_TIME) / 1000000000 ))
echo "Seed ${SEED} done at $(date) — wall clock: ${ELAPSED_S}s"
echo "{\"binder\": \"${BINDER_MODEL}\", \"mode\": \"${MODE}\", \"seed\": ${SEED}, \"wall_s\": ${ELAPSED_S}}" \
    >> "${LOG_LOC}/${DATE}_wall_clock_ba_refit.jsonl"

conda deactivate