#!/usr/bin/env bash

: <<'END'


sh run.sh --stage 2 --stop_stage 2 --system_version centos --file_folder_name file_dir --final_model_name nx-mpnet-aishell \
--noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise" \
--speech_dir "/data/tianxing/HuggingDatasets/aishell/data_aishell/wav/train" \
--max_epochs 100 \
--duration 2 \


END


# params
system_version="windows";
verbose=true;
stage=0 # start from 0 if you need to start from data preparation
stop_stage=9

work_dir="$(pwd)"
file_folder_name=file_folder_name
final_model_name=final_model_name
config_file="yaml/config.yaml"
limit=10

noise_dir=/data/tianxing/HuggingDatasets/nx_noise/data/noise
speech_dir=/data/tianxing/HuggingDatasets/aishell/data_aishell/wav/train
duration=2

nohup_name=nohup.out

# model params
batch_size=64
max_epochs=200
save_top_k=10
patience=5


# parse options
while true; do
  [ -z "${1:-}" ] && break;  # break if there are no arguments
  case "$1" in
    --*) name=$(echo "$1" | sed s/^--// | sed s/-/_/g);
      eval '[ -z "${'"$name"'+xxx}" ]' && echo "$0: invalid option $1" 1>&2 && exit 1;
      old_value="(eval echo \\$$name)";
      if [ "${old_value}" == "true" ] || [ "${old_value}" == "false" ]; then
        was_bool=true;
      else
        was_bool=false;
      fi

      # Set the variable to the right value-- the escaped quotes make it work if
      # the option had spaces, like --cmd "queue.pl -sync y"
      eval "${name}=\"$2\"";

      # Check that Boolean-valued arguments are really Boolean.
      if $was_bool && [[ "$2" != "true" && "$2" != "false" ]]; then
        echo "$0: expected \"true\" or \"false\": $1 $2" 1>&2
        exit 1;
      fi
      shift 2;
      ;;

    *) break;
  esac
done

file_dir="${work_dir}/${file_folder_name}"
final_model_dir="${work_dir}/../../trained_models/${final_model_name}";
evaluation_audio_dir="${file_dir}/evaluation_audio"

dataset="${file_dir}/dataset.xlsx"
train_dataset="${file_dir}/train.xlsx"
valid_dataset="${file_dir}/valid.xlsx"

$verbose && echo "system_version: ${system_version}"
$verbose && echo "file_folder_name: ${file_folder_name}"

if [ $system_version == "windows" ]; then
  alias python3='D:/Users/tianx/PycharmProjects/virtualenv/nx_denoise/Scripts/python.exe'
elif [ $system_version == "centos" ] || [ $system_version == "ubuntu" ]; then
  #source /data/local/bin/nx_denoise/bin/activate
  alias python3='/data/local/bin/nx_denoise/bin/python3'
fi


if [ ${stage} -le 1 ] && [ ${stop_stage} -ge 1 ]; then
  $verbose && echo "stage 1: prepare data"
  cd "${work_dir}" || exit 1
  python3 step_1_prepare_data.py \
  --file_dir "${file_dir}" \
  --noise_dir "${noise_dir}" \
  --speech_dir "${speech_dir}" \
  --train_dataset "${train_dataset}" \
  --valid_dataset "${valid_dataset}" \
  --duration "${duration}" \

fi


if [ ${stage} -le 2 ] && [ ${stop_stage} -ge 2 ]; then
  $verbose && echo "stage 2: train model"
  cd "${work_dir}" || exit 1
  python3 step_2_train_model.py \
  --train_dataset "${train_dataset}" \
  --valid_dataset "${valid_dataset}" \
  --serialization_dir "${file_dir}" \
  --config_file "${config_file}" \

fi


if [ ${stage} -le 3 ] && [ ${stop_stage} -ge 3 ]; then
  $verbose && echo "stage 3: test model"
  cd "${work_dir}" || exit 1
  python3 step_3_evaluation.py \
  --valid_dataset "${valid_dataset}" \
  --model_dir "${file_dir}/best" \
  --evaluation_audio_dir "${evaluation_audio_dir}" \
  --limit "${limit}" \

fi


if [ ${stage} -le 4 ] && [ ${stop_stage} -ge 4 ]; then
  $verbose && echo "stage 4: collect files"
  cd "${work_dir}" || exit 1

  mkdir -p ${final_model_dir}

  cp "${file_dir}/best"/* "${final_model_dir}"
  cp -r "${file_dir}/evaluation_audio" "${final_model_dir}"

  cd "${final_model_dir}/.." || exit 1;

  if [ -e "${final_model_name}.zip" ]; then
    rm -rf "${final_model_name}_backup.zip"
    mv "${final_model_name}.zip" "${final_model_name}_backup.zip"
  fi

  zip -r "${final_model_name}.zip" "${final_model_name}"
  rm -rf "${final_model_name}"

fi


if [ ${stage} -le 5 ] && [ ${stop_stage} -ge 5 ]; then
  $verbose && echo "stage 5: clear file_dir"
  cd "${work_dir}" || exit 1

  rm -rf "${file_dir}";

fi