# Binning with MetaBAT
from scripts.common import contig_length
from scripts.converters import ProfileFormatter, BinningParser

class MetabatFormatter(ProfileFormatter):
    def __init__(self, var=True):
        self.var = var

    def header(self, first_line):
        if self.var:
            columns = ["contigName", "contigLen", "totalAvgDepth"]
            for i in range(1, len(first_line) // 2 + 1):
                sample = "sample" + str(i)
                columns += [sample + ".cov", sample + ".var"]
        else:
            columns = ["contigName"]
            for i in range(1, len(first_line) + 1):
                sample = "sample" + str(i)
                columns += [sample + ".cov"]
        print(*columns, sep="\t", file=self.out)

    def format(self, contig, profile):
        if self.var:
            total_cov = "{:.2f}".format(sum(map(float, profile[0::2])))
            print(contig, contig_length(contig), total_cov, *profile, sep="\t", file=self.out)
        else:
            print(contig, *profile, sep="\t", file=self.out)

rule metabat_pre:
    input:   "binning/profiles.tsv"
    output:  "binning/metabat/profiles.in"
    message: "Preparing MetaBAT input"
    run:
        MetabatFormatter(PROF_VAR).run(input[0], output[0])

rule metabat:
    input:   contigs="assembly/{}/all.fasta".format(FRAGS), profiles="binning/metabat/profiles.in"
    output:  "binning/metabat/binning.out"
    threads: THREADS
    params:  out="binning/metabat/cluster", flag="" if PROF_VAR else "--cvExt"
    log:     "binning/metabat.log"
    message: "Running MetaBAT clustering"
    shell:   "mkdir -p binning/bins\n"
             "{SOFT}/metabat/metabat -t {threads} -m {MIN_CONTIG_LENGTH} "
             " --minContigByCorr {MIN_CONTIG_LENGTH} --saveCls {params.flag}"
             " -i {input.contigs} -a {input.profiles} -o {output} >{log}"

rule metabat_post:
    input:   "binning/metabat/binning.out"
    output:  "binning/binning.tsv" #Also bins
    message: "Postprocessing MetaBAT results"
    run:
        BinningParser(sep="\t", filter="0").run(input[0], output[0])
        shell(
             "mkdir -p binning/bins\n"
             "for file in binning/metabat/*.fa; do\n"
             "    bin=BIN${{file//[^0-9]}}\n"
             "    mv $file binning/bins/$bin.fasta\n"
             "done")
