#!/usr/bin/env bash
export DIR_ORTHOLOGER="/opt/mambaforge/envs/bioconda/conda-bld/orthologer_1748034094891/_h_env_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_pla/orthologer"
if [[ ! -f "$DIR_ORTHOLOGER/bin/common_base.sh" ]]; then
    echo "ERROR: DIR_ORTHOLOGER does not seem to be correct." 1>&2
    echo "ERROR: DIR_ORTHOLOGER = $DIR_ORTHOLOGER" >&2
    exit 255
fi
source "$DIR_ORTHOLOGER/bin/common_base.sh"

import_all() {
    # get all orthologer/orthomapper functions
    [[ -f "$DIR_ORTHOLOGER/bin/orthomapper_func.sh" ]] || { echo "ERROR: orthomapper_func.sh not found! Is DIR_ORTHOLOGER correctly set?" 1>&2; exit 255; }
    source "$DIR_ORTHOLOGER/bin/orthomapper_func.sh"
    source "$DIR_ORTHOLOGER/bin/orthomapper_busco.sh"
    source "$DIR_ORTHOLOGER/bin/set_orthodb_api.sh"
    set_orthodb_api
    # get clade defs
    local odblevels="$DIR_ORTHOLOGER/etc/$URL_ORTHODB_API_VERSION/orthomapper_levels.sh"
    [[ -f "$odblevels" ]] || {
        echo "ERROR: orthomapper_levels.sh not found! Is DIR_ORTHOLOGER correctly set?" 1>&2
        echo "Here: $odblevels"
        exit 255
    }
    source "$odblevels"

    if [[ -f "$DIR_ORTHOLOGER/etc/orthologer_conf.sh" ]]; then
        echo "source etc conf at $DIR_ORTHOLOGER/etc/orthologer_conf.sh"
        source "$DIR_ORTHOLOGER/etc/orthologer_conf.sh"
    else
        echo "cannot find etc conf at $DIR_ORTHOLOGER/etc/orthologer_conf.sh"
    fi
}

intro() {
    echo
    echo "----------------------------------------------------------------"
    echo "                     OrthoLoger Pipeline                        "
    echo "----------------------------------------------------------------"
    echo
    echo "Orthologer software directory : $DIR_ORTHOLOGER"
    echo "Current version               : $OP_VERSION_SHORT"
    echo
    echo " To setup a running using a local storage, proceed as follows:  "
    echo
    echo " - create an empty directory (e.g odb)     "
    echo "   > mkdir odb"
    echo
    echo " - run setup using (assuming tagged as 'orthologer:latest')     "
    echo "   > docker run -u \$(id -u) -v \$(pwd)/odb:/odbwork orthologer:latest orthologer -c create"
    echo
    echo " - import fasta files"
    echo "   this step requires that /odbdata is mounted in the image"
    echo "   > docker run -u \$(id -u) -v \$(pwd)/odb:/odbwork -v <data>:/odbdata orthologer:latest orthologer -c import"
    echo
    echo " - run orthologer"
    echo "   > docker run -u \$(id -u) -v \$(pwd)/odb:/odbwork orthologer:latest orthologer -c run"
    echo
    echo " - alternatively, run all above "
    echo "   > docker run -u \$(id -u) -v \$(pwd)/odb:/odbwork -d <data>:/odbdata orthologer:latest orthologer -c go"
    echo
    echo " - in order to run scripts directly in /odbwork:"
    echo "   > docker run -u \$(id -u) -v \$(pwd)/odb:/odbwork orthologer:latest ./orthologer.sh <args>"
    echo
    echo " Once setup is complete, all parameters are set in orthologer_conf.sh .  "
    echo " Review these settings, in particular wrt nr of cores available,"
    echo " if a computing cluster is available, etc.                      "
    echo " The most time consuming part is usually the ALIGNMENT step.            "
    echo
    echo " Further help is available running: ./orthologer.sh -h                  "
    echo " In addition, there are two files with instructions in ./docs   "
    echo
}

usage() {
    echo "usage: $0 [-h -d <data dir> -p <project name> -c <command>] | [cmdline]"
    echo
    echo "  -h               : this message"
    echo "  -v               : print version"
    echo "  -f               : force action ('create')"
    echo "  -p <name>        : project name, default = $ORTHOLOGER_NAME"
    echo "  -w <work dir>    : work dir, by default it is the current dir"
    echo "  -d <import>      : fasta import source"
    echo "                     1. a location of fasta files, by default $ORTHOLOGER_DATA"
    echo "                     2. a fasta import description file"
    echo "                        each line contains '+<label> <file>'"
    echo "                        the '+' indicates that the fasta records in the file will be relabeled"
    echo "  -B <nbatch>      : nr of jobs in one batch"
    echo "  -L <nproc>       : nr of jobs to be run in parallel"
    echo "  -c <command>     : what to do, can be one of"
    echo "                     create       create orthologer environment"
    echo "                     import       import fasta files from <data dir>"
    echo "                     run          run the orthologer with the imported fasta"
    echo
    echo "Setup a project in an empty directory by calling within that directory:"
    echo "> orthologer -c create"
    echo ""
    echo "If the project is to be created elsewhere, do:"
    echo "> orthologer -c create -w <trg dir>"
    echo ""
    echo "Fasta files are imported from a directory given by -d option."
    echo "The run can be labeled using -p option. If not given, the default is used."
    echo "> orthologer -c import -d <source dir> -p <name>"
    echo
    echo "Finally to run the project:"
    echo "> orthologer -c run -n <name>"
    echo
    echo "The result(s) will be stored in 'Results'."
    echo
}

# default locations for OrthoDB data and work directory
export ORTHOLOGER_DATA=/odbdata
export DIR_PROJECT=$(pwd)

# this maybe set by user
export ORTHOLOGER_NAME=odbproj

# optional tree file
export TREE_FILE=""

COMMAND=
NPROC=0
NBATCH=0
FORCE=0
DO_CREATE=0
DO_IMPORT=0
DO_RUN=0
process_args() {
    COMMAND=
    DO_CREATE=0
    DO_IMPORT=0
    DO_RUN=0
    NPROC=0
    NBATCH=0
    FORCE=0
    
    local opts="hvfd:w:p:t:c:w:L:B:"
    local OPTIND=1
    local OPTS=""
    local nreq=0
    while getopts $opts OPT; do
        case "$OPT" in
            h) usage
               exit 0
               ;;
            v) import_all
               print_version
               print_busco_version
               exit 0
               ;;
            f) FORCE=1
               ;;
            L) NPROC="$OPTARG"
               ;;
            B) NBATCH="$OPTARG"
               ;;
            d) ORTHOLOGER_DATA="$OPTARG"
               ;;
            w) DIR_PROJECT="$OPTARG"
               ;;
            p) ORTHOLOGER_NAME="$OPTARG"
               ;;
            t) TREE_FILE="$OPTARG"
               ;;
            c) COMMAND="$OPTARG"
               ;;
        esac
    done
    shift $((OPTIND-1))

    [[ -n $NPROC    ]] && OP_NJOBMAX_LOCAL=$NPROC
    [[ -n $NBATCH   ]] && OP_NJOBMAX_BATCH=$NBATCH
    
    # check command
    COMMAND="${COMMAND,,}"
    local optsreq=0
    if [[ ${COMMAND} == run ]]; then
        DO_RUN=1
        TREE_FILE=""
        
    elif [[ ${COMMAND} == create ]]; then
        DO_CREATE=1

    elif [[ ${COMMAND} == import ]]; then
        DO_IMPORT=1
        TREE_FILE=""
        
    elif [[ -n $COMMAND ]]; then
        error_msg "invalid command $COMMAND"
        return 255
    fi
    
    return 0
}

# process arguments and show usage if none are given
[[ $# -eq 0 ]] && { usage; exit 0; }
process_args "$@" || exit 255

if [[ -z $COMMAND ]]; then
    usage
    exit 0
fi

# check existence of directories
[[ -d $DIR_PROJECT ]] || { error_msg "missing orthologer work dir at $DIR_PROJECT - is it mounted?"; exit 255; }

info_msg "moving to work directory $DIR_PROJECT"
cd $DIR_PROJECT || { error_msg "missing work directory at $DIR_PROJECT"; exit 255; }

if (( DO_IMPORT )); then
    if [[ -f "${ORTHOLOGER_DATA}" ]]; then
        info_msg "import data from file $ORTHOLOGER_DATA"
    else
        # check if import is given by directory OR file
        # make sure it contains a trailing /
        ORTHOLOGER_DATA="${ORTHOLOGER_DATA}/"
        ORTHOLOGER_DATA="${ORTHOLOGER_DATA//\/\//\/}"
        [[ -d $ORTHOLOGER_DATA ]] || { error_msg "missing orthologer data dir at $ORTHOLOGER_DATA - is it mounted?"; exit 255; }
        info_msg "import data from directory $ORTHOLOGER_DATA"
    fi
fi

#-----------------------------------------------
cmd_setup_orthologer() {
    info_msg "--- setup project environment at $DIR_PROJECT"
    local nfiles=$(find $DIR_PROJECT/ -type f | wc -l)
    (( FORCE )) && nfiles=0
    if (( nfiles == 0 )); then
        # create empty environment
        (
            export DIR_PROJECT=$DIR_PROJECT
            export NTHREADS=$OP_NJOBMAX_LOCAL
            source $DIR_ORTHOLOGER/bin/setup.sh
        )
    else
        if [[ -f $DIR_PROJECT/orthologer.sh ]] && [[ -f $DIR_PROJECT/orthologer_conf.sh ]] && [[ -d $DIR_PROJECT/PWC ]]; then
            info_msg "info: work directory environment already exists"
        else
            error_msg "non empty work directory is not an orthologer work environment"
            return 255
        fi
    fi
    return 0
}

#-----------------------------------------------
import_fasta() {
    if [[ -d $ORTHOLOGER_DATA ]]; then
        info_msg "--- build import file from directory $ORTHOLOGER_DATA"
        local nfiles=$(find $ORTHOLOGER_DATA -type f | wc -l)
        if (( nfiles == 0 )); then
            error_msg "given data directory seems to be empty"
            return 255
        fi
        if [[ -f "$DIR_PROJECT/${ORTHOLOGER_NAME}.txt" ]]; then
            info_msg "import file $DIR_PROJECT/${ORTHOLOGER_NAME}.txt already exists"
        else
            info_msg "build import file $DIR_PROJECT/${ORTHOLOGER_NAME}.txt"
            # build import file from user given fasta files
            ./sbin/build_import_file.sh -o "$DIR_PROJECT/${ORTHOLOGER_NAME}.txt" "$ORTHOLOGER_DATA" || { error_msg "unable to build import file"; return 255;}
        fi
    else
        info_msg "--- copy given import file from  $ORTHOLOGER_DATA to ${ORTHOLOGER_NAME}.txt"
        cp -f "$ORTHOLOGER_DATA" "$DIR_PROJECT/${ORTHOLOGER_NAME}.txt" || { error_msg "unable to copy $ORTHOLOGER_DATA to ${ORTHOLOGER_NAME}.txt"; return 255; }
    fi

    # move to work dir
    cd $DIR_PROJECT
    
    # import them into this project
    info_msg "import ${ORTHOLOGER_NAME}.txt"
    ./orthologer.sh manage -f ${ORTHOLOGER_NAME}.txt || { error_msg "unable to import fasta files"; return 255;}

    # update orthologer_conf.sh such that it includes the new todo file
    info_msg "update orthologer_conf.sh to include new todo file: $DIR_PROJECT/todo/${ORTHOLOGER_NAME}.todo"
    ./orthologer.sh -t $DIR_PROJECT/todo/${ORTHOLOGER_NAME}.todo -u || { error_msg "unable to update orthologer_conf.sh"; return 255;}

    return 0
}

#-----------------------------------------------
run_orthologer() {
    echo
    info_msg "--- run orthologer"
    echo
    local nproc_args=""
    [[ -z NPROC  ]] || nproc_args="-L $NPROC"
    [[ -z NBATCH ]] || nproc_args="${nproc_args} -B $NBATCH"
    cd $DIR_PROJECT
    local ec=0
    T0=$(date +%s)
    ./orthologer.sh -t ./todo/${ORTHOLOGER_NAME}.todo -r all $nproc_args
    ec=$?
    T1=$(date +%s)
    dT=$((T1-T0))
    echo
    info_msg "--------------------------------------"
    info_msg " Run finished in $dT seconds."
    [[ $ec -eq 0 ]] || info_msg " Error code = $ec"
    info_msg " Final results in 'Results' folder."
    info_msg "--------------------------------------"
    echo
}

#-----------------------------------------------
#RUNLOG=${DIR_PROJECT}run_$(date +%y%m%d%H%M%S).log
#echo
#info_msg "--- Start run" #, log file at $RUNLOG"
#echo
{
    # first setup environment in workdir
    if (( DO_CREATE )); then
        cmd_setup_orthologer          || exit 255
    fi
    # import data and prepare for run
    if (( DO_IMPORT )); then
        import_fasta               || exit 255
    fi
    # run pipeline
    if (( DO_RUN )); then
        run_orthologer             || exit 255
    fi
} ###>& $RUNLOG

info_msg "--- DONE ---"
