class PipelineTask(JobWrapperTask):
    """Wrapper task for predefined pipelines. Adds option
    target_generator_handler which must be defined in order to
    collect targets.
    """
    target_generator_handler = luigi.Parameter(default=None)
    """Target generator handler is a function that returns a list of
    :class:`ISample <ratatosk.experiment.ISample>` objects.""" class JobTask(BaseJobTask):
    def job_runner(self):
        return DefaultShellJobRunner()

    def args(self):
        return [] class JobWrapperTask(JobTask):
    """Wrapper task that adds target by default"""
    def complete(self):
        return all(r.complete() for r in flatten(self.requires()))

    def run(self):
        pass class PrintConfig(JobTask):
    """Print global configuration for all tasks, including all
    parameters (customizations as well as defaults) and absolute path
    names to program executables (thus implicitly in many cases giving
    the program version, although there should be another function for
    getting program version information and inserting that info here).
    This task (or function) should probably be called after the last
    task so that we know the exact parameter settings and programs
    used to run an analysis. Aim for reproducible research :)
    """
    header = """# Created by {program} on {date}
#
# Command: TODO: insert command here 
#
# The ratatosk configuration file collects all configuration settings
# for a run. In principle you could use this file as input to rerun
# the analysis preserving exactly the same settings.
#
""".format(program="ratatosk", date=datetime.today().strftime("at %H:%M on %A %d, %B %Y"))

    def requires(self):
        return []

    def output(self):
        filename = "ratatosk_config_{}.yaml".format(datetime.today().strftime("%Y_%m_%d_%H_%M"))
        with open(filename, "w") as fh:
            fh.write(self.header)
            fh.write(yaml.safe_dump(config_to_dict(backend.__global_config__), default_flow_style=False, allow_unicode=True, width=1000))
        return luigi.LocalTarget(filename)

    def run(self):
        return NotImplemented def make_fastq_links(targets, indir, outdir, fastq_suffix="001.fastq.gz", ssheet="SampleSheet.csv"):
    """Given a set of targets and an output directory, create links
    from targets (source raw data) to an output directory.

    :param targets: list of :class:`ratatosk.experiment.ISample` objects
    :param outdir: (top) output directory
    :param fastq_suffix: fastq suffix
    :param ssheet: sample sheet name

    :returns: new targets list with updated output directory
    """
    newtargets = []
    for tgt in targets:
        fastq = glob.glob("{}*{}".format(tgt.prefix("sample_run"), fastq_suffix))
        if len(fastq) == 0:
            logger.warn("No fastq files for prefix {} in {}".format(tgt.prefix("sample_run"), "make_fastq_links"))
        for f in fastq:
            newpath = os.path.join(outdir, os.path.relpath(f, indir))
            if not os.path.exists(os.path.dirname(newpath)):
                logger.info("Making directories to {}".format(os.path.dirname(newpath)))
                os.makedirs(os.path.dirname(newpath))
                if not os.path.exists(os.path.join(os.path.dirname(newpath), ssheet)):
                    try:
                        os.symlink(os.path.abspath(os.path.join(os.path.dirname(f), ssheet)), 
                                   os.path.join(os.path.dirname(newpath), ssheet))
                    except:
                        logger.warn("No sample sheet found for {}".format())
                        
            if not os.path.exists(newpath):
                logger.info("Linking {} -> {}".format(newpath, os.path.abspath(f)))
                os.symlink(os.path.abspath(f), newpath)
            if not os.path.lexists(os.path.join(os.path.dirname(newpath), ssheet)) and os.path.exists(os.path.abspath(os.path.join(os.path.dirname(f), ssheet))):
                os.symlink(os.path.abspath(os.path.join(os.path.dirname(f), ssheet)), os.path.join(os.path.dirname(newpath), ssheet))
        newsample = Sample(project_id=tgt.project_id(), sample_id=tgt.sample_id(), 
                           project_prefix=outdir, sample_prefix=os.path.join(outdir, os.path.relpath(tgt.prefix("sample"), indir)),
                           sample_run_prefix=os.path.join(outdir, os.path.relpath(tgt.prefix("sample_run"), indir)))
        newtargets.append(newsample)
        # newtargets.append((tgt.sample_id(), 
        #                    os.path.join(outdir, os.path.relpath(tgt.prefix("sample"), indir)),
        #                    os.path.join(outdir, os.path.relpath(tgt.prefix("sample_run"), indir))))
    return newtargets def rreplace(s, old, new, occurrence):
    li = s.rsplit(old, occurrence)
    return new.join(li) def fullclassname(o):
    return o.__module__ + "." + o.__name__ class VariantEval(GATKJobTask):
    sub_executable = "VariantEval"
    options = luigi.Parameter(default=("-ST Filter -l INFO --doNotUseAllStandardModules --evalModule CompOverlap --evalModule CountVariants --evalModule GenotypeConcordance --evalModule TiTvVariantEvaluator --evalModule ValidationReport --stratificationModule Filter",), is_list=True)
    dbsnp = luigi.Parameter(default=None)
    parent_task = luigi.Parameter(default=("ratatosk.lib.tools.gatk.InputVcfFile", ), is_list=True)
    suffix = luigi.Parameter(default=".eval_metrics")

    def opts(self):
        retval = list(self.options)
        # TODO: Sort this one out
        if not self.dbsnp:
            raise Exception("need dbsnp for VariantEval")
        retval += ["--dbsnp", self.dbsnp]
        # TODO: This too
        if self.target_region:
            retval += ["-L", self.target_region]
        return retval
    
    def args(self):
        retval = ["--eval", self.input()[0], "-o", self.output()]
        if not self.ref:
            raise Exception("need reference for VariantEval")
        retval += ["-R", self.ref]
        return retval class UnifiedGenotyper(GATKIndexedJobTask):
    sub_executable = "UnifiedGenotyper"
    options = luigi.Parameter(default=("-stand_call_conf 30.0 -stand_emit_conf 10.0  --downsample_to_coverage 30 --output_mode EMIT_VARIANTS_ONLY -glm BOTH",), is_list=True)
    suffix = luigi.Parameter(default=".vcf")
    dbsnp = luigi.Parameter(default=None)
    #label = luigi.Parameter(default=".RAW")?
    can_multi_thread = True

    def opts(self):
        retval = list(self.options)
        retval += ["-nt {}".format(self.num_threads)]
        if self.target_region:
            retval += ["-L {}".format(self.target_region)]
        if self.dbsnp:
            retval += ["--dbsnp", self.dbsnp]
        return retval

    def args(self):
        retval =  ["-I", self.input()[0], "-o", self.output()]
        if not self.ref:
            raise Exception("need reference for UnifiedGenotyper")
        retval += ["-R", self.ref]
        return retval class UnifiedGenotyperAlleles(UnifiedGenotyper):
    """Force genotype calling at variants defined by second argument.
    
    For background, see :ref:`merging batched call sets
    <http://gatkforums.broadinstitute.org/discussion/46/merging-batched-call-sets>`_.
    The 'master vcf file' is the union of variant calls made for a
    batch of samples, in which samples may have unique variant calls.
    This step is therefore needed to ascertain whether a missing call
    in one sample is due to low coverage, or that it is homozygote wrt
    the reference.
    """
    label = luigi.Parameter(default="-genotype")
    parent_task = luigi.Parameter(default=("ratatosk.lib.tools.gatk.InputBamFile",
                                           "ratatosk.lib.tools.gatk.InputVcfFile"), is_list=True)
        
    def opts(self):
        retval = super(UnifiedGenotyperAlleles, self).opts()
        retval = [x.replace("EMIT_VARIANTS_ONLY", "EMIT_ALL_SITES") for x in retval]
        if not any([re.search("--genotyping_mode", x) for x in retval]):
            retval += ["--genotyping_mode", "GENOTYPE_GIVEN_ALLELES"]
        if not any([re.search("--output_mode", x) for x in retval]):
            retval += ["--output_mode", "EMIT_ALL_SITES"]
        return retval

    def args(self):
        retval = super(UnifiedGenotyperAlleles, self).args() + ['--alleles', self.input()[1]]
        return retval class VariantFiltration(GATKJobTask):
    """Generic VariantFiltration class"""
    sub_executable = "VariantFiltration"
    parent_task = luigi.Parameter(default="ratatosk.lib.tools.gatk.InputVcfFile")
    label = luigi.Parameter(default=".filtered")
    suffix = luigi.Parameter(default=".vcf")
        
    def args(self):
        retval = ["--variant", self.input()[0], "--out", self.output()]
        if not self.ref:
            raise Exception("need reference for VariantFiltration")
        retval += ["-R", self.ref]
        return retval class CombineVariants(GATKJobTask):
    """CombineVariants.

    Generic task for combining variants from variant files. Input can
    be tasks that generate vcf files, or output from
    :func:`.target_generator_handler`.
    """
    sub_executable = "CombineVariants"
    suffix = luigi.Parameter(default=".vcf")
    label = luigi.Parameter(default="-variants")
    parent_task = luigi.Parameter(default=("ratatosk.lib.tools.gatk.InputVcfFile",), is_list=True)
    target_generator_handler = luigi.Parameter(default=None)

    def requires(self):
        if self.target_generator_handler:
            cls = self.parent()[0]
            sources = []
            if "target_generator_handler" not in self._handlers.keys():
                tgf = RatatoskHandler(label="target_generator_handler", mod=self.target_generator_handler)
                register_task_handler(self, tgf)
            sources = self._handlers["target_generator_handler"](self)
            return [cls(target=src) for src in sources]
        else:
            return [cls(target=source) for cls, source in izip(self.parent(), self.source())]

    def args(self):
        retval = []
        for x in self.input():
            retval += ["-V", x]
        retval += ["-o", self.output()]
        if not self.ref:
            raise Exception("need reference for CombineVariants")
        retval += ["-R", self.ref]
        return retval class Bgzip(TabixJobTask):
    sub_executable = luigi.Parameter(default="bgzip")
    parent_task = luigi.Parameter(default="ratatosk.lib.variation.tabix.InputVcfFile")
    suffix = luigi.Parameter(default=".vcf.gz")
    options = luigi.Parameter(default=("-f",))

    def args(self):
        return [self.input()[0]] def get_logger():
    logger = logging.getLogger('ratatosk-interface')
    if not logger.handlers:
        setup_logging()
    return logger
import luigi
import os
import ratatosk.lib.tools.samtools
from ratatosk import backend
from ratatosk.job import PipelineTask, JobTask, JobWrapperTask, PrintConfig
from ratatosk.utils import make_fastq_links, rreplace, fullclassname
from ratatosk.lib.tools.gatk import VariantEval, UnifiedGenotyper, UnifiedGenotyperAlleles, VariantFiltration, CombineVariants
from ratatosk.lib.variation.tabix import Bgzip
from ratatosk.log import get_logger
# Copyright (c) 2013 Per Unneberg
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may not
# use this file except in compliance with the License. You may obtain a copy of
# the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations under
# the License.
"""
The haloplex pipeline implements a workflow recommended by Agilent technologies. 

It adds the following custom tasks:

1. :class:`.RawUnifiedGenotyper`
2. :class:`.VariantHaloFiltration`
3. :class:`.HaloPlexUnifiedGenotyperAlleles`
4. :class:`.CombineAllVariants`

The main pipeline tasks are 

1. :class:`.HaloPlex`
2. :class:`.HaloPlexSummary`
3. :class:`.HaloPlexCombine`

Running only HaloPlexSummary or HaloPlexCombine should invoke all dependencies in HaloPlex.
However, for many samples (>20), it is advisable to first run HaloPlex
with smaller number of samples in batches.


Calling via ratatosk_run.py
----------------------------

.. code-block:: text

   ratatosk_run.py HaloPlex --indir inputdir --custom-config custom_config_file.yaml
   ratatosk_run.py HaloPlexSummary --indir inputdir --custom-config custom_config_file.yaml
   ratatosk_run.py HaloPlexCombine --indir inputdir --custom-config custom_config_file.yaml


Classes
-------
"""


logger = get_logger()

class CombineAllVariants(CombineVariants):
    """Combine all variants generated by :class:`.HaloPlexUnifiedGenotyperAlleles`
    """
    _config_section = "ratatosk.lib.tools.gatk"

    def requires(self):
        targets = backend.__global_vars__["targets"]
        out_targets = list(set(["{}.{}".format(x.prefix("sample"), "trimmed.sync.sort.merge.realign.recal.clip.filtered-genotype.vcf") for x in targets]))
        return [HaloPlexUnifiedGenotyperAlleles(target=tgt, outdir=os.path.dirname(self.target)) for tgt in out_targets]

class HaloPlexUnifiedGenotyperAlleles(UnifiedGenotyperAlleles):
    """ Temporary class that resolves the issue of calling
    :class:`CombineVariants <ratatosk.lib.tools.gatk.CombineVariants>`
    with a predefined target name.
    """
    parent_task = luigi.Parameter(default=("ratatosk.lib.tools.gatk.ClipReads",
                                           "ratatosk.lib.tools.gatk.CombineVariants"), is_list=True)
    _config_section = "ratatosk.lib.tools.gatk"
    outdir = luigi.Parameter(description="Where analysis takes place", default=None)
    def requires(self):
        """Task requirements. In many cases this is a single source
        whose name can be generated following the code below, and
        therefore doesn't need reimplementation in the subclasses."""
        bamcls = self.parent()[0]
        indexcls = ratatosk.lib.tools.samtools.Index
        return [bamcls(target=self.source()[0])]  + [CombineVariants(target=os.path.join(self.outdir, "CombinedVariants.vcf"))] + [indexcls(target=rreplace(self.source()[0], bamcls().sfx(), indexcls().sfx(), 1), parent_task=fullclassname(bamcls))]


class RawUnifiedGenotyper(UnifiedGenotyper):
    """
    RawUnifiedGenotyper is a variant calling class done on merged data
    to generate a list of raw candidates around which realignment is
    done.
    """
    _config_section = "ratatosk.lib.tools.gatk"
    parent_task = luigi.Parameter(default=("ratatosk.lib.tools.picard.MergeSamFiles", "ratatosk.lib.tools.picard.PicardMetrics"), is_list=True)
    options = luigi.Parameter(default=("-stand_call_conf 30.0 -stand_emit_conf 10.0  --downsample_to_coverage 30 --output_mode EMIT_VARIANTS_ONLY -glm BOTH",), is_list=True)
    label = luigi.Parameter(default=".BOTH.raw")

class VariantHaloFiltration(VariantFiltration):
    """VariantHaloFiltration implements settings tailored for
    filtering haloplex variant calls

    """
    _config_section = "ratatosk.lib.tools.gatk"
    # Options from Halo
    options = luigi.Parameter(default=('--clusterWindowSize 10 --clusterSize 3 --filterExpression "MQ0 >= 4 && ((MQ0 / (1.0 * DP)) > 0.1)" --filterName "HARD_TO_VALIDATE" --filterExpression "DP < 10" --filterName "LowCoverage" --filterExpression "QUAL < 30.0" --filterName "VeryLowQual" --filterExpression "QUAL > 30.0 && QUAL < 50.0" --filterName "LowQual" --filterExpression "QD < 1.5" --filterName "LowQD"',), is_list=True)



class HaloPipeline(PipelineTask):