#!/usr/bin/env nextflow
/*
========================================================================================
                         nf-core/neutronstar
========================================================================================
 nf-core/neutronstar Analysis Pipeline.
 #### Homepage / Documentation
 https://github.com/nf-core/neutronstar
----------------------------------------------------------------------------------------
*/

def helpMessage() {
    log.info nfcoreHeader()
    log.info"""

    Usage:

    The typical command for running the pipeline is as follows:


    nextflow run -profile singularity nf-core/neutronstar --id assembly_id --fastqs fastq_path --genomesize 1000000

    Mandatory arguments:
      --id                          [Supernova parameter] A unique run id, used to name output folder [a-zA-Z0-9_-]+.
      --fastqs                      [Supernova parameter] Path of folder created by mkfastq or bcl2fastq.
      --genomesize                  The estimated size of the genome(s) to be assembled. This is mainly used by Quast to compute NGxx statstics, e.g. N50 statistics bound by this value and not the assembly size.
      -profile                      Configuration profile to use
      --busco_data                  The dataset BUSCO should use (e.g. eukaryota_odb9, protists_ensembl)

    Options:
      --sample                      [Supernova parameter] Prefix of the filenames of FASTQs to select.
      --lanes                       [Supernova parameter] Comma-separated lane numbers.
      --indices                     [Supernova parameter] Comma-separated sample index set "SI-001" or sequences.
      --bcfrac                      [Supernova parameter] Fraction of barcodes in the sample to use.
      --project                     [Supernova parameter] Name of the project folder within a mkfastq or bcl2fastq-generated folder to pick FASTQs from.
      --maxreads                    [Supernova parameter] Downsample if more than NUM individual reads are provided or 'all' to use all reads provided (default='all')
      --nopreflight                 [Supernova parameter] Skip preflight checks.
      --no_accept_extreme_coverage  Enables input coverage validation in Supernova (i.e. disables --accept-extreme-coverage)
      --minsize                     [Supernova mkoutput parameter]  Skip FASTA records shorter than NUM. (default=1000)
      --max_cpus                    Max amount of cpu cores for the job scheduler to request. Supernova will use all of them. (default=16)
      --max_memory                  Max amount of memory (in Gb) for the jobscheduler to request. Supernova will use all of it. (default=256)
      --max_time                    Max amount of time for the job scheduler to request (in hours). (default=120)
      --full_output                 Keep all the files that are output from Supernova. By default only the final assembly graph is kept, as it is needed to make the output fasta files.
      --clusterOptions              The options to feed to the HPC job manager. For instance for SLURM --clusterOptions='-A project -C node-type'
      --busco_folder                Path to directory containing BUSCO datasets (default=$baseDir/data)


    Other options:
      --outdir                      The output directory where the results will be saved
      --email                       Set this parameter to your e-mail address to get a summary e-mail with details of the run sent to you when the workflow exits
      -name                         Name for the pipeline run. If not specified, the pipeline will automatically generate one from the uuid (i.e. nonsense).
      -params-file                  Give the arguments for this nextflow run as a structured JSON/YAML file
      --email_on_fail               Same as --email, except only send mail if the workflow is not successful
      --maxMultiqcEmailFileSize     Theshold size for MultiQC report to be attached in notification email. If file generated by pipeline exceeds the threshold, it will not be attached (Default: 25MB)
      -name                         Name for the pipeline run. If not specified, Nextflow will automatically generate a random mnemonic.
    """.stripIndent()
}

/*
 * SET UP CONFIGURATION VARIABLES
 */

// Show help emssage
params.help = false
if (params.help){
    helpMessage()
    exit 0
}

/*
 * SET UP CONFIGURATION VARIABLES
 */

// Configurable variables
params.name = false
params.multiqc_config = "$baseDir/conf/multiqc_config.yaml"
params.email = false
params.plaintext_email = false

multiqc_config = file(params.multiqc_config)
output_docs = file("$baseDir/docs/output.md")
def buscoPath = file("${params.busco_folder}/${params.busco_data}")

// Has the run name been specified by the user?
//  this has the bonus effect of catching both -name and --name
custom_runName = params.name

if( !(workflow.runName ==~ /[a-z]+_[a-z]+/) ){
  custom_runName = workflow.runName
}
else {
  custom_runName = "supernova_assembly_${workflow.sessionId}"
}

// Check workDir/outdir paths to be S3 buckets if running on AWSBatch
// related: https://github.com/nextflow-io/nextflow/issues/813
if( workflow.profile == 'awsbatch') {
    if(!workflow.workDir.startsWith('s3:') || !params.outdir.startsWith('s3:')) exit 1, "Workdir or Outdir not on S3 - specify S3 Buckets for each to run on AWSBatch!"
}

// Common options for both supernova and longranger
def tenx_optional = {sample, lanes, indices, project ->
    def str = ""
    str += sample ? "--sample=${sample} " : ""
    str += lanes ? "--lanes=${lanes} " : ""
    str += indices ? "--indices=${indices} " : ""
    str += project ?  "--project=${project} " : ""
    return str
}
// Now only supernova options
def supernova_optional = {maxreads, bcfrac, nopreflight, no_accept_extreme_coverage ->
    def str = ""
    str += maxreads ?  "--maxreads=${maxreads} " : "--maxreads=all "
    str += bcfrac ? "--bcfrac=${bcfrac} " : ""
    str += nopreflight ? "--nopreflight " : ""
    str += no_accept_extreme_coverage ? "" : "--accept-extreme-coverage "
    return str
}

def samples = []
if (params.samples == null) { //We don't have sample JSON/YAML file, just use cmd-line
    assert params.id : "Missing --id option"
    assert params.fastqs : "Missing --fastqs option"
    s = []
    s += params.id
    s += params.fastqs
    s += tenx_optional(params.sample, params.lanes, params.indices, params.project)
    s += supernova_optional(params.maxreads, params.bcfrac, params.nopreflight, params.no_accept_extreme_coverage)
    samples << s
}

params.samples = []

for (sample in params.samples) {
    assert sample.id : "Error in input parameter file"
    assert sample.fastqs : "Error in input parameter file"
    s = []
    s += sample.id
    s += sample.fastqs
    s += tenx_optional(sample.sample, sample.lanes, sample.indices, sample.project)
    s += supernova_optional(sample.maxreads, sample.bcfrac, sample.nopreflight, sample.no_accept_extreme_coverage)
    samples << s
}

//Extra parameter evaluation
for (sample in params.samples) {
    assert sample.id.size() == sample.id.findAll(/[A-z0-9\\.\-]/).size() : "Illegal character(s) in sample ID: ${sample.id}."
    for(folder in sample.fastqs.split(",")) {
      def f = new File(folder)
      assert f.exists() : "Non-existent folder specified for --fastqs: ${folder}"
    }
}

if (!(workflow.runName ==~ /[a-z]+_[a-z]+/)) {
  custom_runName = workflow.runName
}

if ( workflow.profile == 'awsbatch') {
  // AWSBatch sanity checking
  if (!params.awsqueue || !params.awsregion) exit 1, "Specify correct --awsqueue and --awsregion parameters on AWSBatch!"
  // Check outdir paths to be S3 buckets if running on AWSBatch
  // related: https://github.com/nextflow-io/nextflow/issues/813
  if (!params.outdir.startsWith('s3:')) exit 1, "Outdir not on S3 - specify S3 Bucket to run on AWSBatch!"
  // Prevent trace files to be stored on S3 since S3 does not support rolling files.
  if (workflow.tracedir.startsWith('s3:')) exit 1, "Specify a local tracedir or run without trace! S3 cannot be used for tracefiles."
}

// Stage config files
ch_multiqc_config = file(params.multiqc_config, checkIfExists: true)
ch_output_docs = file("$baseDir/docs/output.md", checkIfExists: true)

def no_samples = 0
for (i in samples) { no_samples += 1 }

// Header log info
log.info nfcoreHeader()
def summary = [:]
if (workflow.revision) summary['Pipeline Release'] = workflow.revision
summary['Run Name']         = custom_runName ?: workflow.runName
summary['# samples']        = no_samples
summary['Genome size']      = params.genomesize
summary['clusterOptions']   = params.clusterOptions
summary['Minsize']          = params.minsize
summary['Full output']      = params.full_output
summary['Max Resources']    = "$params.max_memory memory, $params.max_cpus cpus, $params.max_time time per job"
if (workflow.containerEngine) summary['Container'] = "$workflow.containerEngine - $workflow.container, $params.supernova_container"
summary['Output dir']       = params.outdir
summary['Launch dir']       = workflow.launchDir
summary['Working dir']      = workflow.workDir
summary['Script dir']       = workflow.projectDir
summary['User']             = workflow.userName
if (workflow.profile == 'awsbatch') {
  summary['AWS Region']     = params.awsregion
  summary['AWS Queue']      = params.awsqueue
}
summary['Config Profile'] = workflow.profile
if (params.config_profile_description) summary['Config Description'] = params.config_profile_description
if (params.config_profile_contact)     summary['Config Contact']     = params.config_profile_contact
if (params.config_profile_url)         summary['Config URL']         = params.config_profile_url
if (params.email || params.email_on_fail) {
  summary['E-mail Address']    = params.email
  summary['E-mail on failure'] = params.email_on_fail
  summary['MultiQC maxsize']   = params.maxMultiqcEmailFileSize
}
log.info summary.collect { k,v -> "${k.padRight(18)}: $v" }.join("\n")
log.info "-\033[2m--------------------------------------------------\033[0m-"

// Check the hostnames against configured profiles
checkHostname()

def create_workflow_summary(summary) {
    def yaml_file = workDir.resolve('workflow_summary_mqc.yaml')
    yaml_file.text  = """
    id: 'nf-core-neutronstar-summary'
    description: " - this information is collected when the pipeline is started."
    section_name: 'nf-core/neutronstar Workflow Summary'
    section_href: 'https://github.com/nf-core/neutronstar'
    plot_type: 'html'
    data: |
        <dl class=\"dl-horizontal\">
${summary.collect { k,v -> "            <dt>$k</dt><dd><samp>${v ?: '<span style=\"color:#999999;\">N/A</a>'}</samp></dd>" }.join("\n")}
        </dl>
    """.stripIndent()

   return yaml_file
}

Channel
    .from(samples)
    .set { supernova_input }

if (params.full_output) {
    process supernova_full {
        tag "${id}"
        label "supernova"
        publishDir "${params.outdir}/supernova/", mode: 'copy'

        input:
        set val(id), val(fastqs), val(tenx_options), val(supernova_options) from supernova_input

        output:
        set val(id), file("${id}/*") into supernova_results, supernova_results2

        script:
        """
        supernova run --id=${id} --fastqs=${fastqs} ${tenx_options} ${supernova_options}
        """
    }

} else {
    process supernova {
        tag "${id}"
        label "supernova"
        publishDir "${params.outdir}/supernova/", mode: 'copy'

        input:
        set val(id), val(fastqs), val(tenx_options), val(supernova_options) from supernova_input

        output:
        set val(id), file("${id}_supernova") into supernova_results, supernova_results2

        script:
        // We use rsync to fish out only the important assembly graph files for
        // supernova mkoutput. Otherwise this folder would be approx. ten times the
        // size of the input fastqs.
        """
        supernova run --id=${id} --fastqs=${fastqs} ${tenx_options} ${supernova_options}
        rsync -rav --include="_*" --include="*.tgz" --include="outs/" --include="outs/*.*" \
          --include="assembly/" --include="stats/***" --include="logs/***" --include="a.base/" \
          --include="a.base/" --include="a.hbx" --include="a.inv" --include="final/***" --include="gang" \
          --include="micro"  --include="a.hbx" --include="a.inv" --include="final/***" \
          --exclude="*" "${id}/" ${id}_supernova
        """
    }

}

process supernova_mkoutput {
    tag "${id}"
    label "supernova"
    publishDir "${params.outdir}/assemblies/", mode: 'copy'

    input:
    set val(id), file(supernova_folder) from supernova_results

    output:
    set val(id), file("${id}.fasta") into supernova_asm1, supernova_asm2
    file "${id}.phased.fasta"

    script:
    """
    supernova mkoutput --asmdir=${id}_supernova/outs/assembly --outprefix=${id} --style=pseudohap --minsize=${params.minsize} --nozip
    supernova mkoutput --asmdir=${id}_supernova/outs/assembly --outprefix=${id}.phased --style=megabubbles --minsize=${params.minsize} --nozip
    """
}

process quast {
    tag "${id}"
    publishDir "${params.outdir}/quast/${id}", mode: 'copy'

    input:
    set val(id), file(asm) from supernova_asm1

    output:
    file("quast_results/latest/*") into quast_results

    script:
    def size_parameter = params.genomesize ? "--est-ref-size ${params.genomesize}" : ""
    """
    quast.py ${size_parameter} --threads ${task.cpus} ${asm}
    """
}


process busco {
    tag "${id}"
    publishDir "${params.outdir}/busco/", mode: 'copy'

    input:
    set val(id), file(asm) from supernova_asm2
    file(augustus_archive) from Channel.fromPath("$baseDir/misc/augustus_config.tar.bz2")

    output:
    file ("run_${id}/*.{txt,tsv}") into busco_results

    script:
    """
    tar xfj ${augustus_archive}
    export AUGUSTUS_CONFIG_PATH=augustus_config/
    run_BUSCO.py -i ${asm} -o ${id} -c ${task.cpus} -m genome -l ${buscoPath}
    """
}

process get_supernova_version {
    label "supernova"
    output:
    file("v_supernova.txt") into v_supernova

    script:
    """
    supernova run --version > v_supernova.txt
    """

}

process get_software_versions {
    publishDir "${params.outdir}/pipeline_info", mode: 'copy',
        saveAs: { filename ->
            if (filename.indexOf(".csv") > 0) filename
            else null
        }
    input:
      file "v_supernova.txt" from v_supernova
    output:
      file "software_versions_mqc.yaml" into software_versions_yaml
      file "software_versions.csv"

    script:
    """
    echo $workflow.manifest.version > v_pipeline.txt
    echo $workflow.nextflow.version > v_nextflow.txt
    quast.py -v &> v_quast.txt
    multiqc --version > v_multiqc.txt
    run_BUSCO.py -v > v_busco.txt
    scrape_software_versions.py &> software_versions_mqc.yaml
    """
}

process multiqc {
    publishDir "${params.outdir}/multiqc", mode: 'copy'

    input:
    file ('supernova/') from supernova_results2.collect()
    file ('busco/') from busco_results.collect()
    file ('quast/') from quast_results.collect()
    file ('software_versions/') from software_versions_yaml.toList()
    file(mqc_config) from Channel.fromPath("${params.multiqc_config}")

    output:
    file "*multiqc_report.html"
    file "*_data"

    script:
    """
    multiqc -i ${custom_runName} -f -s  --config ${mqc_config} .
    """
}

process output_documentation {
    publishDir "${params.outdir}/pipeline_info", mode: 'copy'

    input:
    file output_docs from ch_output_docs

    output:
    file "results_description.html"

    script:
    """
    markdown_to_html.r $output_docs results_description.html
    """
}


/*
 * Completion e-mail notification
 */
workflow.onComplete {

    // Set up the e-mail variables
    def subject = "[nf-core/neutronstar] Successful: $workflow.runName"
    if (!workflow.success) {
      subject = "[nf-core/neutronstar] FAILED: $workflow.runName"
    }
    def email_fields = [:]
    email_fields['version'] = workflow.manifest.version
    email_fields['runName'] = custom_runName ?: workflow.runName
    email_fields['success'] = workflow.success
    email_fields['dateComplete'] = workflow.complete
    email_fields['duration'] = workflow.duration
    email_fields['exitStatus'] = workflow.exitStatus
    email_fields['errorMessage'] = (workflow.errorMessage ?: 'None')
    email_fields['errorReport'] = (workflow.errorReport ?: 'None')
    email_fields['commandLine'] = workflow.commandLine
    email_fields['projectDir'] = workflow.projectDir
    email_fields['summary'] = summary
    email_fields['summary']['Date Started'] = workflow.start
    email_fields['summary']['Date Completed'] = workflow.complete
    email_fields['summary']['Pipeline script file path'] = workflow.scriptFile
    email_fields['summary']['Pipeline script hash ID'] = workflow.scriptId
    if (workflow.repository) email_fields['summary']['Pipeline repository Git URL'] = workflow.repository
    if (workflow.commitId) email_fields['summary']['Pipeline repository Git Commit'] = workflow.commitId
    if (workflow.revision) email_fields['summary']['Pipeline Git branch/tag'] = workflow.revision
    if (workflow.container) email_fields['summary']['Docker image'] = workflow.container
    email_fields['summary']['Nextflow Version'] = workflow.nextflow.version
    email_fields['summary']['Nextflow Build'] = workflow.nextflow.build
    email_fields['summary']['Nextflow Compile Timestamp'] = workflow.nextflow.timestamp

    // On success try attach the multiqc report
    def mqc_report = null
    try {
        if (workflow.success) {
            mqc_report = multiqc_report.getVal()
            if (mqc_report.getClass() == ArrayList) {
                log.warn "[nf-core/neutronstar] Found multiple reports from process 'multiqc', will use only one"
                mqc_report = mqc_report[0]
            }
        }
    } catch (all) {
        log.warn "[nf-core/neutronstar] Could not attach MultiQC report to summary email"
    }

    // Check if we are only sending emails on failure
    email_address = params.email
    if (!params.email && params.email_on_fail && !workflow.success) {
        email_address = params.email_on_fail
    }

    // Render the TXT template
    def engine = new groovy.text.GStringTemplateEngine()
    def tf = new File("$baseDir/assets/email_template.txt")
    def txt_template = engine.createTemplate(tf).make(email_fields)
    def email_txt = txt_template.toString()

    // Render the HTML template
    def hf = new File("$baseDir/assets/email_template.html")
    def html_template = engine.createTemplate(hf).make(email_fields)
    def email_html = html_template.toString()

    // Render the sendmail template
    def smail_fields = [ email: email_address, subject: subject, email_txt: email_txt, email_html: email_html, baseDir: "$baseDir", mqcFile: mqc_report, mqcMaxSize: params.maxMultiqcEmailFileSize.toBytes() ]
    def sf = new File("$baseDir/assets/sendmail_template.txt")
    def sendmail_template = engine.createTemplate(sf).make(smail_fields)
    def sendmail_html = sendmail_template.toString()

    // Send the HTML e-mail
    if (email_address) {
        try {
          if ( params.plaintext_email ){ throw GroovyException('Send plaintext e-mail, not HTML') }
          // Try to send HTML e-mail using sendmail
          [ 'sendmail', '-t' ].execute() << sendmail_html
          log.info "[nf-core/neutronstar] Sent summary e-mail to $email_address (sendmail)"
        } catch (all) {
          // Catch failures and try with plaintext
          [ 'mail', '-s', subject, email_address ].execute() << email_txt
          log.info "[nf-core/neutronstar] Sent summary e-mail to $email_address (mail)"
        }
    }

    // Write summary e-mail HTML to a file
    def output_d = new File( "${params.outdir}/pipeline_info/" )
    if (!output_d.exists()) {
      output_d.mkdirs()
    }
    def output_hf = new File( output_d, "pipeline_report.html" )
    output_hf.withWriter { w -> w << email_html }
    def output_tf = new File( output_d, "pipeline_report.txt" )
    output_tf.withWriter { w -> w << email_txt }

    c_reset = params.monochrome_logs ? '' : "\033[0m";
    c_purple = params.monochrome_logs ? '' : "\033[0;35m";
    c_green = params.monochrome_logs ? '' : "\033[0;32m";
    c_red = params.monochrome_logs ? '' : "\033[0;31m";

    if (workflow.stats.ignoredCount > 0 && workflow.success) {
      log.info "${c_purple}Warning, pipeline completed, but with errored process(es) ${c_reset}"
      log.info "${c_red}Number of ignored errored process(es) : ${workflow.stats.ignoredCount} ${c_reset}"
      log.info "${c_green}Number of successfully ran process(es) : ${workflow.stats.succeedCount} ${c_reset}"
    }

    if (workflow.success) {
        log.info "${c_purple}[nf-core/neutronstar]${c_green} Pipeline completed successfully${c_reset}"
    } else {
        checkHostname()
        log.info "${c_purple}[nf-core/neutronstar]${c_red} Pipeline completed with errors${c_reset}"
    }

}


def nfcoreHeader(){
    // Log colors ANSI codes
    c_reset = params.monochrome_logs ? '' : "\033[0m";
    c_dim = params.monochrome_logs ? '' : "\033[2m";
    c_black = params.monochrome_logs ? '' : "\033[0;30m";
    c_green = params.monochrome_logs ? '' : "\033[0;32m";
    c_yellow = params.monochrome_logs ? '' : "\033[0;33m";
    c_blue = params.monochrome_logs ? '' : "\033[0;34m";
    c_purple = params.monochrome_logs ? '' : "\033[0;35m";
    c_cyan = params.monochrome_logs ? '' : "\033[0;36m";
    c_white = params.monochrome_logs ? '' : "\033[0;37m";

    return """    -${c_dim}--------------------------------------------------${c_reset}-
                                            ${c_green},--.${c_black}/${c_green},-.${c_reset}
    ${c_blue}        ___     __   __   __   ___     ${c_green}/,-._.--~\'${c_reset}
    ${c_blue}  |\\ | |__  __ /  ` /  \\ |__) |__         ${c_yellow}}  {${c_reset}
    ${c_blue}  | \\| |       \\__, \\__/ |  \\ |___     ${c_green}\\`-._,-`-,${c_reset}
                                            ${c_green}`._,._,\'${c_reset}
    ${c_purple}  nf-core/neutronstar v${workflow.manifest.version}${c_reset}
    -${c_dim}--------------------------------------------------${c_reset}-
    """.stripIndent()
}

def checkHostname(){
    def c_reset = params.monochrome_logs ? '' : "\033[0m"
    def c_white = params.monochrome_logs ? '' : "\033[0;37m"
    def c_red = params.monochrome_logs ? '' : "\033[1;91m"
    def c_yellow_bold = params.monochrome_logs ? '' : "\033[1;93m"
    if (params.hostnames) {
        def hostname = "hostname".execute().text.trim()
        params.hostnames.each { prof, hnames ->
            hnames.each { hname ->
                if (hostname.contains(hname) && !workflow.profile.contains(prof)) {
                    log.error "====================================================\n" +
                            "  ${c_red}WARNING!${c_reset} You are running with `-profile $workflow.profile`\n" +
                            "  but your machine hostname is ${c_white}'$hostname'${c_reset}\n" +
                            "  ${c_yellow_bold}It's highly recommended that you use `-profile $prof${c_reset}`\n" +
                            "============================================================"
                }
            }
        }
    }
}
