#!/bin/bash

PROJECTDIR="$HOME/federated-conformal-fairness"
CONDAENV=fedcf

ds_folder="$PROJECTDIR/datasets"

export PYTHONPATH=$PROJECTDIR
for DATASET in "ACSIncome" "ACSEducation"; do
    for folktables_partition_type in small large all continental_small continental_large continental_all; do
        for train_frac in 0.3; do
            for val_frac in 0.2; do
                best_param_path="$PROJECTDIR/base_fed_model/configs/base_folktables_config.yaml"
                if [ ! -f $best_param_path ]; then
                    echo "Best parameter file not found for ${DATASET} with train_frac=${train_frac} and val_frac=${val_frac}"
                    continue
                fi
                config_output_dir="$PROJECTDIR/outputs/${DATASET}/split/${train_frac}_${val_frac}_${folktables_partition_type}"
                mkdir -p ${config_output_dir}
                job_id="best_${DATASET}_split_${train_frac}_${val_frac}_${folktables_partition_type}"

sbatch <<EOT
#!/bin/bash
#SBATCH --account general
#SBATCH --partition=cpu_batch
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --cpus-per-task=8
#SBATCH --time=1-05:00:00
#SBATCH -J best_${DATASET}
#SBATCH -o ${PROJECTDIR}/logs/best/${DATASET}/best_${train_frac}_${val_frac}_%A.out
#SBATCH -e ${PROJECTDIR}/logs/best/${DATASET}/best_${train_frac}_${val_frac}_%A.err

echo Job started at `date` on `hostname`
# CONDA SETUP
source ~/.bashrc
conda deactivate
conda activate ${CONDAENV}

cd $PYTHONPATH
python base_fed_model/main.py \
--config_path=${best_param_path} \
--logging_config.use_wandb False \
--output_dir ${config_output_dir} \
--dataset.name ${DATASET} \
--dataset_dir ${ds_folder} \
--dataset_split_fractions.train ${train_frac} \
--dataset_split_fractions.valid ${val_frac} \
--job_id ${job_id} \
--num_server_rounds 50 \
--folktables_partition_type ${folktables_partition_type} \
--fraction_fit 0.5 \
--resource_config.cpus 8 \
--resource_config.gpus 0
EOT
            done
        done
    done
done
