binny

#! /bin/bash -i

DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
VARCONFIG=$DIR/VARIABLE_CONFIG

while IFS=$'\t' read var val; do unset $var ; declare $var="$val" ; done < $VARCONFIG

if [ -z "$MAX_THREADS" ]; then
    MAX_THREADS=50
fi

usage() {
    echo "Usage: $0 [-u|d|c|f|l|i] [-b node] [-t number] [-r] [-n name] /absolute_path/to/config_file " 1>&2
    echo "       -n <name for main job>, only works with -c and -f" 1>&2
    echo "       -r if set, a report is generated (it's recommended to run -c, -f and -l with -r)" 1>&2
    echo "       -d if set, a dry-run is performed" 1>&2
    echo "       -c if set, the whole thing is submitted to the cluster" 1>&2
    echo "       -b if -c is set, -b gives the node name to submit the main instance to" 1>&2
    echo "       -i if set, only the conda environments will be installed, if they don't exist" 1>&2
    echo "       -u if set, the working directory will be unlocked (only necessary for crash/kill recovery)" 1>&2
    echo "       -l if set, the main snakemake thread and indivdual rules are run in the current terminal session" 1>&2
    echo "       -x if -c is set, binny will use tmux for snakemake and submit to the cluster, if -l is set, snakemake and rules are run in tmux" 1>&2
    echo "       -f if set, the main snakemake thread and indivdual rules are run in a tmux session (= -l and -x)" 1>&2
    echo "       -t <max_threads> maximum number of cpus to use for all rules at a time. Defaults to $MAX_THREADS for -c, and to 1 for -l and -f. No effect on -r, -d or -u only." 1>&2

}

while getopts n:t:udxflcb:rhi flag
do
    case $flag in
        i)
            INITIAL=true;;
        u)
            UNLOCK=true;;
        d)
            DRYRUN=true;;
        c)
            CLUSTER=true;;
        n)
            JNAME=$OPTARG;;
        r)
            REPORT=true;;
        b)
            NNAME=$OPTARG;;
        l)
            LAPTOP=true;;
        x)
            TMUX=true;;
        f)
            FRONTEND=true;;
        t)
            THREADS=$OPTARG;;
        h)
            usage
            exit;;
        *)
            echo "Unimplemented option: -$OPTARG" >&2
            usage
            exit 1;;
        :)
            echo "Missing option argument for -$OPTARG" >&2
            usage
            exit 1;;
        ?)
            usage
            exit
             ;;
    esac
done

shift $((OPTIND-1))

if [ -z "$1" ]; then
    echo "missing input"
    usage
    exit 1
else
    CONFIGFILE=$1
fi


#if the config file cannot be found
if [[ !  -e "$1" ]]; then
   echo "Configfile "$1" was not found."
   echo "Provide full path."
   exit 1
fi

# Load modules
eval $LOADING_MODULES

# Ensure conda functionality, even without prior user conda setup.
. "$(conda info --base)/etc/profile.d/conda.sh"

# Env setup
conda_source_config_var=$(grep 'conda_source:' ${CONFIGFILE} | grep '^#' -v | cut -d':' -f2 | tr -d '[:space:]' | sed -e 's/"//g')
if [[ -n ${conda_source_config_var} ]]; then
  CONDA_SOURCE="${conda_source_config_var}"
else
  CONDA_SOURCE="${DIR}/conda"
fi

# Get info on Prokka and Mantis envs
prokka_env_config_var=$(grep 'prokka_env:' ${CONFIGFILE} | grep '^#' -v | cut -d':' -f2 | tr -d '[:space:]' | sed -e 's/"//g')
if [[ -n  ${prokka_env_config_var} ]]; then
  # Get env folder path without extension, if input is yaml/yml else get whole string
  prokka_env_path=${prokka_env_config_var%.@(yaml|yml|YAML|YML)}
  # If env is from yaml/yml add parent dir to envs_dirs to trick snakemake into not rebuilding the env and accepting it as named env
  if [[ ${prokka_env_config_var} != ${prokka_env_path} ]]; then
    prokka_env_name=$(basename ${prokka_env_path})
    prokka_env_parent=$(dirname ${prokka_env_path})
    conda config --append envs_dirs ${prokka_env_parent}
#    # Hack to prevent source file not found error?
#    ln -s ${prokka_env_path} ${DIR}/${prokka_env_name}
  else
    prokka_env_name=${prokka_env_path}
  fi
  echo "Using existing Prokka env: ${prokka_env_name}"
fi

# Same for mantis
mantis_env_config_var=$(grep 'mantis_env:' ${CONFIGFILE} | grep '^#' -v | cut -d':' -f2 | tr -d '[:space:]' | sed -e 's/"//g')
if [[ -n  ${mantis_env_config_var} ]]; then
  mantis_env_path=${mantis_env_config_var%.@(yaml|yml|YAML|YML)}
  # If env is from yaml/yml add parent dir to envs_dirs to trick snakemake into not rebuilding the env and accepting it as named env
  if [[ ${mantis_env_config_var} != ${mantis_env_path} ]]; then
    mantis_env_name=$(basename ${mantis_env_path})
    mantis_env_parent=$(dirname ${mantis_env_path})
  else
    mantis_env_name=${mantis_env_path}
  fi
  mantis_env=${mantis_env_name}
  echo "Using existing Mantis env: ${mantis_env}"
else
  if [[ -d ${CONDA_SOURCE} ]]; then
      shopt -s failglob
      env_yamls=$(ls ${CONDA_SOURCE}/*.yaml)
      if [[ ${env_yamls} ]]; then
        for i in ${CONDA_SOURCE}/*.yaml; do
          env_name=$(head -n 1 ${i} | cut -d' ' -f2)
          if [[ ${env_name} == 'mantis' ]]; then
            mantis_env=${i%.yaml}
            if [[ -n  ${mantis_env_config_var} ]]; then
              echo "Loading mantis env from: ${mantis_env}"
            fi
            mantis_set='true'
          fi
        done
      fi
  else
      shopt -u failglob
#  else
#    mantis_set=''
  fi
fi

# Set conda source path to target path or existing Prokka/Mantis source path if exists
if [[ -n  ${conda_source_config_var} ]]; then
  # Find path to mantis env
  if [[  ${conda_source_config_var}} != ${mantis_env_path}  && -n ${mantis_env_path} ]]; then
    for i in ${CONDA_SOURCE}/*.yaml; do
      env_name=$(grep -l 'mantis' ${i})
      if [[ -n ${env_name} ]]; then
        mantis_env=${i%.yaml}
      fi
    done
  fi
elif  [[ ( -n  ${mantis_env_parent} && -n  ${prokka_env_parent} ) && (${prokka_env_parent} == ${mantis_env_parent}) ]]; then
  CONDA_SOURCE=${mantis_env_parent}
elif [[ ( -n  ${mantis_env_parent} && -n  ${prokka_env_parent} ) ]]; then
  echo "Existing Prokka and Mantis env sources must currently be identical (Snakemake limitation)." \
       " Choose one and have Binny create the other." && exit 1
elif [[ -n  ${mantis_env_parent} ]]; then
  CONDA_SOURCE=${mantis_env_parent}
elif [[ -n  ${prokka_env_parent} ]]; then
  CONDA_SOURCE=${prokka_env_parent}
else
  CONDA_SOURCE="${DIR}/conda"
  mkdir -p $CONDA_SOURCE
fi
echo "Will use conda source path: ${CONDA_SOURCE}"

# Check Snakemake and create env, if needed
if [ "$SNAKEMAKE_VIA_CONDA" = true ]; then
  snakemake_env_config_var=$(grep 'snakemake_env:' ${CONFIGFILE} | grep '^#' -v | cut -d':' -f2 | tr -d '[:space:]' | sed -e 's/"//g')
  if [[ ${snakemake_env_config_var} == "in_path" ]]; then
    echo "Will use snakemake from PATH."
    CONDA_START=""
  elif [[ -n  ${snakemake_env_config_var} ]]; then
    snakemake_env=${snakemake_env_config_var}
    echo "Using existing Snakemake env: ${snakemake_env}"
    CONDA_START="conda activate ${snakemake_env}"
  else
    if [[ ! -x "${CONDA_SOURCE}/snakemake_env/bin/snakemake" ]]; then
      mkdir -p ${CONDA_SOURCE}
      if ! [ -x "$(command -v mamba)" ]; then
        env_manager="conda"
      else
        env_manager="mamba"
      fi
      echo "Creating snakemake environment"
      # snakemake originally was 3.8, 7.16.0, now try python = 3.9, snakemake=7.32.4
      ${env_manager} create --prefix ${CONDA_SOURCE}/snakemake_env snakemake>=7.16.0 mamba>=0.22.1 unzip python=3.8 -c conda-forge -c bioconda --yes
    fi
    CONDA_START="conda activate ${CONDA_SOURCE}/snakemake_env"
  fi
   CONDA_END="conda deactivate"
   CONDA_END_t="conda deactivate;"
else
   CONDA_START=""
   CONDA_END=""
   CONDA_END_t=""
fi

START_TIME=`date +%s`
NAMEHASH=`echo $START_TIME| cksum | awk '{print $1}'`
if [ -z "$JNAME" ]; then
    JNAME="binny_${NAMEHASH}"
else
    JNAME="${JNAME}_${NAMEHASH}"
fi

if [ "$UNLOCK" = true ]; then
    echo "Unlocking working directory."
    eval $CONDA_START
    snakemake $SNAKEMAKE_EXTRA_ARGUMENTS --cores 1 -s $DIR/Snakefile --unlock --configfile $CONFIGFILE
    eval $CONDA_END
elif [ "$DRYRUN" = true ]; then
    echo "Dryrun."
    eval $CONDA_START
    snakemake $SNAKEMAKE_EXTRA_ARGUMENTS --cores 1 -s $DIR/Snakefile --dryrun --config sessionName=$JNAME --configfile $CONFIGFILE
    eval $CONDA_END
elif [ "$INITIAL" = true ]; then
    eval $CONDA_START
    snakemake $SNAKEMAKE_EXTRA_ARGUMENTS --verbose --cores 1 -s $DIR/Snakefile --conda-create-envs-only --use-conda \
              --conda-prefix ${CONDA_SOURCE} --local-cores 1 --configfile $CONFIGFILE

    DB_PATH=$(grep "db_path:" $CONFIGFILE | grep '^#' -v | cut -d':' -f2 | tr -d '[:space:]' | sed -e 's/"//g')
    if [[ -z  ${DB_PATH} ]]; then
      DB_PATH="${DIR}/database"
    fi

    echo "Setting up Mantis with the CheckM databases"
    if [[ ! ${mantis_set} ]]; then
      for i in ${CONDA_SOURCE}/*.yaml; do
        env_name=$(head -n 1 ${i} | cut -d' ' -f2)
        if [[ ${env_name} == 'mantis' ]]; then
          mantis_env=${i%.yaml}
          echo "Loading mantis env from: ${mantis_env}"
        fi
      done
    fi
    conda activate ${mantis_env}
    # Make sure a compiler for cython is available
    if ! [ -x "$(command -v gcc)" ]; then
      conda install -c conda-forge gcc_linux-64 --yes
      $CONDA_PREFIX/etc/conda/activate.d/activate-binutils_linux-64.sh
      $CONDA_PREFIX/etc/conda/activate.d/activate-gcc_linux-64.sh
    fi
    hmmpress -f ${DB_PATH}/hmms/checkm_tf/checkm_filtered_tf.hmm && \
    hmmpress -f ${DB_PATH}/hmms/checkm_pf/checkm_filtered_pf.hmm && \
    sed -e "s|__PATH_TO_DB__|${DIR}|g" ${DIR}/config/binny_mantis_template.cfg > ${DIR}/config/binny_mantis.cfg && \
    mantis setup -mc ${DIR}/config/binny_mantis.cfg --chunk_size 1200 --no_taxonomy && \
    mantis check -mc ${DIR}/config/binny_mantis.cfg --no_taxonomy && \
    conda deactivate && \
    echo "Done." && \
    exit 0
elif [ "$CLUSTER" = true ]; then
    if [ -z "$THREADS" ]; then
        THREADS=$MAX_THREADS
    fi
    if [ -z "$NNAME" ]; then
        NNAME=""
    fi

    if [ "$TMUX" = true ]; then
      echo "Submitting workflow to cluster - snakemake stays on the frontend, via tmux"
      tmux new -s $JNAME -d
      tmux send-keys -t $JNAME "$LOADING_MODULES >> $JNAME.stdout 2>> $JNAME.stderr" C-m
      tmux send-keys -t $JNAME "$CONDA_START >> $JNAME.stdout 2>> $JNAME.stderr" C-m
      if [ "$REPORT" = true ]; then
        tmux send-keys -t $JNAME "snakemake --cores $THREADS --jobs $THREADS -s $DIR/Snakefile --keep-going --local-cores 1 \
                 --cluster-config $DIR/config/$SCHEDULER.config.yaml \
                 --cluster \"{cluster.call} {cluster.runtime}{resources.runtime} {cluster.mem_per_cpu}{resources.mem} \
                  {cluster.nodes} {cluster.qos} {cluster.threads}{threads} {cluster.partition} {cluster.stdout}\" \
                 --configfile $CONFIGFILE --config sessionName=$JNAME \
                 --use-conda --conda-prefix $CONDA_SOURCE >> $JNAME.stdout 2>> $JNAME.stderr; \
                 snakemake --cores 1 \
                  -s $DIR/Snakefile --configfile $CONFIGFILE --use-conda --conda-prefix $CONDA_SOURCE --report report.html >>\
                   $JNAME.stdout 2>> $JNAME.stderr; $CONDA_END_t tmux kill-session" C-m
      else
        tmux send-keys -t $JNAME "snakemake $SNAKEMAKE_EXTRA_ARGUMENTS --cores $THREADS --jobs $THREADS \
                 -s $DIR/Snakefile --keep-going --local-cores 1 --cluster-config $DIR/config/$SCHEDULER.config.yaml \
                 --cluster \"{cluster.call} {cluster.runtime}{resources.runtime} {cluster.mem_per_cpu}{resources.mem} \
                 {cluster.threads}{threads} {cluster.nodes} {cluster.qos} {cluster.partition} {cluster.stdout}\" \
                 --configfile $CONFIGFILE --config sessionName=$JNAME --use-conda --conda-prefix $CONDA_SOURCE >> \
                 $JNAME.stdout 2>> $JNAME.stderr; $CONDA_END_t tmux kill-session" C-m
      fi
    else
      if [ "$REPORT" = true ]; then
        run_cmd="${SUBMIT_COMMAND}$NNAME $DIR/runscripts/runBinny_withReport.sh $CONFIGFILE $VARCONFIG $JNAME $THREADS ${CONDA_START##* } $CONDA_SOURCE"
      else
        run_cmd="${SUBMIT_COMMAND}$NNAME $DIR/runscripts/runBinny_withoutReport.sh $CONFIGFILE $VARCONFIG $JNAME $THREADS ${CONDA_START##* } $CONDA_SOURCE"
      fi
      echo "Submitting workflow to cluster."
      eval $run_cmd
    fi
elif [[ "$FRONTEND" = true || ("$LAPTOP" = true && "$TMUX" = true) ]] ; then
    echo "Running workflow on frontend using tmux - don't use this setting except with small datasets and with no more than one run at a time."
    if [ -z "$THREADS" ]; then
        THREADS=1
    fi
    tmux new -s $JNAME -d
    tmux send-keys -t $JNAME "$LOADING_MODULES >> $JNAME.stdout 2>> $JNAME.stderr" C-m
    tmux send-keys -t $JNAME "$CONDA_START >> $JNAME.stdout 2>> $JNAME.stderr" C-m
    if [ "$REPORT" = true ]; then
        tmux send-keys -t $JNAME "snakemake $SNAKEMAKE_EXTRA_ARGUMENTS --cores $THREADS --jobs $THREADS -s $DIR/Snakefile --keep-going \
                                  --configfile $CONFIGFILE --config sessionName=$JNAME --use-conda --conda-prefix $CONDA_SOURCE >> \
                                  $JNAME.stdout 2>> $JNAME.stderr; snakemake $SNAKEMAKE_EXTRA_ARGUMENTS --cores 1 \
                                  -s $DIR/Snakefile --configfile $CONFIGFILE --use-conda --conda-prefix $CONDA_SOURCE \
                                  --report report.html >> $JNAME.stdout 2>> $JNAME.stderr; $CONDA_END_t tmux kill-session" C-m
    else
        tmux send-keys -t $JNAME "snakemake $SNAKEMAKE_EXTRA_ARGUMENTS --cores $THREADS --jobs $THREADS -s $DIR/Snakefile --keep-going \
                                  --configfile $CONFIGFILE --config sessionName=$JNAME --use-conda --conda-prefix $CONDA_SOURCE >> \
                                  $JNAME.stdout 2>> $JNAME.stderr; $CONDA_END_t tmux kill-session" C-m
    fi
elif [ "$LAPTOP" = true ]; then
    echo "Running workflow in current session - don't use this setting except with small datasets and databases."
    echo ${CONDA_SOURCE}
    JNAME=${JNAME//./_}
    if [ -z "$THREADS" ]; then
        THREADS=1
    fi
    eval $CONDA_START
    if [ "$REPORT" = true ]; then
        snakemake $SNAKEMAKE_EXTRA_ARGUMENTS --cores $THREADS -s $DIR/Snakefile --keep-going --configfile $CONFIGFILE --config sessionName=$JNAME --use-conda --conda-prefix ${CONDA_SOURCE}
        snakemake $SNAKEMAKE_EXTRA_ARGUMENTS --cores $THREADS -s $DIR/Snakefile --configfile $CONFIGFILE --use-conda --conda-prefix ${CONDA_SOURCE} --report report.html
        eval $CONDA_END
    else
        snakemake $SNAKEMAKE_EXTRA_ARGUMENTS --cores $THREADS -s $DIR/Snakefile --keep-going --configfile $CONFIGFILE --config sessionName=$JNAME --use-conda --conda-prefix ${CONDA_SOURCE}
        eval $CONDA_END
    fi
elif [ "$REPORT" = true ]; then
    if [ "$TMUX" = true ]; then
       echo "Preparing report on frontend using tmux."
       tmux new -s $JNAME -d
       tmux send-keys -t $JNAME "$LOADING_MODULES >> $JNAME.stdout 2>> $JNAME.stderr" C-m
       tmux send-keys -t $JNAME "$CONDA_START >> $JNAME.stdout 2>> $JNAME.stderr" C-m
       tmux send-keys -t $JNAME "snakemake $SNAKEMAKE_EXTRA_ARGUMENTS --cores 1 \
                                  -s $DIR/Snakefile --configfile $CONFIGFILE --use-conda --conda-prefix $CONDA_SOURCE \
                                  --report report.html >> $JNAME.stdout 2>> $JNAME.stderr; $CONDA_END_t tmux kill-session" C-m
    else
       echo "Writing report."
       eval $CONDA_START
       snakemake $SNAKEMAKE_EXTRA_ARGUMENTS --cores 1 -s $DIR/Snakefile --report report.html --configfile $CONFIGFILE --use-conda --conda-prefix ${CONDA_SOURCE}
       eval $CONDA_END
    fi
else
    echo "Nothing was done, please give -u, -d, -r, -c, -f, -i, -x, or -l to start anything."
fi