include: "Common.snake"

import os
import os.path

import pandas
from pandas import DataFrame

from scripts.common import gather_refs, dump_dict

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

#Autodetect bins
CAGS, = glob_wildcards("binning/{cag,CAG\d+}/left.fastq")
CAGS.sort()

CAG_EDGES = [c + "_edges" for c in CAGS]

#Detect references
REFS = dict(gather_refs(config.get("REFS", [])))
ALL_REFS = ",".join(path for path in REFS.values())

FRAGMENT_NAMES_BY_TYPE = {"reassembly": CAG_EDGES,
                          "initial_assembly": list(GROUPS.keys())}

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

onstart:
    try:
        os.mkdir("tmp")
    except:
        pass
    print("Detected", SAMPLE_COUNT, "samples in", IN)
    if CAGS:
        print("Detected good (abundant) CAGs:", " ".join(CAGS))
    if REFS:
        print("Detected references:", " ".join(REFS))

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

#---- Single alignments for samples per reference -------------------------------
#TODO: use alignments from meta version instead
rule quast_all_samples:
    input:   ref_fn=ref_path, contigs=expand("assembly/{sample}.fasta", sample=GROUPS)
    output:  summary_tsv="stats/summary/q_{ref}.tsv", report="stats/initial_assembly/{ref}/report.txt"
    params:  "stats/initial_assembly/{ref}"
    log:     "stats/initial_assembly/{ref}/quast.log"
    threads: THREADS
    message: "Aligning all samples on {wildcards.ref}"
    shell:   "{QUAST} -t {threads} -R {input.ref_fn} {input.contigs} -o {params} >/dev/null 2>&1 && "
             "cp {params}/report.tsv {output.summary_tsv}"

rule quast_all_reassemblies:
    input:   ref=ref_path, fragments=expand("profile/{cag_edges}.fasta", cag_edges=CAG_EDGES)
    output:  "stats/reassembly/{ref}/report.txt"
    params:  "stats/reassembly/{ref}"
    log:     "stats/reassembly/{ref}/quast.log"
    threads: THREADS
    message: "Aligning all samples on {wildcards.ref}"
    shell:   "{QUAST} -t {threads} -R {input.ref} {input.fragments} -o {params} >/dev/null 2>&1 && "
             "cp {params}/report.tsv {output}"

#---- Contigs of interest ------------------------------------------------------
rule filter_ref_alignments:
    input:   "{path}/report.txt"
    output:  "{path}/{fragments}.info"
    params:  "{path}/contigs_reports/nucmer_output/{fragments}.coords.filtered"
    shell:   "if [ -f {params} ] ; then {SCRIPTS}/filter_nucmer.py {params} {output} {MIN_CONTIG_LENGTH} 70 ; else touch {output} ; fi"

#---- GF of combined sample ----------------------------------------------------
#rule combine_filtered:
#    input:   contigs=expand("assembly/{sample}.fasta", sample=GROUPS),
#             filters=expand("stats/{{ref}}/{sample}.cont", sample=GROUPS)
#    output:  "stats/{ref}.fasta"
#    message: "Gathering all interesting contigs for {wildcards.ref} into a single assembly"
#    shell:   "{SCRIPTS}/filter_contigs.py {SAMPLE_COUNT} {output} {input.contigs} {input.filters}"

rule quast_combined:
    input:   ref=ref_path, contigs="stats/{ref}.fasta"
    output:  "stats/q_{ref}_all/report.tsv"
    params:  "stats/q_{ref}_all"
    log:     "stats/q_{ref}_all.log"
    threads: THREADS
    message: "Aligning combined sample on {wildcards.ref}"
    shell:   "{QUAST} -t {threads} -R {input.ref} {input.contigs} -o {params} >{log} 2>&1"

# Run this
rule quast_combined_all:
    input:   expand("stats/q_{ref}_all/report.tsv", ref=REFS)
    message: "Calculated QUAST metrics on all combined samples"

#---- Bins of interest ---------------------------------------------------------
rule int_bins:
    input:   "annotation/{sample}.ann", "stats/{ref}/{sample}.info"
    output:  "stats/{ref}/{sample}.bin"
    message: "Filtering interesting bins for {wildcards.sample} aligned to {wildcards.ref}"
    shell:   "{SCRIPTS}/filter_bins.py {input} > {output}"

rule int_bins_all_samples:
    input:   expand("stats/{{ref}}/{sample}.bin", sample=GROUPS)
    output:  "stats/{ref}/total.bin"
    message: "Gathering interesting bins for {wildcards.ref} from all samples"
    run:
        bins = set()
        for in_fn in input:
            with open(in_fn) as infile:
                for line in infile:
                    bins.add(line)
        with open(output[0], "w") as outfile:
            for bin in bins:
                print(bin, file=outfile)

# Run this
rule int_bins_all:
    input:   expand("stats/{ref}/total.bin", ref=REFS)
    message: "Gathered all interesting bins"

#---- GF per bin per reference -------------------------------------------------
#Helper formatters for determining input files from different stages
PROP = {"prelim": ("assembly/{}_splits.fasta",   "annotation/{}.ann"),
        "prop":   ("propagation/{}_edges.fasta", "propagation/{}_edges.ann")}

#TODO: split into different directories per sample
rule split_bins:
    input:   lambda w: PROP[w.prop][0].format(w.sample),
             lambda w: PROP[w.prop][1].format(w.sample)
    output:  touch("binning/{prop}/{sample}.log")
    log:     "binning/{prop}/split_{sample}.log"
    params:  "binning/{prop}"
    message: "Splitting assembly of {wildcards.sample} between {wildcards.prop} bins"
    shell:   "{SCRIPTS}/split_bins.py {input} {params} >{log}"

rule cat_binned_contigs:
    input:   expand("binning/{{prop}}/{sample}.log", sample=SAMPLES)
    output:  "binning/{prop}/{cag,CAG\d+}.fasta"
    params:  "`ls binning/{prop}/*-{cag}.fasta`"
    message: "Combine binned contigs ({wildcards.prop}) for {wildcards.cag}"
    shell:   "cat {params} > {output}"

#Two helpers for determining dependencies of QUAST targets.
#For split contigs and reassemblies, we need only corresponding FASTA.
#For combined contigs, we need to glue their split pieces first.
def stats_input(wildcards):
    if wildcards.stage == "reassembly":
        return expand("reassembly/{cag}.fasta", cag=CAGS)
    w_bin, w_prop = wildcards.stage.split("_", 2)
    if w_bin == "split":
        return expand("binning/{prop}/{sample}.log", prop=w_prop, sample=GROUPS)
    elif w_bin == "bin":
        return expand("binning/{prop}/{cag}.fasta", prop=w_prop, cag=CAGS)

def stats_data(wildcards):
    if wildcards.stage == "reassembly":
        return "`ls reassembly/CAG*.fasta`"
    w_bin, w_prop = wildcards.stage.split("_", 2)
    masks = {"bin": "CAG*", "split": "*-CAG*"}
    return "`ls binning/{}/{}.fasta`".format(w_prop, masks[w_bin])

rule quast_stats:
    input:   stats_input
    output:  "stats/summary/gf_{stage}.tsv"
    params:  data=stats_data, out="stats/q_{stage}"
    log:     "stats/q_{stage}.log"
    threads: THREADS
    message: "Aligning {wildcards.stage} assemblies on all references"
    shell:   "{METAQUAST} -t {threads} -R {ALL_REFS} {params.data} -o {params.out} >{log} 2>&1 && "
             "cp '{params.out}/summary/TSV/Genome_fraction_(%).tsv' {output}"

# Run this AFTER 'all'
rule stats_all:
    input:   expand("stats/summary/gf_{bin}_{prop}.tsv", bin=["bin"], prop=["prelim", "prop"]), 
             "stats/initial_assembly/total.cont"
    message: "Gathered some numbers, deal with them."

#---- Reassembly statistics ----------------------------------------------------

# Run this AFTER 'reassembly_all'
rule stats_reassembly:
    input:   "stats/summary/gf_reassembly.tsv",
             "stats/reassembly/total.cont"
    output:  "stats/summary/reassembly.tsv"
    params:  "stats/q_reassembly"
    message: "Gathered bins stats"
    shell:   "{SCRIPTS}/gather_stats.py {params} > {output}"

#---- Propagator statistics ----------------------------------------------------
rule prop_stats:
    input:   prelim="annotation/{sample}.ann", prop="annotation/{sample}_edges.ann",
             contigs="assembly/{sample}.fasta", edges="assembly/{sample}_edges.fasta",
             ref=REFS.values() #, bins="{sample}/{ref}.bin"
    output:  "stats/prop_{cag}/{sample}.tsv"
    log:     "stats/prop_{cag}/{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.cag} -o {output} >{log}"

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

#---- CheckM stats -------------------------------------------------------------
rule checkm:
    input:   expand("reassembly/{cag}.fasta", cag=CAGS)
    output:  qa="stats/checkm/qa.tsv", tree_qa="stats/checkm/tree_qa.tsv"
    params:  dir="stats/checkm"
    threads: THREADS
    shell:   "set +u; source activate concoct_env; set -u \n"
             "checkm tree -x fasta reassembly {params.dir} \n"
             "checkm tree_qa -o 2 --tab_table -f {output.tree_qa} {params.dir}\n"
             "checkm lineage_set {params.dir} {params.dir}/lineage.ms\n"
             "checkm analyze -x fasta {params.dir}/lineage.ms reassembly {params.dir}\n"
             "checkm qa -o 2 --tab_table -f {output.qa} {params.dir}/lineage.ms {params.dir}"

rule parse_checkm:
    input:   qa=rules.checkm.output.qa, tree_qa=rules.checkm.output.tree_qa
    output:  "stats/summary/checkm.tsv"
    #shell:   "{SCRIPTS}/parse_checkm.py {input.qa} {input.tree_qa} > {output}"
    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")

#---- PCA ----------------------------------------------------------------------
rule pca:
    input:   "profile/canopy.in", "profile/canopy.out", "{sample}.cont"
    output:  "stats/{sample}.png"
    message: "Doing some visualization"
    shell:
        "Rscript {SCRIPTS}/pca.R {input} {output}"

def fragments_info_by_assembly_type(wildcards):
    frags=FRAGMENT_NAMES_BY_TYPE[wildcards.assembly_type]
    return expand("stats/{assembly_type}/{ref}/{fragments}.info", assembly_type=wildcards.assembly_type, ref=wildcards.ref, fragments=frags)

rule combine_fragments_info:
    input:  fragments_info_by_assembly_type 
    output: "stats/{assembly_type}/{ref}/ref.cont"
    shell: "rm -rf {output}; for f in {input}; do name=$(basename $f .info); cat $f | sed 's/^/'$name'-/g' >> {output} ; done"

rule combine_refs_info:
    input:  expand("stats/{{assembly_type}}/{ref}/ref.cont", ref=list(REFS.keys()))
    output:  "stats/{assembly_type}/total.cont"
    run:
        shell("rm -rf {output}")
        for ref in REFS.keys():
            shell("awk '{{print $0 \"\t{ref}\"}}' stats/{wildcards.assembly_type}/{ref}/ref.cont >> {output}")

# Run this
rule pca_total:
    input:   "binning/canopy/profiles.in", "binning/canopy/binning.out", "stats/total.cont"
    output:  "stats/summary/pca.png"
    shell:   "Rscript {SCRIPTS}/pca.R {input} {output}"
