#!/usr/bin/env python
"""

    pipeline_multisample.py
                        [--bamdir PATH]
                        [--groups INT]
                        [--log_file PATH]
                        [--verbose]
                        [--target_tasks]
                        [--jobs]
                        [--just_print]
                        [--flowchart]
                        [--key_legend_in_graph]
                        [--forced_tasks]

"""
import sys
import os
import glob


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#   options


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888


if __name__ == '__main__':
    from optparse import OptionParser
    import StringIO

    parser = OptionParser(version="%prog 1.0", usage = "\n\n    %prog --settings PIPELINE_SETTINGS.CFG --groups NUMBER [more_options]")
    parser.add_option("-s", "--settings", dest="pipeline_settings",
                        metavar="FILE",
                        type="string",
                        help="File containing all the settings for the analysis.")                  
                            
    parser.add_option("-g", "--groups", dest="groups",
                        type="int",
                        default=1,
                        help="Split dataset into smaller groups for multisample snp caling.")



    #
    #   general options: verbosity / logging
    #
    parser.add_option("-v", "--verbose", dest = "verbose",
                      action="count", default=0,
                      help="Print more verbose messages for each additional verbose level.")
    parser.add_option("-L", "--log_file", dest="log_file",
                      metavar="FILE",
                      type="string",
                      help="Name and path of log file")




    #
    #   pipeline
    #
    parser.add_option("-t", "--target_tasks", dest="target_tasks",
                        action="append",
                        default = list(),
                        metavar="JOBNAME",
                        type="string",
                        help="Target task(s) of pipeline.")
    parser.add_option("-j", "--jobs", dest="jobs",
                        default=1,
                        metavar="N",
                        type="int",
                        help="Allow N jobs (commands) to run simultaneously.")
    parser.add_option("-n", "--just_print", dest="just_print",
                        action="store_true", default=False,
                        help="Don't actually run any commands; just print the pipeline.")
    parser.add_option("--flowchart", dest="flowchart",
                        metavar="FILE",
                        type="string",
                        help="Don't actually run any commands; just print the pipeline "
                             "as a flowchart.")

    #
    #   Less common pipeline options
    #
    parser.add_option("--key_legend_in_graph", dest="key_legend_in_graph",
                        action="store_true", default=False,
                        help="Print out legend and key for dependency graph.")
    parser.add_option("--forced_tasks", dest="forced_tasks",
                        action="append",
                        default = list(),
                        metavar="JOBNAME",
                        type="string",
                        help="Pipeline task(s) which will be included even if they are up to date.")
    parser.add_option("--rebuild_mode", dest="rebuild_mode",
                        action="store_false", default=True,
                        help="gnu_make_maximal_rebuild_mode")

    # get help string
    f =StringIO.StringIO()
    parser.print_help(f)
    helpstr = f.getvalue()
    (options, remaining_args) = parser.parse_args()


    #vvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvv
    #                                             #
    #   Change this if necessary                  #
    #                                             #
    #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

    #
    #   Add names of mandatory options,
    #       strings corresponding to the "dest" parameter
    #       in the options defined above
    #
    mandatory_options = ['pipeline_settings']

    #vvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvv
    #                                             #
    #   Change this if necessary                  #
    #                                             #
    #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^


    def check_mandatory_options (options, mandatory_options, helpstr):
        """
        Check if specified mandatory options have been defined
        """
        missing_options = []
        for o in mandatory_options:
            if not getattr(options, o):
                missing_options.append("--" + o)
    
        if not len(missing_options):
            return
    
        raise Exception("Missing mandatory parameter%s: %s.\n\n%s\n\n" %
                        ("s" if len(missing_options) > 1 else "",
                         ", ".join(missing_options),
                         helpstr))
        
    check_mandatory_options (options, mandatory_options, helpstr)


    #vvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvv
    #                                             #
    #   Get pipeline settings from a config file  #
    #                                             #
    #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

    import ConfigParser

    config = ConfigParser.ConfigParser()
    config.read(options.pipeline_settings)
    # inputs 
    input_bams = config.get('Inputs','input-bams')
    # variant calls from other projects to call together with (semicolon separated list)
    try: 
        call_with_gvcfs = [ path.strip() for path in config.get('Inputs','call-with-gvcfs').split(";") ]
	#print 'Calling will be done together with:'
	#for p in call_with_gvcfs:
	#	print '\t',p
    except (ConfigParser.NoOptionError): 
        call_with_gvcfs = [] 
    
    # reference files
    reference = config.get('Resources','reference-genome')
    dbsnp = config.get('Resources','dbsnp-vcf')
    hapmap = config.get('Resources','hapmap-vcf')
    omni = config.get('Resources','1000genomes-omni-vcf')
    snps_1kg = config.get('Resources','1000genomes-snps-vcf')
    indels_1kg = config.get('Resources','1000genomes-indels-vcf')
    mills = config.get('Resources','mills-indels-vcf')
    capture = config.get('Resources','capture-regions-bed')
    capture_qualimap = config.get('Resources','capture-regions-bed-for-qualimap')
    exome = config.get('Resources', 'exome-regions-bed')
    # tools 
    java = config.get('Tools','java-binary')
    java_with_params = java
    picard = config.get('Tools','picard-tools-path')
    qualimap = config.get('Tools','qualimap')
    gatk = config.get('Tools','gatk-jar')
    vcftools = config.get('Tools','vcftools')
    # other
    n_cpus = config.getint('Other','n-cpus')





#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#   imports


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

from ruffus import *
import subprocess
# import drmaa
import resource

#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#   Functions 


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888


def setlimits():
    """Set maximum meomory to be used by a child process"""
    resource.setrlimit(resource.RLIMIT_AS, (100000000000,100000000000))

def run_cmd(cmd_str):
    """
    Throw exception if run command fails
    """
    process = subprocess.Popen(cmd_str, stdout = subprocess.PIPE,
                                stderr = subprocess.PIPE, shell = True, preexec_fn=setlimits)
    stdout_str, stderr_str = process.communicate()
    if process.returncode != 0:
        raise Exception("Failed to run '%s'\n%s%sNon-zero exit status %s" %
                            (cmd_str, stdout_str, stderr_str, process.returncode))

def rename(old_file, new_file):
    """rename file"""
    os.rename(old_file, new_file)

def remove(f):
    """remove file"""
    os.remove(f)
                            
def index_bam(bam):
    """Use samtools to create an index for the bam file"""
    run_cmd('samtools index %s' % bam)
                          
def bam_quality_score_distribution(bam,qs,pdf):
    """Calculates quality score distribution histograms"""
    run_cmd("java -jar {picard}/QualityScoreDistribution.jar \
             CHART_OUTPUT={chart} \
             OUTPUT={output} \
             INPUT={bam} \
             VALIDATION_STRINGENCY=SILENT".format(
                 picard=picard,
                 chart=pdf,
                 output=qs,
                 bam=bam
             ))

def bam_alignment_metrics(bam,metrics):
    """Collects alignment metrics for a bam file"""
    run_cmd("java -jar {picard}/CollectAlignmentSummaryMetrics.jar \
             REFERENCE_SEQUENCE={reference} \
             OUTPUT={output} \
             INPUT={bam} \
             VALIDATION_STRINGENCY=SILENT".format(
                 picard=picard,
                 reference=reference,
                 output=metrics,
                 bam=bam
             ))
    
def bam_coverage_metrics(input_bam, output):
    """ Calculates and outputs bam coverage statistics """
    run_cmd("{java} -Xmx4g -jar {gatk} \
            -R {reference} \
            -T DepthOfCoverage \
            -o {output} \
            -I {input} \
            -L {capture} \
            -ct 8 -ct 20 -ct 30 \
            --omitDepthOutputAtEachBase --omitLocusTable \
            ".format(java=java, gatk=gatk,
                reference=reference,
                output=output,
                input=input_bam,
                capture=capture
            ))

def qualimap_bam(input_bam, output_dir):
    """ Generates Qualimap bam QC report """
    # create necessary folders first
    if not os.path.exists('qc'): os.mkdir('qc')
    if not os.path.exists('qc/qualimap/'): os.mkdir('qc/qualimap')
    if not os.path.exists(output_dir): os.mkdir(output_dir)
    run_cmd("{qualimap} bamqc -bam {bam} \
             -c -outformat PDF \
             -gff {target} \
             -gd HUMAN -os \
             -outdir {dir} &> {dir}/qualimap.err".format(
                qualimap=qualimap,
                bam=input_bam,
                target=capture_qualimap,
                dir=output_dir))

    
    
#
#
# The two functions below seem redundant
#
#    
#            
#def merge_indel_and_snp_vcf(snp,indel, output):
#    """Merges vcf files from the batch run"""
#    run_cmd("{java} -Xmx4g -jar {gatk} \
#            -R {reference} \
#            -T CombineVariants \
#            -V:SNP {snp} \
#            -V:INDEL {indel} \
#            -o {output}".format(
#                java=java,
#                gatk=gatk,
#                reference=reference,
#                snp=snp,
#                indel=indel,
#                output=output
#            )) 
#
#def variant_annotator(bam, vcf, output):
#    run_cmd("{java} -Xmx16g -jar {gatk} \
#            -R {reference} \
#            -T VariantAnnotator \
#            -I {bam} \
#            -o {output} \
#            -A Coverage \
#            --variant {vcf} \
#            --dbsnp {dbsnp} \
#            ".format(java=java, gatk=gatk,
#                bam=bam,
#                reference=reference,
#                output=output,
#                vcf=vcf,
#                dbsnp=dbsnp
#            )) 


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#   Logger


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

if __name__ == '__main__':
    import logging
    import logging.handlers

    MESSAGE = 15
    logging.addLevelName(MESSAGE, "MESSAGE")

    def setup_std_logging (logger, log_file, verbose):
        """
        set up logging using programme options
        """
        class debug_filter(logging.Filter):
            """
            Ignore INFO messages
            """
            def filter(self, record):
                return logging.INFO != record.levelno

        class NullHandler(logging.Handler):
            """
            for when there is no logging
            """
            def emit(self, record):
                pass

        # We are interesting in all messages
        logger.setLevel(logging.DEBUG)
        has_handler = False

        # log to file if that is specified
        if log_file:
            handler = logging.FileHandler(log_file, delay=False)
            handler.setFormatter(logging.Formatter("%(asctime)s - %(name)s - %(levelname)6s - %(message)s"))
            handler.setLevel(MESSAGE)
            logger.addHandler(handler)
            has_handler = True

        # log to stderr if verbose
        if verbose:
            stderrhandler = logging.StreamHandler(sys.stderr)
            stderrhandler.setFormatter(logging.Formatter("    %(message)s"))
            stderrhandler.setLevel(logging.DEBUG)
            if log_file:
                stderrhandler.addFilter(debug_filter())
            logger.addHandler(stderrhandler)
            has_handler = True

        # no logging
        if not has_handler:
            logger.addHandler(NullHandler())


    #
    #   set up log
    #
    module_name = "exome"
    logger = logging.getLogger(module_name)
    setup_std_logging(logger, options.log_file, options.verbose)

    #
    #   Allow logging across Ruffus pipeline
    #
    def get_logger (logger_name, args):
        return logger

    from ruffus.proxy_logger import *
    (logger_proxy,
     logging_mutex) = make_shared_logger_and_proxy (get_logger,
                                                    module_name,
                                                    {})


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#   Pipeline


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#       Put pipeline code here

def get_sample_ids():
    files = glob.glob(input_bams)
    return [ os.path.splitext(os.path.basename(f))[0] for f in files ]

def generate_parameters():
    files = glob.glob(input_bams)
    parameters = []
    for f in files:
        prefix = os.path.splitext(os.path.basename(f))[0]
        parameters.append([None,prefix + '/' + prefix + '.bam', [prefix,os.path.abspath(f)]])
    for job_parameters in parameters:
            yield job_parameters

def get_num_files():
    files = glob.glob(input_bams)
    return len(files)

#
#
# Prepare input directory structure
#

@files(generate_parameters)
def link(none, bam, extra):
    """Make working directory and make symlink to bam file"""
    if not os.path.exists(extra[0]):
        os.mkdir(extra[0])
    if not os.path.exists(bam):
        os.symlink(extra[1], bam) 

#@follows(link)
@transform(link, suffix(".bam"), '.bam.bai')
def index(bam, output):
    """create bam index"""
    index_bam(bam)


#
#
# QC the raw bam files
#

@follows(index)
@transform(link, suffix(".bam"), '.quality_score')
def qc_raw_bam_quality_score_distribution(input_bam, output):
    """docstring for metrics1"""
    bam_quality_score_distribution(input_bam, output, output + '.pdf')

@follows(index)
@transform(link, suffix(".bam"), '.metrics')
def qc_raw_bam_alignment_metrics(input_bam, output):
    """docstring for metrics1"""
    bam_alignment_metrics(input_bam, output)
    
@follows(index)
@transform(link, suffix(".bam"), '.coverage.sample_summary', r'\1.coverage')
def qc_raw_bam_coverage_metrics(input_bam, output, output_format):
    bam_coverage_metrics(input_bam, output_format)

@follows(index)
@transform(link, formatter(".*/(?P<SAMPLE_ID>[^/]+).bam"), '{subpath[0][1]}/qc/qualimap/{SAMPLE_ID[0]}')
def qc_raw_bam_qualimap_report(input_bam, output_dir):
    qualimap_bam(input_bam, output_dir)

#@follows(qc_raw_bam_quality_score_distribution, qc_raw_bam_alignment_metrics, qc_raw_bam_coverage_metrics)
@follows(qc_raw_bam_coverage_metrics, qc_raw_bam_qualimap_report)
def raw_bam_qc():
    """ Aggregates raw bam quality control steps """
    pass





#
#
# Prepare the bam files for variant calling
#

def dup_removal_picard(bam,output):
    """Use Picard to remove duplicates"""
    run_cmd('%s -Xmx4096m -jar %s/CleanSam.jar \
             INPUT=%s \
             OUTPUT=%s \
             VALIDATION_STRINGENCY=LENIENT VERBOSITY=ERROR CREATE_INDEX=TRUE' 
             % (java, picard, bam, output))
   
def dup_mark_picard(bam,output):
    """Use Picard to mark duplicates"""
    run_cmd('%s -Xmx4096m -jar %s/MarkDuplicates.jar \
            TMP_DIR=/export/astrakanfs/stefanj/tmp \
            REMOVE_DUPLICATES=true \
            INPUT=%s \
            OUTPUT=%s \
            METRICS_FILE=%s.dup_metrics \
            VALIDATION_STRINGENCY=LENIENT VERBOSITY=ERROR CREATE_INDEX=true'
            % (java, picard, bam, output, bam))


@follows(index)
@transform(link, suffix(".bam"), '.dedup.bam')
def remove_dups(bam, output):
    """Use samtools for dupremoval"""
    #run_cmd('samtools rmdup %s %s' % (bam,output))
    dup_mark_picard(bam, output)
    # remove(input)


#@follows(remove_dups)
@transform(remove_dups, suffix(".dedup.bam"), '.dedup.bam.bai')
def index_dups(bam, output):
    """create bam index"""
    index_bam(bam)


#@follows(index_dups)
@transform(index_dups, suffix(".dedup.bam.bai"), '.realign.intervals', r'\1.dedup.bam')
def find_realignment_intervals(foo, intervals, input_bam):
    
    run_cmd("%s -Djava.io.tmpdir=/export/astrakanfs/stefanj/tmp -Xmx4g -jar %s \
             -T RealignerTargetCreator \
             -I %s \
             -R %s \
             -known %s \
             -known %s \
             -o %s \
             -nt %d" 
             % (java, gatk, input_bam, reference, indels_1kg, mills, intervals, n_cpus))


#@follows(find_realignment_intervals)
@transform(find_realignment_intervals, suffix(".realign.intervals"), '.realigned.bam', r'\1.dedup.bam')
def indel_realigner(intervals_file, realigned_bam, input_bam):
    """Re-aligns regions around indels"""
    run_cmd("%s -Xmx4g -jar %s \
             -T IndelRealigner \
             -I %s \
             -R %s \
             -targetIntervals %s \
             -known %s \
             -known %s \
             -o %s" 
             % (java, gatk, input_bam, reference, intervals_file, indels_1kg, mills, realigned_bam))


#@follows(indel_realigner)
@transform(indel_realigner, suffix('.realigned.bam'), '.gatk.bam.recal_data.grp')
def recalibrate_baseq1(input_bam, output):
    """Base quality score recalibration in bam file - first pass """
    index_bam(input_bam)
    run_cmd("%s -Xmx6g -jar %s \
            -T BaseRecalibrator \
            -R %s \
            -knownSites %s \
            -I %s \
            -o %s \
            -nct %d"
            % (java_with_params, gatk, reference, dbsnp, input_bam, output, n_cpus))
    

# This custom check ensures that the recalibrate_baseq2 step is not run in --rebuild_mode if the .gatk.bam exists
# This way metric_coverage target can be run after the intermediate files (e.g. .realigned.bam) are removed
def is_gatk_bam_missing(inputs, gatk_bam):
    # ignore input files of this step: input[0] - .realigned.bam, and input[1] - .recal_data.grp file
    if not os.path.exists(gatk_bam):
        return True, "Missing file %s" % gatk_bam
    else:
        return False, "File %s exists" % gatk_bam

@follows(recalibrate_baseq1)
@transform(indel_realigner, suffix('.realigned.bam'), add_inputs(r'\1.gatk.bam.recal_data.grp'),'.gatk.bam')
@check_if_uptodate(is_gatk_bam_missing)
def recalibrate_baseq2(inputs, output_bam):
    """Base quality score recalibration in bam file
        Part 2: rewrite quality scores into a new bam file"""   
    bam = inputs[0]
    recal_data = inputs[1]
    run_cmd("%s -Xmx4g -jar %s \
            -T PrintReads \
            -R %s \
            -I %s \
            --out %s \
            -BQSR %s \
            -nct %d" 
            % (java_with_params, gatk, reference, bam, output_bam, recal_data, n_cpus) )
  
    # remove(inputs[0])


#
#
# gatk.bam-level QC measurements
#

#@follows(recalibrate_baseq2)
@transform(recalibrate_baseq2, suffix('.gatk.bam'), '.quality_score')
def qc_gatk_bam_quality_score_distribution(input_bam, output):
    """docstring for metrics1"""
    bam_quality_score_distribution(input_bam, output, output + '.pdf')

#@follows(recalibrate_baseq2)
@transform(recalibrate_baseq2, suffix('.gatk.bam'), '.metrics')
def qc_gatk_bam_alignment_metrics(input_bam, output):
    """docstring for metrics1"""
    bam_alignment_metrics(input_bam, output)

#@follows(recalibrate_baseq2)
@transform(recalibrate_baseq2, suffix('.gatk.bam'), '.coverage.sample_summary', r'\1.coverage')
def qc_gatk_bam_coverage_metrics(input_bam, output, output_format):
    bam_coverage_metrics(input_bam, output_format)

@transform(recalibrate_baseq2, formatter(".*/(?P<SAMPLE_ID>[^/]+).bam"), '{subpath[0][1]}/qc/qualimap/{SAMPLE_ID[0]}')
def qc_gatk_bam_qualimap_report(input_bam, output_dir):
    qualimap_bam(input_bam, output_dir)


#@follows(qc_gatk_bam_quality_score_distribution, qc_gatk_bam_alignment_metrics, qc_gatk_bam_coverage_metrics, qc_gatk_bam_qualimap_report)
@follows(qc_gatk_bam_coverage_metrics, qc_gatk_bam_qualimap_report)
def gatk_bam_qc():
    """ Aggregates gatk_bam quality control steps """
    pass



#
#
# Bam reduce & (batch) variant calling using UnifiedGenotyper
#


@jobs_limit(6)
#@follows(recalibrate_baseq2)
@transform(recalibrate_baseq2, suffix('.gatk.bam'), '.reduced.bam')
def reduce_bam(bam, output):
    """Reduces the BAM file using read based compression that keeps only essential information for variant calling"""
    run_cmd("%s -Xmx6g -jar %s \
            -T ReduceReads \
            -R %s \
            -I %s \
            -o %s"
            % (java_with_params, gatk, reference, bam, output))


def split_seq(seq, num_pieces):
    """ split a list into pieces passed as param """
    start = 0
    for i in xrange(num_pieces):
        stop = start + len(seq[i::num_pieces])
        yield seq[start:stop]
        start = stop

def multisample_variant_call(bams, output):
    """Perform multi-sample variant calling using GATK"""
    cmd = "nice %s -Xmx8g -jar %s \
            -T UnifiedGenotyper \
            -R %s \
            -o %s \
            -glm BOTH \
            -nt %s \
            --dbsnp %s " % (java_with_params, gatk, reference, output, options.jobs, dbsnp)
    for bam in bams:
        cmd = cmd + '-I {} '.format(bam)
    #log the results
    cmd = cmd + '&> {}.log'.format(output)
    run_cmd(cmd)
    
def merge_batch_vcf(vcfs, output):
    """Merges vcf files from the batch run"""
    if len(vcfs) == 1:
        run_cmd('cp {vcf} {output}'.format(vcf = vcfs[0], output = output))
    else:
        merging = ''
        for i in range(len(vcfs)):
            merging = merging + ' -V:batch{number} {file}'.format(number=i,file=vcfs[i])
        run_cmd("{java} -Xmx4g -jar {gatk} \
                -R {reference} \
                -T CombineVariants \
                -o {output} \
                {files}".format(
                    java=java,
                    gatk=gatk,
                    reference=reference,
                    output=output,
                    files=merging
                )) 

@merge(reduce_bam, 'multisample.gatk.vcf')
def call_variants(infiles, output):
    """Splits the files into s number of batches, calls the variants, and merges them back"""
    #splits the calculations into batches
    counter = 1
    batches = []
    for batch in split_seq(infiles, options.groups):
        multisample_variant_call(batch, 'batch{}.vcf'.format(counter))
        batches.append('batch{}.vcf'.format(counter))
        counter += 1
    merge_batch_vcf(batches,output)
    for batch in batches:
        remove(batch)
        remove(batch + '.idx')


# @follows('call_variants')
# @files('multisample.gatk.vcf', ['multisample.gatk.snp.vcf','multisample.gatk.indel.vcf'])
# def split_snps_and_indels(vcf, outputs):
#     """Separates the vcf file into only snp calls"""
#    run_cmd("{} -Xmx2g -jar {} -R {} -T SelectVariants \
#            --variant {} -selectType SNP\
#            -o {}".format(java, gatk,reference,vcf, outputs[0]))
#    run_cmd("{} -Xmx2g -jar {} -R {} -T SelectVariants \
#            --variant {} -selectType INDEL\
#            -o {}".format(java, gatk,reference,vcf,outputs[1]))


#
#
# Ganotyping using HaplotypeCaller
#

@transform(recalibrate_baseq2, suffix('.gatk.bam'), '.gvcf')
def call_haplotypes(bam, output_gvcf):
    """Perform variant calling using GATK HaplotypeCaller"""
    cmd = "nice %s -Xmx6g -jar %s \
            -T HaplotypeCaller \
            -R %s \
            -I %s \
            -o %s \
            --emitRefConfidence GVCF \
            --variant_index_type LINEAR \
            --variant_index_parameter 128000 \
            -minPruning 4 \
            -L %s \
            --dbsnp %s " % (java_with_params, gatk, reference, bam, output_gvcf, exome, dbsnp)
#             -nct %s \
#             -stand_call_conf 50.0 \

    #log the results
    #cmd = cmd + '&> {}.log'.format(output_gvcf)
    run_cmd(cmd)


@merge(call_haplotypes, 'multisample.gatk.gvcf')
def merge_gvcfs(gvcfs, merged_gvcf):
    """Combine the per-sample GVCF files into one project-wide GVCF""" 
    cmd = "nice %s -Xmx4g -jar %s -T CombineGVCFs \
            -R %s -o %s" % (java_with_params, gatk, reference, merged_gvcf)
       
    for gvcf in gvcfs:
        cmd = cmd + " --variant {}".format(gvcf)
    
    cmd = cmd + '&> {}.log'.format(merged_gvcf)
    run_cmd(cmd)
    

@originate(call_with_gvcfs)
def call_with_gvcfs_as_arguments_task(gvcfs):
    pass

@merge([merge_gvcfs, call_with_gvcfs_as_arguments_task], 'multisample.gatk.vcf')
def genotype_gvcfs(gvcfs, output):
    """ Genotype this project's merged GVCF together with other project-wide GVCF files (provided in settings) """
    cmd = "nice %s -Xmx8g -jar %s -T GenotypeGVCFs \
            -R %s -o %s -nt %s" % (java_with_params, gatk, reference, output, options.jobs)

    # if there are any external gvcfs to call with, include them
    for gvcf in gvcfs:
        cmd = cmd + " --variant {}".format(gvcf)

    cmd = cmd + '&> {}.log'.format(output)
    run_cmd(cmd)


#
#
# Genotyping mtDNA
#


@transform(recalibrate_baseq2, suffix('.gatk.bam'), '.mt.gvcf')
def call_mt_haplotypes(bam, output_gvcf):
    """Perform variant calling using GATK HaplotypeCaller"""
    cmd = "nice %s -Xmx6g -jar %s \
            -T HaplotypeCaller \
            -R %s \
            -I %s \
            -o %s \
            --emitRefConfidence GVCF \
            --variant_index_type LINEAR \
            --variant_index_parameter 128000 \
            -minPruning 4 \
            -L MT \
            --dbsnp %s " % (java_with_params, gatk, reference, bam, output_gvcf, dbsnp)
    run_cmd(cmd)

@merge(call_mt_haplotypes, 'multisample.gatk.mt.gvcf')
def merge_mt_gvcfs(gvcfs, merged_gvcf):
    """Combine the per-sample GVCF files into one project-wide GVCF""" 
    cmd = "nice %s -Xmx4g -jar %s -T CombineGVCFs \
            -R %s -o %s" % (java_with_params, gatk, reference, merged_gvcf)
       
    for gvcf in gvcfs:
        cmd = cmd + " --variant {}".format(gvcf)
    
    cmd = cmd + '&> {}.log'.format(merged_gvcf)
    run_cmd(cmd)
    
@merge([merge_mt_gvcfs], 'multisample.gatk.mt.vcf')
def genotype_mt_gvcfs(gvcfs, output):
    """ Genotype this project's merged GVCF together with other project-wide GVCF files (provided in settings) """
    cmd = "nice %s -Xmx8g -jar %s -T GenotypeGVCFs \
            -R %s -o %s -nt %s" % (java_with_params, gatk, reference, output, options.jobs)

    # if there are any external gvcfs to call with, include them
    for gvcf in gvcfs:
        cmd = cmd + " --variant {}".format(gvcf)

    cmd = cmd + '&> {}.log'.format(output)
    run_cmd(cmd)






#
#
# Variant recalisbration and filtering
# 

#@follows(call_variants)
@follows(genotype_gvcfs)
@files('multisample.gatk.vcf', ['multisample.gatk.snp.model','multisample.gatk.snp.model.tranches'])
def find_snp_tranches_for_recalibration(vcf,outputs):
    """Runs VariantRecalibrator on snps"""
    cmd = "{java} -Xmx16g -jar {gatk} \
            -T VariantRecalibrator \
            -R {reference}  \
            -input {input} \
            -resource:hapmap,known=false,training=true,truth=true,prior=15.0 {hapmap} \
            -resource:omni,known=false,training=true,truth=true,prior=12.0 {omni} \
            -resource:1000G,known=false,training=true,truth=false,prior=10.0 {snps_1kg} \
            -resource:dbsnp,known=true,training=false,truth=false,prior=2.0 {dbsnp} \
            -an QD -an MQ -an MQRankSum -an ReadPosRankSum -an FS -an DP \
            -mode SNP \
            -recalFile {output} \
            -tranchesFile {tranches} \
            -rscriptFile {plots}\
            -nt {num_jobs}".format(
                java=java_with_params,
                gatk=gatk,
                reference=reference,
                hapmap=hapmap,
                omni=omni,
                snps_1kg=snps_1kg,
                dbsnp=dbsnp,
                input=vcf,
                output=outputs[0],
                tranches=outputs[1],
                plots=outputs[0]+'.plots.R',
                num_jobs=options.jobs
            )
    if get_num_files() > 10:
        cmd += " -an InbreedingCoeff"
    run_cmd(cmd)


@follows(find_snp_tranches_for_recalibration)
@files('multisample.gatk.vcf', ['multisample.gatk.indel.model','multisample.gatk.indel.model.tranches'])
def find_indel_tranches_for_recalibration(vcf,outputs):
    """Runs VariantRecalibrator on indels"""
    cmd = "{java} -Xmx16g -jar {gatk} \
            -T VariantRecalibrator \
            -R {reference}  \
            -input {input} \
            --maxGaussians 4 \
            -resource:mills,known=false,training=true,truth=true,prior=12.0 {mills} \
            -resource:dbsnp,known=true,training=false,truth=false,prior=2.0 {dbsnp} \
            -an QD -an DP -an FS -an ReadPosRankSum -an MQRankSum \
            -mode INDEL \
            -recalFile {output} \
            -tranchesFile {tranches} \
            -rscriptFile {plots}\
            -nt {num_jobs}".format(
                java=java_with_params,
                gatk=gatk,
                reference=reference,
                mills=mills,
                dbsnp=dbsnp,
                input=vcf,
                output=outputs[0],
                tranches=outputs[1],
                plots=outputs[0]+'.plots.R',
                num_jobs=options.jobs
            )
    if get_num_files() > 10:
        cmd += " -an InbreedingCoeff"
    run_cmd(cmd)


def apply_recalibration_to_snps_or_indels(vcf,recal,tranches,output,mode='SNP'):
    """Apply the recalibration tranch file to either snps or indels (depending on mode=SNP or mode=INDEL"""
    run_cmd("{java} -Xmx16g -jar {gatk} \
            -T ApplyRecalibration \
            -R {reference} \
            -input {vcf} \
            --ts_filter_level 99.9 \
            -tranchesFile {tranches}  \
            -recalFile {recal} \
            -mode {mode} \
            -nt {num_jobs} \
            -o {output}".format(
                java=java_with_params,
                gatk=gatk,
                reference=reference,
                vcf=vcf,
                tranches=tranches,
                recal=recal,
                mode=mode,
                num_jobs=options.jobs,
                output=output
            ))


@follows(find_indel_tranches_for_recalibration)
@files(['multisample.gatk.vcf','multisample.gatk.snp.model','multisample.gatk.snp.model.tranches'],'multisample.gatk.recalibratedSNPS.rawIndels.vcf')
def apply_recalibration_filter_snps(inputs,output):
    apply_recalibration_to_snps_or_indels(inputs[0],inputs[1],inputs[2],output,mode='SNP')
    # remove(input[0])
    #     remove(input[1])
    #     remove(input[2])


@follows(apply_recalibration_filter_snps)
@files(['multisample.gatk.recalibratedSNPS.rawIndels.vcf','multisample.gatk.indel.model','multisample.gatk.indel.model.tranches'],'multisample.gatk.preHardFiltering.vcf')
def apply_recalibration_filter_indels(inputs,output):
    apply_recalibration_to_snps_or_indels(inputs[0],inputs[1],inputs[2],output,mode='INDEL')


#@follows(apply_recalibration_filter_indels)
#@files('multisample.gatk.preHardFiltering.vcf', 'multisample.gatk.markedHardFiltering.vcf')
@transform(apply_recalibration_filter_indels, suffix('.preHardFiltering.vcf'), '.markedHardFiltering.vcf')
def filter_variants(input_vcf, output_vcf):
    """docstring for apply_indel_filter"""
    run_cmd('{java} -Xmx12g -jar {gatk} \
            -T VariantFiltration \
            -o {output} \
            --variant {input} \
            --filterExpression "QD < 3.0" \
            --filterExpression "DP < 6" \
            --filterName QDFilter   \
            --filterName DPFilter   \
            -R {reference}'.format(
                java=java_with_params,
                gatk=gatk,
                output=output_vcf,
                input=input_vcf,
                reference=reference
            ))
    # remove(input)


#@follows(filter_variants)
#@files('multisample.gatk.markedHardFiltering.vcf', 'multisample.gatk.analysisReady.vcf')
@transform(filter_variants, suffix('.markedHardFiltering.vcf'), '.analysisReady.vcf')
def remove_filtered(input_vcf, output_vcf):
    """Remove filtered variants"""
    run_cmd("{java} -Xmx12g -jar {gatk} \
            -T SelectVariants \
            -R {reference} \
            --variant {input} \
            -o {output} \
            -env -ef".format(
                java=java_with_params,
                gatk=gatk,
                reference=reference,
                input=input_vcf,
                output=output_vcf
            ))
    # remove(input)


#@follows(remove_filtered)
@transform(remove_filtered, suffix('.analysisReady.vcf'),'.analysisReady.exome.vcf')
def final_calls(vcf, output):
    """ Produce the final variant calls in the exome regions """
    output = output[:-10]
    # apply filters to the vcf file to limit calling to exome region    
    run_cmd("%s --vcf %s \
             --out %s \
             --recode \
             --bed %s \
             --keep-INFO-all"
            % (vcftools, vcf, output, exome))
    rename('multisample.gatk.analysisReady.recode.vcf','multisample.gatk.analysisReady.exome.vcf')


def split_snp_parameters():
    exome_vcf = 'multisample.gatk.analysisReady.exome.vcf'
    mtdna_vcf = 'multisample.gatk.mt.vcf'
    for s_id in get_sample_ids():
        yield [exome_vcf, s_id + '/' + s_id + '.exome.vcf', s_id]
        yield [mtdna_vcf, s_id + '/' + s_id + '.mt.vcf', s_id]

def cleanup_files():
    run_cmd("rm -rf */*.recal_data.csv */*.realign* */*.dedup* */*.log *.log \
            *.to_filter* multisample.gatk.snp.recal batch* \
            multisample.gatk.recalibratedSNPS.rawIndels.vcf \
            multisample.gatk.indel.model.* multisample.gatk.snp.model.* \
            multisample.gatk.analysisReady.vcf.vcfidx \
            multisample.gatk.analysisReady.vcf.idx \
            multisample.gatk.recalibratedSNPS.rawIndels.vcf.idx \
            ")
#            multisample.gatk.analysisReady.vcf \


@posttask(cleanup_files)
@follows('final_calls')
@files(split_snp_parameters)
def split_snps(vcf, output, sample):
    """ Split variants by sample, and use sample-specific statistics to filter: AD and DP"""
    AD_threshold=5
    DP_threshold=8
    # what if there is multiple alt alleles??? is the first one most covered      
    run_cmd("{java} -Xmx2g -jar {gatk} -R {ref} \
            -T SelectVariants \
            --variant {vcf} \
            -sn {sample} \
            -select 'vc.getGenotype(\"{sample}\").getAD().1 >= {ad_thr} && vc.getGenotype(\"{sample}\").getDP() >= {dp_thr}' \
            -o {out} \
            ".format(java=java_with_params, gatk=gatk, ref=reference, 
                vcf=vcf, sample=sample, out=output, 
                ad_thr=AD_threshold, 
                dp_thr=DP_threshold))


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#   Main logic


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888
if __name__ == '__main__':
    if options.just_print:
        pipeline_printout(sys.stdout, options.target_tasks, options.forced_tasks,
                            gnu_make_maximal_rebuild_mode = options.rebuild_mode,
                            verbose=options.verbose, #verbose_abbreviated_path=0,
                            checksum_level = 0)

    elif options.flowchart:
        pipeline_printout_graph (   open(options.flowchart, "w"),
                                    # use flowchart file name extension to decide flowchart format
                                    #   e.g. svg, jpg etc.
                                    os.path.splitext(options.flowchart)[1][1:],
                                    options.target_tasks,
                                    options.forced_tasks,
                                        gnu_make_maximal_rebuild_mode = options.rebuild_mode,
                                    no_key_legend   = not options.key_legend_in_graph)
    else:
        pipeline_run(options.target_tasks, options.forced_tasks,
                            multiprocess    = options.jobs,
                            logger          = stderr_logger,
                            verbose         = options.verbose,
                            gnu_make_maximal_rebuild_mode = options.rebuild_mode,
                            checksum_level  = 0)

