# Copyright (c) Microsoft Corporation.
# The file is modified based on the original Graphormer's source code.
# Copyright (c) 2022 Tianyu Wen
# Licensed under the MIT License.

[ -z "${seed}" ] && seed="1"
[ -z "${arch}" ] && arch="--ffn_dim 512 --hidden_dim 512 --intput_dropout_rate 0.0 --attention_dropout_rate 0.0 --dropout_rate 0.0 --weight_decay 0.0 --num_heads 32 --n_layers 20 --edge_type multi_hop --multi_hop_max_dist 5"
[ -z "${batch_size}" ] && batch_size="128"
[ -z "${epoch}" ] && epoch="1000"
[ -z "${peak_lr}" ] && peak_lr="3e-4"
[ -z "${end_lr}" ] && end_lr="1e-9"

echo -e "\n\n"
echo "=====================================ARGS======================================"
echo "arg0: $0"
echo "arch: ${arch}"
echo "batch_size: ${batch_size}"
echo "peak_lr ${peak_lr}"
echo "end_lr ${end_lr}"
echo "seed: ${seed}"
echo "epoch: ${epoch}"
echo "==============================================================================="

n_gpu=1
tot_updates=$((330000*epoch/batch_size/n_gpu))
warmup_updates=$((tot_updates/10))
max_epochs=$((epoch+1))
echo "=====================================ARGS======================================"
echo "tot_updates ${tot_updates}"
echo "warmup_updates: ${warmup_updates}"
echo "max_epochs: ${max_epochs}"
echo "==============================================================================="

default_root_dir=./result
mkdir -p $default_root_dir

python ./qm9_entry.py --num_workers 4 --seed $seed --batch_size $batch_size \
      --dataset_name QM9 \
      --gpus $n_gpu --accelerator ddp --precision 16 $arch \
      --default_root_dir $default_root_dir \
      --tot_updates $tot_updates --warmup_updates $warmup_updates --max_epochs $max_epochs \
      --peak_lr $peak_lr --end_lr $end_lr --progress_bar_refresh_rate 10
