#!/bin/bash
set -euo pipefail

REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")"/.. && pwd)"
cd "$REPO_ROOT"

RESULTS_DIR="./results"
RAW_DIR="$RESULTS_DIR/raw"
CHECKPOINT_DIR="$RESULTS_DIR/checkpoints"
DATASET="cifar10"
DATASET_TAG="cifar"
N_WORKERS=30
SEEDS=(3407 114514)
ALPHAS=(0.05 0.1)
SIGMAS=(1 5)

mkdir -p "$RAW_DIR" "$CHECKPOINT_DIR" "$RESULTS_DIR"

for seed in "${SEEDS[@]}"
do
    for alpha in "${ALPHAS[@]}"
    do
        for sigma in "${SIGMAS[@]}"
        do
            python scripts/run.py --logdir "$RAW_DIR" --modeldir "$CHECKPOINT_DIR" --model=cnn --dataset="$DATASET" --partition='noniid-labeldir' --alpha "$alpha" --device='cuda:0' --alg=proposed --lr 0.01 --tau_mean 1 --tot_step=8000 --comput_dist=normal --tau_std "$sigma" --test_freq 30 --save_net --n_workers="$N_WORKERS" --init_seed "$seed" &

            python scripts/run.py --logdir "$RAW_DIR" --modeldir "$CHECKPOINT_DIR" --model=cnn --dataset="$DATASET" --partition='noniid-labeldir' --alpha "$alpha" --device='cuda:0' --alg=async --lr 0.005 --tau_mean 1 --tot_step=8000 --comput_dist=normal --tau_std "$sigma" --test_freq 30 --save_net --n_workers="$N_WORKERS" --init_seed "$seed" &

            python scripts/run.py --logdir "$RAW_DIR" --modeldir "$CHECKPOINT_DIR" --model=cnn --dataset="$DATASET" --partition='noniid-labeldir' --alpha "$alpha" --device='cuda:0' --alg=buffer --buffer_size 5 --lr 0.005 --tau_mean 1 --tot_step=8000 --comput_dist=normal --tau_std "$sigma" --test_freq 30 --save_net --n_workers="$N_WORKERS" --init_seed "$seed" &

            python scripts/run.py --logdir "$RAW_DIR" --modeldir "$CHECKPOINT_DIR" --model=cnn --dataset="$DATASET" --partition='noniid-labeldir' --alpha "$alpha" --device='cuda:0' --alg=buffer --buffer_size "$N_WORKERS" --lr 0.01 --tau_mean 1 --tot_step=8000 --comput_dist=normal --tau_std "$sigma" --test_freq 30 --save_net --n_workers="$N_WORKERS" --init_seed "$seed" &

            python scripts/run.py --logdir "$RAW_DIR" --modeldir "$CHECKPOINT_DIR" --model=cnn --dataset="$DATASET" --partition='noniid-labeldir' --alpha "$alpha" --device='cuda:1' --alg=shuffle --buffer_size 5 --lr 0.01 --tau_mean 1 --tot_step=8000 --comput_dist=normal --tau_std "$sigma" --test_freq 30 --save_net --n_workers="$N_WORKERS" --init_seed "$seed" &

            python scripts/run.py --logdir "$RAW_DIR" --modeldir "$CHECKPOINT_DIR" --model=cnn --dataset="$DATASET" --partition='noniid-labeldir' --alpha "$alpha" --device='cuda:1' --alg=mifa --buffer_size 5 --lr 0.01 --tau_mean 1 --tot_step=8000 --comput_dist=normal --tau_std "$sigma" --test_freq 30 --save_net --n_workers="$N_WORKERS" --init_seed "$seed" &

            python scripts/run.py --logdir "$RAW_DIR" --modeldir "$CHECKPOINT_DIR" --model=cnn --dataset="$DATASET" --partition='noniid-labeldir' --alpha "$alpha" --device='cuda:1' --alg=concurrency --buffer_size=5 --lr 0.01 --tau_mean 1 --tot_step=8000 --comput_dist=normal --tau_std "$sigma" --test_freq 30 --save_net --n_workers="$N_WORKERS" --init_seed "$seed"
        done
    done
done

wait

python scripts/merge_logs_for_comparison.py \
  --logdir "$RAW_DIR" \
  --output-dir "$RAW_DIR" \
  --dataset "$DATASET" \
  --dataset-tag "$DATASET_TAG" \
  --n-workers "$N_WORKERS" \
  --partition noniid-labeldir \
  --model cnn \
  --seeds "${SEEDS[@]}"
