"""Phix removal pipeline

Author: Thomas Cokelaer
Affiliation: Institut Pasteur @ 2016

This pipeline is part of Sequana software (sequana.readthedocs.io)
"""
import json
import os
import sequana
from sequana import snaketools as sm
sm.init("quality_control.rules", globals())


# This must be defined before the include 
configfile: "config.yaml"
__snakefile__ = srcdir(__snakefile__)

# Generic include of some dynamic modules
exec(open(sequana.modules["bwa_mem_dynamic"], "r").read())
exec(open(sequana.modules["fastqc_dynamic"], "r").read())
exec(open(sequana.modules["fastq_stats_dynamic"], "r").read())

# A convenient manager
manager = sm.PipelineManager("quality_control", config)


# Starting of the pipeline
__rawdata__input = manager.getrawdata()

# We can add this optional sampling of the data
if manager.config.fastq_sampling.do is True:
    __fastq_sampling__input_fastq = __rawdata__input
    # TODO: here 001 is added manually but needs to be taken care
    # of inside the mananger. Used only in this pipeline for illumnina
    # so works for now
    if manager.paired:
        __fastq_sampling__output_fastq = [
            "fastq_sampling/%s_R1_001.fastq.gz" % (manager.sample),
            "fastq_sampling/%s_R2_001.fastq.gz" % (manager.sample)]
    else:
        __fastq_sampling__output_fastq = [
            "fastq_sampling/%s_R1_001.fastq.gz" % (manager.sample)]
    include: sm.modules["fastq_sampling"]
    __data__input = __fastq_sampling__output_fastq
else:
    __data__input = __rawdata__input


# FASTQC on input data set
__fastqc_samples__input_fastq = __data__input
__fastqc_samples__output_done = manager.getname("fastqc_samples", ".done")
__fastqc_samples__wkdir       = manager.getwkdir("fastqc_samples")
include: fastqc_dynamic("samples", manager)
expected_output.extend(expand(__fastqc_samples__output_done, sample=manager.samples))


# FASTQ stats on input data set
__fastq_stats_samples__input_fastq = __data__input
__fastq_stats_samples__output_done = manager.getname("fastq_stats_samples", ".done")
__fastq_stats_samples__wkdir       = manager.getwkdir("fastq_stats_samples")
include: fastq_stats_dynamic("samples", manager)
expected_output.extend(expand(__fastq_stats_samples__output_done, sample=manager.samples))


# Do we want to remove the Phix ?
if manager.config.bwa_mem_phix.do:
    reference = config["bwa_mem_phix"]["reference_file"]
    if not os.path.isfile(reference + ".fai"):
        exec(open(sequana.modules["dynamic_copy"], "r").read())
        __copy_ref__input = reference
        __copy_ref__output = "reference/" + os.path.basename(reference)
        include: dynamic_copy("ref", manager)
        __bwa_mem_phix__reference = __copy_ref__output
    else:
        __bwa_mem_phix__reference = reference
    # Mapping input / output
    __bwa_index_phix__log = "common_logs/bwa_index.log"
    __bwa_mem_phix__fastq = __data__input
    __bwa_mem_phix__bam = manager.getname("bwa_mem_phix", ".sorted.bam")
    __bwa_mem_phix__fai = __bwa_mem_phix__reference + ".fai"
    __bwa_mem_phix__log = manager.getlogdir("bwa_mem_phix")

    include: bwa_mem_dynamic("phix", manager)

    # BWA_BAM_TO_FASTQ
    __bwa_bam_to_fastq__input = __bwa_mem_phix__bam

    if manager.paired is True:
        _params = {"sample": manager.sample}
        __bwa_bam_to_fastq__fastq_output = [
            "%(sample)s/bwa_bam_to_fastq/%(sample)s_R1_.unmapped.fastq.gz" % _params,
            "%(sample)s/bwa_bam_to_fastq/%(sample)s_R1_.mapped.fastq.gz" % _params,
            "%(sample)s/bwa_bam_to_fastq/%(sample)s_R2_.unmapped.fastq.gz" % _params,
            "%(sample)s/bwa_bam_to_fastq/%(sample)s_R2_.mapped.fastq.gz" % _params]
    else:
        __bwa_bam_to_fastq__fastq_output = [
            "%(sample)s/bwa_bam_to_fastq/%(sample)s_R1_.unmapped.fastq.gz" % {"sample":manager.sample},
            "%(sample)s/bwa_bam_to_fastq/%(sample)s_R1_.mapped.fastq.gz" % {"sample":manager.sample}]

    __bwa_bam_to_fastq__stats_output = "%s/bwa_bam_to_fastq/bwa_mem_stats.json" % manager.sample
    __bwa_bam_to_fastq__wkdir = "%s/bwa_bam_to_fastq" % manager.sample

    include: sm.modules['bwa_bam_to_fastq']

    # Run a FASTQC and FASTQ_STATS on the data without phix
    __fastqc_phix__input_fastq = [x for x in __bwa_bam_to_fastq__fastq_output if "unmapped" in x]
    __fastqc_phix__output_done = manager.getname("fastqc_phix", ".fastqc.done")
    __fastqc_phix__wkdir = manager.getwkdir("fastqc_phix")
    include: fastqc_dynamic("phix", manager)
    expected_output.extend(expand(__fastqc_phix__output_done, sample=manager.samples))

    __fastq_stats_phix__input_fastq = __bwa_bam_to_fastq__fastq_output
    __fastq_stats_phix__output_done = manager.getname("fastq_stats_phix", ".done")
    __fastq_stats_phix__wkdir = manager.getwkdir("fastq_stats_phix")
    include: fastq_stats_dynamic("phix", manager)
    expected_output.extend(expand(__fastq_stats_phix__output_done, sample=manager.samples))


# Perform the adapter removal and trimming
if manager.config['cutadapt']['do']:
    adapter_removal = manager.config['cutadapt']['tool_choice']

    from sequana.adapters import _get_registered_adapters as registered
    from sequana.adapters import get_sequana_adapters

    # Users may provide Nextera, PCRFree, Rubicon or other registered adapters
    fwd = manager.config.cutadapt.fwd
    if isinstance(fwd, str) and fwd in registered():
        filename = "file:"+ get_sequana_adapters(fwd, "fwd")
        manager.config.cutadapt.fwd = filename

    rev = manager.config.cutadapt.rev
    if isinstance(rev, str) and rev in registered():
        filename = "file:"+ get_sequana_adapters(rev, "revcomp")
        manager.config.cutadapt.rev = filename

    if adapter_removal in ["cutadapt", "atropos"]:
        adapter_removal = "cutadapt"
        if manager.config.bwa_mem_phix.do:
            __cutadapt__input_fastq = [x for x in __bwa_bam_to_fastq__fastq_output
                                       if "unmapped" in x]
            __cutadapt__output = [x.replace("bwa_bam_to_fastq",
                                  "cutadapt").replace("unmapped","cutadapt")
                                  for x in __cutadapt__input_fastq]
        else:
            # If the fix is not yet performed, __data__input is a wildcard
            # function so the output must be specified by hand
            __cutadapt__input_fastq = __data__input
            if manager.paired:
                __cutadapt__output = [
                    "{sample}/cutadapt/{sample}_R1_.cutadapt.fastq.gz",
                    "{sample}/cutadapt/{sample}_R2_.cutadapt.fastq.gz"]
            else:
                __cutadapt__output = [
                    "{sample}/cutadapt/{sample}_R1_.cutadapt.fastq.gz"]

        __cutadapt__wkdir = "%s/cutadapt" % manager.sample

        # Set the fwd and rev adapters
        __cutadapt__fwd = manager.config.cutadapt.fwd
        __cutadapt__rev = manager.config.cutadapt.rev

        __cutadapt__design = manager.config.cutadapt.design_file
        __cutadapt__design_adapter = manager.config['cutadapt']['adapter_choice']
        __cutadapt__options = manager.config.cutadapt.options
        __cutadapt__mode = manager.config.cutadapt.mode
        __cutadapt__log = "%s/logs/cutadapt/cutadapt.txt" % manager.sample
        __cutadapt__sample = manager.sample
        include: sm.modules["cutadapt"]
    else:
        raise ValueError("trimming must be either cutadapt or atropos or skewer")

    # Now we can perform again a FastQC and FastQ stats 
    __fastqc_cutadapt__input_fastq = __cutadapt__output
    __fastqc_cutadapt__output_done = manager.getname("fastqc_cutadapt", ".done")
    __fastqc_cutadapt__wkdir = manager.getwkdir("fastqc_cutadapt")
    include: fastqc_dynamic(adapter_removal, manager)
    expected_output += expand(__fastqc_cutadapt__output_done, sample=manager.samples)

    __fastq_stats_cutadapt__input_fastq = __cutadapt__output
    __fastq_stats_cutadapt__output_done = manager.getname("fastq_stats_cutadapt", ".done")
    __fastq_stats_cutadapt__wkdir = manager.getwkdir("fastq_stats_cutadapt")
    include: fastq_stats_dynamic(adapter_removal, manager)
    expected_output += expand(__fastq_stats_cutadapt__output_done, sample=manager.samples)


# if no database provided, let us make sure kraken option is off
if manager.config.kraken.database_directory is False:
    manager.config.kraken.do = False

if manager.config.kraken.do:
    if manager.config['cutadapt'].do:
        # todo: handle all adapter removal cases
        __kraken__input = __cutadapt__output
    elif manager.config.bwa_mem_phix.do:
        __kraken__input = [x for x in __bwa_bam_to_fastq__fastq_output if "unmapped" in x]
    else:
        __kraken__input = __data__input

    __kraken__output_wkdir = manager.getwkdir("kraken")
    __kraken__output_html = manager.sample + "/kraken/kraken.html"
    __kraken__output = manager.sample + "/kraken/kraken.out"
    __kraken__output_csv = manager.sample + "/kraken/kraken.csv"
    __kraken__output_json = manager.sample + "/kraken/kraken.json"
    include: sm.modules['kraken']
    expected_output.append(expand(__kraken__output_html, sample=manager.samples))
    expected_output.append(expand(__kraken__output_json, sample=manager.samples))


# Include the rule graph
__rulegraph__input = __snakefile__
__rulegraph__output = "rulegraph/rulegraph.svg"
__rulegraph__mapper = {
    "fastqc_samples": "../fastqc_samples.html",
    "fastqc_phix": "../fastqc_phix.html",
    "fastqc_cutadapt": "../fastqc_cutadapt.html",
    "cutadapt": "../cutadapt.html",
    "kraken": "../kraken/kraken.html",
    "kraken_translate": "../kraken/kraken.html",
    "kraken_to_krona": "../kraken/kraken.html",
}
include: sm.modules['rulegraph']
expected_output.extend([__rulegraph__output])


# Add Conda
__conda__output = "requirements.txt"
include: sm.modules['conda']   # Create requirements.txt(dependencies)
expected_output.extend([__conda__output])


# create a json file that summarise information of your pipeline
__summary_pipeline__inputs = __rawdata__input
if manager.config['cutadapt'].do:
    # todo: handle all adapter removal cases
    __summary_pipeline__outputs = [ __cutadapt__output ]
elif manager.config.bwa_mem_phix.do:
    __summary_pipeline__outputs = [x for x in __bwa_bam_to_fastq__fastq_output if "unmapped" in x]
else:
    __summary_pipeline__outputs = []

__summary_pipeline__html = []
__summary_pipeline__rulegraph = __rulegraph__output
__summary_pipeline__requirements = "requirements.txt"
__summary_pipeline__snakefile = __snakefile__
__summary_pipeline__config = "config.yaml"
__summary_pipeline__name = "Quality Control"
__summary_pipeline__json_output = manager.getname("summary_pipeline", ".json")
include: sm.modules["summary_pipeline"]
expected_output.append(expand(__summary_pipeline__json_output,
                              sample=manager.samples))

# Those rules takes a couple of seconds so no need for a cluster
localrules: conda, rulegraph

rule pipeline_quality_control:
    input: expected_output


onsuccess:
    from sequana.modules_report.fastq_stats import FastQStatsModule
    from sequana.modules_report.bwa_bam_to_fastq import BWABAMtoFastQModule
    from sequana.modules_report.kraken import KrakenModule
    from sequana.modules_report.phix import PhixModule
    from sequana.modules_report.cutadapt import CutadaptModule
    from sequana.modules_report.fastqc import FastQCModule
    from sequana.modules_report.summary import SummaryModule

    # the stats image
    snake_parser = snakemake.get_argument_parser().parse_args()
    json_list = expand(__summary_pipeline__json_output, sample=manager.samples)
    sm.add_stats_summary_json(json_list, snake_parser)

    # Main directory
    from sequana.utils import config as conf
    report_dir_format = "%(proj)s/report_qc_%(proj)s"

    for proj in manager.samples.keys():
        sample_summary = {}

        report_dir = report_dir_format % {"proj": proj}
        conf.output_dir = report_dir   # ensure files are stored in the correct location

        # Create the 3 FastQC HTML files (independent)
        for this in ["fastqc_phix", "fastqc_samples", "fastqc_cutadapt"]:
            FastQCModule("{}.html".format(this),
                         proj + "/{}/*_fastqc.html".format(this))

        # add all sections in addition to standard summary
        conf.summary_sections = []

        # First the samples section ------------------------------
        indir = proj + "/fastq_stats_samples"

        fqmod = FastQStatsModule(indir, "fastqc_samples")
        sample_summary["fastq_stats_samples_json"] = json.loads(fqmod.get_stats().to_json())

        conf.summary_sections.append({
            "name": "Stats (input data)",
            "title_links": '<a href="fastqc_samples.html">FastQC</a>',
            "anchor'": "stats",
            "content": fqmod._get_stats_section()
        })

        # the phix section ---------------------------------
        if manager.config.bwa_mem_phix.do:
            phixmod = PhixModule(proj)
            sample_summary["phix_section_json"] = json.loads(phixmod._get_stats().to_json())
            sample_summary["phix_section"] = phixmod._get_summary()
            conf.summary_sections.append({
                "name": "Phix ",
                "title_links": '<a href="fastqc_phix.html">FastQC</a>',
                "anchor'": "phix",
                "content": phixmod._get_html()
            })

        # the cutadapt section ----------------------------
        if manager.config.cutadapt.do:
            cutadapt_mod = CutadaptModule(proj + "/logs/cutadapt/cutadapt.txt", proj)
            sample_summary["cutadapt_json"] = json.loads(cutadapt_mod._get_stats().to_json())
            html = cutadapt_mod._get_stat_section()
            html += FastQStatsModule(proj + "/fastq_stats_cutadapt", "fastqc_cutadapt")._get_stats_section(
                tablename="cutadapt2")
            conf.summary_sections.append({
                "name": "Adapter ",
                "title_links": '<a href="fastqc_cutadapt.html">FastQC</a>|' +
                    '<a href="cutadapt.html">Cutadapt</a>|',
                "anchor'": "phix",
                "content": html
            })

        # The Kraken section
        if manager.config.kraken.do:
            krakenmod = KrakenModule(proj + "/kraken")
            sample_summary['kraken_json'] = json.loads(krakenmod._get_stats().to_json())
            conf.summary_sections.append({
                "name": "Kraken ",
                "anchor'": "kraken",
                "content": krakenmod._get_summary_section()
            })
            print("Copying files")
            shell('cp -r %s/kraken/ %s' % (proj, report_dir))

        # The cutadapt report (independent)
        if manager.config['cutadapt'].do:
            from sequana.modules_report.cutadapt import CutadaptModule
            filename = proj + "/logs/cutadapt/cutadapt.txt"
            CutadaptModule(filename, proj, "cutadapt.html")

        shell('cp -r %s/fastqc_*/ %s' % (proj, report_dir))

        filename = report_dir + "/../summary_pipeline/%s.json" % proj
        #
        intro = """
            <b>Sample:</b> {} <br><b>Paired data:</b> {}<br>
            <b>Data Type:</b> {}<br><b>Pipeline:</b>
<a target="_blank" href="http://sequana.readthedocs.io/en/master/pipelines.html#{}">{}</a><br>
            """.format(
                    proj,
                    manager.paired,
                    config.get('datatype', "unspecified"),
                    "quality_control",
                     __summary_pipeline__name)
        SummaryModule(json.loads(open(filename).read()), intro=intro)

        # save sample summary
        sample_summary['project'] = proj
        with open(report_dir + "/summary.json", "w") as fp:
            data = json.dumps(sample_summary, indent=4, sort_keys=True)
            fp.write(data)
        sm.create_cleanup(proj)

    try:
        print("Create multi summary file")
        shell("sequana_summary -m")
    except: pass

    from sequana.snaketools import OnSuccess
    sm.OnSuccess()() # create instance and execute to create Makefile and cleanup

onerror:
    print("An error occurred. See message above.")
