#!/usr/bin/env bash
# Search Stages
#1. fednas_search -> workshop paper
#2. backbone_cell_search -> backbone + fednas cell
#3. personalized_search -> backbone + personalized cell
#4. train -> resnet18 fedavg

# search (ours)
# sh run_fednas_cifar10.sh 1 4 20 4 resnet18 hetero 2 1 4 0.3 3e-4 88888 nas 1 personalized_search 1 1 0

# train ResNet-18 (baseline) + fine tunning
# sh run_fednas_cifar10.sh 1 4 20 4 resnet18 hetero 2 1 4 0.3 3e-4 88888 nas 1 train 1 1 0 True

# train ResNet-18 (baseline) + no fine tunning
# sh run_fednas_cifar10.sh 1 4 20 4 resnet18 hetero 2 1 4 0.3 3e-4 88888 nas 1 train 1 1 0 False

# FedNAS(workshop paper) + fine tunning
# sh run_fednas_cifar10.sh 1 4 20 4 resnet18 hetero 2 1 32 0.3 3e-4 99999 nas 8 search 1 1 0 True

# FedNAS(workshop paper) + no fine tunning
# sh run_fednas_cifar10.sh 1 4 20 4 resnet18 lda 2 1 8 0.3 3e-4 99999 nas 8 fednas_search 1 0 True
#sh run_fednas_cifar10.sh 1 4 10 4 resnet18 hetero 10 1 32 0.01 3e-4 99999 nas 8 fednas_search 1 0 False

SERVER_NUM=$1
GPU_NUM_PER_SERVER=$2
CLIENT_NUM=$3
WORKER_NUM=$4
MODEL=$5
DISTRIBUTION=$6
ROUND=$7
EPOCH=$8
BATCH_SIZE=$9
W_LR=${10}
ALPHA_LR=${11}
RUN_ID=${12}
DESIGN=${13}
Layers=${14}
STAGE=${15}
TRAIN_EPOCH=${16}
GPU_STARTING_FROM=${17}
LOCAL_FINETUNE=${18}

hostname > mpi_host_file

PROCESS_NUM=`expr $WORKER_NUM + 1`

echo $PROCESS_NUM

mpirun -np $PROCESS_NUM -hostfile ./mpi_host_file python3 ./main_fednas.py \
  --gpu_server_num $SERVER_NUM \
  --gpu_num_per_server $GPU_NUM_PER_SERVER \
  --run_id $RUN_ID \
  --model $MODEL \
  --dataset cifar10 \
  --design $DESIGN \
  --partition_method $DISTRIBUTION  \
  --client_num_per_round $WORKER_NUM \
  --client_num_in_total $CLIENT_NUM \
  --comm_round $ROUND \
  --epochs $EPOCH \
  --batch_size $BATCH_SIZE \
  --learning_rate $W_LR \
  --lr $W_LR \
  --gamma 0 \
  --beta 0 \
  --epochs_for_train $TRAIN_EPOCH\
  --classes_per_client 5\
  --arch_learning_rate $ALPHA_LR \
  --layers $Layers \
  --local_finetune $LOCAL_FINETUNE\
  --stage $STAGE \
  --starting_gpu_id $GPU_STARTING_FROM \
  --data_dir './../../../data/cifar10'
#  --data_dir './../../../../../../data/cifar10'
#  --data_dir '/data/gld'
#  --data_dir './../../../data/gld'
#  --data_dir '/data/gld'
#  --data_dir './../../../data/gld'
#  --data_dir './../../../data/gld'
#  --data_dir '/data/gld'
#  --data_dir './../../../data/gld'
#  --data_dir '/data/gld'
#  --data_dir './../../../data/gld'
#  --data_dir '/data/gld'
#  --data_dir './../../../data/gld'
#  --data_dir '/data/gld'
#  --data_dir "/data/gld/" \
#  --client_num_in_total $CLIENT_NUM
#   --data_dir "/data/gld/" \
#  --data_dir './../../../data/gld/'

#  --data_dir "/data/gld/"
#  --data_dir './../../../data/gld/'
#  --data_dir "/data/gld/"
#  --data_dir './../../../data/gld/'
#  --data_dir "/data/gld/"
#  --data_dir './../../../data/gld/'  # if cpu
#  --data_dir "/data/gld/"
#  --data_dir '/home/chaoyanghe/sourcecode/fedml.ai/data/gld/'
#  --data_dir './../../../../../../../data/gld/images' # if server
#  --data_dir './../../../data/gld/images'  # if cpu



#mpirun -np 3 -hostfile ./mpi_host_file python3 ./main_fednas.py \
#--gpu_server_num 1 \
#--gpu_num_per_server 4 \
#--model "transformer" \
#--dataset cifar10 \
#--partition_method homo  \
#--client_number 4 \
#--comm_round 10 \
#--epochs 1 \
#--batch_size 64