shell:
    """
    gunzip -c {input.fq1} > {output.fq1}
    gunzip -c {input.fq2} > {output.fq2}

    metawrap binning \
        -o {params.out} \
        -t {threads} \
        -l {config[mincontiglength]} \
        -m {resources.mem_GB} \
        -a {input.scaffolds} \
        --concoct --metabat2 --maxbin2 {output.fq1} {output.fq2}

    # rename bins with name
    cd {params.out}concoct_bins
    for f in *.fa ; do mv -- "$f" "{wildcards.project}_{wildcards.sample}_concoct_$f" || true ; done
    cd {params.out}metabat2_bins
    for f in *.fa ; do mv -- "$f" "{wildcards.project}_{wildcards.sample}_metabat2_$f" || true ; done
    cd {params.out}maxbin2_bins
    for f in *.fa ; do mv -- "$f" "{wildcards.project}_{wildcards.sample}_maxbin2_$f" || true ; done

    # remove unbinned bins
    cd {params.out}
    mkdir unbinned
    mv *_bins/*unbinned.fa unbinned/ || true

    # remove temp files
    cd {params.out}
    rm -r work_files
    """

shell:
    """
    FNA=$(realpath {input.fa})
    DB=$(realpath {input.db})
    mkdir -p {params.d}
    if [ ! -z "$(ls -A {params.d})" ]; then
      echo "Folder not empty. Deleting"
      rm -rf {params.d}
    fi
    mkdir -p {params.d}
    cd {params.d}
    busco  --offline \
          -i $FNA \
          -m genome \
          -o out \
          --auto-lineage-euk \
          --download_path $DB \
          -c {threads} || touch out/short_summary.specific_failed.out.txt

    cp out/short_summary.specific*.out.txt short_summary.specific.txt
    rm -r out
    """

shell:
    """
      minimap2 -ax sr -t {threads} {input.fasta} {input.fq1} {input.fq2}  | 
              samtools view -q 20 -Sb - | \
              samtools sort -@ {threads} -O bam - -o {output.bam}
              samtools index {output.bam}

      # cmseq
      breadth_depth.py \
          --combine \
          --mincov {params.mincov} \
          {output.bam} > {output.cov}
    """

shell:
    """
    mkdir -p {params.out}
    rm -r  {params.out}
    mkdir -p {params.out}
    BD=$(pwd)
    # prepare drep quality file
    echo "genome,completeness,contamination" > {output.quality}
    cat {input.concoct_eukcc}  {input.metabat2_eukcc}|\
        awk '{{if($2 - 5*$3 >=50){{print $0}}}}' |\
        sort -k 2,3 -n |\
        tr '\\t' ',' >> {output.quality}
    if [[ $(wc -l <{output.quality}) -lt 2  ]];
    then
        touch {output.quality}
        touch {output.genomes}
        mkdir -p {output.input_bins} || true
        mkdir -p {output.bins} || true
        exit 0
    fi

    cat {output.quality}
    mkdir {output.input_bins}
    cd {output.input_bins}
    ln -s {input.concoct}/*.fa . || true
    ln -s {input.concoct_merged}/*.fa . || true
    ln -s {input.metabat2}/*.fa . || true
    ln -s {input.metabat2_merged}/*.fa . || true
    ls
    ls > ../genomes.txt
    cd ..
    cat quality.csv | grep -v "genome,completeness,contaminatio" |\
        cut -f 1 -d , > .f
    grep -f .f genomes.txt | sed -e 's/^/input_genomes\//'  > .g
    mv .g genomes.txt
    rm .f
    cd $BD

    mkdir -p {params.out}
    cd {params.out}

    cat genomes.txt
    if [[ $(wc -l <genomes.txt) -lt 2  ]];
    then
        echo "No genomes"
        cp -r input_genomes dereplicated_genomes
    else
      dRep dereplicate -p {threads} \
          . -g genomes.txt \
          -pa 0.80 -sa 0.99 -nc 0.40 \
          -cm larger \
          --genomeInfo quality.csv\
          -comp 49 -con 21
    fi
    """

shell:
    """
    rm -r {params.out}workdir || true
    # if no bins
    if [[ $( ls {input.bindir} | wc -l ) -lt 1 ]]; then
      touch {output.csv}
      mkdir -p {output.bins}
      exit 0
    fi

    eukcc --debug folder \
        --improve_percent 10 \
        --n_combine 1 \
        --threads {threads} \
        --improve_ratio     5 \
        --links {input.links} \
        --min_links 100 \
        --suffix .fa \
        --db {input.db} \
        --out {params.out} \
        --prefix {wildcards.project}_{wildcards.sample}_{wildcards.binner}_merged. \
        {input.bindir}
    rm -r {params.out}/refine_workdir || true
    """

import pysam
from Bio import SeqIO
from collections import defaultdict
import os
import argparse
import logging
import csv


def is_in(read, contig_map, within=1000):
    if read.reference_name not in contig_map.keys():
        return False
    if read.reference_start <= within or read.reference_end <= within:
        return True
    elif read.reference_start > (
        contig_map[read.reference_name] - within
    ) or read.reference_end > (contig_map[read.reference_name] - within):
        return True
    else:
        return False


def keep_read(read, contig_map, within=1000, min_ANI=98, min_cov=0):
    ani = (
        (read.query_alignment_length - read.get_tag("NM"))
        / float(read.query_alignment_length)
        * 100
    )
    cov = read.query_alignment_length / float(read.query_length) * 100

    if ani >= min_ANI and cov >= min_cov and is_in(read, contig_map, within) is True:
        return True
    else:
        return False


def contig_map(bindir, suffix=".fa"):
    m = {}
    for f in os.listdir(bindir):
        if f.endswith(suffix) is False:
            continue
        path = os.path.join(bindir, f)
        with open(path, "r") as handle:
            for record in SeqIO.parse(handle, "fasta"):
                m[record.name] = len(record.seq)
    return m


def bin_map(bindir, suffix=".fa"):
    contigs = defaultdict(str)
    contigs_per_bin = defaultdict(int)
    for f in os.listdir(bindir):
        if f.endswith(suffix) is False:
            continue
        path = os.path.join(bindir, f)
        binname = os.path.basename(f)
        with open(path, "r") as handle:
            for record in SeqIO.parse(handle, "fasta"):
                contigs[record.name] = binname
                contigs_per_bin[binname] += 1
    return contigs, contigs_per_bin


def read_pair_generator(bam):
    """
    Generate read pairs in a BAM file or within a region string.
    Reads are added to read_dict until a pair is found.
    From: https://www.biostars.org/p/306041/
    """
    read_dict = defaultdict(lambda: [None, None])
    for read in bam.fetch():
        if not read.is_paired or read.is_secondary or read.is_supplementary:
            continue
        qname = read.query_name
        if qname not in read_dict:
            if read.is_read1:
                read_dict[qname][0] = read
            else:
                read_dict[qname][1] = read
        else:
            if read.is_read1:
                yield read, read_dict[qname][1]
            else:
                yield read_dict[qname][0], read
            del read_dict[qname]


def main():
    # set arguments
    # arguments are passed to classes
    parser = argparse.ArgumentParser(
        description="Evaluate completeness and contamination of a MAG."
    )
    parser.add_argument("bindir", type=str, help="Run script on these bins")
    parser.add_argument(
        "bam",
        type=str,
        help="Bam with allr eads aligned against all contigs making up the bins",
    )
    parser.add_argument(
        "--out",
        "-o",
        type=str,
        required=False,
        help="Path to output table (Default: links.csv)",
        default="links.csv",
    )
    parser.add_argument(
        "--ANI", type=float, required=False, help="ANI of matching read", default=99
    )
    parser.add_argument(
        "--within",
        type=int,
        required=False,
        help="Within this many bp we need the read to map",
        default=1000,
    )
    parser.add_argument(
        "--contigs",
        "-c",
        action="store_true",
        default=False,
        help="Instead of bins print contigs",
    )
    parser.add_argument(
        "--quiet",
        "-q",
        dest="quiet",
        action="store_true",
        default=False,
        help="Silcence most output",
    )
    parser.add_argument(
        "--debug",
        "-d",
        action="store_true",
        default=False,
        help="Debug and thus ignore safety",
    )
    args = parser.parse_args()

    # define logging
    logLevel = logging.INFO
    if args.quiet:
        logLevel = logging.WARNING
    elif args.debug:
        logLevel = logging.DEBUG
    logging.basicConfig(
        format="%(asctime)s %(message)s",
        datefmt="%d-%m-%Y %H:%M:%S: ",
        level=logLevel,
    )

    bindir = args.bindir
    samfile = pysam.AlignmentFile(args.bam, "rb")

    cm = contig_map(bindir)
    bm, contigs_per_bin = bin_map(bindir)

    link_table = defaultdict(lambda: defaultdict(int))
    bin_table = defaultdict(lambda: defaultdict(int))

    # generate link table
    logging.info("Parsing Bam file. This can take a few moments")
    for read, mate in read_pair_generator(samfile):
        if keep_read(read, cm, args.within, min_ANI=args.ANI) and keep_read(
            mate, cm, args.within, min_ANI=args.ANI
        ):
            # fill in the table
            link_table[read.reference_name][mate.reference_name] += 1
            if read.reference_name != mate.reference_name:
                link_table[mate.reference_name][read.reference_name] += 1

    # generate bin table
    for contig_1, dic in link_table.items():
        for contig_2, links in dic.items():
            bin_table[bm[contig_1]][bm[contig_2]] += 1
    out_data = []
    logging.debug("Constructing output dict")
    if args.contigs:
        for contig_1, linked in link_table.items():
            for contig_2, links in linked.items():
                out_data.append(
                    {
                        "bin_1": bm[contig_1],
                        "bin_2": bm[contig_2],
                        "contig_1": contig_1,
                        "contig_2": contig_2,
                        "links": links,
                        "bin_1_contigs": contigs_per_bin[bm[contig_1]],
                        "bin_2_contigs": contigs_per_bin[bm[contig_2]],
                    }
                )
    else:
        for bin_1, dic in bin_table.items():
            for bin_2, links in dic.items():
                out_data.append({"bin_1": bin_1, "bin_2": bin_2, "links": links})
    # results
    logging.info("Writing output")
    with open(args.out, "w") as fout:
        cout = csv.DictWriter(fout, fieldnames=list(out_data[0].keys()))
        cout.writeheader()
        for row in out_data:
            cout.writerow(row)


if __name__ == "__main__":
    main()