#!/bin/python

intro_string = r'''
        _______       _____      _____              _     _           
     /\|__   __|/\   / ____|    |  __ \            | |   (_)          
    /  \  | |  /  \ | |   ______| |  | | ___ ______| |__  _  __ _ ___ 
   / /\ \ | | / /\ \| |  |______| |  | |/ _ \______| '_ \| |/ _` / __|
  / ____ \| |/ ____ \ |____     | |__| |  __/      | |_) | | (_| \__ \
 /_/    \_\_/_/    \_\_____|    |_____/ \___|      |_.__/|_|\__,_|___/
                                                                      
Pipeline to debias, aggregate, and analyze ATAC-seq data.
Author: Allen Lynch
Date: 10/10/2020
'''

print(intro_string)
import os
import sys
from snakemake.utils import validate

configfile: "./config/config.json"

if not 'skiptest' in config or not config['skiptest'] == True:
    print("Executing with parameters:\n" + '\n'.join([str(k) + ': ' + str(v) for k, v in config.items()]))
    file_params = ['cutsite_model','fragment_model','reference']
    for file_param in file_params:
        if not os.path.isfile(config[file_param]):
            print('\nERROR: path "{path}" provided for param "{param}" is invalid.'\
                .format(path = str(config[file_param]), param = file_param), file = sys.stderr)
            exit(1)

wildcard_constraints:
    chrom="chr[a-zA-Z0-9_]+"

rule all:
    input:
        "{dataset}/{sample}.done"

rule sort_fragments:
    input:
        "{dataset}/{sample}.bed"
    output:
        temp("{dataset}/{sample}.sorted.bed")
    threads: 8
    resources: cpus=config["sort_cpus"], mem_mb=config["sort_mem"], time_min=300
    shell:
        "sort -k1,1 -k2,2n {input} > {output}"

rule call_peaks:
    input:
        "{dataset}/{sample}.bed"
    output:
        temp("{dataset}/{sample}_peaks.bed")
    resources: cpus=1, mem_mb=8000, time_min=180
    conda: "envs/macs2.yaml"
    shell:
        "macs2 callpeak -t \"{input}\" -f BEDPE -g {config[genomesize]} --keep-dup all -B -q 0.05 --nomodel --extsize=50 --SPMR -n {wildcards.sample} --outdir \"{wildcards.dataset}/{wildcards.sample}_peak_calling/\" && "
        "cat {wildcards.dataset}/{wildcards.sample}_peak_calling/{wildcards.sample}_peaks.narrowPeak | cut -f1-3 > {output} &&"
        "rm -rf -r {wildcards.dataset}/{wildcards.sample}_peak_calling"

rule sort_peaks:
    input:
        ancient("{dataset}/{sample}_peaks.bed")
    output:
        "{dataset}/sorted_{sample}_peaks.bed"
    threads: 4
    resources: cpus=4
    shell:
        "sort -k1,1 -k2,2n {input} > {output}"

rule count_barcodes:
    input:
        fragments="{dataset}/{sample}.sorted.bed",
        peaks="{dataset}/sorted_{sample}_peaks.bed"
    output:
        temp("{dataset}/{sample}_barcodes.txt")
    threads: 8
    resources: 
        cpus=config["sort_cpus"], mem_mb=config["sort_mem"], time_min=300
    params:
        peak_frac=int(1/config['min_peak_proportion'])
    conda: "envs/bedtools.yaml"
    shell:
        "bedtools intersect -a {input.fragments} -b {input.peaks} -wa -wb -sorted | cut -f4,6,7,8 | sort -k1,1 -k2,2 -k3,3n -k4,4n | uniq | cut -f1 | uniq -c | "
        "sort -k1,1n | awk -v min_peaks=$(( $(wc -l {input.peaks} | awk '{{print $1}}') / {params.peak_frac} )) '$1 >= min_peaks' "
        "| tail -n {config[num_cells]} | awk '{{print $2}}' > {output}"

checkpoint split_bedfile:
    input: 
        bedfile="{dataset}/{sample}.sorted.bed",
        barcodes="{dataset}/{sample}_barcodes.txt"
    output:
        temp(directory("{dataset}/{sample}_processing/chromosome_fragments/"))
    threads: 1
    conda: "envs/bedtools.yaml"
    shell:
        "cat {input.bedfile} | cut -f1-5 | python workflow/scripts/filter_by_barcode.py - {input.barcodes} | python workflow/scripts/split_bed.py -d {output} --pattern {{chrom}}-{wildcards.sample}.sorted.bed"

#__APPLY MODELS________

rule mark_duplicate_fragments:
    input:
        "{dataset}/{sample}_processing/chromosome_fragments/{chrom}-{sample}.sorted.bed"
    output:
        temp("{dataset}/{sample}_processing/dupmarked_chromosome_fragments/{chrom}-{sample}.sorted.mkdups.bed")
    threads: 1
    shell:
        "python workflow/scripts/fast_bulk_deduplicate.py {input} --mark > {output}"

rule get_plus_strand_nucleotide_sequences:
    input:
        "{dataset}/{sample}_processing/dupmarked_chromosome_fragments/{chrom}-{sample}.sorted.mkdups.bed"
    output:
        temp("{dataset}/{sample}_processing/nucleotide_sequences/{chrom}-{sample}_plus_strand_nucs.txt")
    threads: 1
    conda: "envs/pyfaidx.yaml"
    shell:
        "workflow/scripts/get_cut_centers.sh {input} + | workflow/scripts/expand_cut_centers.sh - 4 9 |"
        "faidx {config[reference]} --bed - --no-name --lazy -s N > {output}"

rule get_minus_strand_nucleotide_sequences:
    input:
        "{dataset}/{sample}_processing/dupmarked_chromosome_fragments/{chrom}-{sample}.sorted.mkdups.bed"
    output:
        temp("{dataset}/{sample}_processing/nucleotide_sequences/{chrom}-{sample}_minus_strand_nucs.txt")
    threads: 1
    conda: "envs/pyfaidx.yaml"
    shell:
        "workflow/scripts/get_cut_centers.sh {input} - | workflow/scripts/expand_cut_centers.sh - 4 9 |"
        "faidx {config[reference]} --bed - --no-name --reverse --complement --lazy -s N > {output}"

rule get_cutsite_bias:
    input:
        "{dataset}/{sample}_processing/nucleotide_sequences/{chrom}-{sample}_{strand}_strand_nucs.txt"
    output:
        temp("{dataset}/{sample}_processing/fragment_features/{chrom}-{sample}_{strand}_strand_biases.txt")
    threads: 4
    resources: cpus=4, time_min=300, mem_mb=4000
    conda: "envs/sklearn.yaml"
    shell:
        "python workflow/scripts/predict_cutsite_bias.py -s {input} -c {resources.cpus} -m {config[cutsite_model]} -e 390 > {output}"

rule get_fragment_gc:
    input:
        "{dataset}/{sample}_processing/dupmarked_chromosome_fragments/{chrom}-{sample}.sorted.mkdups.bed"
    output:
        temp("{dataset}/{sample}_processing/fragment_features/{chrom}-{sample}.gc_content.tsv")
    threads: 1
    conda: "envs/pyfaidx.yaml"
    shell:
        "faidx {config[reference]} --bed {input} --transform nucleotide --lazy -s N | cut -f 4-9 | "
        "awk '{{OFS=\"\t\"}} NR>1 {{print ($1+$2+$3+$4 > 0) ? ($3+$4)/($1+$2+$3+$4) : \"nan\", $1+$2+$3+$4+$5}}' > {output}"

rule get_fragment_duprate:
    input:
        bedfile="{dataset}/{sample}_processing/dupmarked_chromosome_fragments/{chrom}-{sample}.sorted.mkdups.bed",
        gc_content="{dataset}/{sample}_processing/fragment_features/{chrom}-{sample}.gc_content.tsv",
        plus_strand="{dataset}/{sample}_processing/fragment_features/{chrom}-{sample}_plus_strand_biases.txt",
        minus_strand="{dataset}/{sample}_processing/fragment_features/{chrom}-{sample}_minus_strand_biases.txt",
        peaks="{dataset}/sorted_{sample}_peaks.bed"
    output:
        temp("{dataset}/{sample}_processing/chromosome_duprates/{chrom}-{sample}.sorted.mkdups.bias.bed")
    conda: "envs/fragment_bias_env.yaml"
    shell:
        "paste {input.bedfile} {input.plus_strand} {input.minus_strand} {input.gc_content} | "
        "python workflow/scripts/predict_fragment_level_bias.py -b - -m {config[fragment_model]} --fillnan |"
        "bedtools intersect -a - -b {input.peaks} -loj -sorted > {output}"

#__REAGGREGATE SUBDAGS_____
def aggregate_chromosomes(wildcards):
    checkpoint_output = checkpoints.split_bedfile.get(**wildcards).output[0]
    return expand("{dataset}/{sample}_processing/chromosome_duprates/{chrom}-{sample}.sorted.mkdups.bias.bed",
           sample=wildcards.sample,
           dataset=wildcards.dataset,
           chrom=sorted(glob_wildcards(os.path.join(checkpoint_output, "{chrom}-{sample}.sorted.bed")).chrom))

rule concatenate_chromosomes:
    input:
        aggregate_chromosomes
    output: 
        "{dataset}/{sample}.sorted.mkdups.bias.bed"
    shell:
        "echo {input} | xargs cat > {output}"

rule filter_peak_fragments:
    input:
        "{dataset}/{sample}.sorted.mkdups.bias.bed"
    output:
        temp("{dataset}/{sample}.sorted.mkdups.bias.filtered.bed")
    shell:
        "awk '$11!=\".\"' {input} > {output}"

rule calculate_lambda_bar:
    input:
        "{dataset}/{sample}.sorted.mkdups.bias.filtered.bed"
    output:
        temp("{dataset}/{sample}_processing/average_bias.txt")
    shell:
        "awk 'BEGIN {{sum=0; num=0}} $10!=\"nan\" {{sum=sum+$10; num=num+1}} END {{print sum/num}}' {input} > {output}"

rule correct_bias:
    input:
        bedfile="{dataset}/{sample}.sorted.mkdups.bias.filtered.bed",
        lambda_bar="{dataset}/{sample}_processing/average_bias.txt"
    output:
        "{dataset}/{sample}.sorted.mkdups.bias.corrected.bed"
    shell:
        "awk -v lambda=$(cat {input.lambda_bar}) '{{OFS=\"\t\"}} {{print $0,($10!=\"nan\") ? exp(lambda)/exp($10) : \"1\"}}' {input.bedfile} > {output}"

#__AGGREGATION AND ANALYSIS PREPARATION______

rule get_barcode_stats:
    input:
        "{dataset}/{sample}.sorted.mkdups.bias.filtered.bed"
    output:
        temp("{dataset}/{sample}_analysis/{sample}_barcode_stats.tsv")
    threads: 4
    resources:
        cpus=config["sort_cpus"], mem_mb=config["sort_mem"], time_min=300
    conda: "envs/sparse.yaml"
    shell:
        "cat {input} | cut -f4,10 | sort -k1,1 | "
        "python workflow/scripts/memefficient_agg.py -g 1 -c 2 2 2 -o mean count sample | "
        "sort -k3,3n | tail -n{config[num_cells]} > {output}"

rule get_peak_stats:
    input:
        "{dataset}/{sample}.sorted.mkdups.bias.filtered.bed"
    output:
        temp("{dataset}/{sample}_analysis/{sample}_peak_stats.tsv")
    threads: 4
    resources:
        cpus=config["sort_cpus"], mem_mb=config["sort_mem"], time_min=300
    conda: "envs/sparse.yaml"
    shell:
        "cat {input} | cut -f10-14 | sort -k2,2 -k3,3n -k4,4n | awk '{{OFS=\"\t\"}} {{print ($2 \"_\" $3 \"_\" $4),$1}}' | "
        "python workflow/scripts/memefficient_agg.py -g 1 -c 2 2 2 -o median count sample --ignore_warnings > {output}"

def set_cnt_matrix_flag(wildcards, *args):
    if wildcards.corrected == 'corrected':
        return '--corrected'
    elif wildcards.corrected == 'dedup':
        return '--dedup'
    elif wildcards.corrected == 'raw':
        return ''
    else:
        raise Exception()

rule make_count_matrix:
    input:
        barcodes="{dataset}/{sample}_analysis/{sample}_barcode_stats.tsv",
        peaks="{dataset}/{sample}_analysis/{sample}_peak_stats.tsv",
        fragments="{dataset}/{sample}.sorted.mkdups.bias.corrected.bed"
    output:
        temp("{dataset}/{sample}_analysis/{sample}_{corrected}_counts.npz")
    conda: "envs/sparse.yaml"
    params: 
        corrected=set_cnt_matrix_flag
    resources: mem_mb=8000
    shell:
        "python workflow/scripts/make_sparse_counts.py -b <(cat {input.barcodes} | cut -f1) -p <(cat {input.peaks} | cut -f1) -f {input.fragments} -o {output} {params.corrected}"

rule count_duplicate_groups:
    input:
        fragments="{dataset}/{sample}.sorted.mkdups.bias.corrected.bed"
    output:
        "{dataset}/{sample}_analysis/{sample}_dupgroup_counts.txt"
    shell:
        "echo \"group\tduplicate_counts\n-1\t1\" > {output} &&"
        "cat {input.fragments} | cut -f5 | awk '$1!=\"-1\"' | uniq -c | awk '{{OFS=\"\t\"}} {{print $2,$1}}' >> {output}"

rule sample_analysis_fragments:
    input:
        "{dataset}/{sample}.sorted.mkdups.bias.corrected.bed"
    output:
        "{dataset}/{sample}_analysis/{sample}_fragments.bed"
    shell:
        "shuf -n5000 {input} > {output}"

msg_str = '''For this step to complete properly, install/load R, then run:
if (!requireNamespace("BiocManager", quietly = TRUE))
    install.packages("BiocManager")

BiocManager::install("ChIPseeker")
BiocManager::install("TxDb.Hsapiens.UCSC.hg38.knownGene")'''

rule join_tad_data:
    input:
        "{dataset}/sorted_{sample}_peaks.bed"
    output:
        temp("{dataset}/temp_{sample}_peaks.bed")
    params:
        tads="./resources/TADs.bed"
    conda: "envs/bedtools.yaml"
    shell:
        "bedtools intersect -a {input} -b {params.tads} -loj -f 0.5 | awk '{{OFS=\"\t\"}} {{print $1,$2,$3,$7,\".\",\"*\"}}' > {output}"

rule annotate_peaks:
    input:
        "{dataset}/temp_{sample}_peaks.bed"
    output:
        "{dataset}/{sample}_analysis/peak_annotations.tsv"
    envmodules:
        "R"
    message: msg_str
    shell:
        "Rscript ./workflow/scripts/annotate.r {input} {output}"

rule detect_doublets:
    input:
        counts="{dataset}/{sample}_analysis/{sample}_raw_counts.npz"
    output:
        "{dataset}/{sample}_analysis/doublet_analysis.tsv"
    conda: "envs/doublets.yaml"
    shell:
        "python ./workflow/scripts/find_doublets.py {input.counts} {output}"

rule convert_to_anndata:
    input:
        barcodes="{dataset}/{sample}_analysis/{sample}_barcode_stats.tsv",
        peaks="{dataset}/{sample}_analysis/{sample}_peak_stats.tsv",
        counts="{dataset}/{sample}_analysis/{sample}_{treatment}_counts.npz"
    conda:
        "envs/report.yaml"
    output:
        "{dataset}/{sample}_analysis/{sample}_{treatment}.h5ad"
    shell:
        "python ./workflow/scripts/convert_to_h5.py --obs {input.barcodes} --var {input.peaks} --X {input.counts} -o {output}"

'''
rule cluster_bin_based:
    input:
        fragments="{dataset}/{sample}.sorted.mkdups.bias.corrected.bed"
    output:
    
    shell:
        "bedtools makewindows -g {config[chrom_sizes]} -w {params.bin_size} | bedtools "
'''

#___ REPORT GENERATION _____

rule generate_report:
    input:
        peak_annotations="{dataset}/{sample}_analysis/peak_annotations.tsv",
        corrected_counts="{dataset}/{sample}_analysis/{sample}_corrected.h5ad",
        counts="{dataset}/{sample}_analysis/{sample}_raw.h5ad",
        dedup_counts="{dataset}/{sample}_analysis/{sample}_dedup.h5ad",
        dupcounts="{dataset}/{sample}_analysis/{sample}_dupgroup_counts.txt",
        sample_fragments="{dataset}/{sample}_analysis/{sample}_fragments.bed",
        fragments="{dataset}/{sample}.sorted.mkdups.bias.corrected.bed",
        doublets="{dataset}/{sample}_analysis/doublet_analysis.tsv"
    output:
        paths="{dataset}/{sample}_analysis/paths.pkl",
        diffpeaks=directory("{dataset}/{sample}_analysis/diffpeaks/"),
        delta_diffpeaks="{dataset}/{sample}_analysis/delta_diffpeaks.bed",
        report_finished=touch("{dataset}/{sample}.report")
    params:
        min_peak_proportion=config['min_peak_proportion']
    resources:
        cpus=1, mem_mb=config["notebook_mem"], time_min=300
    message: "Generating report: " + os.path.abspath("{wildcards.dataset}/{wildcards.sample}_analysis/{wildcards.sample}_report.ipynb")
    log:
        notebook="{dataset}/{sample}_report.ipynb"
    conda:
        "envs/report.yaml"
    notebook:
        "report_template.py.ipynb"

#___VISUALIZATION_____

rule prepare_peaks:
    input:
        "{dataset}/{sample}_analysis/peak_annotations.tsv"
    output:
        "{dataset}/{sample}_visualization/peak_regions.bed"
    shell:
        "tail -n+2 {input} | awk '{{OFS=\"\t\"}} {{print $1,$2,$3,$9,$6,\".\"}}' > {output}"

def get_awk_script(wildcards, *args):
    if wildcards.cutsite == 'plus':
        return "{print $1,$2,$2+1}"
    elif wildcards.cutsite == 'minus':
        return "{print $1,$3,$3+1}"
    else:
        raise Exception()

rule fragments_to_bedgraph:
    input:
        fragments="{dataset}/{sample}.sorted.mkdups.bias.bed",
        genome=config['chrom_sizes']
    output:
        temp("{dataset}/{sample}_visualization/{cutsite}_signal.bedgraph")
    params:
        awk_script=get_awk_script
    conda: "envs/bedtools.yaml"
    threads: 1
    shell:
        "awk '{{OFS=\"\t\"}} {params.awk_script}' {input.fragments} | bedtools slop -g {input.genome} -b 8 | bedtools genomecov -g {input.genome} -i - -bg > {output}"

rule get_cutsite_bias_background:
    input:
        "{dataset}/sorted_{sample}_peaks.bed"
    output:
        temp("{dataset}/{sample}_visualization/plus_background.bedgraph"),
        temp("{dataset}/{sample}_visualization/minus_background.bedgraph")
    params:
        slop_peak=config['slop_peak']
    threads: 8
    conda: "envs/background_track.yaml"
    resources: cpus=16
    shell:
        "cat {input} | bedtools slop -b {params.slop_peak} -g {config[chrom_sizes]} | bedtools merge -i - | "
        "python ./workflow/scripts/build_background_bias.py -m {config[cutsite_model]} -w 50 -g {config[reference]} -c {threads} -o {wildcards.dataset}/{wildcards.sample}_visualization/"

rule bedgraph_to_bigwig:
    input:
        bedgraph="{dataset}/{sample}_visualization/{fragment_type}_{signal}.bedgraph",
        genome=config['chrom_sizes']
    output:
        "{dataset}/{sample}_visualization/{fragment_type}_{signal}.bigwig"
    threads: 1
    conda: "envs/bigwig.yaml"
    shell:
        "bedGraphToBigWig {input.bedgraph} {input.genome} {output}"

rule generate_vis_data:
    input:
        "{dataset}/{sample}_visualization/peak_regions.bed",
        "{dataset}/{sample}_visualization/plus_signal.bigwig",
        "{dataset}/{sample}_visualization/minus_signal.bigwig",
        "{dataset}/{sample}_visualization/plus_background.bigwig",
        "{dataset}/{sample}_visualization/minus_background.bigwig"
    output:
        touch("{dataset}/{sample}.vis")