# coding: utf-8
#
#  This file is part of Sequana software
#
#  Copyright (c) 2016 - Sequana Development Team
#
#  File author(s):
#      Dimitri Desvillechabrol <dimitri.desvillechabrol@pasteur.fr>,
#          <d.desvillechabrol@gmail.com>
#
#  Distributed under the terms of the 3-clause BSD license.
#  The full license is in the LICENSE file, distributed with this software.
#
#  website: https://github.com/sequana/sequana
#  documentation: http://sequana.readthedocs.io
#
##############################################################################
"""
Author: Dimitri Desvillechabrol
Affiliation: Institut Pasteur
Aim: Variant calling
Data: paired end or single reads
Run: snakemake -s variant_calling.rules
"""
import os

import sequana
from sequana import snaketools as sm
from sequana.snaketools import SequanaConfig


sm.init("variant_calling.rules", globals())

# This must be defined before the include
configfile: "config.yaml"
manager = sm.PipelineManager("variant_calling", config)

# Clean template and replace None with empty string
config = manager.config
__snakefile__ = srcdir(__snakefile__)
__rawdata__input = manager.getrawdata()
__report_dir__ = 'report_vc_{0}'.format(manager.sample)

#
# Preprocessing for variant calling
#
# - Mapping: BWA
# - Sort: sambamba
# - MarkDuplicates: sambamba 
# - Quality filter: samtools
#
##############################################################################

reference = config["bwa_mem_ref"]["reference"]
# Add locus in FASTA file for snpEff
if config["snpeff"]["do"]:
    from sequana.snpeff import SnpEff
    __snpeff_add_locus_in_fasta__log = "common_logs/snpeff_add_locus_in_fasta.log"
    __snpeff_add_locus_in_fasta__input = reference
    __snpeff_add_locus_in_fasta__output = "reference/{0}".format(
        os.path.basename(reference))
    include: sm.modules["snpeff_add_locus_in_fasta"]
    __bwa_mem_ref__reference = __snpeff_add_locus_in_fasta__output
# Copy the reference index if it exists
elif 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_ref__reference = __copy_ref__output
else:
    __bwa_mem_ref__reference = reference

# The pipeline can be started with sorted BAM files
if config['input_extension'].endswith('bam') or \
    config['input_pattern'].endswith('bam'):
    __bwa_mem_ref__fai = __bwa_mem_ref__reference + ".fai"
    __bwa_mem_ref__bam = __rawdata__input
else:
    # Mapping with BWA MEM
    exec(open(sequana.modules["bwa_mem_dynamic"], "r").read())
    __bwa_mem_ref__fastq = __rawdata__input
    __bwa_mem_ref__fai = __bwa_mem_ref__reference + ".fai"
    __bwa_mem_ref__bam = manager.getname("bwa_mem_ref", ".sorted.bam")
    __bwa_mem_ref__log = manager.getlogdir("bwa_mem_ref")
    __bwa_index_ref__log = "common_logs/bwa_index.log"
    include: bwa_mem_dynamic("ref", manager)

# Add read groups of sample
__add_read_group__input = __bwa_mem_ref__bam
__add_read_group__output = manager.getname("add_read_group/", ".rg.sorted.bam")
__add_read_group__log_err = manager.getlogdir("add_read_group.err")
__add_read_group__log_std = manager.getlogdir("add_read_group.std")
__add_read_group__rg = "ID=%s LB=%s PL=%s PU=%s SM=%s" % (
    manager.sample, manager.sample, manager.config.sequencing.platform,
    manager.config.sequencing.flowcell, manager.sample)
include: sm.modules["add_read_group"]

# Init all input of next rules. These rules are optional and can be inactivate
# in the config file with the "do" option.
__sambamba_markdup__input = __add_read_group__output
__sambamba_filter__input = __add_read_group__output
__samtools_depth__input = __add_read_group__output
__freebayes__input = __add_read_group__output

# Mark duplicates with sambamba markdup
if config["sambamba_markdup"]["do"]:
    __sambamba_markdup__output = manager.getname("sambamba_markdup",
                                                ".rmdup.sorted.bam")
    __sambamba_markdup__log_err = manager.getlogdir("sambamba_markdup.err")
    __sambamba_markdup__log_std = manager.getlogdir("sambamba_markdup.std")
    include: sm.modules["sambamba_markdup"]
    __sambamba_filter__input = __sambamba_markdup__output
    __freebayes__input = __sambamba_markdup__output
    __samtools_depth__input = __sambamba_markdup__output

# bam quality filter with sambamba
if config["sambamba_filter"]["do"]:
    __sambamba_filter__output = manager.getname("sambamba_filter",
                                                ".filter.sorted.bam")
    __sambamba_filter__log = manager.getlogdir("sambamba_filter")
    include: sm.modules["sambamba_filter"]
    __freebayes__input = __sambamba_filter__output
    __samtools_depth__input = [
        __sambamba_filter__output, 
        __sambamba_filter__input,
    ]

#
# Coverage analysis
#
##############################################################################

# Sequana_coverage analysis
if config["sequana_coverage"]["do"]:
    __samtools_depth__output = manager.getname("samtools_depth", ".bed")
    __samtools_depth__log = manager.getlogdir("samtools_depth")
    include: sm.modules["samtools_depth"]
    __sequana_coverage__bed = __samtools_depth__output
    __sequana_coverage__fasta = __bwa_mem_ref__reference
    if config['snpeff']['do']:
        __sequana_coverage__gbk = config['snpeff']['reference']
    else:
        __sequana_coverage__gbk = []
    __sequana_coverage__csv = manager.getname("sequana_coverage", ".csv")
    __sequana_coverage__report_dir = __report_dir__
    __sequana_coverage__html = os.sep.join([
        __report_dir__, 'sequana_coverage.html'
    ])
    include: sm.modules["sequana_coverage"]
    expected_output.append(expand(__sequana_coverage__csv,
                                  sample=manager.samples))


#
# Variant calling
# - Calling: freebayes
# - annotation: snpEff
# - filter: sequana
#
##############################################################################

# Variant calling with Freebayes
# bai file is setup in freebayes rule for pipeline summary
__freebayes__reference = __bwa_mem_ref__reference
__freebayes__output = manager.getname("freebayes", ".raw.vcf")
__freebayes__log = manager.getlogdir("freebayes")
include: sm.modules["freebayes"]

# Annotate detected variants with snpEff
if config["snpeff"]["do"]:
    __snpeff__input = __freebayes__output
    __snpeff__output = manager.getname("snpeff", ".ann.vcf")
    __snpeff__html = manager.getname("snpeff", ".snpeff.html")
    __snpeff__log = manager.getlogdir("snpeff")
    include: sm.modules["snpeff"]
    __freebayes_vcf_filter__input = __snpeff__output
else:
    __freebayes_vcf_filter__input = __freebayes__output

# Freebayes filter
__freebayes_vcf_filter__output = manager.getname("freebayes_vcf_filter",
                                                 ".filter.vcf")
__freebayes_vcf_filter__csv = manager.getname("freebayes_vcf_filter", ".csv")
__freebayes_vcf_filter__report_dir = __report_dir__
__freebayes_vcf_filter__html = os.sep.join([
    __report_dir__, 'variant_calling.html'
])
include: sm.modules["freebayes_vcf_filter"]


#
# Joint variant calling with freebayes
#
##############################################################################

# if there are more than one sample lets do a joint calling with all samples
if len(manager.samples) == 1:
    config['joint_freebayes']['do'] = False
if config['joint_freebayes']['do']:
    __joint_freebayes__input = expand(__freebayes__input,
                                      sample=manager.samples)
    __joint_freebayes__reference = __freebayes__reference
    __joint_freebayes__output = 'joint_calling/joint_calling.raw.vcf'
    __joint_freebayes__log = 'joint_calling/logs/joint_freebayes.log'
    __joint_freebayes__ploidy = config['freebayes']['ploidy']
    include: sm.modules['joint_freebayes']

    if config['snpeff']['do']:
        __snpeff_joint__input = __joint_freebayes__output
        __snpeff_joint__log = 'joint_calling/logs/snpeff_joint.log'
        __snpeff_joint__output = 'joint_calling/joint_calling.ann.vcf'
        __snpeff_joint__html = 'joint_calling/snpeff_joint.html'
        __snpeff_joint__gbk = config['snpeff']['reference']
        __snpeff_joint__options = config['snpeff']['options']
        exec(open(sequana.modules["snpeff_dynamic"], "r").read())
        include: snpeff_dynamic('joint', manager)
        __joint_freebayes_vcf_filter__input = __snpeff_joint__output
    else:
        __joint_freebayes_vcf_filter__input = __joint_freebayes__input
    
    __joint_freebayes_vcf_filter__html = 'joint_calling/joint_calling.html'
    __joint_freebayes_vcf_filter__output = 'joint_calling/joint_calling.filter.vcf'
    __joint_freebayes_vcf_filter__report_dir = 'joint_calling'
    include: sm.modules['joint_freebayes_vcf_filter']


#
# Utils
#
##############################################################################

# Create requirements.txt(dependencies)
include: sm.modules["conda"] 

# Create rulegraph
__rulegraph__input = __snakefile__
__rulegraph__output = "rulegraph/rulegraph.svg"
__rulegraph__mapper = {
    "sequana_coverage": "../sequana_coverage.html",
    "freebayes_vcf_filter": "../variant_calling.html",
}
if len(manager.samples) > 1:
    __rulegraph__mapper['joint_freebayes_vcf_filter'] = '../joint_calling.html'
include: sm.modules["rulegraph"]
expected_output.extend(["requirements.txt", __rulegraph__output])

# create a json file that summarise information of your pipeline
# they must be complete in the onsuccess block
__summary_pipeline__inputs = __rawdata__input
__summary_pipeline__outputs = [
    __bwa_mem_ref__reference,
    __freebayes__input,
    __freebayes__output,
    __freebayes_vcf_filter__input,
    __freebayes_vcf_filter__output,
]
if config['joint_freebayes']['do']:
    __summary_pipeline__outputs += [
        __joint_freebayes_vcf_filter__output,
        __joint_freebayes_vcf_filter__input,
    ]
if os.path.isfile(__bwa_mem_ref__fai):
    __summary_pipeline__outputs.append(__bwa_mem_ref__fai)
if config["snpeff"]["do"]:
    __summary_pipeline__html = [__snpeff__html]
else:
    __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 = "Variant Calling"
__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))


# these rules don't need to be submit on a node.
# snpeff_download_database needs an internet connection
localrules: conda, rulegraph

rule pipeline_variant:
    input:
        expected_output

onsuccess:
    import os
    import shutil
    import json

    from sequana.modules_report.summary import SummaryModule
    from sequana.utils import config as conf

    # add file name for snakemake stats in json
    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)

    # create summary pipeline for each samples
    report_dir_format = 'report_vc_{0}'
    for proj in manager.samples.keys():
        report_dir = report_dir_format.format(proj)
        conf.output_dir = report_dir
        filename = os.sep.join([
            proj,
            'summary_pipeline',
            '{0}.json'.format(proj),
        ])
        SummaryModule(json.loads(open(filename).read()))
        try:
            shutil.copy(__joint_freebayes_vcf_filter__html, report_dir)
        except (FileNotFoundError, NameError):
            pass
