include: "Common.snake"

import os
import os.path

import pandas
from pandas import DataFrame
from snakemake.utils import report
import textwrap

from scripts.common import gather_refs

#Autodetect bins and stage
if config["reassembly"]["enabled"]:
    LAST_STAGE = "reassembly"
    BINS, = glob_wildcards("reassembly/{{bin,{}}}.info".format(NAME_TEMPLATE))
else:
    if config["propagation"]["enabled"]:
        LAST_STAGE = "propagation"
    else:
        LAST_STAGE = "binning"
    with open("binning/filtered_bins.tsv") as input:
        BINS = [line.split()[0] for line in input]
if not BINS:
    raise WorkflowError("No bins detected!")
BINS.sort()

#Additional config parameters
try:
    stats_config = config["stats"]
    QUAST_DIR = stats_config["quast"]
    QUAST_PARAMS = stats_config.get("params", "")
    QUAST = os.path.join(QUAST_DIR, "quast.py")
    METAQUAST = os.path.join(QUAST_DIR, "metaquast.py")
except KeyError:
    QUAST = "quast"
    METAQUAST = "metaquast"

#Detect references
REFS = dict(gather_refs(config["stats"].get("refs", [])))
REFS_STR = ",".join(path for path in sorted(REFS.values()))
METAQUAST_PARAMS = "--scaffolds --no-plots --no-icarus --no-snps --no-gc --no-sv"

def ref_path(wildcards):
    return REFS[wildcards.ref]

onstart:
    print("Detected", SAMPLE_COUNT, "samples in", IN)
    if BINS:
        print("Detected good (abundant) bins:", " ".join(BINS))
    if REFS:
        print("Detected references:", " ".join(REFS))

enabled_stats = []
stats_targets = {"checkm": "{}_checkm.tsv",
                 "pca":    "{}_pca.png",
                 "tsne":   "{}_tsne.png",
                 "report": "{}_report.html"}

if config["stats"].get("refs"):
    ref_stats = {"f1": "txt", "summary": "tsv", "nga": "tsv", "lost": "tsv", "group": "tsv"}
    for stat, ext in ref_stats.items():
        stats_targets[stat] = "{{}}_{}.{}".format(stat, ext)

def unsupported_target(stage, stat):
    if stage == "assembly":
        return stat != "summary"
    elif stat == "lost":
        return stage != "binning"
    elif stage == "reassembly":
        return stat == "group"
    else:
        return stat == "nga" and len(GROUPS) > 1

def stat_path(template, stage):
    return os.path.join("stats/summary", template.format(stage))

STAT_DIR = "stats/summary"
STAGES = ["assembly", "binning", "propagation", "reassembly"]
for stage in STAGES:
    if not config.get(stage, {}).get("enabled", True):
        continue
    stage_stats = config["stats"].get(stage, [])
    if stage_stats == "all":
        stats = [target for target in stats_targets if not unsupported_target(stage, target)]
    else:
        stats = [target for target in config["stats"].get("all", []) if not unsupported_target(stage, target)]
        for stat in stage_stats:
            if unsupported_target(stage, stat):
                print(stat, "is not supported for", stage, "and disabled")
            else:
                stats.append(stat)
    if stage == LAST_STAGE:
        stats.extend(config["stats"].get("last", []))
    enabled_stats.extend(os.path.join(STAT_DIR, stats_targets[st].format(stage)) for st in stats)

misc_targets = {"ref_profiles": "ref_profiles.tsv"}
enabled_stats.extend(os.path.join(STAT_DIR, misc_targets[st]) for st in config["stats"].get("misc", []))

rule all_stats:
    input:   enabled_stats
    message: "Gathered stats: {input}"

#===============================================================================
#---- Statistics section -------------------------------------------------------
#===============================================================================

#---- Binning of full contigs based on majority of their splits ----------------
rule binning_full:
    input:   "assembly/{}/all.fasta".format(ASSEMBLER), "binning/binning.tsv"
    output:  expand("binning/full/{bin}.fasta", bin=BINS)
    log:     "binning/full.log"
    message: "Distributing full contigs between bins"
    shell:   "mkdir -p binning/full && {SCRIPTS}/split_bins.py {input} binning/full -g >{log}"

#---- F1 stats -----------------------------------------------------------------
rule combine_bins:
    input:   expand("{{stage}}/bins/{bin}.fasta", bin=BINS)
    output:  "{stage,(binning|propagation|reassembly)}/bins/all.fasta"
    message: "Combine all {wildcards.stage} bins"
    shell:   "{SCRIPTS}/combine_contigs.py {input} > {output}"

rule filter_ref_alignments:
    input:   "stats/q_{stage}.done"
    output:  "stats/q_{stage}/{ref}.tsv"
    params:  bins=" ".join(BINS),
             prefix=lambda w: "| awk -v b=$bin '{{print b \"-\" $0}}'" if w.stage == "reassembly" else ""
    message: "Filtering alignments from {wildcards.stage} onto {wildcards.ref}"
    shell:   "rm -f {output}\n"
             "for bin in {params.bins}\ndo\n"
             "    file=stats/q_{wildcards.stage}/runs_per_reference/{wildcards.ref}/contigs_reports/contigs_report_$bin.stdout\n"
             "    if [ -f $file ] ; then\n"
             "        {SCRIPTS}/filter_nucmer.py $file {MIN_CONTIG_LENGTH} 70 {params.prefix} >> {output}\n"
             "    fi\n"
             "done\n"
             "if [ -n {output} ] ; then touch {output} ; fi\n"

rule combine_refs_info:
    input:   expand("stats/q_{{stage}}/{ref}.tsv", ref=REFS)
    output:  "stats/q_{stage}/total.tsv"
    params:  " ".join(REFS)
    message: "Combining good contigs of {wildcards.stage} for all bins"
    shell:   "rm -f {output}\n"
             "for ref in {params}\ndo\n"
             "    file=stats/q_{wildcards.stage}/$ref.tsv\n"
             "    if [ -f $file ] ; then\n"
             "        awk -v r=$ref '{{print $0 \"\\t\" r}}' $file >> {output}\n"
             "    fi\n"
             "done"

ruleorder: combine_refs_info > filter_ref_alignments

rule calc_f1:
    input:   binning="{stage}/binning.tsv", contigs="{stage}/bins/all.fasta", etalons="stats/q_{stage}/total.tsv"
    output:  "stats/summary/{stage}_f1.txt"
    log:     "stats/f1_{stage}.log"
    message: "Calculating F1 metrics for {wildcards.stage}"
    shell:   "{SCRIPTS}/validate.pl --ffile={input.contigs} --cfile={input.binning}"
             " --sfile={input.etalons} --ofile={log} >{output} 2>&1"

#---- Reference profiles -------------------------------------------------------
rule combine_refs:
    input:   list(REFS.values())
    output:  "stats/refs.fasta"
    message: "Combining all references in a single file"
    shell:   "cat {input} > {output}"

rule ref_profiles:
    input:   "stats/refs.fasta"
    output:  "stats/summary/ref_profiles.tsv"
    log:     "stats/ref_profiles.log"
    message: "Calculating etalon profiles of references"
    shell:   "{BIN}/contig_abundance_counter -k {PROFILE_K} -w tmp -c {input}"
             " -n {SAMPLE_COUNT} -m profile/mts/kmers -o {output} >{log} 2>&1"

#---- Genome fraction etc. -----------------------------------------------------
def quast_input(w):
    if w.stage == "assembly":
        return expand("assembly/{assembler}/{group}.fasta", assembler=ASSEMBLER, group=GROUPS)
    frag = "full" if w.stage == "binning" else "bins"
    return expand("{stage}/{frag}/{bin}.fasta", stage=w.stage, frag=frag, bin=BINS)

rule quast_all:
    input:   quast_input
    output:  touch("stats/q_{stage,\w+}.done")
    params:  out="stats/q_{stage}"
    threads: THREADS
    message: "Aligning all of {wildcards.stage} on all references"
    shell:   "{METAQUAST} {QUAST_PARAMS} {METAQUAST_PARAMS} -t {threads} -R {REFS_STR}"
             " {input} -o {params.out} >/dev/null 2>&1"

rule gather_stats:
    input:   "stats/q_{stage}.done"
    output:  "stats/summary/{stage}_summary.tsv", "stats/summary/{stage}_nga.tsv"
    params:  "--gf", "--summary", "--problematic", "--nga", "--plot", "stats/q_{stage}", "stats/summary/{stage}"
    message: "Gathering {wildcards.stage} stats"
    shell:   "{SCRIPTS}/gather_stats.py {params}"

#---- Per-assembly stats -------------------------------------------------------
def group_input(w):
    dirs = {"binning": "assembly/splits", "propagation": "propagation/edges"}
    return expand("{dir}/{group}.fasta", dir=dirs[w.stage], group=w.group)

rule bin_group:
    input:   contigs=group_input, ann="{stage}/annotation/{group}.ann"
    output:  touch("{stage}/{group,\w+\d+}/group.done")
    params:  "{stage}/{group}"
    message: "Distributing {wildcards.stage} contigs from {wildcards.group} between bins"
    shell:   "{SCRIPTS}/split_bins.py {input.contigs} {input.ann} {params}"

rule group_all:
    input:   expand("{{stage}}/{group}/group.done", group=GROUPS)
    output:  touch("{stage,\w+}/group.done")
    message: "Redistributed all {wildcards.stage} contigs"

rule grouped_quast:
    input:   "{stage}/group.done"
    output:  touch("stats/q_{stage,\w+}/q_{bin,\w+\d+}.done")
    params:  input="{stage}/sample*/{bin}.fasta", output="stats/q_{stage}/q_{bin}"
    threads: THREADS
    message: "Aligning {wildcards.stage} grouped contigs from {wildcards.bin} on all references"
    shell:   "{METAQUAST} {QUAST_PARAMS} {METAQUAST_PARAMS} -t {threads} -R {REFS_STR}"
             " {params.input} -o {params.output} >/dev/null 2>&1"

rule grouped_stat:
    input:   expand("stats/q_{{stage}}/q_{bin}.done", bin=BINS)
    output:  "stats/summary/{stage}_group.tsv"
    params:  gf=lambda w: expand("'stats/q_{stage}/q_{bin}/summary/TSV/Genome_fraction_(%).tsv'", stage=w.stage, bin=BINS)
    message: "Gathering grouped {wildcards.stage} stats"
    run:
        bins = set(BINS)
        for bin in BINS:
            dir = "stats/q_{}/q_{}".format(wildcards.stage, bin)
            if not os.path.isdir(os.path.join(dir, "summary")):
                bins.remove(bin)
            else:
                shell("{SCRIPTS}/gather_stats.py --gf {dir} {dir}/res")
        all = expand("'stats/q_{stage}/q_{bin}/res_gf.tsv'", stage=wildcards.stage, bin=sorted(list(bins)))
        shell("cat {all} >{output}")

#---- Propagator statistics ----------------------------------------------------
rule prop_annotation:
    input:   expand("propagation/annotation/{group}.ann", group=GROUPS)
    output:  "propagation/binning.tsv"
    message: "Combine propagated annotation"
    run:
        shell("rm -f {output}")
        for sample_ann in input:
            sample, _ = os.path.splitext(os.path.basename(sample_ann))
            shell("sed -e 's/^/{sample}-/' {sample_ann} >> {output}")

#Redistribute sample pieces to corresponding bins
rule extract_prop_bins:
    input:   "propagation/edges/all.fasta", "propagation/binning.tsv"
    output:  expand("propagation/bins/{bin}.fasta", bin=BINS)
    message: "Distribute propagated contigs between bins"
    shell:   "{SCRIPTS}/split_bins.py {input} propagation/bins/"

# rule prop_stats:
#     input:   prelim="binning/annotation/{sample}.ann", prop="propagation/annotation/{sample}.ann",
#              contigs="assembly/{sample}.fasta", edges="assembly/edges/{sample}.fasta",
#              ref=REFS.values() #, bins="{sample}/{ref}.bin"
#     output:  "stats/prop_{bin}/{sample}.tsv"
#     log:     "stats/prop_{bin}/{sample}.log"
#     message: "Calculating propagation statistics for {wildcards.sample}"
#     shell:   "{BIN}/stats -k {K} -s {wildcards.sample}/assembly/{SAVES} -r {input.ref}"
#              " -c {input.contigs} -a {input.prelim} -e {input.edges} -p {input.prop}"
#              " -b {wildcards.bin} -o {output} >{log}"

# Run this
# rule prop_stats_all:
#     input:   expand("stats/prop_{cag}/{sample}.tsv", sample=GROUPS, cag=BINS)
#     message: "Calculated propagation statistics"

#---- Lost contigs -------------------------------------------------------------
rule binned_list:
    input:   "binning/binning.tsv"
    output:  "binning/binned.info"
    message: "Filtering list of contigs which were binned"
    shell:   "awk '{{print $1}}' {input} > {output}"

rule lost_contigs:
    input:   "assembly/splits/all.fasta", "binning/binned.info"
    output:  "binning/bins/lost.fasta"
    message: "Filtering contigs which were not binned"
    shell:   "{SCRIPTS}/contig_name_filter.py {input} {output} remove"

rule lost_gf:
    input:   "binning/bins/lost.fasta"
    output:  "stats/summary/binning_lost.tsv"
    params:  out="stats/lost_binning"
    message: "Aligning lost binned contigs on all references"
    shell:   "{METAQUAST} --fast -t {THREADS} -R {REFS_STR} {input} -o {params.out} >/dev/null 2>&1 && "
             "cp {params.out}/summary/Genome_fraction_\(%\).tsv {output}"

#===============================================================================
#---- Reference-free stats and metrics -----------------------------------------
#===============================================================================

#---- Minimus for pre-reassembly stages ----------------------------------------
rule to_amos:
    input:   "{stage}/bins/{bin}.fasta"
    output:  "{stage}/minimus/{bin}.afg"
    log:     "{stage}/minimus/to_amos_{bin}.log"
    message: "Converting {wildcards.bin} contigs from {wildcards.stage} into AMOS format"
    shell:   "toAmos -s {input} -o {output} >{log} 2>&1"

rule minimus:
    input:   "{stage}/minimus/{bin}.afg"
    output:  "{stage}/minimus/{bin}.fasta"
    params:  "{stage}/minimus/{bin}"
    log:     "{stage}/minimus/{bin}.log"
    message: "Merging {wildcards.bin} from {wildcards.stage} with minimus2"
    shell:   "minimus2 {params} >{log} 2>&1 && cat {params}.singletons.seq >> {output}"

#---- Deduplicate multi-bins via Prodigal and CD-HIT ---------------------------
rule checkm_tree:
    input:   expand("binning/full/{bin}.fasta", bin=BINS)
    output:  touch("stats/checkm_merge.done")
    threads: THREADS
    log:     "stats/checkm_merge.log"
    message: "Running Prodigal on multi-bins"
    shell:   "set +u; source activate mts27; set -u \n"
             "checkm tree -t {threads} --pplacer_threads {threads} -x .fasta --nt"
             " binning/full stats/checkm_merge &>{log}"

rule cdhit:
    input:   "stats/checkm_merge.done"
    output:  "binning/cdhit/{bin}.fasta"
    params:  "stats/checkm_merge/bins/{bin}/genes.fna"
    threads: THREADS
    log:     "binning/cdhit/{bin}.log"
    message: "Clustering {wildcards.bin} with CD-HIT"
    shell:   "{SOFT}/cd-hit/cd-hit-est -s 0.5 -A 0.95 -i {params} -o {output} -T {threads} &>{log}"

#---- CheckM stats -------------------------------------------------------------
def checkm_input_dir(w):
    return w.stage + ("/bins" if (w.stage == "reassembly" or len(GROUPS) == 1) else "/cdhit")

rule checkm:
    input:   lambda w: expand("{dir}/{bin}.fasta", dir=checkm_input_dir(w), bin=BINS)
    output:  qa="stats/checkm_{stage,binning|reassembly}/qa.tsv",
             tree_qa="stats/checkm_{stage}/tree_qa.tsv"
    params:  dir=checkm_input_dir, out="stats/checkm_{stage}"
    threads: THREADS
    log:     "stats/checkm_{stage}.log"
    message: "Running CheckM for results of {wildcards.stage}"
    run:
        ext = ".fasta"
        shell("set +u; source activate mts27; set -u \n"
             "checkm tree -t {THREADS} --pplacer_threads {THREADS} -x {ext} {params.dir} {params.out} >{log} 2>&1\n"
             "checkm tree_qa -o 2 --tab_table -f {output.tree_qa} {params.out} >>{log} 2>&1\n"
             "checkm lineage_set {params.out} {params.out}/lineage.ms >>{log} 2>&1\n"
             "checkm analyze -t {THREADS} -x fasta {params.out}/lineage.ms {params.dir} {params.out} >>{log} 2>&1\n"
             "checkm qa -t {THREADS} -o 2 --tab_table -f {output.qa} {params.out}/lineage.ms {params.out} >>{log} 2>&1")

rule parse_checkm:
    input:   qa="stats/checkm_{stage}/qa.tsv", tree_qa="stats/checkm_{stage}/tree_qa.tsv"
    output:  "stats/summary/{stage}_checkm.tsv"
    message: "Parse CheckM results for {wildcards.stage}"
    run:
        table = pandas.read_table(input.qa, dtype="str")
        tree_table = pandas.read_table(input.tree_qa, dtype="str", na_filter=False)
        all_table = pandas.merge(table, tree_table, on="Bin Id")
        res_table = all_table[["Bin Id", "Taxonomy (contained)", "Taxonomy (sister lineage)", "Genome size (Mbp)", "Completeness", "Contamination"]].copy()
        def extract_taxon(taxonomy):
            return str(taxonomy).split(";")[-1]
        for column in ["Taxonomy (contained)", "Taxonomy (sister lineage)"]:
            res_table[column] = res_table[column].apply(extract_taxon)
        res_table.to_csv(output[0], index=False, sep="\t")

rule reassembly_mock_binning:
    output:  "reassembly/binning.tsv"
    message: "Preparing reassembly mock binning"
    run:
        shell("rm -f {output}")
        for bin in BINS:
            shell("grep '>' reassembly/bins/{bin}.fasta | cut -c 2- | awk '{{print \"{bin}-\" $0 \"\\t{bin}\"}}' >> {output}")

# To be moved into main pipeline
rule split_groups:
    input:   "binning/bins/{bin}.fasta"
    output:  touch("binning/groups/{bin}.done")
    message: "Splitting {wildcards.bin} into groups"
    shell:   "{SCRIPTS}/split_groups.py {input} $(dirname {output})"

rule checkm_bin:
    input:   "binning/groups/{bin}.done"
    output:  qa="stats/checkm_{bin,BIN\d+}/qa.tsv",
             tree_qa="stats/checkm_{bin}/tree_qa.tsv"
    params:  dir="binning/groups/{bin}/", out="stats/checkm_{bin}"
    threads: THREADS
    log:     "stats/checkm_{bin}.log"
    message: "Running CheckM for groups from {wildcards.bin}"
    run:
        ext = ".fasta"
        shell("set +u; source activate mts27; set -u \n"
             "checkm tree -t {THREADS} --pplacer_threads {THREADS} -x {ext} {params.dir} {params.out} >{log} 2>&1\n"
             "checkm tree_qa -o 2 --tab_table -f {output.tree_qa} {params.out} >>{log} 2>&1\n"
             "checkm lineage_set {params.out} {params.out}/lineage.ms >>{log} 2>&1\n"
             "checkm analyze -t {THREADS} -x fasta {params.out}/lineage.ms {params.dir} {params.out} >>{log} 2>&1\n"
             "checkm qa -t {THREADS} -o 2 --tab_table -f {output.qa} {params.out}/lineage.ms {params.out} >>{log} 2>&1")

#Fake rule to test
rule checkm_all_bins:
    input:   expand("stats/checkm_{bin}/qa.tsv", bin=BINS)
    output:  touch("stats/checkm_bins.done")
    message: "FInished CheckM for all regrouped bins"

#---- PCA ----------------------------------------------------------------------
# FRAGMENT_NAMES_BY_STAGE = {"reassembly": CAG_EDGES,
#                            "assembly": list(GROUPS.keys())}

def fragments_info_by_stage(wildcards):
    fragments=FRAGMENT_NAMES_BY_STAGE[wildcards.stage]
    return expand("stats/q_{stage}/runs_per_reference/{ref}/{fs}.info", stage=wildcards.stage, ref=wildcards.ref, fs=fragments)

rule pca:
    input:   "binning/profiles.tsv", "binning/binning.tsv",
             #"stats/q_{stage}/total.info"
    output:  "stats/summary/pca_{stage}.png"
    message: "Drawing PCA visualisation for {wildcards.stage}"
    shell:   "Rscript {SCRIPTS}/pca.R {input} {output}"

#---- TSNE ---------------------------------------------------------------------

rule tsne:
    input:   "binning/profiles.tsv", "binning/binning.tsv",
    output:  "stats/summary/tsne_{stage}.png"
    message: "Drawing BH-TSNE visualisation for {wildcards.stage}"
    shell:   "python2 {SCRIPTS}/run_tsne.py {input} {output}"
             " --percent 1.0 --iteration 2000 --perplexity 50"

#---- HTML summary report ------------------------------------------------------

report_targets = {"p1": "reassembly_gf.png",
                  "t1": "f1_reassembly.txt",
                  "p2": "reassembly_nga.png"
                 }

def render(key, value):
    if key[0] == "p": #This is a Picture
        return "        .. figure:: {}\n\n           {}".format(value, os.path.splitext(value)[0])
    elif key[0] == "t": #This is a Text
        return "\n        {}\n".format(value)
    return None

def render_all(targets):
    return ("\n" + "\n\n".join(render(key, value) for key, value in targets.items()))

rule report:
    input:   "stats/summary/reassembly_summary.tsv"
    output:  "stats/summary/report_reassembly.html"
    message: "Printing the report"
    run:
        report("""
        =============
        Stats summary
        =============
        Data: {SAMPLE_COUNT} samples from {IN}

        Reassembly stats
        ----------------
        """
        #
        # .. figure:: reassembly_gf.png
        #
        #    reassembly_gf
        #"""
        + render_all(report_targets)
        , output[0])
