include: "Common.snake"

import yaml

reassembly_config = config.get("reassembly", dict())
CUTOFF_COEFF = reassembly_config.get("cutoff_coeff", 0.3)
READ_LENGTH = config.get("read_length", 100) #TODO: auto-detect default value

#Autodetecting bins for reassembly based on the choose_samples result in the main pipeline
BINS, = glob_wildcards("reassembly/{{bin,{}}}.info".format(NAME_TEMPLATE))
if not BINS:
    raise WorkflowError("No bins detected!")

rule reassemble_all:
    input:   expand("reassembly/bins/{bin}.fasta", bin=BINS)
    message: "Reassembly finished."

rule reassembly_config:
    input:   "reassembly/{bin}.info"
    output:  "reassembly/{bin}.yaml"
    message: "Generating config file for reassembly of {wildcards.bin}"
    run:
        with open(output[0], "w") as outfile:
            conf = {"k": PROFILE_K, "sample_cnt": SAMPLE_COUNT,
                    "kmer_mult": "profile/mts/kmers",
                    "bin": wildcards.bin, "bin_prof": "binning/{}/bins.prof".format(BINNER),
                    "edges_sqn": "reassembly/{}/edges.fasta".format(wildcards.bin),
                    "edges_mpl": "reassembly/{}/edges.mpl".format(wildcards.bin),
                    "edge_fragments_mpl": "reassembly/{}/edges_frag.mpl".format(wildcards.bin),
                    "frag_size": SPLIT_LENGTH, "min_len": 100}
            yaml.dump(conf, outfile)

rule subsample:
    input:   "reassembly/{bin}.info"
    output:  "reassembly/{bin}/reads.info"
    params:  "reads/{bin}", "reassembly/{bin}"
    log:     "reassembly/{bin}/subsample.log"
    message: "Subsampling reads for {wildcards.bin}"
    shell:   "{SCRIPTS}/subsample.py {input} {params} >{log}"

rule reassemble:
    input:   info="reassembly/{bin}/reads.info", config="reassembly/{bin}.yaml"
    output:  "reassembly/bins/{bin}.fasta"
    params:  out="reassembly/{bin}"
    log:     "reassembly/{bin}.log"
    threads: THREADS
    message: "Reassembling reads for {wildcards.bin}"
    run:
        with open(input.info) as reads_info:
            sampled_cov = float(reads_info.readline())
            files = [line.split() for line in reads_info]
        cov_cutoff = CUTOFF_COEFF * sampled_cov * (READ_LENGTH - ASSEMBLY_K) / (READ_LENGTH - PROFILE_K)
        reads_param = " ".join(["--pe1-1 {} --pe1-2 {}".format(left, right) for left, right in files])
        series_analysis = "--series-analysis " + input.config if PROFILER == "mts" else ""
        shell("{REASSEMBLER_DIR}/spades.py --meta --only-assembler -t {threads} {series_analysis}"
             " {reads_param} --pe1-ff --hidden-cov-cutoff {cov_cutoff} -o {params.out}"
             " >{log} 2>&1 && cp {params.out}/scaffolds.fasta {output}")
