#!/bin/bash
#SBATCH --account=p32222
#SBATCH --partition=gengpu
#SBATCH --time=48:00:00
#SBATCH --gres=gpu:2
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=8
#SBATCH --mem=150G
#SBATCH --job-name=adamson_st_se_250_2gpu
#SBATCH --output=logs/adamson_st_se_250_2gpu.%j.out
#SBATCH --error=logs/adamson_st_se_250_2gpu.%j.err

module purge

cd /projects/p32222/aTester/state_parse_repro
source .venv/bin/activate
export PATH="$HOME/.local/bin:$PATH"

export WANDB_MODE=disabled

export OMP_NUM_THREADS=1
export MKL_NUM_THREADS=1
export OPENBLAS_NUM_THREADS=1
export NUMEXPR_NUM_THREADS=1

export ADAMSON_ROOT="$PWD/data/adamson_state_compatible"
export ADAMSON_TOML="$ADAMSON_ROOT/tomls/adamson_fewshot.toml"
export ADAMSON_RUN_ROOT="$PWD/models/adamson_st_se_runs"
export ADAMSON_RUN_NAME="adamson_st_se_250steps_2gpu_8core_nw4"

mkdir -p "$ADAMSON_RUN_ROOT" logs

echo "Job started on $(hostname)"
echo "Date: $(date)"
echo "CUDA_VISIBLE_DEVICES=$CUDA_VISIBLE_DEVICES"
echo "SLURM_JOB_GPUS=$SLURM_JOB_GPUS"
echo "SLURM_GPUS=$SLURM_GPUS"
echo "SLURM_CPUS_PER_TASK=$SLURM_CPUS_PER_TASK"
echo "ADAMSON_TOML=$ADAMSON_TOML"
echo "ADAMSON_RUN_ROOT=$ADAMSON_RUN_ROOT"
echo "ADAMSON_RUN_NAME=$ADAMSON_RUN_NAME"

nvidia-smi

state tx train \
  data.kwargs.toml_config_path="$ADAMSON_TOML" \
  data.kwargs.embed_key=X_state \
  data.kwargs.pert_col=gene \
  data.kwargs.cell_type_key=cell_line \
  data.kwargs.batch_col=gem_group \
  data.kwargs.control_pert=non-targeting \
  data.kwargs.num_workers=4 \
  trainer.accelerator=gpu \
  trainer.devices=2 \
  training.max_steps=250 \
  model=state \
  output_dir="$ADAMSON_RUN_ROOT" \
  name="$ADAMSON_RUN_NAME"

echo "Training finished at $(date)"

export ADAMSON_RUN_DIR="$ADAMSON_RUN_ROOT/$ADAMSON_RUN_NAME"

echo "Run dir: $ADAMSON_RUN_DIR"
find "$ADAMSON_RUN_DIR" -maxdepth 4 -name "*.ckpt" -print

state tx predict \
  --output_dir "$ADAMSON_RUN_DIR" \
  --checkpoint final.ckpt \
  --profile minimal

echo "Prediction finished at $(date)"

find "$ADAMSON_RUN_DIR" -maxdepth 5 -name "*agg_results.csv" -print -exec cat {} \;
find "$ADAMSON_RUN_DIR" -maxdepth 5 -name "*results.csv" -print -exec head -20 {} \;