#!/bin/bash

PROJECTDIR="$HOME/federated_dir/federated-conformal-fairness"
CONDAENV=myenv
TRAIN_FRAC=0.3
VAL_FRAC=0.2
ALPHA=0.1
DATASET="Fitzpatrick"
base_output_dir="$PROJECTDIR/outputs"
ds_folder="/fs/ess/PAS2030/fedcf/datasets/finalfitz17k"

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


for num_clients in 8; do # 2 4 8
# Build [1,0,1,0,...] with length = $num_clients (odd_pos=1, even_pos=0) for hybrid client formulation
    # arr=()
    # for ((i=1; i<=num_clients; i++)); do arr+=($((i%2))); done
    # client_formulation="[$(IFS=,; echo "${arr[*]}")]"
    for client_formulation in 0; do
        for closeness_measure in 0.1; do #0.15 0.2
            for sketch_method in tdigest; do # mean tdigest ddsketch
                config_output_dir="$PROJECTDIR/outputs/${DATASET}/split/${TRAIN_FRAC}_${VAL_FRAC}_${num_clients}"
                job_id="best_${DATASET}_split_${TRAIN_FRAC}_${VAL_FRAC}_${num_clients}_faster"

                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_fitzpatrick_config.yaml"
                            cd ${PROJECTDIR}
sbatch <<EOT
#!/bin/bash
#SBATCH --account PAS2030
#SBATCH --partition=cpu
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --cpus-per-task=40
#SBATCH --time=1-00:00:00
#SBATCH -J c_${DATASET}_${TRAIN_FRAC}_${VAL_FRAC}_${METHOD}
#SBATCH -e logs/fedcf/${DATASET}/${TRAIN_FRAC}_${VAL_FRAC}_${sketch_method}/${METRIC}/closeness_${closeness_measure}_%A.err
#SBATCH -o logs/fedcf/${DATASET}/${TRAIN_FRAC}_${VAL_FRAC}_${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 \
    --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=${job_id} \
    --dataset.name=${DATASET} \
    --output_dir=${config_output_dir} \
    --dataset_dir ${ds_folder} \
    --dataset.csv_file ${PROJECTDIR}/fitzpatrick17k.csv \
    --closeness_measure ${closeness_measure} \
    --quantile_method ${sketch_method} \
    --num_clients ${num_clients} \
    --resource_config.cpus 10 \
    --resource_config.gpus 0 \
    --client_formulation ${client_formulation} \
    --cf_opt.use_mle ${USE_MLE}

EOT
                        done
                    done
                done
            done
        done
    done
done