#!/bin/bash

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


function set_defaults()
{
    JOBS_DB="$PBGAPS/results/jobs.lst"
    STAGES=(
        setup
        mapping
        support
        extraction
        assembly
        output
    )
}


function error()
{
    if (( $# > 0 ));
    then
        echo "$PROG: error:" "$@"
    else
        echo
    fi
} >&2


function bail_out()
{
    error "$@"
    exit 1
}


function bail_out_usage()
{
    error "$@"
    error
    print_usage
    exit 1
}


function log()
{
    echo "--" "$@"
} >&2


function fetch_job_ids()
{
    sed -nE '/^[0-9]+$/ p; s/[^0-9]*Submitted batch job ([0-9]+)[^0-9]*/\1\n/gp' | \
        grep -oE '\b[0-9]+$'
}


function comma_list()
{
    local IFS=','

    echo "$*"
}


function print_usage()
{
    echo "USAGE:  $PROG [-h] <dataset>"
} >&2


function print_help()
{
    print_usage
    echo
    echo 'This runs the whole PBjelly pipeline doing some checks to see if stages are ok.'
    echo
    echo 'Positional arguments:'
    echo ' <dataset>       The name of the dataset, e.i. a path to the folder containing'
    echo '                 the protocol.xml and cluster.json files.'
    echo
    echo 'Optional arguments:'
    echo ' --dry-run, -n   Do not run any commands but list output only.'
    echo ' --help, -h      Prints this help.'
    echo ' --usage         Print a short command summary.'
    echo ' --version       Print software version.'
} >&2


function print_version()
{
    echo "$PROG v0"
    echo
    echo "Copyright © 2019, Arne Ludwig <arne.ludwig@posteo.de>"
} >&2


function parse_args()
{
    ARGS=()
    for ARG in "$@";
    do
        if [[ "${ARG:0:1}" == - ]];
        then
            case "$ARG" in
                -n|--dry-run)
                    DRY_RUN=1
                    ;;
                -h|--help)
                    print_help

                    exit
                    ;;
                --usage)
                    print_usage

                    exit
                    ;;
                --version)
                    print_version

                    exit
                    ;;
                *)
                    bail_out_usage "unkown option $ARG"
                    ;;
            esac
        else
            ARGS+=( "$ARG" )
        fi
    done

    (( ${#ARGS[*]} == 1 )) || bail_out_usage "<dataset> is missing"
}


function dry_run()
{
    [[ -v DRY_RUN ]]
}


function setup_variables()
{
    DATASET_PATH="${ARGS[0]}"
    [[ -d "$DATASET_PATH" ]] || bail_out "<dataset> is not a directory: $DATASET_PATH"

    PROTOCOL="$DATASET_PATH/protocol.xml"
    [[ -f "$PROTOCOL" ]] || bail_out "protocol.xml does not exist: $PROTOCOL"

    CLUSTER_CONFIG="$DATASET_PATH/cluster.json"
    [[ -f "$CLUSTER_CONFIG" ]] || bail_out "cluster.json does not exist: $CLUSTER_CONFIG"

    DATASET="$(basename "$(realpath "$DATASET_PATH")")"
    [[ -d "$PBGAPS/data/$DATASET" ]] || bail_out "invalid dataset $DATASET: not a directory: $PBGAPS/data/$DATASET"

    WORKDIR="$PBGAPS/workdir/$DATASET/pbjelly"
    LOGDIR="$WORKDIR/logs"

    ASSEMBLY="$PBGAPS/data/$DATASET/assembly-test.fasta"
    [[ -f "$ASSEMBLY" ]] || bail_out "missing assembly FASTA: $ASSEMBLY"

    if [[ "$DATASET" =~ _(pacbio|real)$ ]];
    then
        READS_TYPE="real-pb"
    else
        READS_TYPE="simulated-pb"
    fi

    READS="$PBGAPS/data/$DATASET/reads-$READS_TYPE.fasta"
    [[ -f "$READS" ]] || bail_out "missing reads FASTA: $READS"
}


function get_num_reads_parts()
{
    grep -oF '<job>' "$PROTOCOL" | wc -l
}


function count_split_reads()
{
    if [[ -d "$SPLIT_READS_DIR" ]]
    then
        find "$SPLIT_READS_DIR" -name '*.part_*.fasta' | wc -l
    else
        echo 0
    fi
}


function split_reads()
{
    seqkit split -O"$SPLIT_READS_DIR" -p"$NUM_READS_PARTS" "$READS"
}


function record_split_reads()
{
    export SPLIT_READS
    SPLIT_READS="$(find "$SPLIT_READS_DIR" -name '*.fasta' | sort | head -c-1)"

    export NUM_SPLIT_READS
    NUM_SPLIT_READS="$(wc -l <<<"$SPLIT_READS")"
}


function fix_read_headers()
{
    local JOB_IDS=( $(
        IFS="$IFS" sbatch --array=0-$(( NUM_SPLIT_READS - 1 )) --wrap='
            FILES=( $SPLIT_READS )
            FILE="${FILES[$SLURM_ARRAY_TASK_ID]}"
            sed -i -E "s/(^>[^ ]+).*/\1/" "$FILE"
        ' | fetch_job_ids
    ) )

    log "waiting for jobs: ${JOB_IDS[@]}"
    swait "${JOB_IDS[@]}"
}


function count_split_reads_quals()
{
    if [[ -d "$SPLIT_READS_DIR" ]]
    then
        find "$SPLIT_READS_DIR" -name '*.part_*.qual' | wc -l
    else
        echo 0
    fi
}


function need_fix_headers()
{
    if [[ -v SPLIT_READS ]]
    then
        builtin read -ra FILES <<<"$SPLIT_READS"
        head -q "${FILES[@]}" | grep -qF 'RQ='
    else
        dry_run || bail_out "unreachable"
    fi
}


function fake_read_quals()
{
    local JOB_IDS=( $(
        IFS="$IFS" sbatch --array=0-$(( NUM_SPLIT_READS - 1 )) --wrap='
            FILES=( $SPLIT_READS )
            FILE="${FILES[$SLURM_ARRAY_TASK_ID]}"
            fakeQuals.py "$FILE" "${FILE%.fasta}.qual" -s8
        ' | fetch_job_ids
    ) )

    log "waiting for jobs: ${JOB_IDS[@]}"
    swait "${JOB_IDS[@]}"
}


function prepare_workdir()
{
    if [[ ! -d "$WORKDIR" || ! -d "$LOGDIR" ]]
    then
        log "creating working directory structure"
        dry_run || mkdir -p "$WORKDIR" "$LOGDIR"
    fi

    WORK_ASSEMBLY="$WORKDIR/$(basename "$ASSEMBLY")"
    if [[ "$ASSEMBLY" -nt "$WORK_ASSEMBLY" ]]
    then
        log "copying assembly to working directory"
        dry_run || cp "$ASSEMBLY" "$WORKDIR"
    fi

    SPLIT_READS_DIR="$WORKDIR/$(basename "$READS" .fasta)"
    NUM_READS_PARTS=$(( $(get_num_reads_parts) ))

    if [[ ! -d "$SPLIT_READS_DIR" ]] || (( $(count_split_reads) != NUM_READS_PARTS ));
    then
        if [[ -d "$SPLIT_READS_DIR" ]]
        then
            log "deleting stale split reads directory"
            dry_run || rm -rf "$SPLIT_READS_DIR"
        fi

        log "splitting reads"
        dry_run || split_reads

    fi

    [[ ! -d "$SPLIT_READS_DIR" ]] || record_split_reads || bail_out "read splitting failed"

    if need_fix_headers
    then
        log "fixing read headers"
        dry_run || fix_read_headers
    fi

    if (( $(count_split_reads_quals) != NUM_READS_PARTS ))
    then
        log "faking read qualities"
        dry_run || fake_read_quals
    fi
}


function submit_stage()
{
    STAGE="$1"
    LOG="$LOGDIR/$STAGE.log"

    Jelly.py "$STAGE" "$PROTOCOL" --debug |& tee "$LOG"
}


function stage_failed()
{
    if [[ "${1:-}" == --silent ]]
    then
        return 1
    else
        bail_out "$STAGE failed: " "$@"
    fi
}


function enforce_stage_file()
{
    if [[ "$1" == --silent ]]
    then
        local FILE="$2"
        set -- "$1"
    else
        local FILE="$1"
    fi

    [[ -e "$FILE" ]] || stage_failed "$@" "missing output file: $FILE"
}


function verify_stage()
{
    case "$STAGE" in
        setup)
            enforce_stage_file "$@" "$WORK_ASSEMBLY.original"
            enforce_stage_file "$@" "$WORK_ASSEMBLY.sa"
            enforce_stage_file "$@" "${WORK_ASSEMBLY%.fasta}.gapInfo.bed"
            ;;
        mapping)
            local MAPPING_DIR="$WORKDIR/mapping"
            [[ -d "$MAPPING_DIR" ]] || stage_failed "$@" "mapping directory not created: $MAPPING_DIR" || return 1

            local NUM_MAPPINGS=$(( $(find "$MAPPING_DIR" -name '*.m4' | wc -l) ))

            (( NUM_MAPPINGS == NUM_READS_PARTS )) || stage_failed "$@" "missing read mappings: $NUM_MAPPINGS < $NUM_READS_PARTS"
            ;;
        support)
            local SUPPORT_DIR="$WORKDIR/support"
            [[ -d "$SUPPORT_DIR" ]] || stage_failed "$@" "support directory not created: $SUPPORT_DIR" || return 1

            local NUM_GRAPHS=$(( $(find "$SUPPORT_DIR" -name '*.gml' | wc -l) ))

            (( NUM_GRAPHS == NUM_READS_PARTS )) || stage_failed "$@" "missing support graphs: $NUM_GRAPHS < $NUM_READS_PARTS"
            ;;
        extraction)
            local ASSEMBLY_DIR="$WORKDIR/assembly"
            [[ -d "$ASSEMBLY_DIR" ]] || stage_failed "$@" "support directory not created: $ASSEMBLY_DIR" || return 1
            local NUM_EXTRACTED=$(( $(find "$ASSEMBLY_DIR" -name 'ref*' | wc -l) ))

            (( NUM_EXTRACTED > 50 )) || stage_failed "$@" "insufficient number of extracted sequences: $NUM_EXTRACTED"
            ;;
        assembly)
            local ASSEMBLY_DIR="$WORKDIR/assembly"
            [[ -d "$ASSEMBLY_DIR" ]] || stage_failed "$@" "support directory not created: $ASSEMBLY_DIR" || return 1

            local NUM_EXTRACTED=$(( $(find "$ASSEMBLY_DIR" -name 'ref*' | wc -l) ))
            local NUM_METRICS=$(( $(find "$ASSEMBLY_DIR" -name 'fillingMetrics.json' | wc -l) ))
            local THRESHOLD=$(( NUM_EXTRACTED - NUM_EXTRACTED/10 ))

            (( NUM_METRICS >= THRESHOLD )) || stage_failed "$@" "missing fillingMetrics: $NUM_METRICS < $THRESHOLD"
            ;;
        output)
            enforce_stage_file "$@" "$WORKDIR/jelly.out.fasta"
            enforce_stage_file "$@" "$WORKDIR/jelly.out.qual"
            enforce_stage_file "$@" "$WORKDIR/liftOverTable.json"
            enforce_stage_file "$@" "$WORKDIR/gap_fill_status.txt"
            ;;
    esac
}


function __invalidate_stage()
{
    case "$STAGE" in
        setup)
            rm -f \
                "$WORK_ASSEMBLY" \
                "$WORK_ASSEMBLY.original" \
                "$WORK_ASSEMBLY.sa" \
                "${WORK_ASSEMBLY%.fasta}.gapInfo.bed"
            ;;
        mapping)
            rm -rf "$WORKDIR/mapping"
            ;;
        support)
            rm -rf "$WORKDIR/support"
            ;;
        extraction)
            find "$WORKDIR/assembly" -name 'ref*' -delete
            ;;
        assembly)
            find "$WORKDIR/assembly" -name 'fillingMetrics.json' -delete
            ;;
        output)
            rm -f \
                "$WORKDIR/jelly.out.fasta" \
                "$WORKDIR/jelly.out.qual" \
                "$WORKDIR/liftOverTable.json" \
                "$WORKDIR/gap_fill_status.txt"
            ;;
    esac
}


function invalidate_stage()
{
    for LIST_STAGE in "${STAGES[@]}"
    do
        if [[ "$LIST_STAGE" == "$STAGE" ]]
        then
            IS_INVALIDATING=1
        fi

        [[ -v IS_INVALIDATING ]] && __invalidate_stage "$LIST_STAGE"
    done
}


function record_jobs()
{
    JOB_IDS="$(comma_list "$@")" "$PBGAPS/source/record-jobs.sh" --tool=pbjelly --run-id=auto
}


function delete_jobs()
{
    "$PBGAPS/source/delete-jobs.sh" "$@"
}


function current_stage_job_ids()
{
    export RUN_ID
    RUN_ID="$("$PBGAPS/source/record-jobs.sh" --tool=pbjelly --run-id=auto --print-run-id)"
    awk -F'\t' '
        BEGIN {
            dataset = ENVIRON["DATASET"];
            stage = ENVIRON["STAGE"];
            run_id = ENVIRON["RUN_ID"];
        }

        ($1 == dataset && $3 == "pbjelly" && $4 == stage && $5 == run_id) {
            print $6
        }
    ' "$JOBS_DB"
}


function is_running_stage()
{
    (( $(current_stage_job_ids | wc -l) > 0 ))
}


function run_stage()
{
    if ! verify_stage --silent
    then
        log "running $STAGE"

        if ! dry_run
        then
            if is_running_stage
            then
                log "resuming $STAGE"
                JOB_IDS=( $(current_stage_job_ids) )
            else
                submit_stage "$STAGE"
                JOB_IDS=( $(fetch_job_ids < "$LOGDIR/$STAGE.log") )

                (( ${#JOB_IDS[*]} > 0 )) || bail_out "submission failed"

                record_jobs "${JOB_IDS[@]}"
            fi

            log "waiting for $STAGE: ${JOB_IDS[@]}"
            if ! swait "${JOB_IDS[@]}"
            then
                delete_jobs "${JOB_IDS[@]}"
                bail_out "$STAGE failed; please investigate"
            fi
            log "$STAGE finished"

            if ! verify_stage
            then
                delete_jobs "${JOB_IDS[@]}"
                bail_out "$STAGE failed; please investigate"
            fi

            log "validated stage $STAGE"
        fi
    else
        log "skipping $STAGE: already done"
    fi
}


function main()
{
    set_defaults
    parse_args "$@"
    dry_run && log "this is a dry run; not performing any changes"
    setup_variables
    export DATASET
    export STAGE
    log "running pipeline for $DATASET with $READS_TYPE reads"
    prepare_workdir

    for STAGE in "${STAGES[@]}"
    do
        run_stage
    done
}


PIPELINE_LOG="run.$(date -Iseconds).log"
main "$@" |& tee "$PIPELINE_LOG"
