shell.executable("/bin/bash")
shell.prefix("source ~/.bashrc; ")

from collections import defaultdict
import os
from os.path import join
import glob
from riboraptor.helpers import mkdir_p

ADAPTER = 'default'

include:
    config['config_path']

workdir: OUT_DIR


mkdir_p(join(OUT_DIR, 'slurm-logs'))
ALL_SRA_FILES = glob.glob('{}/**/*.sra'.format(RAWDATA_DIR), recursive=True)
SRX_ID = defaultdict(list)

for sample in ALL_SRA_FILES:
    srx, srr = sample.replace('{}'.format(RAWDATA_DIR),'').lstrip('/').rstrip('/').split('/')
    SRX_ID[srx].append(srr.replace('.sra', ''))

SRR_ID = list(SRX_ID.values())
SRX_SAMPLES = list(SRX_ID.keys())

ALL_SRR = [item for sublist in SRR_ID for item in sublist]


def merge_bams_input(wildcards):
    return ['bams_srr/{}.bam'.format(srr) for srr in SRX_ID[wildcards.sample]]


def merge_fastq_input(wildcards):
    return ['sratofastq/{}.fastq.gz'.format(srr) for srr in SRX_ID[wildcards.sample]]


def sra_to_fastq_input(wildcards):
    srr_id = wildcards.sample
    for srx_id in list(SRX_ID.keys()):
        value = SRX_ID[srx_id]
        if srr_id in list(value):
            return ancient(str(join(RAWDATA_DIR, srx_id, srr_id+'.sra')))
    print("WRONG encodeterend: {}".format(srr_id))


def get_wrapper(wrapper_name):
    path = os.path.dirname(os.path.abspath(os.path.realpath(workflow.snakefile)))
    return 'file://' + os.path.join(path,
                                    'wrappers',
                                    wrapper_name + '.py')


def get_multiqc_report_input(wildcards):
    return glob.glob('multiqc_report/multiqc_plots/png/*.png')


rule all:
    input:
        expand('reports/{sample}_riboraptor_report.html', sample=SRX_SAMPLES),
        #expand('qc_post_trimming/{sample}_trimmed_trimmed_fastqc.zip', sample=ALL_SRR),
        #expand('qc/{sample}_fastqc.zip', sample=SRX_SAMPLES),
        expand('plots/periodicity/{sample}.png', sample=SRX_SAMPLES),
        expand('bigWig_lengthwise_normalized/{sample}/{fragment_length}/{orientation}_{strand}.bw',
               sample=SRX_SAMPLES,
               fragment_length=FRAGMENT_LENGTHS,
               orientation=ORIENTATIONS,
               strand=STRANDS),
        expand('bams_unique/{sample}.bam', sample=SRX_SAMPLES),
        dynamic('multiqc_report/multiqc_plots/png/{plot}.png'),
        'featureCounts/fcounts.tsv',
        'multiqc_report/aggregated_report.html',
        'reports/riboraptor_report_summarized.html',


rule sra_to_fastq:
    input: sra_to_fastq_input
    benchmark: 'benchmarks/sra_to_fastq/{sample}.txt'
    output: 'sratofastq/{sample}.fastq.gz'
    threads: 1
    params:
        prefix_sra='sratofastq/{sample}.sra.fastq',
        prefix='sratofastq/{sample}.fastq',
        temp_dir='/tmp/{sample}_sratofastq',
    shell:
        r'''
        fastq-dump --split-3 \
        -O sratofastq {input} \
        && gzip {params.prefix}
        '''

#    fasterq-dump --split-3 \
#    -O sratofastq {input} -t {params.temp_dir}\
#    -e {threads} \
#    && mv {params.prefix_sra} {params.prefix} \
#    && gzip {params.prefix}

rule merge_fastq:
    input:
        all_fastq = expand('sratofastq/{srr}.fastq.gz', srr=ALL_SRR),
        dynamic_input = merge_fastq_input
    benchmark: 'benchmarks/merge_fastq/{sample}.txt'
    output: temp('fastq_merged/{sample}.fastq.gz')
    wrapper:
        get_wrapper('merge_fastq_wrapper')


rule perform_qc:
    input:
        'fastq_merged/{sample}.fastq.gz'
    benchmark: 'benchmarks/perform_qc/{sample}.txt'
    output:
       'qc/{sample}_fastqc.html',
       'qc/{sample}_fastqc.zip',
    params:
        out_dir = 'qc'
    resources:
        mem_mb=10000
    shell:
        r'''
        fastqc -o {params.out_dir} \
        -f fastq {input}
        '''


rule perfom_trimming_merged:
    input:
        R1 = 'fastq_merged/{sample}.fastq.gz',
    benchmark: 'benchmarks/perform_trimming_merged/{sample}.txt'
    output: 'merged_preprocessed/{sample}_trimmed.fq.gz'
    params:
        out_dir = 'merged_preprocessed',
        phred_cutoff = 5,
        adaptor_type = '--illumina'
    shell:
        r'''
        trim_galore -o {params.out_dir} \
        -q {params.phred_cutoff} {input.R1}
        '''


rule perform_trimming:
    input:
        R1 = 'sratofastq/{sample}.fastq.gz'
    benchmark: 'benchmarks/perform_trimming/{sample}.txt'
    output:
        pass1_fq = 'preprocessed_step1/{sample}_trimmed.fq.gz',
        pass2_fq = 'preprocessed/{sample}_trimmed_trimmed.fq.gz',
        pass1_fq_report = 'preprocessed_step1/{sample}.fastq.gz_trimming_report.txt',
        #pass2_fq_report = 'preprocessed/{sample}_trimmed.fq.gz_trimming_report.txt',
    params:
        pass1_dir = 'preprocessed_step1',
        pass2_dir = 'preprocessed',
        phred_cutoff = 5,
        min_length = 18, # Discard sequences shorter than this length
        max_length = 38, # Discard sequences longer than this length
        adapter = ADAPTER
        # --max_length for SmallRNA: Should let to 37?
        # --illumina: For adapter trimming?
    #shell:
    #    r'''
    #    trim_galore -o {params.out_dir} \
    #    -q {params.phred_cutoff} {input.R1}
    #    '''
    wrapper:
        get_wrapper('trimming_wrapper')

rule perform_qc_post_trimming:
    input:
        'preprocessed/{sample}_trimmed_trimmed.fq.gz'
    benchmark: 'benchmarks/perform_qc_post_trimming/{sample}.txt'
    output:
       'qc_post_trimming/{sample}_trimmed_trimmed_fastqc.html',
       'qc_post_trimming/{sample}_trimmed_trimmed_fastqc.zip',
    params:
        out_dir = 'qc_post_trimming'
    resources:
        mem_mb=10000
    shell:
        r'''
        fastqc -o {params.out_dir} \
        -f fastq {input}
        '''


rule map_star:
    input:
        R1 = 'preprocessed/{sample}_trimmed_trimmed.fq.gz',
        index = STAR_INDEX
    benchmark: 'benchmarks/map_star/{sample}.txt'
    output:
        bam = 'bams_srr/{sample}.bam',
        txbam = 'bams_srr_tx/{sample}.bam',
        counts = 'STARcounts/{sample}.counts',
        starlogs = 'starlogs/{sample}Log.final.out'
    params:
        chrom_sizes = CHROM_SIZES,
        intron_bed = INTRON_BED,
        tmp_dir = '/tmp'
    threads: 16
    wrapper:
        get_wrapper('star_wrapper')


rule merge_bams:
    input: merge_bams_input
    benchmark: 'benchmarks/merge_bams/{sample}.txt'
    output: temp('bams/{sample}.bam')
    params:
        tmp_dir = '/tmp'
    threads: 1
    wrapper:
        get_wrapper('merge_bams_wrapper')


rule extract_uniq_mapping:
    input: 'bams/{sample}.bam'
    benchmark: 'benchmarks/extract_uniq_mapping/{sample}.txt'
    output: 'bams_unique/{sample}.bam'
    params:
        tmp_dir = '/tmp'
    threads: 16
    wrapper:
        get_wrapper('uniq_mapping_wrapper')


rule bam_to_bedgraph:
    input: 'bams_unique/{sample}.bam'
    benchmark: 'benchmarks/bam_to_bedgraph/{sample}.txt'
    output: temp('bedGraphs/{sample}.bg')
    shell:
        r'''
        riboraptor bam-to-bedgraph \
        --bam {input} \
        --end_type 5prime \
        --strand + \
        --saveto {output} \
        && bedSort {output} {output}
        '''


rule bigwig_from_uniq_bedgraph_raw:
    input: 'bedGraphs/{sample}.bg'
    benchmark: 'benchmarks/bigwig_from_uniq_bedgraph_raw/{sample}.txt'
    output: 'bigWigs/{sample}.bw',
    shell:
        r'''bedGraphToBigWig {input} {CHROM_SIZES} {output}'''


rule export_utr5_coverage:
    input: 'bigWigs/{sample}.bw'
    benchmark: 'benchmarks/export_utr5_coverage/{sample}.txt'
    output: 'gene_coverages/UTR5/{sample}_gene_coverages.tsv'
    params:
        offset_5p = 10
    shell: r'''riboraptor export-gene-coverages \
           --bw {input} \
           --bed {UTR5_BED} \
           --saveto {output} \
           --offset_5p {params.offset_5p}'''


rule export_cds_coverage:
    input: 'bigWigs/{sample}.bw'
    benchmark: 'benchmarks/export_cds_coverage/{sample}.txt'
    output: 'gene_coverages/CDS/{sample}_gene_coverages.tsv'
    params:
        offset_5p = 10
    shell: r'''riboraptor export-gene-coverages \
           --bw {input} \
           --bed {CDS_BED} \
           --saveto {output} \
           --offset_5p {params.offset_5p}'''


rule export_utr3_coverage:
    input: 'bigWigs/{sample}.bw'
    benchmark: 'benchmarks/export_utr3_coverage/{sample}.txt'
    output: 'gene_coverages/UTR3/{sample}_gene_coverages.tsv'
    params:
        offset_5p = 10
    shell: r'''riboraptor export-gene-coverages \
           --bw {input} \
           --bed {UTR3_BED} \
           --saveto {output} \
           --offset_5p {params.offset_5p}'''


rule export_read_length:
    input: 'bams_unique/{sample}.bam'
    benchmark: 'benchmarks/export_read_length/{sample}.txt'
    output: 'read_lengths/{sample}.tsv'
    shell:
        r'''
        riboraptor export-read-length \
        --bam {input} \
        --saveto {output}
        '''


rule plot_read_length:
    input: 'read_lengths/{sample}.tsv'
    benchmark: 'benchmarks/plot_read_length/{sample}.txt'
    output: 'plots/read_length/{sample}.png'
    shell:
        r'''
        riboraptor plot-read-length \
        --millify_labels \
        --read-lengths {input} --saveto {output}
        '''


rule export_metagene:
    input: 'bigWigs/{sample}.bw'
    benchmark: 'benchmarks/export_metagene/{sample}.txt'
    output: 'metagene_coverages/{sample}.tsv'
    params:
        orientation = '5prime'
    shell:
        r'''
        riboraptor export-metagene-coverage \
        --bw {input} \
        --bed {CDS_BED} --saveto {output} \
        --orientation {params.orientation}
        '''


rule plot_metagene:
    input: 'metagene_coverages/{sample}.tsv'
    benchmark: 'benchmarks/plot_metagene/{sample}.txt'
    output: 'plots/metagene/{sample}.png'
    shell:
        r'''
        riboraptor plot-metagene \
        --counts {input} \
        --saveto {output} \
        --positions -60:100
        '''


rule metagene_coverage_cds:
    input: 'bams_unique/{sample}.bam'
    benchmark: 'benchmarks/metagene_coverage_cds/{sample}.txt'
    output: 'genewise_counts_CDS/{sample}.tsv'
    shell:
        r'''
        riboraptor count-reads-bed \
        --bam {input} \
        --bed {CDS_BED} \
        --saveto {output}
        '''


rule infer_protocol:
    input: 'bams_unique/{sample}.bam'
    benchmark: 'benchmarks/infer_protocol/{sample}.txt'
    output: 'inferred_protocol/{sample}.txt'
    shell:
        r'''
        riboraptor infer-protocol \
        --bam {input} \
        --refseq {GENE_BED} > {output}
        '''


rule create_hdf:
    input: 'bams_unique/{sample}.bam'
    benchmark: 'benchmarks/create_hdf/{sample}.txt'
    output:
        hdf = 'hdf/{sample}.hdf5',
        tsv = 'hdf/{sample}.tsv',
    params:
        prefix = 'hdf/{sample}'
    shell:
        r'''riboraptor bam-coverage \
        --bam {input} \
        --genebed {GENE_BED} \
        --outprefix {params.prefix}
        '''


rule hdf_to_bw:
    input: 'hdf/{sample}.hdf5'
    benchmark: 'benchmarks/hdf_to_bw/{sample}_{fragment_length}.txt'
    output:
        prime5_pos = 'bigWig_lengthwise/{sample}/{fragment_length}/5prime_pos.bw',
        prime3_pos = 'bigWig_lengthwise/{sample}/{fragment_length}/3prime_pos.bw',
        prime5_neg = 'bigWig_lengthwise/{sample}/{fragment_length}/5prime_neg.bw',
        prime3_neg = 'bigWig_lengthwise/{sample}/{fragment_length}/3prime_neg.bw',
        prime5_combined = 'bigWig_lengthwise/{sample}/{fragment_length}/5prime_combined.bw',
        prime3_combined = 'bigWig_lengthwise/{sample}/{fragment_length}/3prime_combined.bw',
    params:
        prefix = 'bigWig_lengthwise/{sample}',
    shell:
        r'''riboraptor hdf-to-bw \
        --hdf {input} \
        --prefix {params.prefix} \
        --readlength {wildcards.fragment_length} \
        && touch {output}'''


rule export_metagene_individual:
    input: 'bigWig_lengthwise/{sample}/{fragment_length}/{orientation}_{strand}.bw'
    benchmark: 'benchmarks/hdf_to_bw/{sample}_{fragment_length}_{orientation}_{strand}.txt'
    output: 'metagene_coverage_lengthwise/{sample}/{fragment_length}/{orientation}_{strand}.tsv'
    params:
        orientation = '{orientation}'
    run:
        if os.stat(str(input)).st_size:
            shell(r'''
            riboraptor export-metagene-coverage \
            --bw {input} \
            --bed {CDS_BED} \
            --orientation {params.orientation} \
            --saveto {output} \
            --offset_5p 60 \
            --offset_3p 60
            ''')
        else:
            # Just touch the file
            shell(r'''touch {output}''')

rule normalize_bigwig:
    input:
        bw = 'bigWig_lengthwise/{sample}/{fragment_length}/{orientation}_{strand}.bw',
        hdf = 'hdf/{sample}.hdf5'
    benchmark: 'benchmarks/normalize_bigwig/{sample}_{fragment_length}_{orientation}_{strand}.txt'
    output: 'bigWig_lengthwise_normalized/{sample}/{fragment_length}/{orientation}_{strand}.bw'
    wrapper:
        get_wrapper('normalize_bigwig_wrapper')


rule plot_metagene_individual_length:
    input: 'metagene_coverage_lengthwise/{sample}/{fragment_length}/{orientation}_{strand}.tsv'
    benchmark: 'benchmarks/plot_metagene_individual_length/{sample}_{fragment_length}_{orientation}_{strand}.txt'
    output: 'plots/metagene_lengthwise/{sample}/{fragment_length}/{orientation}_{strand}.png'
    wrapper:
        get_wrapper('plot_metagene_wrapper')


rule featurecounts:
    input:
        bams = expand('bams/{sample}.bam',
                      sample=SRX_SAMPLES),
        hdfs = expand('hdf/{sample}.hdf5',
                      sample=SRX_SAMPLES)
    benchmark: 'benchmarks/featurecounts/featurecounts.txt'
    output: 'featureCounts/fcounts.tsv'
    params:
        annotation  = GTF
    threads: 16
    wrapper:
        get_wrapper('featurecounts_wrapper')


rule multiqc:
    input:
        #fastqc = expand('qc/{sample}_fastqc.html',
        #                sample=SRX_SAMPLES),
        #trimmed = expand('merged_preprocessed/{sample}_trimmed.fq.gz',
        #                 sample=SRX_SAMPLES),
        bams = expand('bams/{sample}.bam',
                      sample=SRX_SAMPLES),
        counts = 'featureCounts/fcounts.tsv',
    output:
        dynamic('multiqc_report/multiqc_plots/png/{plot}.png')
    resources:
        mem_mb=10000
    shell:
        r'''
        multiqc --export -f --outdir multiqc_report starlogs featureCounts bams
        '''


rule aggregate_multiqc_plots:
    input:
        dynamic_input = get_multiqc_report_input,
        counts = 'featureCounts/fcounts.tsv',
    benchmark: 'benchmarks/aggregate_multiqc_plots/multiqc.txt'
    output:
        html='multiqc_report/aggregated_report.html'
    wrapper:
        get_wrapper('multiqc_report_wrapper')



rule periodicity_matrix:
    input: expand('metagene_coverage_lengthwise/{{sample}}/{fragment_length}/5prime_combined.tsv', fragment_length=FRAGMENT_LENGTHS)
    output: 'periodicity/{sample}.tsv'
    wrapper:
        get_wrapper('periodicity_wrapper')


rule plot_periodicity_matrix:
    input: 'periodicity/{sample}.tsv'
    output: 'plots/periodicity/{sample}.png'
    shell:
        r'''riboraptor plot-periodicity-df --df {input} --saveto {output}'''


rule qc_report:
    input:
      metagene = 'plots/metagene/{sample}.png',
      fragment_length = 'plots/read_length/{sample}.png',
      prime5_pos = expand('plots/metagene_lengthwise/{{sample}}/{fragment_length}/5prime_pos.png',
                          fragment_length=FRAGMENT_LENGTHS),
      prime3_pos = expand('plots/metagene_lengthwise/{{sample}}/{fragment_length}/3prime_pos.png',
                          fragment_length=FRAGMENT_LENGTHS),
      prime5_neg = expand('plots/metagene_lengthwise/{{sample}}/{fragment_length}/5prime_neg.png',
                          fragment_length=FRAGMENT_LENGTHS),
      prime3_neg = expand('plots/metagene_lengthwise/{{sample}}/{fragment_length}/3prime_neg.png',
                          fragment_length=FRAGMENT_LENGTHS),
      prime5_combined = expand('plots/metagene_lengthwise/{{sample}}/{fragment_length}/5prime_combined.png',
                               fragment_length=FRAGMENT_LENGTHS),
      prime3_combined = expand('plots/metagene_lengthwise/{{sample}}/{fragment_length}/3prime_combined.png',
                               fragment_length=FRAGMENT_LENGTHS),
    benchmark: 'benchmarks/qc_report/{sample}.txt'
    output:
        html='reports/{sample}_riboraptor_report.html'
    wrapper:
        get_wrapper('qc_report_wrapper')


rule qc_report_summarized:
    input:
      metagene = expand('plots/metagene/{sample}.png', sample=SRX_SAMPLES),
      fragment_length = expand('plots/read_length/{sample}.png', sample=SRX_SAMPLES),
      periodicity = expand('plots/periodicity/{sample}.png', sample=SRX_SAMPLES),
    benchmark: 'benchmarks/qc_report_summarized/benchmark.txt'
    output:
        html='reports/riboraptor_report_summarized.html'
    wrapper:
        get_wrapper('qc_report_summarizer_wrapper')
