#!/bin/bash

# Optional arguments [MODEL_NAME] [NUM_GPUS] [NUM_NODES]
NUM_GPUS=${1:-4}                # default 4 GPUs per node
NUM_NODES=${2:-1}               # default 1 node

# Clear first 2 args to avoid hydra issues
shift 2

# CPUs and memory automatically determined 
CPUS_PER_TASK=8                             # 8 CPUs per GPU
TOTAL_CPUS=$((NUM_GPUS * CPUS_PER_TASK))
TOTAL_MEM=$((NUM_GPUS * 32))                # 32gb ram per GPU

# Extract model name from hydra overrides
MODEL_NAME="transformer"
for arg in "$@"; do
  if [[ "$arg" == model=* ]]; then
    MODEL_NAME="${arg#model=}"
    break
  fi
done

JOB_NAME="${MODEL_NAME}_train_medium"

DATE=$(date +%F)

echo "Launching training job: $JOB_NAME"
echo "Nodes: ${NUM_NODES}, GPUs per node: ${NUM_GPUS}"
echo "Resources per node: ${TOTAL_CPUS} CPUs, ${TOTAL_MEM} GB memory"

SBATCH_ARGS="--job-name=${JOB_NAME} \
             --nodes=${NUM_NODES} \
             --gres=gpu:h200:${NUM_GPUS} \
             --ntasks-per-node=${NUM_GPUS} \
             --cpus-per-task=${CPUS_PER_TASK} \
             --mem=${TOTAL_MEM}G \
             --output=../../logs/${MODEL_NAME}/${DATE}/%x-%j.out \
             --error=../../logs/${MODEL_NAME}/${DATE}/%x-%j.err"

# bash /work/jf381/code/lm-research/scripts/training/train_medium_bash_resume.sh model=fst;bash /work/jf381/code/lm-research/scripts/training/train_medium_bash_resume_transformer.sh model=transformer;

# bash  /work/jf381/code/lm-research/scripts/jobs/run_job_h200_bash_resume.sh \
#     size=medium \
#     training=medium \
#     dataset=tinygsm_resume \
#     "$@"

# bash  /work/jf381/code/lm-research/scripts/jobs/run_job_h200_bash_resume_1_3b.sh \
#     size=1_3b \
#     training=1_3b \
#     dataset=tinygsm_resume \
#     "$@" \
#     model=fst_1_3B

# cp /work/jf381/code/lm-research/config/config_gpt2.yaml /work/jf381/code/lm-research/config/config.yaml

# cp /work/jf381/code/lm-research/src/lmr/models/fst/fst_353M.py /work/jf381/code/lm-research/src/lmr/models/fst/fst.py
# cp /work/jf381/code/lm-research/src/lmr/tokenizer/tokenizer_gpt2.py /work/jf381/code/lm-research/src/lmr/tokenizer/tokenizer.py  


# cp /work/jf381/code/lm-research/src/lmr/models/fst/fst_1_3B.py /work/jf381/code/lm-research/src/lmr/models/fst/fst.py
# cp /work/jf381/code/lm-research/src/lmr/tokenizer/tokenizer_gpt2.py /work/jf381/code/lm-research/src/lmr/tokenizer/tokenizer.py  


# bash /work/jf381/code/lm-research/scripts/jobs/run_job_h200_bash_resume_353M.sh \
#     size=medium \
#     training=medium \
#     dataset=tinygsm_resume \
#     model=fst_353M \
#     "$@"
    
bash  /work/jf381/code/lm-research/scripts/jobs/run_job_h200_bash_resume_transformer_353M.sh \
    size=medium \
    training=medium \
    dataset=tinygsm_resume \
    model=transformer_353M \
    "$@"


# bash  /work/jf381/code/lm-research/scripts/jobs/run_job_h200_bash_transformer_resume_353M.sh \
#     size=medium \
#     training=medium \
#     dataset=tinygsm_resume \
#     model=transformer_353M \
#     "$@"

# bash  /work/jf381/code/lm-research/scripts/jobs/run_job_h200_bash_resume_1_3B.sh \
#     size=medium \
#     training=medium \
#     dataset=tinygsm_resume \
#     model=fst_1_3B \
#     "$@" 

# bash  /work/jf381/code/lm-research/scripts/jobs/run_job_h200_bash_resume_transformer_353M.sh \
#     size=medium \
#     training=medium \
#     dataset=tinygsm_resume \
#     model=transformer_353M \
#     "$@" 

# bash  /work/jf381/code/lm-research/scripts/jobs/run_job_h200_bash_resume_transformer_1_3B.sh \
#     size=medium \
#     training=medium \
#     dataset=tinygsm_resume \
#     model=transformer_1_3B \
#     "$@" 



# bash  /work/jf381/code/lm-research/scripts/jobs/run_job_h200_bash_resume_353M.sh \
#     size=1_3b \
#     training=1_3b \
#     dataset=tinygsm_resume \
#     "$@" \
#     model=fst_1_3B

# bash  /work/jf381/code/lm-research/scripts/jobs/run_job_h200_bash_resume_generate.sh \
#     size=medium \
#     training=medium \
#     dataset=tinygsm_resume \
#     "$@"

# bash /work/jf381/code/lm-research/scripts/jobs/run_job_h200_bash_resume_1_3b_generate.sh \
#     size=1_3b \
#     training=1_3b \
#     dataset=tinygsm_resume \
#     "$@"

# bash /work/jf381/code/lm-research/scripts/training/train_medium_bash_resume.sh model=fst;

# bash /work/jf381/code/lm-research/scripts/training/train_medium_bash_resume.sh model=transformer;