#!/bin/bash

## [Goal] Run AQUAMIS in batches; Wrapper for aquamis.py, direct Snakemake calls or Docker container
## [Author] Holger Brendebach

## Shell Behaviour --------------------------------------------------
set -Euo pipefail # to exit immediately, add e
[[ -z ${email+x} ]] || trap error_handler SIGINT SIGTERM ERR

## (1) Argument Logic

error_handler() {
  trap - SIGINT SIGTERM ERR
  ## send message on failure
  cat << MSG | mail -s "[${script_name:-}] Script Error" ${email:-${EMAIL:-NGSadmin@bfr.bund.de}}  # ${email:-${EMAIL:-NGSadmin@bfr.bund.de}}
Date:       $(date +"%F %T")
Script:     ${script_real:-}
Arguments:  ${*:-${args:-"No arguments provided"}}
User/Host:  $(whoami)@$(hostname)
Error:      A SIGINT/SIGTERM/ERR signal was caught.

Please review the logs:
* ${logfile:-}
* ${logfile_global:-}
MSG
  exit 1
}

usage() {
  cat <<- USE

${bold}Basic usage:${normal}
bash $script_name

${bold}Description:${normal}
Run AQUAMIS in batch mode based on entries in the aquamis_runs.tsv table.
Any script arguments that do not match the ones in the next section will be appended to the run command, e.g. "--unlock".
This wrapper may evaluate several run strategies in its run_modules() function: aquamis.py, direct Snakemake calls or Docker container runs.

${bold}Optional Named Arguments:${normal}
-p, --profile                   Set a Snakemake profile, e.g. workstation (default=bfr.hpc)
-o, --onerun                   Deactivate batch mode and use hardcoded single run parameters (default=false)
-s, --snakemake                 Skip the aquamis.py wrapper and use a customized Snakemake command
--dryrun                        Perform a dryrun and include a Snakemake dryrun for DAG calculation (default=false)
-n, --norun                     Perform a dryrun without Snakemake commands (default=false)
-a, --auto                      Do not ask before starting the program (default=false), e.g. for tmux-wrapped script calls

${bold}Notes:${normal}
Check the manual_override() function to tweak script arguments for repetitive use.
Check the tweak_settings() function to review potential system-specific expert settings.
If you like script fail notifications by email, set the environment variable "email" to activate the error_handler() function.

${bold}Example:${normal}
bash $script_name -p bfr.hpc -a -n

USE
  exit 0
}

parse_args() {
  local args=("$@")

  ## default values of variables
  repo_path=$(dirname $script_path)
  profile="bfr.hpc"
  conda_recipe="$repo_path/envs/aquamis.yaml"
  dag_search="$repo_path/profiles/smk_directives.yaml"
  docker_image="bfrbioinformatics/aquamis:1.3.12"

  ## default values of switches
  batch_mode=true
  snakemake=false
  norun=false
  dryrun=false
  interactive=true

  while :; do
    case "${1-}" in
      -h | --help) usage ;;
      --dev) dev_notes "$script_real" ;; # from helper_functions.sh
      -v | --verbose)
        set -x
        shift
        ;;
      -p | --profile)
        profile="$2"
        shift 2
        ;;
      -o | --onerun)
        batch_mode=false
        shift
        ;;
      -s | --snakemake)
        snakemake=true
        shift
        ;;
      -n | --norun)
        norun=true
        dryrun=true
        shift
        ;;
      --dryrun)
        dryrun=true
        user_args+="--dryrun "
        shift
        ;;
      -a | --auto)
        interactive=false
        shift
        ;;
      -?*)
        user_args+="${1-} "
        shift
        ;;
      *) break ;;
    esac
  done

  ## checks
  [[ -f $conda_recipe ]] || die "The conda recipe does not exist: $conda_recipe"

  [[ $dryrun == true ]] && echo -e "\n${bold}${green}-->  This is a dryrun with no log  <--${normal}\n"
  [[ $dryrun == false ]] && echo -e "\n${bold}${magenta}-->  This is NOT a dryrun!  <--${normal}\n"

  return 0
}

## Source Helper Functions --------------------------------------------------
script_real=$(realpath -P "${BASH_SOURCE[0]}")
script_path=$(dirname $script_real)
script_name=$(basename $script_real)
[[ ! -f "$script_path/helper_functions.sh" ]] && echo "[ERROR] File not found: $script_path/helper_functions.sh" && kill -s TERM $$
source "$script_path/helper_functions.sh"

## Parse Arguments --------------------------------------------------
args="$*"
parse_args "$@"

## (2) Script Logic

manual_override() {
  ## (1) Hardcoded Parameters that will override defaults, comment-in to activate

  ## Code Base
#  repo_path="/home/$USER/anaconda3/envs/aquamis_mamba/opt/aquamis"  # BioConda version
#  repo_path="/cephfs/abteilung4/NGS/Pipelines/aquamis_akkreditierung"  # tagged version: Accreditation

  ## Snakemake Profiles (see config.yaml within)
#  profile="denbi.user"
#  profile="denbi.large"
#  profile="nris.saga.slurm.dev"
#  profile="nris.saga.slurm.normal"
#  profile="bfr.hpc"
#  profile="bfr.hpc.greedy"
#  profile="workstation"

  ## DAG Job Prioritization
#  dag_search=$repo_path/profiles/smk_directives_depthfirst.yaml
#  dag_search=$repo_path/profiles/smk_directives_breadthfirst.yaml
#  dag_search=$repo_path/profiles/smk_directives_breadthfirst_groups.yaml
#  dag_search=$repo_path/profiles/smk_directives_depthfirst_groups.yaml

  ## Docker Images
#  docker_image="bfrbioinformatics/aquamis:latest"
#  docker_image="bfrbioinformatics/aquamis:1.3.7"

  true
}

## Define Paths --------------------------------------------------

define_paths() {
  ## [Rationale] Manual Run Parameters Definition; 

  ## Conda Environment
  conda_recipe_env_name=$(head -n1 $conda_recipe | cut -d' ' -f2)
  #[[ "$profile" == "workstation" ]] && conda_recipe_env_name="aquamis_mamba"  # NOTE: fix for non-standard env names

  ## Resources
  [[ "$profile" == "bfr.hpc.greedy"         ]] && threads_sample=6
  [[ "$profile" == "bfr.hpc"                ]] && threads_sample=4
  [[ "$profile" == "nris.saga.slurm.normal" ]] && threads_sample=2
  [[ "$profile" == "nris.saga.slurm.dev"    ]] && threads_sample=2
  [[ "$profile" == "denbi.large"            ]] && threads_sample=2
  [[ "$profile" == "denbi.user"             ]] && threads_sample=1
  [[ "$profile" == "workstation"            ]] && threads_sample=1

  ## Logging
  logfile_name=${script_name/.sh/.log}
  [[ $dryrun == false ]] && logfile="$script_path/$logfile_name"
  [[ $dryrun == false ]] && logfile_global=/dev/null
  
  ## Run variables; if in batch mode {run_name,dir_samples,dir_data} will be read from run_table
  run_table="$script_path/aquamis_runs.tsv"
  run_name="test_data_$(hostname)"
  dir_samples="$repo_path/test_data"
  dir_data="$repo_path/test_data/analysis"
}

set_run_paths() {
  [[ -d $dir_data                     ]] || logexec mkdir -p $dir_data
  [[ -z ${run_name+x}                 ]] && run_name=$(date +"%Y-%m-%d %T")
  [[ -f "$dir_samples/samples.tsv"    ]] && sample_list=$dir_samples/samples.tsv
  [[ -f "$dir_data/samples_fixed.tsv" ]] && sample_list=$dir_data/samples_fixed.tsv
  logexec cd $dir_data || true
}

## Info --------------------------------------------------

show_info() {
  logheader "General Settings:"
  echo "
Repository:|$repo_path
Conda Recipe:|$conda_recipe
Threads per Sample:|$threads_sample
Profile:|$profile
DAG Priority:|$dag_search
Snakemake User Args:|${user_args:-}
" | column -t -s="|" | tee -a $logfile $logfile_global
  echo
}

show_info_on_run() {
  ## print current run infos
  cat <<- TXT | tee -a $logfile $logfile_global

${yellow}#################################################################### Processing:${normal}

Run:        ${inverted}$run_name${normal}
SampleList: ${bold}${sample_list:-not yet created}${normal}
WorkDir:    ${bold}$dir_data${normal}

${blue}To monitor the Snakemake log use:${normal}
tail -F -s3 \$(find $dir_data/.snakemake/log -name "*.snakemake.log" -print0 | xargs -r -0 ls -1 -t | head -1)
TXT
}

show_info_on_tmux() {
  ## Prerequisite for tmux wrapping: add `set-option -g default-shell "/bin/bash"` to ~/.tmux.conf
  cat <<- TXT | tee -a $logfile $logfile_global

${blue}Best experience:${normal}
tmux new -d -s aquamis_\$USER
tmux send-keys -t aquamis_\${USER}.0 'bash $script_real --auto' ENTER

All steps were logged to: $logfile and $logfile_global (global)
TXT
}

## Checks --------------------------------------------------

run_checks() {
  if [[ $batch_mode == true ]]; then
    [[ -f $run_table ]] || die "The run table does not exist: $run_table"
  fi

  # TODO: add more for novices
}

## Modules --------------------------------------------------

tweak_settings() {
  ## more privacy towards NCBI
  declare -x BLAST_USAGE_REPORT=false # see https://github.com/tseemann/mlst/issues/115

  ## append environment libraries to LD path, if dependencies are not found
  [[ $profile =~ "saga" ]] && declare -x LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$conda_env_path/lib/ && logentry "Conda libraries appended to LD_LIBRARY_PATH"
  # NOTE: this may cause a "/bin/bash: <conda-prefix>/envs/aquamis/lib/libtinfo.so.6: no version information available (required by /bin/bash)"
  # NOTE: checkout Issues #13 and #15 at https://gitlab.com/bfr_bioinformatics/AQUAMIS/-/issues/?state=all

  ## use krakenDB in ramdisk
  [[ -f /mnt/ramdisk/kraken/hash.k2d ]] && krakenDB="--kraken2db /mnt/ramdisk/kraken" && echo "Using RamDisk copy of kraken2db"
  [[ -f /mnt/ramdisk/kraken/hash.k2d ]] || krakenDB=""

  ## keep referenceDB in sync with server
#  sync_referenceDB
}

activate_conda() {
  ## activate environment
  if [[ "$conda_env_active" != "$conda_env_name" ]]; then
    first_conda_env=$(head -n1 <<< $conda_env_path)  # use first entry if multiple envs with the same name exist, e.g. from other users
    echo "Conda environment $conda_env_name has not been activated yet. Activating Conda environment \"$first_conda_env\""
    set -u && source $conda_base/etc/profile.d/conda.sh && set +u
    conda activate $first_conda_env
  else
    echo "Conda environment $conda_env_name was already activated."
  fi

  ## check
  logentry "Using $(conda info | grep "active env location" | sed -e "s/^\s*//")"
  echo
}

sync_referenceDB() {
  rsync -avuP /home/DenekeC/Snakefiles/aquamis/reference_db/confindr/ $repo_path/reference_db/confindr/
}

delete_previousRun() {
  echo "The script function 'delete_previousRun()' is active. Are you sure?"
  select yn in "Yes, I want to delete previous AQUAMIS calculations!" "No, I fear consequences."; do
    case $yn in
      "Yes, I want to delete previous AQUAMIS calculations!" )
        echo "Starting Cleanup..."
#        [ -d "$dir_data/json_bak"     ] && mv $dir_data/json_bak $dir_data/json_old
#        [ -d "$dir_data/json"         ] && mv $dir_data/json $dir_data/json_9dab9ee
#        [ -d "$dir_data/json"         ] && rm -f -R $dir_data/json
#        [ -d "$dir_data/reports_bak"  ] && mv $dir_data/reports_bak $dir_data/reports_old
#        [ -d "$dir_data/reports"      ] && mv $dir_data/reports $dir_data/reports_9dab9ee
#        [ -d "$dir_data/reports"      ] && rm -f -R $dir_data/reports
#        [ -d "$dir_data/Assembly_bak" ] && mv $dir_data/Assembly_bak $dir_data/Assembly_old
#        [ -d "$dir_data/Assembly"     ] && mv  $dir_data/Assembly $dir_data/Assembly_bak
#        [ -d "$dir_data/Assembly"     ] && rm -f -R $dir_data/Assembly
#        [ -d "$dir_data/.snakemake"   ] && rm -f -R $dir_data/.snakemake
#        [ -d "$dir_data/confindr"     ] && rm -f -R $dir_data/confindr
#        [ -d "$dir_data/kraken2"      ] && rm -f -R $dir_data/kraken2
#        [ -d "$dir_data/logs"         ] && rm -f -R $dir_data/logs
#        [ -d "$dir_data/trimmed"      ] && rm -f -R $dir_data/trimmed
#        [ -f "$dir_data/config.yaml"         ] && mv $dir_data/config.yaml $dir_data/config_9dab9ee.yaml
#        [ -f "$dir_data/config.yaml"         ] && rm -f $dir_data/config.yaml
#        [ -f "$dir_data/pipeline_status.txt" ] && rm -f $dir_data/pipeline_status.txt
#        [ -f "$dir_data/Assembly/fails.txt"  ] && rm -f $dir_data/Assembly/fails.txt
#        [ -f "$dir_data/samples_fixed.tsv"   ] && rm -f $dir_data/samples_fixed.tsv
        echo "Cleanup Finished."
        break
        ;;
      "No, I fear consequences." )
        echo "Cleanup Aborted."
        break
        ;;
    esac
  done < /dev/tty
}

run_aquamis() {
  logheader "Running AQUAMIS via aquamis.py:"
  read -r -d '' wrapper_cmd <<- EOF
  python3 $repo_path/aquamis.py
    --working_directory $dir_data
    --snakefile $repo_path/Snakefile
    --profile $profile
    --rule_directives $dag_search
    --threads_sample $threads_sample
    --sample_list $sample_list
    --run_name $run_name
    $krakenDB
    --fix_fails
    ${user_args:-}
EOF
#    --shovill_ram 32
#    --json
#    --profile $profile
#    --shovill_tmpdir /tmp/shovill_$USER  # add this for accreditation version
#    --ephemeral
#    --threads 60
#    --dryrun
#    --iontorrent --shovill_extraopts="--iontorrent"
#    --single_end
#    --shovill_depth 0
#    --remove_temp

#    --mashdb "/cephfs/abteilung4/Projects_NGS/aquamis_contamination/repo/AQUAMIS/reference_db/mash/mashDB.msh"
#    --taxonkit_db "/cephfs/abteilung4/Projects_NGS/aquamis_contamination/repo/AQUAMIS/reference_db/taxonkit"
#    --kraken2db "/cephfs/abteilung4/Projects_NGS/aquamis_contamination/repo/AQUAMIS/reference_db/kraken"
#    --confindr_database "/cephfs/abteilung4/Projects_NGS/aquamis_contamination/repo/AQUAMIS/reference_db/confindr"

#    --mashdb "/home/DenekeC/Snakefiles/aquamis/reference_db/mash/refseq_bacteria_viruses_plasmids.k21s1000.msh"
#    --taxonkit_db "/home/DenekeC/Snakefiles/aquamis/reference_db/taxonkit"
#    --kraken2db "/home/DenekeC/Snakefiles/aquamis/reference_db/krakenDB/minikraken2_v1_8GB"
#    --confindr_database "/cephfs/abteilung4/deneke/Projects/confindr_database_development/Campylobacter_2021-01-20"

#    --force all_trimming_only
#    --version
#    --conda_frontend
#    --condaprefix $CONDA_BASE/envs
  run_command_array $wrapper_cmd
}

run_snakemake() {
  logheader "Running AQUAMIS via direct Snakemake command:"
  # for CLI options, see https://snakemake.readthedocs.io/en/stable/executing/cli.html
  [[ -f $dir_data/config_aquamis.yaml ]] || die "The config yaml was not created yet in the working directory"

  read -r -d '' wrapper_cmd <<- EOF
  snakemake
    --profile $repo_path/profiles/$profile
    --configfile $dir_data/config_aquamis.yaml
    --snakefile $repo_path/Snakefile
    --notemp
    --dryrun
    --quiet progress
    ${user_args:-}
EOF
#    --rerun-triggers mtime params input software-env code
#    --forcerun download_mash_reference
#    --until download_mash_reference
#    --force all_ephemeral
#    --summary
#    --keep-incomplete
#    --cleanup-metadata  $(snakemake --configfile $dir_data/config.yaml --snakefile $repo_path/Snakefile --list-params-changes)
#    --list-input-changes
#    --list-params-changes
#    --list-code-changes
#    --list-version-changes
#    --list-untracked
#    --cleanup-shadow
#    --conda-prefix $CONDA_BASE/envs
#    --use-conda
  run_command_array $wrapper_cmd
}

run_snakemake_report() {
  logheader "Generating Snakemake Report:"
  read -r -d '' wrapper_cmd <<- EOF
  snakemake
    --profile $repo_path/profiles/$profile
    --configfile $dir_data/config_aquamis.yaml
    --snakefile $repo_path/Snakefile
    --forceall
    --rulegraph | dot -Tpdf -o $dir_data/rulegraph-all.pdf
EOF
#    --dag | dot -Tpdf -o $dir_data/dag-all.pdf
#    --report $dir_data/reports/smk_report.html
  run_command_array $wrapper_cmd
}

run_docker() {
  logheader "Running AQUAMIS via DOCKER:"

  # create sample sheet
  [[ -f $dir_data/samples.tsv ]] || logheader "Creating AQUAMIS Sample Sheet via DOCKER:"
  read -r -d '' wrapper_cmd <<- EOF
  docker run --rm
    -e LOCAL_USER_ID=$(id -u $USER)
    -v ${dir_data}:/AQUAMIS/analysis
    $docker_image
    /AQUAMIS/scripts/create_sampleSheet.sh
    --mode illumina
    --fastxDir /AQUAMIS/analysis/fastq
    --outDir /AQUAMIS/analysis
    ${user_args:-}
EOF
  [[ -f $dir_data/samples.tsv ]] || run_command_array $wrapper_cmd

  # run pipeline
  read -r -d '' wrapper_cmd <<- EOF
  docker run --rm
    -e LOCAL_USER_ID=$(id -u $USER)
    -v $dir_data:/AQUAMIS/analysis
    -v $(realpath $repo_path/reference_db/confindr):/AQUAMIS/reference_db/confindr
    -v $(realpath $repo_path/reference_db/kraken):/AQUAMIS/reference_db/kraken
    -v $(realpath $repo_path/reference_db/mash):/AQUAMIS/reference_db/mash
    -v $(realpath $repo_path/reference_db/taxonkit):/AQUAMIS/reference_db/taxonkit
    -v $(realpath $repo_path/reference_db/trimmomatic):/AQUAMIS/reference_db/trimmomatic
    $docker_image
    /AQUAMIS/aquamis.py
    --docker $dir_data
    --working_directory /AQUAMIS/analysis
    --snakefile /AQUAMIS/Snakefile
    --profile $profile
    --rule_directives /AQUAMIS/profiles/$(basename $dag_search)
    --threads_sample $threads_sample
    --sample_list /AQUAMIS/analysis/samples.tsv
    --fix_fails
    ${user_args:-}
EOF
  # legacy incompatible options
#    --profile $profile
#    --rule_directives /AQUAMIS/profiles/$(basename $dag_search)
  # for database mapping when using only bfrbioinformatics/aquamis:conda
#    -v $(realpath $repo_path/reference_db/confindr):/AQUAMIS/reference_db/confindr
#    -v $(realpath $repo_path/reference_db/kraken):/AQUAMIS/reference_db/kraken
#    -v $(realpath $repo_path/reference_db/mash):/AQUAMIS/reference_db/mash
#    -v $(realpath $repo_path/reference_db/taxonkit):/AQUAMIS/reference_db/taxonkit
#    -v $(realpath $repo_path/reference_db/trimmomatic):/AQUAMIS/reference_db/trimmomatic
  run_command_array $wrapper_cmd

  # clean-up
  [[ -f $dir_data/core ]] && rm $dir_data/core
}

run_docker_compose() {
  logheader "Running AQUAMIS via DOCKER-compose:"

  # git image
  app_name="aquamis_app"
  compose_yaml="$repo_path/docker/docker-compose.yaml"

  # bioconda image
  app_name="aquamis_bioconda_app"
  compose_yaml="$repo_path/docker/docker-compose-bioconda.yaml"

  # fire up docker containers
  declare -x LOCAL_USER_ID=$(id -u $USER)
  declare -x AQUAMIS_HOST_PATH=$dir_data
  grep -q $app_name <(docker ps --filter "name=$app_name" --format '{{.Names}}') || logheader "Docker compose UP:"
  grep -q $app_name <(docker ps --filter "name=$app_name" --format '{{.Names}}') || logexec docker compose -f $compose_yaml up --no-build --pull missing --detach

  # create sample sheet
  [[ -f $dir_data/samples.tsv ]] || logheader "Creating AQUAMIS Sample Sheet via DOCKER:"
  read -r -d '' wrapper_cmd <<- EOF
  docker exec --user $LOCAL_USER_ID $app_name
    /AQUAMIS/scripts/create_sampleSheet.sh
    --force
    --mode ncbi
    --fastxDir /AQUAMIS/test_data/fastq
    --outDir /AQUAMIS/analysis
EOF
  [[ -f $dir_data/samples.tsv ]] || run_command_array $wrapper_cmd

  # run pipeline
  read -r -d '' wrapper_cmd <<- EOF
  docker exec --user $LOCAL_USER_ID $app_name
    micromamba run -n aquamis
    /AQUAMIS/aquamis.py
    --docker $dir_data
    --working_directory /AQUAMIS/analysis
    --snakefile /AQUAMIS/Snakefile
    --profile $profile
    --rule_directives /AQUAMIS/profiles/$(basename $dag_search)
    --threads_sample $threads_sample
    --sample_list /AQUAMIS/analysis/samples.tsv
    --fix_fails
    ${user_args:-}
EOF
  # legacy incompatible options
#    --profile $profile
#    --rule_directives /AQUAMIS/profiles/$(basename $dag_search)
  # for database mapping when using only bfrbioinformatics/aquamis:conda
#    -v $(realpath $repo_path/reference_db/confindr):/AQUAMIS/reference_db/confindr
#    -v $(realpath $repo_path/reference_db/kraken):/AQUAMIS/reference_db/kraken
#    -v $(realpath $repo_path/reference_db/mash):/AQUAMIS/reference_db/mash
#    -v $(realpath $repo_path/reference_db/taxonkit):/AQUAMIS/reference_db/taxonkit
#    -v $(realpath $repo_path/reference_db/trimmomatic):/AQUAMIS/reference_db/trimmomatic
  run_command_array $wrapper_cmd

  # clean-up
  [[ -f $dir_data/core ]] && rm $dir_data/core
}

run_modules() {
  ## [Rationale] Define per-run module set - choose between Python wrapper, Snakemake or Docker
  set_run_paths
  show_info_on_run
#  delete_previousRun
  [[ $snakemake == false ]] && run_aquamis
  [[ $snakemake == true  ]] && run_snakemake
#  run_snakemake_report
#  run_docker_sample
#  run_docker
#  run_docker_compose
}

## (3) Main

manual_override
define_paths
logentry "[START] $script_real"
logentry "Arguments: ${*:-"No arguments provided"}"
show_info
run_checks

## Workflow
check_conda $conda_recipe_env_name
activate_conda
[[ $interactive == true ]] && interactive_query
tweak_settings

## Switch between single-run or batch mode
if [[ $batch_mode == true ]]; then
  logheader "Processing all active runs in run table: $run_table"
  while IFS=$'\t' read -r -a runArray; do  # loop through run_table and read as array
    run_name="${runArray[0]}"
    [[ "$run_name" =~ ^#.* ]] && continue  # skip commented samples
    dir_samples=$(eval "echo ${runArray[1]}")
    dir_data=$(eval "echo ${runArray[2]}")
    run_modules
  done < $run_table
else
  logheader "Processing a single run:"
  run_modules
fi

#show_info_on_tmux
logentry "[STOP] $script_real"
