# 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: Denovo assembly
Data: paired end illumina
Run: snakemake -s denovo_assembly.rules
"""
import os

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


sm.init('denovo_assembly.rules', globals())


# This must be defined before the include
configfile: 'config.yaml'
manager = sm.PipelineManager('denovo_assembly', config)
config = manager.config
__snakefile__ = srcdir(__snakefile__)
__report_dir__ = 'report_da_{0}'.format(manager.sample)

__rawdata__input = manager.getrawdata()

# Initiate variable that will be modified during the pipeline
__summary_pipeline__outputs = []

#
# Normalisation of the coverage - Khmer (digital normalisation)
# 
##############################################################################

if config['digital_normalisation']['do']:
    __digital_normalisation__input = __rawdata__input
    __digital_normalisation__prefix = manager.getname('digital_normalisation')
    __digital_normalisation__output = [
        manager.getname('digital_normalisation', '_R%i_.dn.fastq') % i
        for i in (1, 2)]
    __digital_normalisation__log = manager.getlogdir('digital_normalisation')
    include: sm.modules['digital_normalisation']
    __spades__fastq = __digital_normalisation__output
else:
    __spades__fastq = __rawdata__input

# 
# De Novo Assembly - Spades
# 
##############################################################################

__spades__outdir = manager.getwkdir('spades')
__spades__contigs = manager.getname('spades', '.contigs.fasta')
__spades__scaffolds = manager.getname('spades', '.scaffolds.fasta')
__spades__log = manager.getlogdir('spades')
include: sm.modules['spades']

# get corrected reads
if not config['spades']['only_assembler']:
    __spades_get_fastq__input = __spades__scaffolds
    __spades_get_fastq__yaml = __spades__outdir + 'corrected/corrected.yaml'
    __spades_get_fastq__output = [
        manager.getname('spades_get_fastq', 'R1.fastq.gz'),
        manager.getname('spades_get_fastq', 'R2.fastq.gz')]
    include: sm.modules['spades_get_fastq']
    __bwa_mem_assembly__fastq = __spades_get_fastq__output
    __summary_pipeline__outputs.append(__spades_get_fastq__output)
else:
    __bwa_mem_assembly__fastq = __rawdata__input

exec(open(sequana.modules['bwa_mem_dynamic'], 'r').read())

#
# Assembly assessment
#
# - Quast
# - BUSCO
#
##############################################################################

__quast__genes = []
if config['quast']['reference']:
    __quast__reference = config['quast']['reference']
    if config['quast']['genes_file']:
        __quast__genes = config['quast']['genes_file']
else:
    __quast__reference = []
__quast__input = [__spades__contigs, __spades__scaffolds]
__quast__outdir = manager.getwkdir('quast')
__quast__done = manager.getname('quast', '.done')
__quast__log = manager.getlogdir('quast')
include: sm.modules['quast']
expected_output.append(expand(__quast__done, sample=manager.samples))

# reduce names of contigs
__format_contigs__input = __spades__scaffolds
__format_contigs__output = manager.getname('format_contigs', '.fasta')
include: sm.modules['format_contigs']

#
# Annotation - Prokka
#
##############################################################################

if config['prokka']['do']:
    __prokka__input = __format_contigs__output
    __prokka__gbk = manager.getname('prokka', '.gbk')
    __prokka__prefix = manager.sample
    __prokka__outdir = manager.getwkdir('prokka')
    __prokka__log = manager.getlogdir('prokka')
    include: sm.modules['prokka']
    expected_output.append(expand(__prokka__gbk, sample=manager.samples))

#
# Remapping of reads on assembly:
#  
#  - Variant calling: Freebayes
#  - Coverage analysis: Sequana Coverage
#
##############################################################################

# assess de novo quality with coverage analysis and variant calling
__bwa_mem_assembly__reference = __format_contigs__output
__bwa_mem_assembly__fai = __bwa_mem_assembly__reference + '.fai'
__bwa_mem_assembly__bam = manager.getname('bwa_mem_assembly', '.sorted.bam')
__bwa_mem_assembly__log = manager.getlogdir('bwa_mem_assembly')
__bwa_index_assembly__log = 'common_logs/bwa_index.log'
include: bwa_mem_dynamic('assembly', manager)

# initiate variables
__sambamba_markdup__input = __bwa_mem_assembly__bam
__sambamba_filter__input = __bwa_mem_assembly__bam
__samtools_depth__input = __bwa_mem_assembly__bam
__freebayes__input = __bwa_mem_assembly__bam

# 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]

# 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_assembly__reference
    if config['prokka']['do']:
        __sequana_coverage__gbk = __prokka__gbk
    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 with Freebayes
if config['freebayes']['do']:
    __freebayes__bai = __freebayes__input + '.bai'
    __freebayes__reference = __bwa_mem_assembly__reference
    __freebayes__output = manager.getname('freebayes', '.raw.vcf')
    __freebayes__log = manager.getlogdir('freebayes')
    include: sm.modules['freebayes']

    # Freebayes filter
    __freebayes_vcf_filter__input = __freebayes__output
    __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']
    expected_output.append(expand(__freebayes_vcf_filter__csv,
                                  sample=manager.samples))

# 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',
                       'quast': '../quast/icarus.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.append(__format_contigs__output)
__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 = "De Novo Assembly"
__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.
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_da_{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),
        ])

        # Quast section
        quast_src = os.sep.join([proj, 'quast'])
        quast_dst = os.sep.join([conf.output_dir, 'quast'])
        if os.path.isdir(quast_dst):
            shutil.rmtree(quast_dst)
        elif os.path.isfile(quast_dst):
            os.remove(quast_dst)
        shutil.copytree(quast_src, quast_dst)

        # Prokka section
        if manager.config.prokka.do:
            # copy prokka directory
            prokka_src = os.sep.join([proj, 'prokka'])
            prokka_dst = os.sep.join([conf.output_dir, 'prokka'])
            if os.path.isdir(prokka_dst):
                shutil.rmtree(prokka_dst)
            elif os.path.isfile(prokka_dst):
                os.remove(prokka_dst)
            shutil.copytree(prokka_src, prokka_dst)
            
            ref = '<li><a href="prokka/%(f)s" download="%(f)s">%(f)s</a></li>'
            html = '\n'.join(ref % {'f': f} for f in os.listdir(prokka_dst))
            conf.summary_sections.append({
                'name': 'Prokka',
                'anchor': 'prokka',
                'content': '<ul>%s</ul>' % html,
            })

        SummaryModule(json.loads(open(filename).read()))
