#!/bin/bash

PROJECTDIR="$HOME/federated-conformal-fairness"
CONDAENV=fedcf
TRAIN_FRAC=0.3
VAL_FRAC=0.2
ALPHA=0.1

base_output_dir="$PROJECTDIR/outputs"
ds_folder="$PROJECTDIR/datasets"

client_formulation=0 # 0 for low overhead, 1 for more private

for closeness_measure in 0.1 0.15 0.2
do
for folktables_partition_type in continental_small continental_large continental_all all small large
do
    for DATASET in ACSEducation ACSIncome
    do
        for sketch_method in tdigest ddsketch
        do
            base_job_id="best_${DATASET}_split_${TRAIN_FRAC}_${VAL_FRAC}_${folktables_partition_type}"
            config_output_dir="$base_output_dir/${DATASET}/split/${TRAIN_FRAC}_${VAL_FRAC}_${folktables_partition_type}"

            for METRIC in "Demographic_Parity" "Equal_Opportunity" "Predictive_Equality"
            do
                for METHOD in "aps" "raps"
                do
                    for USE_MLE in false true
                    do
                        config_path="$PROJECTDIR/configs/fair_fedcf_acs_config.yaml"
                        cd ${PROJECTDIR}
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-00:00:00
#SBATCH -J c_${DATASET}_${TRAIN_FRAC}_${VAL_FRAC}_${folktables_partition_type}_${METHOD}
#SBATCH -e logs/fedcf/${DATASET}/${TRAIN_FRAC}_${VAL_FRAC}_${folktables_partition_type}_${sketch_method}/${METRIC}/closeness_${closeness_measure}_%A.err
#SBATCH -o logs/fedcf/${DATASET}/${TRAIN_FRAC}_${VAL_FRAC}_${folktables_partition_type}_${sketch_method}/${METRIC}/closeness_${closeness_measure}_%A.out

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

export DGLBACKEND=pytorch

cd ${PROJECTDIR}
python run_fed_cf.py --config_path=${config_path} \
    --logging_config.use_wandb False \
    --dataset.name=${DATASET} \
    --fairness_metric=${METRIC} \
    --conformal_method=${METHOD} \
    --job_id=c_${DATASET}_${TRAIN_FRAC}_${VAL_FRAC}_${METHOD}_${sketch_method} \
    --dataset_split_fractions.train ${TRAIN_FRAC} \
    --dataset_split_fractions.valid ${VAL_FRAC}  \
    --alpha $ALPHA \
    --base_job_id=${base_job_id} \
    --dataset.name=${DATASET} \
    --output_dir=$config_output_dir \
    --dataset_dir ${ds_folder} \
    --folktables_partition_type=${folktables_partition_type} \
    --closeness_measure ${closeness_measure} \
    --quantile_method ${sketch_method} \
    --resource_config.cpus 8 \
    --resource_config.gpus 0 \
    --client_formulation ${client_formulation} \
    --cf_opt.use_mle ${USE_MLE}
EOT
                        done
                    done
                done
            done
        done
    done
done
