#!/bin/bash

# Config file location
cfg_file=

effect1="denoiser"
effect2="superres"
gpu="t4"
ulimit_val=0
sample_rate1=16
sample_rate2=16
output_sample_rate1=16
output_sample_rate2=48
batch_size=1
reset=
multi_input=""
frame_size=10;
input_file=

args=()

declare -A gpu_map=( ["a100"]="sm_80"
                     ["a30"]="sm_80"
                     ["a2"]="sm_86"
                     ["a10"]="sm_86"
                     ["a16"]="sm_86"
                     ["a40"]="sm_86"
                     ["t4"]="sm_75"
                     ["v100"]="sm_70"
                    )
declare -a supported_effects=( "denoiser_16k_superres_16k"
                               "dereverb_16k_superres_16k"
                               "dereverb_denoiser_16k_superres_16k"
                               "superres_8k_denoiser_16k"
                               "superres_8k_dereverb_16k"
                               "superres_8k_dereverb_denoiser_16k"
                               "superres_8k_dereverb_denoiser_16k"
                             )

gpu_list=$(IFS=/; printf "%s" "${!gpu_map[*]}")

while [[ "$#" -gt 0 ]]; do
    args+=($1)
    case $1 in
        -e1|--effect1) effect1="$2"; args+=($2); shift ;;
        -e2|--effect2) effect2="$2"; args+=($2); shift ;;
        -g|--gpu) gpu="$2"; args+=($2); shift ;;
        -u|--ulimit) ulimit_val="$2"; args+=($2); shift ;;
        -s1|--sample_rate) sample_rate1="$2"; args+=($2); shift ;;
        -s2|--sample_rate2) sample_rate2="$2"; args+=($2); shift ;;
        -o1|--output_sample_rate1) output_sample_rate1="$2"; args+=($2); shift ;;
        -o2|--output_sample_rate2) output_sample_rate2="$2"; args+=($2); shift ;;
        -b|--batch_size) batch_size="$2"; args+=($2); shift ;;
        -c|--cfg-file) cfg_file="$2"; args+=($2); shift;;
        -r|--reset) multi_input="0,1,2,3,4" reset="reset 2 3 5";;
        -f|--frame_size) frame_size="$2"; args+=($2); shift;;
        -i|--input_file) input_file="$2"; unset 'args[${#args[@]}-1]'; shift;;
        -h|--help) echo -e "-g, --gpu : $gpu_list [default=t4]"
                   echo -e "-e1, --effect1 : First effect to use: denoiser/dereverb/dereverb_denoiser/superres [default=denoiser]"
                   echo -e "-e2, --effect2 : Second effect to use: denoiser/dereverb/dereverb_denoiser/superres [default=denoiser]"
                   echo -e "-u, --ulimit : raise files open ulimit, this is used when running larger batch sizes that exceed the kernel maximum open files limit. Warning: Use this option with caution. [default=0(unusued)]"
                   echo -e "-s1, --sample_rate1 : sample rate to use 8/16/48 (in kHz) for first effect. [default=16]"
                   echo -e "-s2, --sample_rate2 : sample rate to use 8/16/48 (in kHz) for second effect. [default=48]"
                   echo -e "-o1, --output_sample_rate1 : output sample rate to use 16/48 (in kHz) for first effect. [default=48]"
                   echo -e "-o2, --output_sample_rate2 : output sample rate to use 16/48 (in kHz) for second effect. [default=48]"
                   echo -e "-b, --batch_size : batch size to use. [default=1, max=1024]"
                   echo -e "-c, --cfg-file: configuration file to write to"
                   echo -e "-f, --frame_size: frame size to be used 10/20ms. [default=10]"
                   echo -e "-i, --input_file: file to be used as input, instead of using sample input files (file must be in correct format)"
                   echo -e "-h, --help: print this help menu";exit;;
        *) echo "Unknown parameter passed: $1"; exit 1 ;;
    esac
    shift
done

effect_check="\b${effect1}_${sample_rate1}k_${effect2}_${sample_rate2}k\b"
if ! [[ "${supported_effects[@]}" =~ ${effect_check} ]]; then
  echo "Chain (${effect1} ${sample_rate1}k + ${effect2} ${sample_rate2}k) is not supported. Supported chains are:"
  for e in "${supported_effects[@]}"; do
    v=${e//k_/k + }
    echo "- ${v//_/ }"
  done
  exit 1
fi

OUT_DIR="chained/${effect1}${sample_rate1}k_${effect2}${sample_rate2}k"
mkdir -p "$OUT_DIR"
input_list=""
out_file=""
if [ ! -z "$input_file" ]; then
  if [[ ! (-f "$input_file" || -d "$input_file") ]]; then
    echo "Error: File/folder '$input_file' does not exist"
    exit 1
  fi
  if [ $(($batch_size)) -ne 1 ]; then
    echo "Error: Batch size must be 1 when input file/folder is specified (got $batch_size)"
    exit 1
  fi
  if [ -d "$input_file" ]; then
    if [ ! -z "$cfg_file" ]; then
      echo "Error: config file cannot be specified if input is a folder"
      exit 1
    fi
    # If folder, process file by file for all files in that folder
    for i in "${input_file}"/*.wav; do
      ./$0 "${args[@]}" -i "$i"
      if [ $? -ne 0 ]; then
        echo "Error processing file '$i'"
        exit 1
      fi
    done
    exit 0
  else
    input_list="input_wav_list $input_file"
    out_file="$OUT_DIR/$(basename $input_file)"
  fi
else
    folder="${effect1}_${sample_rate1}k_${effect2}_${sample_rate2}k"
    wav=$(./list.sh ../input_files/chaining/${folder}/ ${batch_size} ${multi_input})
    if [ $? -ne 0 ]; then
       echo "Error generating input files list"
       echo "Command returned: $wav"
       exit 1
    fi
    input_list="input_wav_list $wav"
    for i in $(seq 1 $batch_size); do
        out_file="$out_file ${OUT_DIR}/${effect1}${sample_rate1}k_${effect2}${sample_rate2}k_$i.wav"
    done
fi

if [ -z "$cfg_file" ]; then
  cfg_file="/tmp/tmp_cfg.txt"
fi

if [ -z "${gpu_map[$gpu]:-}" ]; then
  echo "Invalid GPU (got $gpu, supported $gpu_list)"
  exit 1
fi
arch="${gpu_map[$gpu]}"

model1="../../models"
model2="../../models"
if [ "$effect1" == "superres" ]; then
  model1="${model1}/${arch}/${effect1}_${sample_rate1}k_to_${output_sample_rate1}k.trtpkg"
else
  model1="${model1}/${arch}/${effect1}_${sample_rate1}k.trtpkg"
fi

if [ "$effect2" == "superres" ]; then
  model2="${model2}/${arch}/${effect2}_${sample_rate2}k_to_${output_sample_rate2}k.trtpkg"
else
  model2="${model2}/${arch}/${effect2}_${sample_rate2}k.trtpkg"
fi

echo "effect ${effect1} ${effect2}
sample_rate ${sample_rate1}000 ${sample_rate2}000
model ${model1} ${model2}
real_time 0
$input_list
output_wav_list $out_file
frame_size ${frame_size}" > $cfg_file

if [ ! -z "$reset" ]; then
  if [ $batch_size -lt 5 ]; then
    echo "Batch size must be at least 5 for reset configs"
    exit 1
  fi
  echo "$reset" >> $cfg_file
fi

echo "Using the config file $cfg_file"

if [ "$ulimit_val" -ne "0" ]; then
   if [ "$ulimit_val" -le "1024" ]; then
      echo "You probably dont want to do this. ulimit is too low.";
      exit;
   fi
   echo "ulimit set to ${ulimit_val}";
   ulimit -n $ulimit_val
   exit;
fi

./effects_demo -c $cfg_file