
# abort if no config is set
if len(config) == 0:

    raise ValueError('config not set, make sure to run snakemake with config file')

snakemake_dir = sys.path[0]

output_dir = config['output_dir']
alignment_dir = output_dir + '/alignment'
log_dir = output_dir + '/logs'
benchmark_dir = log_dir + '/benchmarks'
methylation_dir = output_dir + '/methylation'
qc_dir = output_dir + '/qc'
mbias_dir = qc_dir + '/mbias'
segmentation_dir = output_dir + '/segmentation'
dmr_dir = output_dir + '/dmr'
metilene_dir = dmr_dir + '/metilene'
bsseq_dir = dmr_dir + '/bsseq'
camel_dir = dmr_dir + '/camel'

cgi_annotation_file = snakemake_dir + '/' + config['cgi_annotation_file']
gene_annotation_file = snakemake_dir + '/' + config['gene_annotation_file']
repeat_masker_annotation_file = snakemake_dir + '/' + config['repeat_masker_annotation_file']
transcript_start_site_file = snakemake_dir + '/' + config['transcript_start_site_file']

num_threads = config['computing_threads']
io_threads = config['io_threads']

ref_fasta = config['ref']

target_files = [ output_dir + '/' + outfile for outfile in config['target_files']]
runtime_report = qc_dir + '/runtime.pdf'
mbias_ob_reports = expand(mbias_dir + '/{sample}_OB.svg', sample = config['samples'])
mbias_ot_reports = expand(mbias_dir + '/{sample}_OT.svg', sample = config['samples'])

rule all:
    input:
        target_files +
        [ runtime_report ] +
        mbias_ob_reports +
        mbias_ot_reports

### ALIGNMENT

rule find_fqs:
    output:
        first  = output_dir +'/raw/{sample}.first.txt',
        second = output_dir +'/raw/{sample}.second.txt'
    params:
        rawdir = config['rawdir'],
        first_regex = config['rawsuffixregex']['first'],
        second_regex = config['rawsuffixregex']['second']
    shell:
        """
        find {params.rawdir} | grep -E '.*{wildcards.sample}.*{params.first_regex}'  | sort > {output.first}
        find {params.rawdir} | grep -E '.*{wildcards.sample}.*{params.second_regex}' | sort > {output.second}
        """

rule bwameth_indexing:
    input:
        ref = ref_fasta
    output:
        indices = [
            ref_fasta + '.bwameth.c2t',
            ref_fasta + '.bwameth.c2t.amb',
            ref_fasta + '.bwameth.c2t.ann',
            ref_fasta + '.bwameth.c2t.bwt',
            ref_fasta + '.bwameth.c2t.pac',
            ref_fasta + '.bwameth.c2t.sa',
            ref_fasta + '.fai'
        ]
    log:
        log_dir + '/bwameth_indexing.log'
    shell:
        """
        bwameth.py index {input} 2> {log}
        samtools faidx {input} 2> {log}
        """

rule bwameth_align:
    input:
        ref = ref_fasta,
        indices = [
            ref_fasta + '.bwameth.c2t',
            ref_fasta + '.bwameth.c2t.amb',
            ref_fasta + '.bwameth.c2t.ann',
            ref_fasta + '.bwameth.c2t.bwt',
            ref_fasta + '.bwameth.c2t.pac',
            ref_fasta + '.bwameth.c2t.sa',
            ref_fasta + '.fai'
        ],
        first  = output_dir +'/raw/{sample}.first.txt',
        second = output_dir +'/raw/{sample}.second.txt'
    output:
        temp(alignment_dir + '/{sample}.unsorted.bam')
    log:
        log_dir + '/{sample}.align.log'
    threads:
        num_threads
    benchmark:
        benchmark_dir + '/alignment/{sample}.tsv'
    shell:
        """
        FQ1=$(cat {input.first} | tr '\n' ',')
        FQ2=$(cat {input.second} | tr '\n' ',')
        bwameth.py -t {threads} --reference {input.ref} $FQ1 $FQ2 2> {log} | samtools view -b - > {output}
        """


rule sort_bam:
    input:
        rules.bwameth_align.output
    output:
        temp(alignment_dir + '/{sample}.sorted.bam')
    log:
        log_dir + '/{sample}.sorting.log'
    threads:
        num_threads
    benchmark:
        benchmark_dir + '/sorting/{sample}.tsv'
    shell:
        'samtools sort -o {output} -@ {threads} {input} 2> {log}'


rule mark_duplicates:
    input:
        rules.sort_bam.output
    output:
        bam = alignment_dir + '/{sample}.bam',
        metrics = temp(alignment_dir + '/{sample}-dup-metrics.txt')
    log:
        log_dir + '/{sample}.mark_duplicates.log'
    benchmark:
        benchmark_dir + '/deduplication/{sample}.tsv'
    shell:
        'picard MarkDuplicates I={input} O={output.bam} M={output.metrics} &> {log}'


rule index_bam:
    input:
        bam = alignment_dir + '/{sample}.bam'
    output:
        alignment_dir + '/{sample}.bai'
    benchmark:
        benchmark_dir + '/indexing/{sample}.tsv'
    shell:
        'samtools index {input} {output}'


### QC

rule fastqc:
    input:
        bam = alignment_dir + '/{sample}.bam'
    output:
        qc_dir + '/fastqc/{sample}_fastqc.html'
    log:
        log_dir + '/{sample}.fastqc.log'
    benchmark:
        benchmark_dir + '/fastqc/{sample}.tsv'
    shell:
        'fastqc -o $(dirname {output}) {input.bam} &> {log}'

rule picard_metrics:
    input:
        ref = ref_fasta,
        bam = alignment_dir + '/{sample}.bam'
    output:
        alignment = qc_dir +'/picard-metrics/{sample}-alignment.txt',
        insert    = qc_dir +'/picard-metrics/{sample}-insert-size.txt',
        hist      = temp(qc_dir +'/picard-metrics/{sample}-hist.pdf')
    log:
        log_dir + '/{sample}.picard_metrics.log'
    benchmark:
        benchmark_dir + '/picard_metrics/{sample}.tsv'
    shell:
        """
        picard CollectAlignmentSummaryMetrics R={input.ref} I={input.bam} O={output.alignment} &> {log}
        picard CollectInsertSizeMetrics I={input.bam} O={output.insert} H={output.hist} &> {log}
        """

rule qualimap:
    input:
        alignment_dir + '/{sample}.bam'
    output:
        qc_dir + '/qualimap/{sample}/qualimapReport.html'
    threads:
        io_threads
    log:
        log_dir + '/{sample}.qualimap.log'
    benchmark:
        benchmark_dir + '/qualimap/{sample}.tsv'
    params:
        memory = config['qualimap_memory_gb']
    shell:
        'qualimap bamqc -bam {input} -outdir $(dirname {output}) -nt {threads} --collect-overlap-pairs --skip-duplicated --java-mem-size={params.memory}G &> {log} '

rule mbias:
    input:
        fasta_ref = ref_fasta,
        fasta_index = ref_fasta + '.fai',
        bam_file = alignment_dir + '/{sample}.bam',
        bam_index = alignment_dir + '/{sample}.bai'
    log:
        log_dir + '/{sample}_mbias.log'
    benchmark:
        benchmark_dir + '/{sample}_mbias.tsv'
    output:
        mbias_ot = mbias_dir + '/{sample}_OT.svg',
        mbias_ob = mbias_dir + '/{sample}_OB.svg'
    shell:
        'MethylDackel mbias {input.fasta_ref} {input.bam_file} {mbias_dir}/{wildcards.sample} &> {log}'

rule multiqc:
    input:
        fastqc           = expand(qc_dir + '/fastqc/{sample}_fastqc.html', sample = config['samples']),
        alnMetrics       = expand(qc_dir + '/picard-metrics/{sample}-alignment.txt', sample = config['samples']),
        insertMetrics    = expand(qc_dir + '/picard-metrics/{sample}-insert-size.txt', sample = config['samples']),
        qualimap         = expand(qc_dir + '/qualimap/{sample}/qualimapReport.html', sample = config['samples'])
    output:
        qc_dir + '/multiqc_report.html'
    log:
        log_dir + '/multiqc.log'
    benchmark:
        benchmark_dir + '/multiqc.tsv'
    shell:
        'multiqc -f -o {qc_dir} {qc_dir}/fastqc {qc_dir}/picard-metrics {qc_dir}/qualimap &> {log} '

rule methylation_metrics:
    input:
        bed_graphs = expand(methylation_dir + '/{sample}_CpG.bedGraph', sample = config['samples'])
    output:
        methylation_metrics = qc_dir + '/methylation_metrics.csv'
    log:
        log_dir + '/methylation_metrics.log'
    params:
        methylation_rate_on_chromosomes = config['methylation_rate_on_chromosomes']
    benchmark:
        benchmark_dir + '/methylation_metrics.tsv'
    script:
        'scripts/methylationMetrics.R'

### DMR CALLING

rule methyl_dackel:
    input:
        ref = ref_fasta,
        bai = alignment_dir + '/{sample}.bai',
        bam = alignment_dir + '/{sample}.bam'
    output:
        methylation_dir + '/{sample}_CpG.bedGraph'
    params:
        min_cov  = config['min_cov'],
        methylation_dir = methylation_dir
    benchmark:
        benchmark_dir + '/methylation_calling/{sample}.tsv'
    shell:
        'MethylDackel extract --mergeContext -o {params.methylation_dir}/{wildcards.sample} {input.ref} {input.bam}'

rule bedgraph_to_methylation_ratio:
    input:
        bedGraph = methylation_dir + '/{sample}_CpG.bedGraph'
    output:
        bedGraph = temp(methylation_dir + '/{sample}_CpG_ratio.bedGraph')
    log:
        log_dir + '/{sample}.bedgraph_to_methylation_ratio.log'
    benchmark:
        benchmark_dir + '/metilene/{sample}.transformation.tsv'
    script:
        'scripts/transformBedGraph.R'

rule metilene_input:
    input:
        expand(methylation_dir + '/{sample}_CpG_ratio.bedGraph', sample = config['group1']),
        expand(methylation_dir + '/{sample}_CpG_ratio.bedGraph', sample = config['group2'])
    output:
        metilene_dir + '/metilene-input.bedGraph'
    params:
        g1_header = ' '.join([ 'g1_' + s for s in config['group1']]),
        g2_header = ' '.join([ 'g2_' + s for s in config['group2']])
    benchmark:
        benchmark_dir + '/metilene/input.tsv'
    shell:
        'bedtools unionbedg -filler NA -header -names {params.g1_header} {params.g2_header} -i {input} > {output}'

rule metilene:
    input:
        metilene_dir + '/metilene-input.bedGraph'
    output:
        metilene_dir + '/dmrs.csv'
    params:
        min_cpg = config['min_cpg'],
        min_diff = config['min_diff']
    threads:
        1
    log:
        log_dir + '/metilene.log'
    benchmark:
        benchmark_dir + '/metilene/calling.tsv'
    shell:
        'metilene -m {params.min_cpg} -d {params.min_diff} -t {threads} {input} > {output} 2> {log}'

rule camel_index:
    input:
        ref = ref_fasta,
    output:
        camel_h5 = ref_fasta + '.h5'
    params:
        snakemake_dir = snakemake_dir
    shell:
        'python {params.snakemake_dir}/submodules/camel/modules/index.py {input.ref} {output}'

rule camel_call:
    input:
        index = ref_fasta + '.h5',
        bam = alignment_dir + '/{sample}.bam',
        bai = alignment_dir + '/{sample}.bai'
    output:
        camel_dir + '/{sample}.h5'
    params:
        snakemake_dir = snakemake_dir
    log:
        log_dir + '/{sample}.camel_call.log'
    benchmark:
        benchmark_dir + '/camel/{sample}.tsv'
    shell:
        'python {params.snakemake_dir}/submodules/camel/modules/call.py {input.bam} {input.index} {output} 2> {log}'

rule camel_dmr:
    input:
        index   = ref_fasta + '.h5',
        control = expand(camel_dir + '/{sample}.h5', sample = config['group1']),
        case    = expand(camel_dir + '/{sample}.h5', sample = config['group2'])
    output:
        camel_dir + '/dmrs.csv'
    params:
        min_cpg  = config['min_cpg'],
        min_diff = config['min_diff'],
        min_cov  = config['min_cov'],
        snakemake_dir = snakemake_dir
    log:
        log_dir + '/camel.log'
    benchmark:
        benchmark_dir + '/camel/dmr.tsv'
    shell:
        'python {params.snakemake_dir}/submodules/camel/modules/dmr.py {input.index} --case {input.case} --control {input.control} --min_diff {params.min_diff} --min_cpg {params.min_cpg} --min_cov {params.min_cov} > {output} 2> {log}'

rule bsseq:
    input:
        meth = expand(methylation_dir + '/{sample}_CpG.bedGraph', sample = config['samples'])
    output:
        rdata = bsseq_dir + '/bsseq.Rdata',
        csv   = bsseq_dir + '/dmrs.csv',
        pdf   = bsseq_dir + '/top100.pdf'
    threads:
        config['io_threads']
    params:
        local_correct = config['bsseq_local_correct']
    log:
        log_dir + '/bsseq.log'
    benchmark:
        benchmark_dir + '/bsseq/bsseq.tsv'
    script:
        'scripts/bsseq.R'

### DMR comparison

def expected_tool_files():
    return {tool: dmr_dir + '/' + tool + '/dmrs.csv' for tool in config['dmr_tools']}

rule dmr_combination:
    input:
        **expected_tool_files()
    output:
        csv = dmr_dir + '/combined-dmrs.csv',
        bed = dmr_dir + '/dmr-coverage/combined-dmrs.bed'
    log:
        log_dir + '/dmr_combination.log'
    benchmark:
        benchmark_dir + '/dmr_annotation/combination.tsv'
    script:
        'scripts/dmrCombination.R'

rule dmr_coverage:
    input:
        bed = dmr_dir + '/dmr-coverage/combined-dmrs.bed',
        bam = alignment_dir + '/{sample}.bam',
        bai = alignment_dir + '/{sample}.bai'
    output:
        dmr_dir + '/dmr-coverage/{sample}.regions.bed.gz'
    params:
        min_mapq = config['annotation_min_mapq']
    threads: config['computing_threads']
    benchmark:
        benchmark_dir + '/dmr_coverage/{sample}.tsv'
    shell:
        'mosdepth --threads {threads} --no-per-base --mapq {params.min_mapq} --by {input.bed} $(dirname {output})/{wildcards.sample} {input.bam}'

rule repeat_masker_download:
    output:
        gzipped_csv = snakemake_dir + '/' + config['repeat_masker_annotation_file']
    params:
        link = config['repeat_masker_links'][config['repeat_masker_annotation_file']]
    log:
        log_dir + '/repeat_masker_download.log'
    shell:
        'wget -O {output} {params.link} 2> {log}'

rule dmr_annotation:
    input:
        cgi_annotation_file = cgi_annotation_file,
        gene_annotation_file = gene_annotation_file,
        repeat_masker_annotation_file = repeat_masker_annotation_file,
        transcript_start_site_file = transcript_start_site_file,
        coverages = expand(dmr_dir + '/dmr-coverage/{sample}.regions.bed.gz', sample = config['samples']),
        combined_dmrs = dmr_dir + '/combined-dmrs.csv'
    output:
        annotated_dmrs = dmr_dir + '/annotated-dmrs.csv'
    params:
        biotypes = config['annotation_allowed_biotypes'],
        tss_distances = config['promoter_tss_distances']
    log:
        log_dir + '/dmr_annotation.log'
    benchmark:
        benchmark_dir + '/dmr_annotation/annotation.tsv'
    script:
        'scripts/dmrAnnotation.R'

### SEGMENTATION

rule methylseekr:
    input:
        ref = ref_fasta,
        methylation_table = expand(methylation_dir + '/{sample}_CpG.bedGraph', sample = config['samples']),
        cgi_annotation_file = cgi_annotation_file,
        gene_annotation_file = gene_annotation_file,
        repeat_masker_annotation_file = repeat_masker_annotation_file,
        transcript_start_site_file = transcript_start_site_file
    output:
        pmd_all = segmentation_dir + '/pmd-all.csv',
        umr_lmr_all = segmentation_dir + '/umr-lmr-all.csv',
        pmd_segments = expand(segmentation_dir + '/{sample}/pmd-segments.csv', sample = config['samples']),
        umr_lmr_with_pmd = expand(segmentation_dir + '/{sample}/LMRUMRwithPMD/umr-lmr.csv', sample = config['samples']),
        umr_lmr_without_pmd = expand(segmentation_dir + '/{sample}/LMRUMRwithoutPMD/umr-lmr.csv', sample = config['samples'])
    params:
        target_dir = segmentation_dir,
        samples = config['samples'],
        methylation_dir = methylation_dir,
        calibration_chr = config['methylseekr_pmd_chromosome'],
        genome_build = config['methylseekr_cgi_genome'],
        biotypes = config['annotation_allowed_biotypes'],
        tss_distances = config['promoter_tss_distances'],
        min_coverage = config['min_cov'],
        methylation_cutoff = config['methylseekr_methylation_cutoff'],
        fdr_cutoff = config['methylseekr_fdr_cutoff']
    log:
        log_dir + '/methylseekr.log'
    benchmark:
        benchmark_dir + '/methylseekr.tsv'
    threads:
        num_threads
    script:
        'scripts/methylseekRSegmentation.R'

### BENCHMARKING

rule benchmark_plot:
    input:
        target_files = target_files
    output:
        runtime_report = runtime_report
    params:
        benchmark_dir = benchmark_dir
    script:
        'scripts/benchmarkPlot.R'
