"""RNASeq pipeline

Affiliation: Institut Pasteur @ 2016

This pipeline is part of Sequana software (sequana.readthedocs.io)
"""
import os
import sequana
from os.path import join
from sequana import snaketools as sm
sm.init("rnaseq.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["bowtie1_mapping_dynamic"], "r").read())
exec(open(sequana.modules["bowtie1_index_dynamic"], "r").read())
exec(open(sequana.modules["fastqc_dynamic"], "r").read())
exec(open(sequana.modules["dynamic_unpigz"], "r").read())

manager = sm.PipelineManager("rnaseq", config)

__data__input = manager.getrawdata()


TODO = "todo"

__prefix_name__ = join(config["genome"]["genome_directory"], config["genome"]["name"])

if manager.config.genome.do:
    # create sequence dict file
    __create_sequence_dictionary__reference = config["genome"]["fasta_file"]
    __create_sequence_dictionary__output = config["genome"]["fasta_file"] + ".dict"
    __create_sequence_dictionary__log = "logs/indexing/create_sequence_dictionary.log"
    include: sm.modules["create_sequence_dictionary"]
    # indexing for bowtie1
    # whole genome
    __bowtie1_index_ref__fasta = config["genome"]["fasta_file"]
    __bowtie1_index_ref__output_done = __prefix_name__ + ".1.ebwt"
    __bowtie1_index_ref__output_prefix = __prefix_name__
    __bowtie1_index_ref__log = "logs/indexing/bowtie_genome.log"
    include: bowtie1_index_dynamic("ref")
    # rRNA genome
    if manager.config.genome.rRNA_file:
        __bowtie1_index_rna__fasta = config["genome"]["rRNA_file"]
        __bowtie1_index_rna__output_done = __prefix_name__ + "_rRNA.1.ebwt"
        __bowtie1_index_rna__output_prefix = __prefix_name__ + "_rRNA"
        __bowtie1_index_rna__log = "logs/indexing/bowtie_rRNA.log"
    else :
        __extract_fasta_from_bed__input = config["genome"]["fasta_file"]
        __extract_fasta_from_bed__gff = config["genome"]["gff_file"]
        __extract_fasta_from_bed__feature = config["genome"]["rRNA_feature"]
        __extract_fasta_from_bed__output = __prefix_name__ + "_rRNA.fa"
        __extract_fasta_from_bed__output_features = __prefix_name__ + "_rRNA.gff"
        __extract_fasta_from_bed__log = "logs/indexing/get_rRNA.log"
        include: sm.modules["extract_fasta_from_bed"]
        __bowtie1_index_rna__fasta = __extract_fasta_from_bed__output
        __bowtie1_index_rna__output_done = __prefix_name__ + "_rRNA.1.ebwt"
        __bowtie1_index_rna__output_prefix = __prefix_name__ + "_rRNA"
        __bowtie1_index_rna__log = "logs/indexing/bowtie_rRNA.log"
    include: bowtie1_index_dynamic("rna")


    # indexing for bowtie2
    __bowtie2_index__fasta = config["genome"]["fasta_file"]
    __bowtie2_index__output_done = __prefix_name__ + ".1.bt2"
    __bowtie2_index__output_prefix = __prefix_name__
    __bowtie2_index__log = "logs/indexing/bowtie2_genome.log"
    include: sm.modules["bowtie2_index"]

    # indexing for star
    __star_index__fasta = config["genome"]["fasta_file"]
    __star_index__output_done = config["genome"]["genome_directory"] +  "/SAindex"
    __star_index__output_dir = config["genome"]["genome_directory"]
    __star_index__log = "logs/indexing/star_genome.log"
    include: sm.modules["star_index"]

    #output index
    __star_dir__ = __star_index__output_dir
    __bowtie2_index__ = __bowtie2_index__output_prefix
    __bowtie1_index__ = __bowtie1_index_ref__output_prefix
    __RNA_index__ = __bowtie1_index_rna__output_prefix
    __sequence_dictionary__ = __create_sequence_dictionary__output
    expected_output.extend([__create_sequence_dictionary__output,__bowtie1_index_ref__output_done,
                            __bowtie1_index_rna__output_done,__bowtie2_index__output_done,__star_index__output_done])

else:
    __star_dir__ = config["genome"]["genome_directory"]
    __star_index__output_done = config["genome"]["genome_directory"] +  "/SAindex"
    __bowtie2_index__ =__prefix_name__
    __bowtie2_index__output_done = __prefix_name__ + ".1.bt2"
    __bowtie1_index__ = __prefix_name__
    __bowtie1_index_ref__output_done = __prefix_name__+ ".1.ebwt"
    __RNA_index__ = __prefix_name__ + "_rRNA"
    __bowtie1_index_rna__output_done = __prefix_name__ + "_rRNA.1.ebwt"
    __sequence_dictionary__ = config["genome"]["fasta_file"] + ".dict"



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



if manager.config.cutadapt.do:
    adapter_tool = manager.config.cutadapt.tool_choice

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

    # Users may provide TruSeq, Nextera, PCRFree 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

    # Clean NGS
    if adapter_tool == "clean_ngs":
        __clean_ngs__input = __data__input
        __clean_ngs__output = [manager.getname("clean_ngs", "_R1_.clean_ngs.fastq.gz")]
        if manager.paired:
            __clean_ngs__output += [manager.getname("clean_ngs", "_R2_.clean_ngs.fastq.gz")]
        __clean_ngs__logs = manager.getname("clean_ngs", ".logs")
        __clean_ngs__rej_files = [manager.getname("clean_ngs", "_rejected.fastq.gz")]
        include: sm.modules["clean_ngs"]

        #expected_output.extend(expand(__clean_ngs__output, sample=manager.samples))
    elif adapter_tool in ["cutadapt", "atropos"]:
        adapter_tool = "cutadapt"
        __cutadapt__input_fastq = __data__input
        __cutadapt__wkdir = manager.getwkdir("cutadapt")
        __cutadapt__output = [manager.getname("cutadapt", "_R1_.cutadapt.fastq.gz")]
        if manager.paired:
            __cutadapt__output += [manager.getname("cutadapt", "_R2_.cutadapt.fastq.gz")]

        # 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
        __clean_ngs__output = __cutadapt__output
        include: sm.modules["cutadapt"]

    else:
        raise ValueError("Invalid choice of cutadapt:tool in config file. Use either clean_ngs or cutadapt")
else:
    __clean_ngs__output = __data__input

if manager.config.fastq_screen.do:
    # Fastq Screen
    __fastq_screen__input = __clean_ngs__output
    __fastq_screen__logs = manager.getname("fastq_screen", ".logs")
    __fastq_screen__output = manager.getname("fastq_screen", ".done")
    __fastq_screen__outdir = manager.getwkdir("fastq_screen")
    include: sm.modules["fastq_screen"]
    if manager.config.fastq_screen.pf2_report:
        __fastq_screen_report__input = __fastq_screen__outdir
        __fastq_screen_report__sample = manager.sample
        __fastq_screen_report__input_done = __fastq_screen__output
        __fastq_screen_report__figure = manager.getname("fastq_screen", "_report.svg")
        __fastq_screen_report__logs =  manager.getname("fastq_screen", "pf2_report.logs")
        include: sm.modules["fastq_screen_report"]
        expected_output.extend(expand(__fastq_screen_report__figure, sample=manager.samples))
    else:
        expected_output.extend(expand(__fastq_screen__output, sample=manager.samples))


# FASTQC on input data set
__fastqc_filtered__input_fastq = __clean_ngs__output
__fastqc_filtered__output_done = manager.getname("fastqc_filtered", ".done")
__fastqc_filtered__wkdir       = manager.getwkdir("fastqc_filtered")
include: fastqc_dynamic("filtered", manager)
expected_output.extend(expand(__fastqc_filtered__output_done, sample=manager.samples))



# Decompress fastq.gz file before to run bowtie1
__unpigz_R1__input = manager.getname("cutadapt", "_R1_.cutadapt.fastq.gz")
__unpigz_R1__output = manager.getname("cutadapt", "_R1_.cutadapt.fastq")
include: dynamic_unpigz("R1", manager)
__unpigz__output = [__unpigz_R1__output]
if manager.paired:
    __unpigz_R2__input = manager.getname("cutadapt", "_R2_.cutadapt.fastq.gz")
    __unpigz_R2__output = manager.getname("cutadapt", "_R2_.cutadapt.fastq")
    include: dynamic_unpigz("R2", manager)
    __unpigz__output += [__unpigz_R2__output]



if manager.config.bowtie1_mapping_rna.do:
    # rRNA
    __bowtie1_mapping_rna__input = __unpigz__output
    __bowtie1_mapping_rna__index_done = __bowtie1_index_rna__output_done
    __bowtie1_mapping_rna__bam = manager.getname("bowtie1_mapping_rna", ".bam")
    __bowtie1_mapping_rna__sort = manager.getname("bowtie1_mapping_rna", ".sorted.bam")
    __bowtie1_mapping_rna__prefix_index = __RNA_index__
    __bowtie1_mapping_rna__stdout = manager.getname("bowtie1_mapping_rna", "_rRNA.out")
    __bowtie1_mapping_rna__stderr = manager.getname("bowtie1_mapping_rna", "_rRNA.err")

    include: bowtie1_mapping_dynamic("rna", manager)
    expected_output.extend(expand(__bowtie1_mapping_rna__sort, sample=manager.samples))

if manager.config.bowtie1_mapping_ref.do:
    # Mapper bowtie 1
    __bowtie1_mapping_ref__input = __unpigz__output
    __bowtie1_mapping_ref__index_done = __bowtie1_index_ref__output_done
    __bowtie1_mapping_ref__bam = manager.getname("bowtie1_mapping_ref", ".bam")
    __bowtie1_mapping_ref__sort = manager.getname("bowtie1_mapping_ref", ".sorted.bam")
    __bowtie1_mapping_ref__prefix_index = __bowtie1_index__
    __bowtie1_mapping_ref__stdout = manager.getname("bowtie1_mapping_ref", ".out")
    __bowtie1_mapping_ref__stderr = manager.getname("bowtie1_mapping_ref", ".err")
    include: bowtie1_mapping_dynamic("ref", manager)
    expected_output.extend(expand(__bowtie1_mapping_ref__sort, sample=manager.samples))
    __mapping_output = __bowtie1_mapping_ref__sort
if manager.config.star_mapping.do:
    # Mapper rna-star
    __star_mapping__input = __clean_ngs__output
    __star_mapping__index_done = __star_index__output_done
    __star_mapping__index = __star_dir__
    __star_mapping__ref = config["genome"]["fasta_file"]
    __star_mapping__logs = manager.getname("star_mapping", ".logs")
    __star_mapping__output_prefix1 = manager.getname("star_mapping", "_init_")
    __star_mapping__output_prefix2 = manager.getname("star_mapping", "_")
    __star_mapping__read_groups = ""
    __star_mapping__sort = manager.getname("star_mapping", "_Aligned.sortedByCoord.out.bam")
    __star_mapping__genome_dir = manager.getname("star_mapping", "_star_2nd_pass")
    __star_mapping__splice_file = manager.getname("star_mapping", "_init_SJ.out.tab")
    include: sm.modules["star_mapping"]
    expected_output.extend(expand(__star_mapping__sort, sample=manager.samples))
    __mapping_output = __star_mapping__sort


# generating of bigwig files
if manager.config.coverage.do:
    __bamCoverage__input = __mapping_output
    __bamCoverage__log = manager.getname("bamCoverage", ".logs")
    __bamCoverage__output = manager.getname("bamCoverage", ".norm.bw")
    include: sm.modules["bamCoverage"]
    expected_output.extend(expand(__bamCoverage__output, sample=manager.samples))
else :
    __igvtools__input = __mapping_output
    __igvtools__output = manager.getname("igvtools", ".tdf")
    __igvtools__log = manager.getname("igvtools", ".logs")
    include: sm.modules["igvtools"]
    #expected_output.extend(expand(__igvtools__output, sample=manager.samples))

# Feature counts from subread suite
if manager.config.star_mapping.do:
    __feature_counts__input = __star_mapping__sort
else :
    __feature_counts__input = __bowtie1_mapping_ref__sort
if manager.config.feature_counts.do:
	__feature_counts__output_count = manager.getname("feature_counts", "_feature.out")
	__feature_counts__gff = config["genome"]["gff_file"]
	__feature_counts__output_gene_count = manager.getname("feature_counts", "_count.out")
	__feature_counts__log = manager.getname("feature_counts", ".logs")
	include: sm.modules["feature_counts"]
	expected_output.extend(expand(__feature_counts__output_gene_count, sample=manager.samples))


# Add Read group on BAM files
__add_read_group__input = __mapping_output
__add_read_group__output = manager.getname("add_read_group/", ".bam")
__add_read_group__log_err = "%s/logs/AddOrReplaceReadGroups/stderr.logs" % manager.sample
__add_read_group__log_std ="%s/logs/AddOrReplaceReadGroups/stdout.logs" % manager.sample
__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"]



# Mark duplicates
if config["mark_duplicates"]["do"]:
    __mark_duplicates__input = __add_read_group__output
    __mark_duplicates__output = manager.getname("mark_duplicates", ".bam")
    __mark_duplicates__metrics = manager.getname("mark_duplicates", ".metrics")
    __mark_duplicates__log_std = "%s/logs/mark_duplicates/stdout.logs" % manager.sample
    __mark_duplicates__log_err =  "%s/logs/mark_duplicates/stderr.logs" % manager.sample
    include: sm.modules["mark_duplicates"]
    expected_output.extend(expand(__mark_duplicates__output, sample=manager.samples))


if config["RNAseQC"]["do"]:
    __reorderSam__input = __mark_duplicates__output
    __reorderSam__logs = manager.getname("reorderSam", ".logs")
    __reorderSam__output = manager.getname("reorderSam", "_reorder.bam")
    include: sm.modules["reorderSam"]
    __RNAseQC__input_bam = expand(__mark_duplicates__output, sample=manager.samples)
    __RNAseQC__logs = "RNAseQC/RNAseQC.log"
    __RNAseQC__output = "RNAseQC"
    __RNAseQC__conf = "RNAseQC/sample_index.txt"
    __RNAseQC__figure =  "RNAseQC/Transcript-associated_Reads_metrics.svg"
    include: sm.modules["RNAseQC"]
    expected_output.extend([__RNAseQC__figure])

# !Reset expected_output variable after multiqc
# Multiqc rule
__multiqc__input = expected_output
__multiqc__input_dir = "."
__multiqc__logs = "multiqc/multiqc.log"
__multiqc__output = config['multiqc']['output-directory'] + "/multiqc_report.html"
include: sm.modules["multiqc"]
expected_output = [__multiqc__output]



# Include rule graph for each sample
__rulegraph__input = __snakefile__
__rulegraph__output = "rulegraph/rulegraph.svg"
__rulegraph__mapper = {"fastqc_samples": "fastqc_samples/"}
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 = __data__input
#if manager.config.cutadapt.do:
#    __summary_pipeline__outputs = [ __cutadapt__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 = "RNA-seq"
#__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 rnaseq:
    input: expected_output


onsuccess:
    import os
    # Create plots about stats
    sm.plot_stats(N=len(manager.samples))

    # Main directory
    report_dir_format = "%(proj)s/report_rnaseq_%(proj)s"
    for proj in manager.samples.keys():
        report_dir = report_dir_format % {"proj": proj}
        try:os.mkdir(report_dir)
        except:pass

        shell("cp %s %s" % (__snakefile__, report_dir))
        #shell("cp rulegraph.svg %s/rulegraph.svg" % (report_dir))
        shell("cp config.yaml %s" % report_dir)
        shell("cp requirements.txt %s" % report_dir)
        shell("cp snakemake_stats.png %s" % report_dir)
        try: os.mkdir("cluster_logs")
        except:pass

        try: shell("mv slurm* cluster_logs/")
        except: pass

        # Create a cleanup python file to clean a sub-directory
        sm.create_cleanup(proj)

    sm.OnSuccess()() # create instance to create main cleanup


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


