#!/bin/bash
#SBATCH --cpus-per-task=8
#SBATCH --mem=16G
#SBATCH --time=01:00:00

# Unofficial Bash Strict Mode (http://redsymbol.net/articles/unofficial-bash-strict-mode/)
set -euo pipefail
IFS=$'\n\t'

if (( $# == 0));
then
    echo "USAGE: $0 [-g] TOOL SPECIES[/READ_SET] [COMMAND]" >&2
    echo >&2
    echo "Flags:" >&2
    echo "  -g                Run in debugging mode (interactive session with gdb)" >&2
    echo >&2
    echo "Commands:" >&2
    echo "  all               Run the complete pipeline; this is the default command." >&2
    echo "  create-result-db  Create DB of the result if not existing." >&2
    echo "  cache             Create the contig alignments cache. Does not recreate the" >&2
    echo "                    cache of it already exists." >&2
    echo "  check             Run checks in block-wise in parallel." >&2
    echo "  merge             Merge the result from all blocks." >&2
    echo "  mk-stats          Create the result summary as TSV for use with Gnuplot." >&2
    echo "  clear             Clear the contig alignments cache and remove the result DB." >&2
    echo "  debug             Output dentist command line for debugging." >&2

    exit 1
fi


if [[ "$1" == "-g" ]];
then
    DEBUG="$1"
    shift
fi

if (( $# < 2));
then
    echo "$0: missing SPECIES" >&2

    exit 1
fi

TOOL="$1"
DATASET="$2"
SPECIES="${DATASET%/*}"
[[ "$SPECIES" == "$DATASET" ]] || READ_SET="${DATASET##*/}"
RUN_ID="${READ_SET#*.}"
if [[ "$RUN_ID" == "$READ_SET" ]]
then
    unset RUN_ID
else
    READ_SET="${READ_SET%%.*}"
fi

COMMAND="all"
[[ "${3:--}" =~ ^- ]] || COMMAND="$3"

if [[ "$COMMAND" == "${3:-}" ]];
then
    shift 3
else
    shift 2
fi

declare -A RESULT_DB=(
    [dentist]=gap-closed.dam
    [pbjelly]=jelly.out.dam
    [arrow]=arrow.dam
    [finisher_sc]=improved3.dam
    [lr_gapcloser]=gapclosed.dam
)

if [[ ! ${RESULT_DB[$TOOL]+set} ]];
then
    echo "$0: unkown TOOL"

    exit 1
fi

TOOLDIR="$TOOL${READ_SET+/$READ_SET}${RUN_ID+.$RUN_ID}"
DATADIR="$PBGAPS/data/$SPECIES"
WORKDIR="$PBGAPS/workdir/$SPECIES/$TOOLDIR"
LOGDIR="$PBGAPS/workdir/$SPECIES/$TOOLDIR/logs"
STATSDIR="$PBGAPS/workdir/$SPECIES/$TOOLDIR/stats"
RESULTSDIR="$PBGAPS/results/$SPECIES/$TOOLDIR"

ASSEMBLY_TRUE="${ASSEMBLY_TRUE:-$DATADIR/${ASSEMBLY_TRUE_FILE:-assembly-reference.dam}}"
MAPPED_CONTIGS="${MAPPED_CONTIGS:-mapped}"
ASSEMBLY_TEST="${ASSEMBLY_TEST:-$DATADIR/${ASSEMBLY_TEST_FILE:-assembly-test.dam}}"
ASSEMBLY_RESULT="$WORKDIR/${RESULT_DB[$TOOL]}"

declare -A RESULT_FASTA=(
    [dentist]="$RESULTSDIR/gap-closed.fasta"
    [pbjelly]="$WORKDIR/jelly.out.fasta"
    [arrow]="$RESULTSDIR/arrow.fasta"
    [finisher_sc]="$WORKDIR/improved3.fasta"
    [lr_gapcloser]="$WORKDIR/workdir/iteration-3/gapclosed.fasta"
)

ASSEMBLY_RESULT_FASTA="${RESULT_FASTA[$TOOL]}"
CONTIG_ALIGNMENTS_CACHE="$WORKDIR/contig-alignments-cache.json"

BATCH_SIZE="${BATCH_SIZE:-250}"
(( BATCH_SIZE > 0 ))
NUM_CONTIGS=$( DBdump "$ASSEMBLY_TEST" | awk '/^+ R/ { print $3 }')
NUM_BATCHES=$(( 1 + ((NUM_CONTIGS - 1) / BATCH_SIZE) ))
GAP_DETAILS_CONTEXT=2500
CROP_AMBIGUOUS=$(( ${CROP_AMBIGUOUS:-300} ))


function dentist()
{
    if [[ -v DEBUG ]];
    then
        local IFS=" "
        gdb "$(command -v dentist)" -ex 'break main' -ex "run $*"
    else
        /usr/bin/env dentist "$@" > "$STDOUT" 2> "$STDERR"
    fi
}


function notify_mail()
{
    local EXIT=$?

    if [[ "${1:-}" == '--stdin' ]];
    then
        mail -s "$TOOL $DATASET check/$COMMAND ($EXIT)" "$USER@mpi-cbg.de"
    else
        echo "Job id: $JOBID" | mail -s "$TOOL $DATASET check/$COMMAND ($EXIT)" "$USER@mpi-cbg.de"
    fi

    return $EXIT
}


function assert_result_db()
{
    if [[ ! -f "$ASSEMBLY_RESULT" ]];
    then
        trap 'DBrm "$ASSEMBLY_RESULT"' err
        fold -w500 "$ASSEMBLY_RESULT_FASTA" | fasta2DAM -i "$ASSEMBLY_RESULT"
        DBsplit -x1000 "$ASSEMBLY_RESULT"

        (( $(wc -l < "$ASSEMBLY_RESULT") > 1 ))

        echo "Result DB created"
    elif [[ "${1:-}" == -v ]];
    then
        echo "Result DB already exists"
    fi
}


function cancel_current_job()
{
    scancel "$JOBID" || true

    echo "Batch job $JOBID cancelled"
}


function submit_to_slurm()
{
    if [[ -v DEBUG ]];
    then
        if [[ -v BLOCK_IDS ]];
        then
            export SLURM_ARRAY_TASK_ID="$BLOCK_IDS"
        fi

        srun --pty "$@"
    else
        JOBNAME="$(basename "$0") $TOOL $DATASET${READ_SET:+/$READ_SET}${RUN_ID:+.$RUN_ID}"
        JOBID="$(sbatch -J"$JOBNAME" ${BLOCK_IDS:+--array="$BLOCK_IDS"} "$@")"
        JOBID=${JOBID#Submitted batch job }

        trap cancel_current_job ERR SIGINT SIGTERM
        echo "Submitted batch job $JOBID"
        echo "A notification mail will be sent when the job has finished."

        ( set +e; swait "$JOBID"; notify_mail --stdin <<<"${LOG:-Jobid: $JOBID}" )
    fi
}


function prepare_common_dentist_flags()
{
    if [[ "${DENTIST_FLAGS:+set}" ]];
    then
        echo "$DENTIST_FLAGS" | tr ' ' '\n'
    fi

    echo -v
    echo -v
    echo --cache-contig-alignments="$CONTIG_ALIGNMENTS_CACHE"
    echo --crop-ambiguous=$CROP_AMBIGUOUS

    echo --threads="${SLURM_CPUS_PER_TASK:-1}"
    echo --aux-threads=1

    case "$TOOL" in
        pbjelly)
            echo --crop-alignment=100
            ;;
        lr_gapcloser)
            echo --recover-imperfect-contigs
            echo --crop-alignment=300
            ;;
    esac

    case "${RUN_ID:-}" in
        join-*)
            echo --recover-imperfect-contigs
            echo --max-imperfect-contigs-error=0.05
            ;;
    esac
}


case "$COMMAND" in
    all)
        echo "Running stage 'create-result-db'..."
        "$0" "$TOOL" "$DATASET" create-result-db "$@"
        echo "Running stage 'cache'..."
        "$0" "$TOOL" "$DATASET" cache "$@"
        echo "Running stage 'check'..."
        "$0" "$TOOL" "$DATASET" check "$@"
        echo "Running stage 'merge'..."
        "$0" "$TOOL" "$DATASET" merge "$@"
        echo "Running stage 'mk-stats'..."
        "$0" "$TOOL" "$DATASET" mk-stats "$@" | tee /dev/stderr | notify_mail --stdin
        ;;
    create-result-db)
        assert_result_db -v
        ;;
    cache)
        if [[ ! -e "$CONTIG_ALIGNMENTS_CACHE" ]]
        then
            submit_to_slurm "$@" "$0" ${DEBUG:-} "$TOOL" "$DATASET" __exec-cache
        fi
        ;;
    __exec-cache)
        mkdir -p "$WORKDIR" "$LOGDIR"
        assert_result_db
        DENTIST_FLAGS=( $(prepare_common_dentist_flags) )
        LOG="$LOGDIR/check.cache.log"

        STDOUT="/dev/null"
        STDERR="$LOG"
        dentist check-results \
            --cache-only \
            "${DENTIST_FLAGS[@]}" \
            "$ASSEMBLY_TRUE" \
            "$ASSEMBLY_TEST" \
            "$ASSEMBLY_RESULT" \
            "$MAPPED_CONTIGS"
        ;;
    check)
        BLOCK_IDS="${BLOCK_IDS:-0-$(( NUM_BATCHES - 1 ))}"

        submit_to_slurm "$@" "$0" ${DEBUG:-} "$TOOL" "$DATASET" __exec-check
        ;;
    __exec-check)
        BATCH_ID=$(printf "%02d" "$SLURM_ARRAY_TASK_ID")
        FROM=$(( BATCH_SIZE * SLURM_ARRAY_TASK_ID ))
        TO=$(( FROM + BATCH_SIZE ))

        if (( TO > NUM_CONTIGS ));
        then
            TO=$NUM_CONTIGS
        fi

        mkdir -p "$WORKDIR" "$STATSDIR" "$LOGDIR"
        assert_result_db
        DENTIST_FLAGS=( $(prepare_common_dentist_flags) )
        LOG="$LOGDIR/check.$BATCH_ID.log"

        STDOUT="$STATSDIR/stats.$BATCH_ID.json"
        STDERR="$LOG"
        dentist check-results \
            --batch="$FROM..$TO" \
            --gap-details="$STATSDIR/gap-details.$BATCH_ID.json" \
            --gap-details-context=$GAP_DETAILS_CONTEXT \
            --json \
            "${DENTIST_FLAGS[@]}" \
            "$ASSEMBLY_TRUE" \
            "$ASSEMBLY_TEST" \
            "$ASSEMBLY_RESULT" \
            "$MAPPED_CONTIGS"
        ;;
    merge)
        mkdir -p "$RESULTSDIR"

        jq -sc add "$STATSDIR/stats".*.json > "$RESULTSDIR/stats.json"
        cat "$STATSDIR/gap-details".*.json > "$RESULTSDIR/gap-details.json"
        ;;
    mk-stats)
        jq -srf "$PBGAPS/source/mk-gap-stats.jq" "$RESULTSDIR"/{gap-details,stats}.json > "$RESULTSDIR/gap-stats.tsv"

        tail -n9 "$RESULTSDIR/gap-stats.tsv"
        echo
        echo "Inspect full gap statistics at $RESULTSDIR/gap-stats.tsv"
        ;;
    clear)
        if [[ -f "$CONTIG_ALIGNMENTS_CACHE" ]];
        then
            if rm "$CONTIG_ALIGNMENTS_CACHE";
            then
                echo "Cached cleared"
            else
                echo "Failed to clear cache"
                echo "Check cache file $CONTIG_ALIGNMENTS_CACHE"
            fi
        else
            echo "No cache to be cleared"
        fi

        if [[ -f "$ASSEMBLY_RESULT" ]];
        then
            if DBrm "$ASSEMBLY_RESULT";
            then
                echo "Result DB cleared"
            else
                echo "Failed to clear result DB"
                echo "Check DB ${RESULT_DB[$TOOL]}"
            fi
        else
            echo "No result DB to be cleared"
        fi

        if (( $(find "$STATSDIR" -name 'gap-details.*.json' | wc -l) > 0 ));
        then
            if find "$STATSDIR" -name 'gap-details.*.json' -delete;
            then
                echo "Block-level gap details cleared"
            else
                echo "Failed to clear block-level gap details"
                echo "Check block-level gap details in $STATSDIR"
            fi
        else
            echo "No block-level gap details to be cleared"
        fi

        if (( $(find "$STATSDIR" -name 'stats.*.json' | wc -l) > 0 ));
        then
            if find "$STATSDIR" -name 'stats.*.json' -delete;
            then
                echo "Block-level stats cleared"
            else
                echo "Failed to clear block-level stats"
                echo "Check block-level stats in $STATSDIR"
            fi
        else
            echo "No block-level stats to be cleared"
        fi
        ;;
    debug)
        mkdir -p "$WORKDIR"
        assert_result_db
        DENTIST_FLAGS=( $(prepare_common_dentist_flags) )

        echo dentist check-results "${DENTIST_FLAGS[@]}" "$ASSEMBLY_TRUE" "$ASSEMBLY_TEST" "$ASSEMBLY_RESULT" "$MAPPED_CONTIGS"
        ;;
    *)
        echo "$0: illegal command: $COMMAND" >&2

        exit 1
        ;;
esac
