# vim: set syntax=python

import random
import sys
import csv
import time
import json
import math

import pandas as pd
import numpy as np
import pysam
import vcf
import matplotlib
matplotlib.use("agg")
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid.axislines import Subplot
from scipy.stats import kde, binned_statistic
import seaborn as sns
sns.set_style("ticks")

savefig = partial(plt.savefig, bbox_inches="tight")


cm_blue_alpha = matplotlib.colors.LinearSegmentedColormap.from_list("BlueAlpha", [(0.24, 0.47, 0.71, 0.0), (0.24, 0.47, 0.71, 1.0)])
cm_red_alpha = matplotlib.colors.LinearSegmentedColormap.from_list("RedAlpha", [(0.67, 0.06, 0.086, 0.3), (0.67, 0.06, 0.086, 1.0)])


shell.prefix("set -o pipefail; ")


######################################### Config ###############################

# see also https://github.com/broadinstitute/somatic-benchmark
config = {
    "raw_samples": {
        "NA12878": {  # daughter
            "gold_variants": (
                "ftp://ftp-trace.ncbi.nih.gov/giab/ftp/data/NA12878/variant_calls/"
                "NIST/NISTIntegratedCalls_14datasets_131103_allcall_UGHapMerge_HetHomVarPASS"
                "_VQSRv2.18_all_nouncert_excludesimplerep_excludesegdups_"
                "excludedecoy_excludeRepSeqSTRs_noCNVs.vcf.gz"
            ),
            "giab_gold_regions": (
                "ftp://ftp-trace.ncbi.nih.gov/giab/ftp/data/NA12878/variant_calls/NIST/"
                "union13callableMQonlymerged_addcert_nouncert_excludesimplerep_"
                "excludesegdups_excludedecoy_excludeRepSeqSTRs_noCNVs"
                "_v2.18_2mindatasets_5minYesNoRatio.bed.gz"
            ),
            "reads": (
                "ftp://ftp-trace.ncbi.nih.gov/1000genomes/ftp/technical/working/20120117_ceu_trio_b37_decoy/CEUTrio.HiSeq.WEx.b37_decoy.NA12878.clean.dedup.recal.20120117.bam"
            )
        },
        "NA12891": {  # father
            "gold_variants": (
                "ftp://platgene:G3n3s4me@ussd-ftp.illumina.com/NA12891_S1.genome.vcf.gz"
            ),
            "reads": (
                "ftp://ftp-trace.ncbi.nih.gov/1000genomes/ftp/technical/working/"
                "20120117_ceu_trio_b37_decoy/"
                "CEUTrio.HiSeq.WEx.b37_decoy.NA12891.clean.dedup.recal.20120117.bam"
            )
        },
        "NA12892": {  # mother
            "gold_variants": (
                "ftp://platgene:G3n3s4me@ussd-ftp.illumina.com/NA12892_S1.genome.vcf.gz"
            ),
            "reads": (
                "ftp://ftp-trace.ncbi.nih.gov/1000genomes/ftp/technical/working/"
                "20120117_ceu_trio_b37_decoy/"
                "CEUTrio.HiSeq.WEx.b37_decoy.NA12892.clean.dedup.recal.20120117.bam"
            )
        }
    },
    "fpr": {
        "queries": [
            "A0-A1", "A0-(A1+A2)", "(A0+A1)-A2", "(A0+A1)-(A2+A3)"
        ]
    },
    "tpr": {
        "sample": "NA12878",
        "filter_sample": "NA12892",
        "queries": [
            "A0-B0", "A0-(B0+B1)", "(A0+A1)-B0", "(A0+A1)-(B0+B1)"
        ],
        "exclude": [
            'GL000207.1', 'GL000226.1', 'GL000229.1', 'GL000231.1', 'GL000210.1', 'GL000239.1', 'GL000235.1', 'GL000201.1', 'GL000247.1', 'GL000245.1', 'GL000197.1', 'GL000203.1', 'GL000246.1', 'GL000249.1', 'GL000196.1', 'GL000248.1', 'GL000244.1', 'GL000238.1', 'GL000202.1', 'GL000234.1', 'GL000232.1', 'GL000206.1', 'GL000240.1', 'GL000236.1', 'GL000241.1', 'GL000243.1', 'GL000242.1', 'GL000230.1', 'GL000237.1', 'GL000233.1', 'GL000204.1', 'GL000198.1', 'GL000208.1', 'GL000191.1', 'GL000227.1', 'GL000228.1', 'GL000214.1', 'GL000221.1', 'GL000209.1', 'GL000218.1', 'GL000220.1', 'GL000213.1', 'GL000211.1', 'GL000199.1', 'GL000217.1', 'GL000216.1', 'GL000215.1', 'GL000205.1', 'GL000219.1', 'GL000224.1', 'GL000223.1', 'GL000195.1', 'GL000212.1', 'GL000222.1', 'GL000200.1', 'GL000193.1', 'GL000194.1', 'GL000225.1', 'GL000192.1', 'MT', 'NC_012920', 'NC_012920', 'hs37d5']
    },
    "samples": {
        "A0": "NA12878.subsample-0.25-0",
        "A1": "NA12878.subsample-0.25-1",
        "A2": "NA12878.subsample-0.25-2",
        "A3": "NA12878.subsample-0.25-3",
        "B0": "NA12892.subsample-0.25-0",
        "B1": "NA12892.subsample-0.25-1"
    },
    "max_pileup_depth": 250,
    "caller_names": {
        "gatk": "GATK",
        "alpaca": "ALPACA",
        "freebayes": "FreeBayes",
        "samtools": "SAMtools"
    },
    "parameter_space": {
        "min_qual": list(range(10, 80, 10))
    }
}


def get_sample_name(
    wildcards,
    sample_to_name={sample: name for name, sample in config["samples"].items()}
):
    return sample_to_name[wildcards.sample]


def get_samples_from_query(query, filter_only=False, call_only=False):
    if filter_only:
        try:
            query = query.split("-")[1]
        except IndexError:
            return []
    if call_only:
        query = query.split("-")[0]
    return [s.strip(" ()") for t in query.split("+") for s in t.split("-")]


def expand_samples_from_query(pattern, filter_only=False, call_only=False, orig_name=False):
    def apply(wildcards):
        samples = get_samples_from_query(
            wildcards.query, filter_only=filter_only, call_only=call_only
        )
        if orig_name:
            samples = [config["samples"][s] for s in samples]
        return expand(
            pattern,
            sample=samples
        )
    return apply


include:
    "https://bitbucket.org/johanneskoester/snakemake-workflows/raw/master/bio/ngs/rules/mapping/samfiles.rules"
include:
    "https://bitbucket.org/johanneskoester/snakemake-workflows/raw/master/bio/ngs/rules/variant_calling/bcftools.rules"


target_fpr_vs_tpr =  expand(expand(
    "plots/fpr_vs_tpr/{{depth}}/{query_fpr}.{query_tpr}.{layout}.pdf",
    zip,
    query_fpr=config["fpr"]["queries"],
    query_tpr=config["tpr"]["queries"],
    layout=["y", "l", "xy", "x"]
), depth=[4, 6, 8])


########################### Target rules #######################################


rule all:
    input:
        target_fpr_vs_tpr,
        expand(
            "data/reads/{sample}.readcount.txt",
            sample="NA12878 NA12892".split()
        )


######################################### Sample stats #########################


rule read_count:
    input:
        "{prefix}.bam"
    output:
        "{prefix}.readcount.txt"
    shell:
        "samtools idxstats {input} | awk '{{s+=$3+$4}} END {{print s}}' "
        "> {output}"


rule coverage_dist:
    input:
        "{prefix}.bam"
    output:
        "{prefix}.coverage.txt"
    resources: benchmark=1
    shell:
        "samtools mpileup {input} | cut -f4 | sort -n | uniq -c | "
        r"awk '{{print $2,$1}}' OFS='\t' > {output}"


rule plot_coverages:
    input:
        expand("data/reads/{sample}.coverage.txt", sample=config["samples"].values())
    output:
        "plots/coverage.pdf"
    resources: benchmark=1
    run:
        plt.figure(figsize=(3.86, 2.4), dpi=90)
        for f in input:
            cov = np.loadtxt(f, dtype=np.int64)[1:]
            depth, freq = cov[:,0], cov[:,1]
            plt.semilogy(depth, freq, alpha=0.5, color=sns.color_palette()[0])
            sites = freq.sum()
            mean = (depth * freq).sum() / sites
            print(mean)
            #plt.semilogy([mean, mean], [0, 10 ** 6], "--")
        plt.xlabel("depth")
        plt.ylabel("count")
        sns.despine()
        plt.gca().xaxis.set_major_locator(matplotlib.ticker.MaxNLocator(nbins=5))
        plt.savefig(output[0], bbox_inches="tight")


############################### FPR vs TPR #####################################

fpr_vs_tpr_caller = "alpaca/dep0.5/1 gatk freebayes samtools strelka".split()
fpr_vs_tpr_labels = "ALPACA GATK FreeBayes SAMtools Strelka".split()

def fpr_calls(pattern):
    def apply(wildcards):
        return expand(
            pattern,
            query=wildcards.query,
            caller=fpr_caller[wildcards.query]
        )
    return apply


def extract_depth_qual(calls, output, query):
    call_samples = get_samples_from_query(query, call_only=True)
    filter_samples = get_samples_from_query(query, filter_only=True)
    with open(calls) as f, open(output, "w") as out:
        print("chrom", "pos", "depth_call", "depth_filter", "qual", file=out, sep="\t")
        for record in vcf.Reader(f):
            def dp(samples):
                calls = [
                    record.genotype(config["samples"][sample])
                    for sample in samples
                ]
                return sum(
                    call.data.DP
                    for call in calls
                    if call.data.DP is not None
                )
            if record.QUAL is not None:
                qual = record.QUAL
            else:
                qual = 0
                if "QSS" in record.INFO:
                    qual = record.INFO["QSS"]
                elif "QSI" in record.INFO:
                    qual = record.INFO["QSI"]
            print(record.CHROM, record.POS, dp(call_samples), dp(filter_samples), qual, file=out, sep="\t")


def get_bams(filter_only=False, call_only=False):
    _get_bams = lambda wildcards: expand(
        "data/reads/{sample}.bam",
        sample=[
            config["samples"][sample]
            for sample in get_samples_from_query(wildcards.query, filter_only=filter_only, call_only=call_only)
        ]
    )
    return _get_bams


def pileup_depths_cmd(bams, positions=""):
    if positions:
        positions = "--positions '{}'".format(positions)
    cols = [str(4 + k * 3) for k in range(len(bams))]
    return (
        r"samtools mpileup -B -q1 -Q0 {positions} {bams} | "
        r"cut -f{cols}"
    ).format(positions=positions, bams=bams, cols=",".join(cols))


def format_positives(input, output, wildcards):
    depths = pileup_depths_cmd(input.bams, input.vcf)
    cols = r"\t".join(get_samples_from_query(wildcards.query))
    shell(r"echo -e 'qual\t{cols}' > '{output}'")
    shell("paste <(grep -v '^#' '{input.vcf}' | "
          "awk '!_[$1\" \"$2]++' | cut -f6) <({depths}) >> '{output}'")


rule fpr_fp:
    input:
        bams=get_bams(),
        vcf="{caller}/{query}.vcf"
    output:
        "fpr/fp/{caller}/{query}.txt"
    run:
        format_positives(input, output, wildcards)


rule fpr_covered_depths:
    input:
        bams=get_bams()
    output:
        "fpr/tn/covered_depths/{query}.txt"
    resources: benchmark=1
    run:
        depths = pileup_depths_cmd(input.bams)
        cols = r"\t".join(get_samples_from_query(wildcards.query))
        shell(r"echo -e '{cols}' > '{output}'")
        shell("{depths} >> '{output}'")


rule tpr_gold_variants:
    input:
        vcf=expand("data/gold_variants/{sample}.vcf", sample=config["tpr"]["sample"]),
        gtf="ref.genes.gtf"
    output:
        "tpr/gold_variants.vcf"
    params:
        exclude=" --not-chr ".join(config["tpr"]["exclude"])
    resources: benchmark=1
    shell:
        "bedtools intersect -header -u -a {input.vcf} -b {input.gtf} | "
        "vcftools --vcf - --stdout --not-chr {params.exclude} "
        "--recode > {output}"


rule tpr_exclusive_gold_variants:
    input:
        filter_vcf=expand("data/gold_variants/{sample}.fixed.vcf", sample=config["tpr"]["filter_sample"]),
        vcf="tpr/gold_variants.vcf"
    output:
        "tpr/exclusive_gold_variants.vcf"
    resources: benchmark=1
    shell:
        "grep '#' {input.vcf} > {output} && "
        "bedtools subtract -A -a {input.vcf} -b {input.filter_vcf} >> {output}"


rule tpr_tp_calls:
    input:
        vcf="{caller}/{query}.vcf",
        gold="tpr/exclusive_gold_variants.vcf"
    output:
        "tpr/tp/{caller}/{query}.vcf"
    resources: benchmark=1
    shell:
        "bedtools intersect -wa -header -a '{input.vcf}' -b {input.gold} > '{output}'"


rule tpr_tp:
    input:
        bams=get_bams(),
        vcf="tpr/tp/{caller}/{query}.vcf"
    output:
        "tpr/tp/{caller}/{query}.txt"
    run:
        format_positives(input, output, wildcards)


rule tpr_covered_depths:
    input:
        bams=get_bams(),
        vcf="tpr/exclusive_gold_variants.vcf"
    output:
        "tpr/fn/covered_depths/{query}.txt"
    resources: benchmark=1
    run:
        depths = pileup_depths_cmd(input.bams, input.vcf)
        cols = r"\t".join(get_samples_from_query(wildcards.query))
        shell(r"echo -e '{cols}' > '{output}'")
        shell("{depths} >> '{output}'")


def filter_positives(positives, min_depth=4, depth_samples=None):
    positives = positives[positives[depth_samples].sum(axis=1) >= min_depth]
    positives = positives[positives["qual"] >= 1]
    return positives


def filter_negatives(negatives, min_depth=4, depth_samples=None):
    negatives = negatives[negatives[depth_samples].sum(axis=1) >= min_depth]
    return negatives


def rate(positives, negative_depths, label, min_depth=4, depth_samples=None):
    positives = filter_positives(positives, min_depth=min_depth, depth_samples=depth_samples)
    negative_total = filter_negatives(negative_depths, min_depth=min_depth, depth_samples=depth_samples).shape[0]

    positive_cum = np.bincount(np.asarray(np.rint(positives["qual"]), dtype=np.int64))[::-1].cumsum()[::-1]
    negative_cum = negative_total - positive_cum

    total_cum = positive_cum + negative_cum
    rate_cum = positive_cum / total_cum

    return pd.DataFrame({"qual": np.arange(len(rate_cum)), label: rate_cum})


def fpr(fp, covered_depths, min_depth=4, depth_samples=None):
    return rate(fp, covered_depths, "fpr", min_depth=min_depth, depth_samples=depth_samples)


def tpr(tp, covered_depths, min_depth=4, depth_samples=None):
    return rate(tp, covered_depths, "tpr", min_depth=min_depth, depth_samples=depth_samples)


rule plot_fpr_vs_tpr:
    input:
        tp=lambda wildcards: expand("tpr/tp/{caller}/{query}.txt", query=wildcards.query_tpr, caller=fpr_vs_tpr_caller if wildcards.query_tpr == "A0-B0" else fpr_vs_tpr_caller[:-1]),
        fp=lambda wildcards: expand("fpr/fp/{caller}/{query}.txt", query=wildcards.query_fpr, caller=fpr_vs_tpr_caller if wildcards.query_fpr == "A0-A1" else fpr_vs_tpr_caller[:-1]),
        fn_cov="tpr/fn/covered_depths/{query_tpr}.txt",
        tn_cov="fpr/tn/covered_depths/{query_fpr}.txt"
    output:
        "plots/fpr_vs_tpr/{depth,[0-9]+}/{query_fpr}.{query_tpr}.{layout,[nlxy]+}.pdf"
    run:
        tn_cov = pd.read_table(input.tn_cov)
        fn_cov = pd.read_table(input.fn_cov)
        depth = int(wildcards.depth)

        plt.figure(figsize=(2.8, 2.5), dpi=90)
        last_line = None
        for tp, fp, label in zip(input.tp, input.fp, fpr_vs_tpr_labels):
            print(label)
            tp = pd.read_table(tp)
            fp = pd.read_table(fp)
            print("fpr")
            _fpr = fpr(fp, tn_cov, depth, depth_samples=get_samples_from_query(wildcards.query_fpr))
            print("tpr")
            _tpr = tpr(tp, fn_cov, depth, depth_samples=get_samples_from_query(wildcards.query_tpr))
            d = pd.merge(_fpr, _tpr)
            if label == "ALPACA(FDR)":
                plt.semilogx(d["fpr"][0], d["tpr"][0], "x", color=last_line.get_color(), mew=2)
            else:
                last_line, = plt.semilogx(d["fpr"], d["tpr"], label=label)
        plt.ylim((0,0.9))
        plt.xlim((1e-9, 1e-3))
        if "x" in wildcards.layout:
            plt.xlabel("FPR")
        if "y" in wildcards.layout:
            plt.ylabel("TPR")
        if "l" in wildcards.layout:
            if len(fpr_vs_tpr_labels) > len(input.tp):
                plt.semilogx([],[], label=fpr_vs_tpr_labels[-1])
            plt.legend(loc="lower right")
        sns.despine()

        savefig(output[0], bbox_inches="tight")


def plot_rate_vs_qual(positives_files, out_file, layout="lxy", min_depth_call=0, min_depth_filter=0):
    plt.figure(figsize=(2.8, 2.5), dpi=90)
    for positives, label in zip(positives_files, fpr_vs_tpr_labels):
        if label == "ALPACA(FDR)":
            continue
        positives = pd.read_table(positives)
        positives = filter_positives(positives, min_depth_call=min_depth_call, min_depth_filter=min_depth_filter)
        sns.kdeplot(positives["qual"].apply(lambda q: 10 ** (-q / 10)), cumulative=True, label=label, legend=False)
    plt.xlim((0,1))
    plt.ylim((0,1))
    if "x" in layout:
        plt.xlabel("posterior")
    if "y" in layout:
        plt.ylabel("CDF")
    if "l" in layout:
        plt.legend(loc="lower right")
    sns.despine()
    plt.locator_params(nbins=5)
    savefig(out_file, bbox_inches="tight")


rule plot_fpr_vs_qual:
    input:
        fp=lambda wildcards: expand("fpr/fp/{caller}/{query}.txt", query=wildcards.query, caller=fpr_vs_tpr_caller if wildcards.query == "A0-A1" else fpr_vs_tpr_caller[:-1]),
    output:
        "plots/fpr_vs_qual/{query}.{layout,[lxyn]+}.pdf"
    run:
        plot_rate_vs_qual(input.fp, output[0], wildcards.layout, min_depth_filter=4)


rule plot_tpr_vs_qual:
    input:
        tp=lambda wildcards: expand("tpr/tp/{caller}/{query}.txt", query=wildcards.query, caller=fpr_vs_tpr_caller if wildcards.query == "A0-B0" else fpr_vs_tpr_caller[:-1])
    output:
        "plots/tpr_vs_qual/{query}.{layout,[lxyn]+}.pdf"
    run:
        plot_rate_vs_qual(input.tp, output[0], wildcards.layout, min_depth_call=4)


############################### estimate FDR ###################################


rule plot_fp_vs_tp_probs:
    input:
        tp="tpr/tp/alpaca/dep0.5/1/{query_tpr}.txt",
        fp="fpr/fp/alpaca/dep0.5/1/{query_fpr}.txt"
    output:
        "plots/fp_tp_dist/{query_fpr}.{query_tpr}.pdf"
    run:
        import random
        qual_to_prob = lambda qual: 10 ** (-qual / 10)

        tp = pd.read_table(input.tp)
        tp["depth"] = tp.iloc[:, 1:].sum(axis="columns")
        tp["fp"] = False
        tp["prob"] = tp["qual"].apply(qual_to_prob)

        fp = pd.read_table(input.fp)
        fp["depth"] = fp.iloc[:, 1:].sum(axis="columns")
        fp["prob"] = fp["qual"].apply(qual_to_prob)
        fp["fp"] = True


        plt.figure(figsize=(3.86, 2.4), dpi=90)
        #plt.figure(figsize=(7.72, 2.4), dpi=90)
        #plt.subplot(121)
        sns.kdeplot(np.log10(tp["depth"]), tp["prob"], label="TP", legend=False, cmap=cm_blue_alpha, shade=True, shade_lowest=False)
        plt.ylabel("posterior probability")
        plt.xlabel("log depth")
        plt.ylim((0,1))
        plt.xlim((0,3))
        #sns.despine()

        #plt.subplot(122)
        sns.kdeplot(np.log10(fp["depth"]), fp["prob"], label="FP", legend=False, cmap=cm_red_alpha, shade=True, shade_lowest=False)
        #plt.xlabel("log depth")
        #plt.ylim((0,1))
        #plt.xlim((0,0.5))
        sns.despine()
        """
        p = pd.concat([tp, fp])
        p.sort("prob", inplace=True)
        p = p.reset_index(drop=True)
        exp_fdr = p["prob"].cumsum() / np.arange(1, p.shape[0] + 1)
        for fdr in [0.01, 0.05, 0.1]:
            i = np.searchsorted(exp_fdr, fdr, side="left")
            plt.plot([p["prob"][i]] * 2, [0, plt.ylim()[1]], label="FDR={}".format(fdr), ls="--")
            print(p[:i]["fp"].sum() / i)
        """
        #plt.legend()
        savefig(output[0], bbox_inches="tight")



rule plot_fdr:
    input:
        tp="tpr/tp/alpaca/dep0.5/1/{query_tpr}.txt",
        fp="fpr/fp/alpaca/dep0.5/1/{query_fpr}.txt"
    output:
        "plots/fdr/{query_fpr}.{query_tpr}.pdf"
    run:
        import random
        qual_to_prob = lambda qual: 10 ** (-qual / 10)

        tp = pd.read_table(input.tp)
        #tp = tp[tp["depth"] >= 50]
        tp["fp"] = False
        tp["prob"] = tp["qual"].apply(qual_to_prob)

        fp = pd.read_table(input.fp)
        #fp = fp[fp["depth"] >= 50]
        fp["prob"] = fp["qual"].apply(qual_to_prob)
        fp["fp"] = True
        print(fp.shape, tp.shape)


        plt.figure(figsize=(3.86, 2.4), dpi=90)
        for total_fdr in [0.0001, 0.001, 0.01, 0.05, 0.10, 0.25]:
            d = []
            for _ in range(100):
                tp_sample = 100
                fp_sample = int(tp_sample * total_fdr / (1 - total_fdr))
                sample = pd.concat([
                    tp.ix[np.random.choice(tp.index, tp_sample, replace=False)],
                    fp.ix[np.random.choice(fp.index, fp_sample, replace=False)]
                ])
                exp_fdr = sample["prob"].sum() / sample.shape[0]
                d.append(exp_fdr)
            sns.kdeplot(np.asarray(d), cumulative=True, label="true FDR={}".format(total_fdr))
        #plt.xlim((0.15, 0.25))
        plt.xlabel("expected FDR")
        plt.ylabel("CDF")
        sns.despine()
        plt.savefig(output[0], bbox_inches="tight")



############################### run time performance ###########################

performance_queries = config["fpr"]["queries"] + config["tpr"]["queries"]
performance_sample_names = sorted(config["samples"])
performance_samples = [config["samples"][name] for name in performance_sample_names]
performance_callers = "alpaca gatk freebayes samtools".split()


def performance_write_runtimes(samples, merge, queries, out):
    def get_runtime(f):
        if f is None:
            return "-"
        with open(f) as f:
            s = json.load(f)["wall_clock_times"]["s"][0]
            m = s // 60
            s = s % 60
            return "{:.0f}:{:02.0f}".format(m, s)

    with open(out, "w") as out:
        json.dump(
            [get_runtime(f) for f in samples] +
            [get_runtime(merge)] +
            [get_runtime(f) for f in queries],
            out
        )


rule performance_alpaca:
    input:
        samples=expand("benchmarks/alpaca/{sample}.index.gh.json", sample=performance_samples),
        merge="benchmarks/alpaca/merge.gh.json",
        calls=expand(
            "benchmarks/alpaca/nodep/1/{query}.call.json", query=performance_queries
        )
    output:
        "performance/alpaca.json"
    run:
        performance_write_runtimes(input.samples, input.merge, input.calls, output[0])


rule performance_gatk:
    input:
        samples=expand("benchmarks/gatk/{sample}.haplotype_caller.json", sample=performance_sample_names),
        calls=expand("benchmarks/gatk/{query}.genotyping.json", query=performance_queries)
    output:
        "performance/gatk.json"
    run:
        performance_write_runtimes(input.samples, None, input.calls, output[0])


rule performance_freebayes:
    input:
        calls=expand("benchmarks/freebayes/{query}.json", query=performance_queries)
    output:
        "performance/freebayes.json"
    run:
        performance_write_runtimes([None] * len(performance_samples), None, input.calls, output[0])


rule performance_samtools:
    input:
        calls=expand("benchmarks/samtools/{query}.json", query=performance_queries)
    output:
        "performance/samtools.json"
    run:
        performance_write_runtimes([None] * len(performance_samples), None, input.calls, output[0])


rule performance:
    input:
        expand("performance/{caller}.json", caller=performance_callers)
    output:
        "tables/performance.csv"
    run:
        table = []
        for caller, f in zip(performance_callers, input):
            with open(f) as f:
                table.append(json.load(f))
        table = np.array(table).T
        with open(output[0], "w") as out:
            writer = csv.writer(out, delimiter="\t")
            writer.writerow(
                ["task"] +
                [config["caller_names"][caller] for caller in performance_callers]
            )
            row_labels = performance_sample_names + ["merging"] + performance_queries
            for label, row in zip(row_labels, table):
                writer.writerow([label] + list(row))


################################ Bayesian Calling ##############################


rule plot_bayesian_calling:
    output:
        "plots/bayesian_calling/single_sample.pdf"
    run:
        from alpaca.prototype import reference_genotype_probability, Pileup, PHRED_TO_LOG_FACTOR

        quals = list(range(ord("("), ord("J"), 10))
        qual_chars = [chr(qual) for qual in quals]
        depths = np.arange(11)

        bases = lambda n, type: (b"AC" * n if type == "heterozygous" else b"CC" * n)
        types = ["heterozygous", "homozygous"]

        figure()
        styles = ["-", "--"]
        handles = [None, None]
        for i, type in enumerate(types):
            for qual in qual_chars:
                pileups = [
                    Pileup(bases(n, type), qual.encode() * n * 2)
                    for n in depths
                ]
                probs = [math.exp(reference_genotype_probability([pileup])) for pileup in pileups]
                handles[i], = plt.semilogy(depths * 2, probs, styles[i], label=type, clip_on=False)

        plt.xlabel("read depth")
        plt.ylabel("reference genotype probability")
        plt.legend(handles, types, "lower left", handlelength=2.5)
        print("QUALS:", [math.exp((qual - 33) * PHRED_TO_LOG_FACTOR) for qual in quals])

        plt.savefig(output[0])


rule table_bayesian_calling_sample_depth:
    output:
        "tables/bayesian_calling/sample_depth.csv"
    run:
        from alpaca.prototype import reference_genotype_probability, Pileup

        setups = [
            [(18, 18), (2, 2)],
            [(10, 10), (10, 10)],
            [(15, 15), (5, 5)],
            [(20, 20), (40, 0)],
            [(20, 20), (0, 0)]
        ]

        with open(output[0], "w") as out:
            print("A", "B", sep="\t", file=out)
            for setup in setups:
                pileups = []
                for ref, alt in setup:
                    pileups.append(Pileup(b"A" * ref + b"C" * alt, b"<" * (ref + alt)))
                    print(ref, alt, sep="+", end="\t", file=out)
                print("{:.4e}".format(math.exp(reference_genotype_probability(pileups))), file=out)


rule plot_algebraic_calling:
    output:
        "plots/algebraic_calling/depth.pdf"
    run:
        from alpaca.prototype import reference_genotype_probability, Pileup, difference
        def calc_prob(n):
            pileup_a = Pileup(b"A" * 20 + b"C" * 20, b"<" * 40)
            pileup_b = Pileup(b"A" * n + b"C" * n, b"<" * n * 2)
            prob_a = reference_genotype_probability([pileup_a])
            prob_b = reference_genotype_probability([pileup_b])
            return difference(prob_a, prob_b), prob_b

        depths = np.arange(6)
        figure()
        probs = np.array([calc_prob(n) for n in depths])
        query_prob = np.exp(probs[:, 0])
        sample_prob = np.exp(probs[:, 1])
        plt.semilogy(depths * 2, query_prob, "-", clip_on=False, label="query probability")
        plt.semilogy(depths * 2, sample_prob, "--", clip_on=False, label="filter sample probability")
        plt.xlabel("filter sample depth")
        plt.legend(loc="lower left", handlelength=2.5)
        plt.savefig(output[0])


rule plot_priors:
    output:
        "plots/bayesian_calling/priors.pdf"
    run:
        het = 0.001
        def prob(m, sample_count):
            if m == 0:
                return 1 - math.fsum(prob(i, sample_count) for i in range(1, 2 * sample_count))
            return het / m
        figure()
        x = np.arange(1, 1001)
        y = [prob(0, n) for n in x]
        plt.plot(x, y, "-")
        plt.xlabel("number of samples")
        plt.ylabel("$Pr(M = 0)$")
        plt.savefig(output[0])


rule plot_dependency_model:
    output:
        "plots/algebraic_calling/dependency.pdf"
    run:
        het = 0.001
        k = 20 # sampling rate
        likelihoods = np.array([0, 0, 0])
        def het_prior(): # Pr(het)
            return 1 / k

        def prior(m, n): # Pr(M=m)
            if m == 0:
                return 1 - math.fsum(prior(i, n) for i in range(1, 2 * n))
            return het / m

        def prob_given_het(het): # Pr(M=0 | D, het)
            return

        def pep_to_het(pep, n=1):
            den = (1 - pep)
            nom = math.fsum(1 / i for i in range(1, 2 * n))
            return den / nom
        plt.figure(figsize=(3.86, 2.4), dpi=90)
        peps = pd.Series(np.linspace(0, 1, 20))
        for n in range(1, 5):
            line = plt.plot(peps, peps.apply(pep_to_het, n=n), label="n={}".format(n))
            halfdep = peps ** 0.5
            plt.plot(peps, halfdep.apply(pep_to_het, n=n), "--", color=line[0].get_color(), label="n={}".format(n))
        sns.despine()
        plt.legend()
        plt.savefig(output[0], bbox_inches="tight")

##################################### Alpaca ###################################


rule alpaca_preprocess:
    input:
        "data/reads/{sample}.bam.bai",
        bam="data/reads/{sample}.bam",
        ref="ref.fasta",
        refindex="ref.fasta.fai"
    output:
        "alpaca/{sample}.preprocessed.bcf"
    params:
        sample=get_sample_name
    threads: 12
    resources: benchmark=100
    log:
        "logs/alpaca/{sample}.log"
    benchmark:
        "benchmarks/alpaca/{sample}.preprocess.json"
    shell:
        "alpaca preprocess --no-BAQ --threads {threads} "
        "{input.ref} {input.bam} > {output} 2> {log}"


rule alpaca_merge:
    input:
        ref="ref.fasta",
        refindex="ref.fasta.fai",
        bcf=lambda wildcards: expand(
            "alpaca/{sample}.preprocessed.bcf",
            sample=config["samples"].values()
        ),
        csi=lambda wildcards: expand(
            "alpaca/{sample}.preprocessed.bcf.csi",
            sample=config["samples"].values()
        )
    output:
        "alpaca/all.bcf"
    threads: 12
    resources: benchmark=100
    log:
        "logs/alpaca/merge.log"
    benchmark:
        "benchmarks/alpaca/merge.json"
    shell:
        "alpaca merge -t {threads} {input.ref} {input.bcf} > {output} 2> {log}"


rule alpaca_merge_two:
    input:
        ref="ref.fasta",
        refindex="ref.fasta.fai",
        bcf=lambda wildcards: expand(
            "alpaca/{sample}.preprocessed.bcf",
            sample=[config["samples"]["A0"], config["samples"]["A1"]]
        ),
        csi=lambda wildcards: expand(
            "alpaca/{sample}.preprocessed.bcf.csi",
            sample=[config["samples"]["A0"], config["samples"]["A1"]]
        )
    output:
        "alpaca/A0+A1.bcf"
    resources: benchmark=100
    threads: 12
    log:
        "logs/alpaca/A0+A1.merge.log"
    benchmark:
        "benchmarks/alpaca/A0+A1.merge.json"
    shell:
        "alpaca merge -t {threads} {input.ref} {input.bcf} > {output} 2> {log}"


def alpaca_filter(wildcards):
    try:
        int(wildcards.filter)
    except ValueError:
        # FDR filtering
        return "--fdr " + wildcards.filter
    return "--min-qual " + wildcards.filter


def alpaca_query(wildcards):
    query = wildcards.query
    return re.sub("[AB][0-9]", lambda m: " {} ".format(config["samples"][m.group()]), query).replace("-(", "- (").replace(")-", ") -")


rule alpaca_call:
    input:
        "alpaca/all.bcf"
    output:
        "alpaca/dep{dependency,(0|0\.5)}/{filter}/{query,[^\.]+}.vcf"
    params:
        filter=alpaca_filter, query=alpaca_query
    threads: 12
    resources: benchmark=100
    log:
        "logs/alpaca/dep{dependency}/{filter}/{query}.call.log"
    benchmark:
        "benchmarks/alpaca/dep{dependency}/{filter}/{query}.call.json"
    shell:
        "alpaca call --threads {threads} --dep {wildcards.dependency} "
        "'{params.query}' {params.filter} < {input} | bcftools view - > '{output}' 2> '{log}'"


##################################### GATK #####################################


rule gatk_haplotype_caller:
    input:
        lambda wildcards: expand("data/reads/{sample}.bam.bai", sample=config["samples"][wildcards.sample]),
        "ref.dict",
        bam=lambda wildcards: expand("data/reads/{sample}.bam", sample=config["samples"][wildcards.sample]),
        ref="ref.fasta"
    output:
        "gatk/{sample}.gvcf"
    log:
        "logs/gatk/{sample}.log"
    benchmark:
        "benchmarks/gatk/{sample}.haplotype_caller.json"
    threads: 2
    resources: benchmark=100
    shell:
        "gatk -T HaplotypeCaller -nct {threads} -R {input.ref} -I {input.bam} "
        "--emitRefConfidence GVCF --variant_index_type LINEAR "
        "-variant_index_parameter 128000 "
        "-o {output} &> {log}"


rule gatk_genotyping:
    input:
        expand_samples_from_query("gatk/{sample}.gvcf"),
        ref="ref.fasta"
    output:
        "gatk/{query,[^\.]+}.raw.vcf"
    params:
        gvcfs=expand_samples_from_query("--variant gatk/{sample}.gvcf")
    log:
        "logs/gatk/{query}.call.log"
    benchmark:
        "benchmarks/gatk/{query}.genotyping.json"
    threads: 12
    resources: benchmark=100
    shell:
        "gatk -T GenotypeGVCFs {params.gvcfs} -nt {threads} -R {input.ref} "
        "-stand_call_conf 1 -stand_emit_conf 1 "
        "-o '{output}' &> '{log}'"


def gatk_query_to_select(query):
    translate = config["samples"].get

    get_sample = lambda sample: 'vc.getGenotype("{}")'

    filter_samples = " && ".join(map(
        '(vc.getGenotype("{0}").isHomRef() || !vc.getGenotype("{0}").isCalled())'.format,
        map(translate, get_samples_from_query(query, filter_only=True))
    ))
    call_samples = " || ".join(map(
        '(vc.getGenotype("{0}").isCalled() && !vc.getGenotype("{0}").isHomRef())'.format,
        map(translate, get_samples_from_query(query, call_only=True))
    ))

    select = "-select '{}'".format(call_samples)
    if filter_samples:
        select += " -select '{}'".format(filter_samples)
    return select

rule gatk_filter_by_query:
    input:
        ref="ref.fasta",
        vcf="{caller}/{query}.raw.vcf"
    output:
        "{caller,(freebayes|gatk|samtools)}/{query,[^\.]+}.vcf"
    params:
        select=lambda wildcards: gatk_query_to_select(wildcards.query)
    log:
        "logs/{caller}/{query}.select.log"
    benchmark:
        "benchmarks/{caller}/{query}.select.json"
    threads: 2
    resources: benchmark=1
    shell:
        "gatk -T SelectVariants -R {input.ref} --variant '{input.vcf}' "
        "-nt {threads} {params.select} -o '{output}' &> '{log}'"


##################################### Freebayes ################################

rule freebayes:
    input:
        expand_samples_from_query("data/reads/{sample}.bam.bai", orig_name=True),
        bams=expand_samples_from_query("data/reads/{sample}.bam", orig_name=True),
        ref="ref.fasta",
        fai="ref.fasta.fai"
    output:
        "freebayes/{query,[^\.]+}.raw.vcf"
    log:
        "logs/freebayes/{query}.log"
    benchmark:
        "benchmarks/freebayes/{query}.json"
    threads: 12
    resources: benchmark=100
    shell:
        "freebayes-parallel <(fasta_generate_regions.py {input.fai} 1000000) "
        "{threads} -f {input.ref} --no-mnps --no-complex "
        "{input.bams} > '{output}' 2> '{log}'"


################################# MuTect #######################################


# MuTect crashes with the used datasets
rule mutect:
    input:
        expand_samples_from_query("data/reads/{sample}.bam.bai", orig_name=True),
        bams=expand_samples_from_query("data/reads/{sample}.bam", orig_name=True),
        ref="ref.fasta",
        fai="ref.fasta.fai"
    output:
        "mutect/{query,(A0-A1|A0-B0)}.vcf"
    log:
        "logs/mutect/{query}.log"
    benchmark:
        "benchmarks/mutect/{query}.json"
    threads: 12
    resources: benchmark=100
    shell:
        "mutect -nt {threads} --analysis_type MuTect --reference_sequence {input.ref} "
        "--input_file:normal {input.bams[1]} --input_file:tumor {input.bams[0]} "
        "--vcf {output} &> {log}"



################################### Strelka ####################################


rule strelka:
    input:
        expand_samples_from_query("data/reads/{sample}.bam.bai", orig_name=True),
        bams=expand_samples_from_query("data/reads/{sample}.bam", orig_name=True),
        ref="ref.fasta",
        fai="ref.fasta.fai",
        config="strelka_config.ini"
    output:
        "strelka/{query,(A0-A1|A0-B0)}/results/all.somatic.snvs.vcf",
        "strelka/{query,(A0-A1|A0-B0)}/results/all.somatic.indels.vcf"
    params:
        dir="strelka/{query,(A0-A1|A0-B0)}"
    benchmark:
        "benchmarks/strelka/{query}.json"
    threads: 12
    resources: benchmark=100
    shell:
        "rm -r {params.dir}; "
        "configureStrelkaWorkflow --normal={input.bams[1]} --tumor={input.bams[0]} "
        "--config={input.config} --ref={input.ref} --output-dir={params.dir}; "
        "cd {params.dir}; make -j {threads}"


rule strelka_merge:
    input:
        "strelka/{query,(A0-A1|A0-B0)}/results/all.somatic.snvs.vcf",
        "strelka/{query,(A0-A1|A0-B0)}/results/all.somatic.indels.vcf"
    output:
        "strelka/{query,(A0-A1|A0-B0)}.merged.vcf"
    params:
        normal=lambda wildcards: config["samples"]["A1" if wildcards.query == "A0-A1" else "B0"],
        tumor=lambda wildcards: config["samples"]["A0"]
    shell:
        "for f in {input}; do bgzip -c $f > $f.gz; tabix -p vcf $f.gz; done; "
        "bcftools concat -a -Ov {input[0]}.gz {input[1]}.gz | "
        "sed 's/TUMOR/{params.tumor}/' | sed 's/NORMAL/{params.normal}/' > {output}"


rule strelka_qual:
    input:
        "strelka/{query}.merged.vcf"
    output:
        "strelka/{query,(A0-A1|A0-B0)}.vcf"
    run:
        with open(input[0]) as f, open(output[0], "w") as out:
            reader = vcf.Reader(f)
            writer = vcf.Writer(out, template=reader)
            for record in reader:
                if "QSI" in record.INFO:
                    record.QUAL = record.INFO["QSI"]
                elif "QSS" in record.INFO:
                    record.QUAL = record.INFO["QSS"]
                else:
                    record.QUAL = 0
                writer.write_record(record)


#################################### Samtools ##################################


rule samtools:
    input:
        expand_samples_from_query("data/reads/{sample}.bam.bai", orig_name=True),
        bams=expand_samples_from_query("data/reads/{sample}.bam", orig_name=True),
        ref="ref.fasta",
        fai="ref.fasta.fai"
    output:
        "samtools/{query,[^\.]+}.raw.vcf"
    log:
        "logs/samtools/{query}.log"
    benchmark:
        "benchmarks/samtools/{query}.json"
    threads: 12
    resources: benchmark=100
    shell:
        "cut -f1 {input.fai} | parallel -j {threads} "
        "'samtools mpileup -t DP -gu -B -C50 -Q0 -q1 -f {input.ref} "
        "-r {{}} {input.bams} | "
        "bcftools call -m -v - > \"{output}.{{}}.bcf\"'; "
        "cut -f1 {input.fai} | "
        "xargs -I '{{}}' bcftools concat -Ov '{output}.{{}}.bcf' > '{output}'; "
        "cut -f1 {input.fai} | xargs -I '{{}}' rm '{output}.{{}}.bcf'"


##################################### Utils ####################################



def get_subsample_param(wildcards):
    seed = int(wildcards.sample[2:]) + int(wildcards.run)
    return "{}{}".format(seed, wildcards.fraction)


rule subsample:
    input:
        "data/reads/{sample}.bam"
    output:
        temp("data/reads/{sample}.subsample-{fraction}-{run,\d+}.pre.bam")
    params:
        subsample=get_subsample_param
    resources: benchmark=1
    shell:
        "samtools view -bs {params.subsample} {input} > {output}"


rule set_read_group:
    input:
        "data/reads/{sample}.subsample-{fraction}-{run}.pre.bam"
    output:
        "data/reads/{sample}.subsample-{fraction}-{run,\d+}.bam"
    params:
        rgsm="{sample}.subsample-{fraction}-{run}"
    log:
        "logs/reads/{sample}.subsample-{fraction}-{run}.set_read_group.log"
    resources: benchmark=1
    shell:
        "picard-tools AddOrReplaceReadGroups INPUT={input} OUTPUT={output} "
        "RGID={params.rgsm} RGLB={params.rgsm} RGPU={params.rgsm} "
        "VALIDATION_STRINGENCY=SILENT "
        "RGPL=illumina RGSM={params.rgsm} &> {log}"


rule subtract_variants:
    input:
        a="data/gold_variants/{a}.vcf",
        b="data/gold_variants/{b}.vcf"
    output:
        "data/gold_variants/{a}-{b}.vcf"
    resources: benchmark=1
    shell:
        "bedtools subtract -A -a {input.a} -b {input.b} > {output}"


rule download_reads:
    output:
        "data/reads/{sample}.bam"
    params:
        url=lambda wildcards: config["raw_samples"][wildcards.sample]["reads"]
    shell:
        "wget -O {output} {params.url}"


rule download_variants:
    output:
        "data/gold_variants/{sample}.vcf"
    params:
        url=lambda wildcards: config["raw_samples"][wildcards.sample]["gold_variants"]
    shell:
        """
        wget -O - {params.url} > {output}.gz
        gzip -d {output}.gz || true
        """

rule fix_variants:
    input:
        "data/gold_variants/{sample}.vcf"
    output:
        "data/gold_variants/{sample}.fixed.vcf"
    resources: benchmark=1
    shell:
        "sed 's/chr//' {input} > {output}"


rule download_ref:
    output:
        "ref.fasta"
    shell:
        """
        wget -O {output}.gz ftp://ftp.1000genomes.ebi.ac.uk/vol1/ftp/technical/reference/phase2_reference_assembly_sequence/hs37d5.fa.gz
        gzip -d --quiet {output}.gz || true  # ignore trailing garbage
        """


rule download_track:
    output:
        "ref.genes.gtf"
    shell:
        "wget -O - ftp://ftp.ensembl.org/pub/release-75/gtf/homo_sapiens/Homo_sapiens.GRCh37.75.gtf.gz | zcat --quiet > {output}"


def figure(right=False, top=False, left=True, bottom=True, figsize=None):
    fig = plt.figure(figsize=figsize)
    ax = Subplot(fig, 111)
    ax.axis["right"].set_visible(right)
    ax.axis["top"].set_visible(top)
    ax.axis["left"].set_visible(left)
    ax.axis["bottom"].set_visible(bottom)
    fig.add_subplot(ax)


############### Heterozygosity distribution ####################


rule dbsnp:
    output:
        "data/dbsnp.vcf"
    shell:
        "wget -O - ftp://ftp.ncbi.nih.gov/snp/organisms/human_9606_b142_GRCh38/VCF/All.vcf.gz | gzip -c -d > {output}"
