mode=$1
dataset=$2

if [ "$mode" = "train" ]; then
    if [ "$dataset" = "HCC" ]; then
        echo "Warning: This dataset can be access unless you have IRB approval."
        
        model="RF"
        echo ""
        echo "Train random forest (${model}) on ${dataset}."
        echo ""
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "train" -c "./config/overall_result/HCC/random_forest.json"
        
        model="XGB"
        echo ""
        echo "Train XGBoost (${model}) on ${dataset}."
        echo ""
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "train" -c "./config/overall_result/HCC/xgboost.json"
        
        model="GRU"
        echo ""
        echo "Train GRU (${model}) on ${dataset}."
        echo ""
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "train" -c "./config/overall_result/HCC/gru.json"
        
        model="GRUD"
        echo ""
        echo "Train GRUD (${model}) on ${dataset}."
        echo ""
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "train" -c "./config/overall_result/HCC/grud.json"
        
        model="TE"
        echo ""
        echo "Train original transformer encoder (${model}) on ${dataset}."
        echo ""
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "train" -c "./config/overall_result/HCC/te.json"
        
        model="TranSCANE"
        echo ""
        echo "Train TranSCANE (${model}) on ${dataset}."
        echo ""
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "train" -c "./config/overall_result/HCC/tesne.json"

    elif [ "$dataset" = "MI3" ]; then
        echo "Warning: You have to get the approval before accessing the dataset. More detail in: https://physionet.org/content/mimiciii/"
        
        model="RF"
        echo ""
        echo "Train random forest (${model}) on ${dataset}."
        echo ""
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "train" -c "./config/overall_result/MI3/random_forest.json"
        
        model="XGB"
        echo ""
        echo "Train XGBoost (${model}) on ${dataset}."
        echo ""
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "train" -c "./config/overall_result/MI3/xgboost.json"
        
        model="GRU"
        echo ""
        echo "Train GRU (${model}) on ${dataset}."
        echo ""
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "train" -c "./config/overall_result/MI3/gru.json"
        
        model="GRUD"
        echo ""
        echo "Train GRUD (${model}) on ${dataset}."
        echo ""
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "train" -c "./config/overall_result/MI3/grud.json"
        
        model="TE"
        echo ""
        echo "Train original transformer encoder (${model}) on ${dataset}."
        echo ""
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "train" -c "./config/overall_result/MI3/te.json"
        
        model="TranSCANE"
        echo ""
        echo "Train TranSCANE (${model}) on ${dataset}."
        echo ""
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "train" -c "./config/overall_result/MI3/tesne.json"

    elif [ "$dataset" = "P12" ]; then
        
        model="RF"
        echo ""
        echo "Train random forest (${model}) on ${dataset}."
        echo ""
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "train" -c "./config/overall_result/P12/random_forest.json"
        
        model="XGB"
        echo ""
        echo "Train XGBoost (${model}) on ${dataset}."
        echo ""
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "train" -c "./config/overall_result/P12/xgboost.json"
        
        model="GRU"
        echo ""
        echo "Train GRU (${model}) on ${dataset}."
        echo ""
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "train" -c "./config/overall_result/P12/gru.json"
        
        model="GRUD"
        echo ""
        echo "Train GRUD (${model}) on ${dataset}."
        echo ""
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "train" -c "./config/overall_result/P12/grud.json"
        
        model="TE"
        echo ""
        echo "Train original transformer encoder (${model}) on ${dataset}."
        echo ""
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "train" -c "./config/overall_result/P12/te.json"
        
        model="TranSCANE"
        echo ""
        echo "Train TranSCANE (${model}) on ${dataset}."
        echo ""
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "train" -c "./config/overall_result/P12/tesne.json"
    else
        echo "dataset should be \"P12\" , \"MI3\", or \"HCC\"."
    fi
elif [ "$mode" = "train_test" ]; then
    echo "This mode will test all models after all models is trained."
    if [ "$dataset" = "HCC" ]; then
        echo "Warning: This dataset can be access unless you have IRB approval."

        model="RF"
        now=$(date)
        rf_train_date=$(date +"%Y%m%d")
        echo ""
        echo "Train random forest (${model}) on ${dataset}."
        echo ""
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "train" -c "./config/overall_result/HCC/random_forest.json"

        model="XGB"
        now=$(date)
        xgb_train_date=$(date +"%Y%m%d")
        echo ""
        echo "Train XGBoost (${model}) on ${dataset}."
        echo ""
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "train" -c "./config/overall_result/HCC/xgboost.json"

        model="GRU"
        now=$(date)
        gru_train_date=$(date +"%Y%m%d")
        echo ""
        echo "Train GRU (${model}) on ${dataset}."
        echo ""
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "train" -c "./config/overall_result/HCC/gru.json"

        model="GRUD"
        now=$(date)
        grud_train_date=$(date +"%Y%m%d")
        echo ""
        echo "Train GRUD (${model}) on ${dataset}."
        echo ""
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "train" -c "./config/overall_result/HCC/grud.json"

        model="TE"
        now=$(date)
        te_train_date=$(date +"%Y%m%d")
        echo ""
        echo "Train original transformer encoder (${model}) on ${dataset}."
        echo ""
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "train" -c "./config/overall_result/HCC/te.json"

        model="TranSCANE"
        now=$(date)
        tesne_train_date=$(date +"%Y%m%d")
        echo ""
        echo "Train TranSCANE (${model}) on ${dataset}."
        echo ""
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "train" -c "./config/overall_result/HCC/tesne.json"

        model="RF"
        echo "Test random forest (${model}) on ${dataset}."
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "test" -c "./exp_output/overall_result/${rf_train_date}_${dataset}_${model}/training_result/config.json" -r "./exp_output/overall_result/${rf_train_date}_${dataset}_${model}/training_result/checkpoint/model.txt" -b 1

        model="XGB"
        echo "Test XGBoost (${model}) on ${dataset}."
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "test" -c "./exp_output/overall_result/${xgb_train_date}_${dataset}_${model}/training_result/config.json" -r "./exp_output/overall_result/${xgb_train_date}_${dataset}_${model}/training_result/checkpoint/model.txt" -b 1

        model="GRU"
        echo "Test GRU (${model}) on ${dataset}."
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "test" -c "./exp_output/overall_result/${gru_train_date}_${dataset}_${model}/training_result/config.json" -r "./exp_output/overall_result/${gru_train_date}_${dataset}_${model}/training_result/checkpoint/ckt_ep_30.pth" -b 1

        model="GRUD"
        echo "Test GRUD (${model}) on ${dataset}."
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "test" -c "./exp_output/overall_result/${grud_train_date}_${dataset}_${model}/training_result/config.json" -r "./exp_output/overall_result/${grud_train_date}_${dataset}_${model}/training_result/checkpoint/ckt_ep_75.pyh" -b 1

        model="TE"
        echo "Test original transformer encoder (${model}) on ${dataset}."
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "test" -c "./exp_output/overall_result/${te_train_date}_${dataset}_${model}/training_result/config.json" -r "./exp_output/overall_result/${te_train_date}_${dataset}_${model}/training_result/checkpoint/ckt_ep_85.pth" -b 1

        model="TranSCANE"
        echo "Test TranSCANE (${model}) on ${dataset}."
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "test" -c "./exp_output/overall_result/${tesne_train_date}_${dataset}_${model}/training_result/config.json" -r "./exp_output/overall_result/${tesne_train_date}_${dataset}_${model}/training_result/checkpoint/ckt_ep_100.pth"

    elif [ "$dataset" = "MI3" ]; then
        echo "Warning: You have to get the approval before accessing the dataset. More detail in: https://physionet.org/content/mimiciii/"
        model="RF"
        now=$(date)
        rf_train_date=$(date +"%Y%m%d")
        echo ""
        echo "Train random forest (${model}) on ${dataset}."
        echo ""
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "train" -c "./config/overall_result/MI3/random_forest.json"

        model="XGB"
        now=$(date)
        xgb_train_date=$(date +"%Y%m%d")
        echo ""
        echo "Train XGBoost (${model}) on ${dataset}."
        echo ""
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "train" -c "./config/overall_result/MI3/xgboost.json"

        model="GRU"
        now=$(date)
        gru_train_date=$(date +"%Y%m%d")
        echo ""
        echo "Train GRU (${model}) on ${dataset}."
        echo ""
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "train" -c "./config/overall_result/MI3/gru.json"

        model="GRUD"
        now=$(date)
        grud_train_date=$(date +"%Y%m%d")
        echo ""
        echo "Train GRUD (${model}) on ${dataset}."
        echo ""
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "train" -c "./config/overall_result/MI3/grud.json"

        model="TE"
        now=$(date)
        te_train_date=$(date +"%Y%m%d")
        echo ""
        echo "Train original transformer encoder (${model}) on ${dataset}."
        echo ""
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "train" -c "./config/overall_result/MI3/te.json"

        model="TranSCANE"
        now=$(date)
        tesne_train_date=$(date +"%Y%m%d")
        echo ""
        echo "Train TranSCANE (${model}) on ${dataset}."
        echo ""
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "train" -c "./config/overall_result/MI3/tesne.json"

        model="RF"
        echo "Test random forest (${model}) on ${dataset}."
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "test" -c "./exp_output/overall_result/${rf_train_date}_${dataset}_${model}/training_result/config.json" -r "./exp_output/overall_result/${rf_train_date}_${dataset}_${model}/training_result/checkpoint/model.txt" -b 1

        model="XGB"
        echo "Test XGBoost (${model}) on ${dataset}."
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "test" -c "./exp_output/overall_result/${xgb_train_date}_${dataset}_${model}/training_result/config.json" -r "./exp_output/overall_result/${xgb_train_date}_${dataset}_${model}/training_result/checkpoint/model.txt" -b 1

        model="GRU"
        echo "Test GRU (${model}) on ${dataset}."
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "test" -c "./exp_output/overall_result/${gru_train_date}_${dataset}_${model}/training_result/config.json" -r "./exp_output/overall_result/${gru_train_date}_${dataset}_${model}/training_result/checkpoint/ckt_ep_280.pth" -b 1

        model="GRUD"
        echo "Test GRUD (${model}) on ${dataset}."
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "test" -c "./exp_output/overall_result/${grud_train_date}_${dataset}_${model}/training_result/config.json" -r "./exp_output/overall_result/${grud_train_date}_${dataset}_${model}/training_result/checkpoint/ckt_ep_400.pth" -b 1

        model="TE"
        echo "Test original transformer encoder (${model}) on ${dataset}."
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "test" -c "./exp_output/overall_result/${te_train_date}_${dataset}_${model}/training_result/config.json" -r "./exp_output/overall_result/${te_train_date}_${dataset}_${model}/training_result/checkpoint/ckt_ep_125.pth" -b 1

        model="TranSCANE"
        echo "Test TranSCANE (${model}) on ${dataset}."
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "test" -c "./exp_output/overall_result/${tesne_train_date}_${dataset}_${model}/training_result/config.json" -r "./exp_output/overall_result/${tesne_train_date}_${dataset}_${model}/training_result/checkpoint/ckt_ep_380.pth" -b 1

    elif [ "$dataset" = "P12" ]; then
        model="RF"
        now=$(date)
        rf_train_date=$(date +"%Y%m%d")
        echo ""
        echo "Train random forest (${model}) on ${dataset}."
        echo ""
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "train" -c "./config/overall_result/P12/random_forest.json"

        model="XGB"
        now=$(date)
        xgb_train_date=$(date +"%Y%m%d")
        echo ""
        echo "Train XGBoost (${model}) on ${dataset}."
        echo ""
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "train" -c "./config/overall_result/P12/xgboost.json"

        model="GRU"
        now=$(date)
        gru_train_date=$(date +"%Y%m%d")
        echo ""
        echo "Train GRU (${model}) on ${dataset}."
        echo ""
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "train" -c "./config/overall_result/P12/gru.json"

        model="GRUD"
        now=$(date)
        grud_train_date=$(date +"%Y%m%d")
        echo ""
        echo "Train GRUD (${model}) on ${dataset}."
        echo ""
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "train" -c "./config/overall_result/P12/grud.json"

        model="TE"
        now=$(date)
        te_train_date=$(date +"%Y%m%d")
        echo ""
        echo "Train original transformer encoder (${model}) on ${dataset}."
        echo ""
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "train" -c "./config/overall_result/P12/te.json"

        model="TranSCANE"
        now=$(date)
        tesne_train_date=$(date +"%Y%m%d")
        echo ""
        echo "Train TranSCANE (${model}) on ${dataset}."
        echo ""
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "train" -c "./config/overall_result/P12/tesne.json"

        model="RF"
        echo "Test random forest (${model}) on ${dataset}."
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "test" -c "./exp_output/overall_result/${rf_train_date}_${dataset}_${model}/training_result/config.json" -r "./exp_output/overall_result/${rf_train_date}_${dataset}_${model}/training_result/checkpoint/model.txt" -b 1

        model="XGB"
        echo "Test XGBoost (${model}) on ${dataset}."
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "test" -c "./exp_output/overall_result/${xgb_train_date}_${dataset}_${model}/training_result/config.json" -r "./exp_output/overall_result/${xgb_train_date}_${dataset}_${model}/training_result/checkpoint/model.txt" -b 1

        model="GRU"
        echo "Test GRU (${model}) on ${dataset}."
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "test" -c "./exp_output/overall_result/${gru_train_date}_${dataset}_${model}/training_result/config.json" -r "./exp_output/overall_result/${gru_train_date}_${dataset}_${model}/training_result/checkpoint/ckt_ep_100.pth" -b 1

        model="GRUD"
        echo "Test GRUD (${model}) on ${dataset}."
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "test" -c "./exp_output/overall_result/${grud_train_date}_${dataset}_${model}/training_result/config.json" -r "./exp_output/overall_result/${grud_train_date}_${dataset}_${model}/training_result/checkpoint/ckt_ep_480.pth" -b 1

        model="TE"
        echo "Test original transformer encoder (${model}) on ${dataset}."
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "test" -c "./exp_output/overall_result/${te_train_date}_${dataset}_${model}/training_result/config.json" -r "./exp_output/overall_result/${te_train_date}_${dataset}_${model}/training_result/checkpoint/ckt_ep_75.pth" -b 1

        model="TranSCANE"
        echo "Test TranSCANE (${model}) on ${dataset}."
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "test" -c "./exp_output/overall_result/${tesne_train_date}_${dataset}_${model}/training_result/config.json" -r "./exp_output/overall_result/${tesne_train_date}_${dataset}_${model}/training_result/checkpoint/ckt_ep_350.pth" -b 1

    else
        echo "dataset should be \"P12\" , \"MI3\", or \"HCC\"."
    fi
elif [ "$mode" = "test" ]; then
    echo "This mode will use the provided checkpoint (trained by the author) to test on the dataset."
    if [ "$dataset" = "HCC" ]; then
        echo "Warning: This dataset can be access unless you have IRB approval."
        model="RF"
        echo "Test random forest (${model}) on ${dataset}."
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "test" -c "./exp_output/model_ckt_HCC/rf/config.json" -r "./exp_output/model_ckt_HCC/rf/ckt.pth" -b 1

        model="XGB"
        echo "Test XGBoost (${model}) on ${dataset}."
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "test" -c "./exp_output/model_ckt_HCC/xgb/config.json" -r "./exp_output/model_ckt_HCC/xgb/ckt.pth" -b 1

        model="GRU"
        echo "Test GRU (${model}) on ${dataset}."
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "test" -c "./exp_output/model_ckt_HCC/gru/config.json" -r "./exp_output/model_ckt_HCC/gru/ckt.pth" -b 1

        model="GRUD"
        echo "Test GRUD (${model}) on ${dataset}."
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "test" -c "./exp_output/model_ckt_HCC/grud/config.json" -r "./exp_output/model_ckt_HCC/grud/ckt.pth" -b 1

        model="TE"
        echo "Test original transformer encoder (${model}) on ${dataset}."
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "test" -c "./exp_output/model_ckt_HCC/te/config.json" -r "./exp_output/model_ckt_HCC/te/ckt.pth" -b 1

        model="TranSCANE"
        echo "Test TranSCANE (${model}) on ${dataset}."
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "test" -c "./exp_output/model_ckt_HCC/tesne/config.json" -r "./exp_output/model_ckt_HCC/tesne/ckt.pth" -b 1

    elif [ "$dataset" = "MI3" ]; then
        echo "Warning: You have to get the approval before accessing the dataset. More detail in: https://physionet.org/content/mimiciii/"
        model="RF"
        echo "Test random forest (${model}) on ${dataset}."
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "test" -c "./exp_output/model_ckt_MI3/rf/config.json" -r "./exp_output/model_ckt_MI3/rf/ckt.pth" -b 1

        model="XGB"
        echo "Test XGBoost (${model}) on ${dataset}."
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "test" -c "./exp_output/model_ckt_MI3/xgb/config.json" -r "./exp_output/model_ckt_MI3/xgb/ckt.pth" -b 1

        model="GRU"
        echo "Test GRU (${model}) on ${dataset}."
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "test" -c "./exp_output/model_ckt_MI3/gru/config.json" -r "./exp_output/model_ckt_MI3/gru/ckt.pth" -b 1

        model="GRUD"
        echo "Test GRUD (${model}) on ${dataset}."
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "test" -c "./exp_output/model_ckt_MI3/grud/config.json" -r "./exp_output/model_ckt_MI3/grud/ckt.pth" -b 1

        model="TE"
        echo "Test original transformer encoder (${model}) on ${dataset}."
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "test" -c "./exp_output/model_ckt_MI3/te/config.json" -r "./exp_output/model_ckt_MI3/te/ckt.pth" -b 1

        model="TranSCANE"
        echo "Test TranSCANE (${model}) on ${dataset}."
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "test" -c "./exp_output/model_ckt_MI3/tesne/config.json" -r "./exp_output/model_ckt_MI3/tesne/ckt.pth" -b 1

    elif [ "$dataset" = "P12" ]; then
        model="RF"
        echo "Test random forest (${model}) on ${dataset}."
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "test" -c "./exp_output/model_ckt_P12/rf/config.json" -r "./exp_output/model_ckt_P12/rf/ckt.pth" -b 1

        model="XGB"
        echo "Test XGBoost (${model}) on ${dataset}."
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "test" -c "./exp_output/model_ckt_P12/xgb/config.json" -r "./exp_output/model_ckt_P12/xgb/ckt.pth" -b 1

        model="GRU"
        echo "Test GRU (${model}) on ${dataset}."
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "test" -c "./exp_output/model_ckt_P12/gru/config.json" -r "./exp_output/model_ckt_P12/gru/ckt.pth" -b 1

        model="GRUD"
        echo "Test GRUD (${model}) on ${dataset}."
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "test" -c "./exp_output/model_ckt_P12/grud/config.json" -r "./exp_output/model_ckt_P12/grud/ckt.pth" -b 1

        model="TE"
        echo "Test original transformer encoder (${model}) on ${dataset}."
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "test" -c "./exp_output/model_ckt_P12/te/config.json" -r "./exp_output/model_ckt_P12/te/ckt.pth" -b 1

        model="TranSCANE"
        echo "Test TranSCANE (${model}) on ${dataset}."
        python3 main.py -e "overall_result" -a "${dataset}_${model}" -m "test" -c "./exp_output/model_ckt_P12/tesne/config.json" -r "./exp_output/model_ckt_P12/tesne/ckt.pth" -b 1

    else
        echo "dataset should be \"P12\" , \"MI3\", or \"HCC\"."
    fi
else
    echo "mode shoud be \"train\", \"train_test\", or \"test\""
fi
