"""Small RNASeq pipeline

Affiliation: Institut Pasteur @ 2016

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

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

__data__input = manager.getrawdata()


if manager.config.genome.do:


    # indexing for bowtie1
    # whole genome
    __bowtie1_index_hairpin__fasta = config["genome"]["fasta_hairpin_file"]
    __bowtie1_index_hairpin__output_done = config["genome"]["fasta_hairpin_file"] + ".1.ebwt"
    __bowtie1_index_hairpin__output_prefix = config["genome"]["fasta_hairpin_file"]
    __bowtie1_index_hairpin__log = "logs/indexing/bowtie_hairpin.log"
    include: bowtie1_index_dynamic("hairpin")
    # rRNA genome
    __bowtie1_index_mature__fasta = config["genome"]["fasta_mature_file"]
    __bowtie1_index_mature__output_done = config["genome"]["fasta_mature_file"] + ".1.ebwt"
    __bowtie1_index_mature__output_prefix = config["genome"]["fasta_mature_file"]
    __bowtie1_index_mature__log = "logs/indexing/bowtie_mature.log"
    include: bowtie1_index_dynamic("mature")


    #output index
    __haipin_index__ = __bowtie1_index_hairpin__output_prefix
    __mature_index__ = __bowtie1_index_mature__output_prefix
    expected_output.extend([__bowtie1_index_hairpin__output_done, __bowtie1_index_mature__output_done])

else:

    __haipin_index__ = config["genome"]["fasta_hairpin_file"]
    __bowtie1_index_hairpin__output_done = config["genome"]["fasta_hairpin_file"] + ".1.ebwt"

    __mature_index__ = config["genome"]["fasta_mature_file"]
    __bowtie1_index_mature__output_done = config["genome"]["fasta_mature_file"] + ".1.ebwt"




# 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

    if 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")]

        # 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__output = __cutadapt__output
        include: sm.modules["cutadapt"]
    else:
        raise ValueError("Invalid choice of cutadapt:tool in config file")



if manager.config.fastq_screen.do:
    # Fastq Screen
    __fastq_screen__input = __clean__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__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__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)



# mirRNA mature
__bowtie1_mapping_mature__input = __unpigz_R1__output
__bowtie1_mapping_mature__index_done = __bowtie1_index_mature__output_done
__bowtie1_mapping_mature__bam = manager.getname("bowtie1_mapping_miRNA", ".bam")
__bowtie1_mapping_mature__sort = manager.getname("bowtie1_mapping_miRNA", ".sorted.bam")
__bowtie1_mapping_mature__prefix_index = __mature_index__
__bowtie1_mapping_mature__stdout = manager.getname("bowtie1_mapping_miRNA", "_miRNA.out")
__bowtie1_mapping_mature__stderr = manager.getname("bowtie1_mapping_miRNA", "_miRNA.err")
include: bowtie1_mapping_dynamic("mature", manager.sample)
expected_output.extend(expand(__bowtie1_mapping_mature__sort, sample=manager.samples))

# counting step
__miRNA_count_mature__input = __bowtie1_mapping_mature__sort
__miRNA_count_mature__output = manager.getname("miRNA_counts", "_miRNA_mature_counts.txt")
__miRNA_count_mature__log = "%s/logs/counts/mature_counts.txt" % manager.sample
include: miRNA_count_dynamic("mature", manager)
expected_output.extend(expand(__miRNA_count_mature__output, sample=manager.samples))

# hairpin mapping
__bowtie1_mapping_hairpin__input = __unpigz_R1__output
__bowtie1_mapping_hairpin__index_done = __bowtie1_index_hairpin__output_done
__bowtie1_mapping_hairpin__bam = manager.getname("bowtie1_mapping_hairpin", ".bam")
__bowtie1_mapping_hairpin__sort = manager.getname("bowtie1_mapping_hairpin", ".sorted.bam")
__bowtie1_mapping_hairpin__prefix_index = __haipin_index__
__bowtie1_mapping_hairpin__stdout = manager.getname("bowtie1_mapping_hairpin", "_hairpin.out")
__bowtie1_mapping_hairpin__stderr = manager.getname("bowtie1_mapping_hairpin", "_hairpin.err")
include: bowtie1_mapping_dynamic("hairpin", manager.sample)
expected_output.extend(expand(__bowtie1_mapping_hairpin__sort, sample=manager.samples))


# counting step
__miRNA_count_hairpin__input = __bowtie1_mapping_hairpin__sort
__miRNA_count_hairpin__output = manager.getname("miRNA_counts", "_miRNA_hairpin_counts.txt")
__miRNA_count_hairpin__log = "%s/logs/counts/hairpin_counts.txt" % manager.sample
include: miRNA_count_dynamic("hairpin", manager)
expected_output.extend(expand(__miRNA_count_hairpin__output, sample=manager.samples))


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

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

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

        os.mkdir("cluster_logs")
        shell("mv slurm* cluster_logs/")
        # from sequana.reporting.report_summary import SequanaSummary
        # #summary = SequanaSummary(proj, directory=report_dir,
        #     output_filename="summary.html",
        #    snakefile=__snakefile__, configfile=report_dir+"/config.yaml",
        #    manager=manager)
        # summary.create_report()

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

    sm.create_recursive_cleanup()
    sm.message("Done. To further cleanup the directory, type \n python .sequana_cleanup.py")


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


