import click
from Bio import SeqIO
import pandas as pd
import os


@click.command()
@click.option('--input', '-i', type=click.Path(exists=True), help='Input FASTA file')
@click.option('--output', '-o', type=click.Path(), help='Output FASTA file')
@click.option('--stats', '-s', type=click.Path(), help='Genome stats outfile')
def main(input, output, stats):
    # read in the input file
    records = list(SeqIO.parse(input, "fasta"))

    # modify the sequence IDs and descriptions
    protcount = 0
    for record in records:
        protcount += 1
        if "|" in record.id:
            record.id = record.id.split("|")[-1]
        record.id = f"{os.path.splitext(os.path.basename(input))[0]}|{record.id}"
        record.description = ""

    # write the reformatted sequences to the output file
    SeqIO.write(records, output, "fasta")

    # calculate statistics and write them to the stats file
    stats_dict = {
        'query': os.path.splitext(os.path.basename(input))[0],
        'contigs': 'no_fna',
        'LENbp': 'no_fna',
        'GCperc': 'no_fna',
        'genecount': protcount,
        'CODINGperc': 'no_fna',
        'ttable': 'no_fna'
    }
    pd.DataFrame.from_dict(stats_dict, orient='index').T.to_csv(stats, sep="\t", index=False)

if __name__ == '__main__':
    main()

import click
import os
import sys
import shutil
import glob
from subprocess import call
from Bio import SeqIO
import pandas as pd


def calc_stats(fnain, faain):
    """
    Calculate coding density, assembly size, GC%
    """
    cumulative_gc = 0
    cumulative_len = 0
    cumulative_len_aa = 0
    gene_count = 0
    contigs = 0

    for seq_record in SeqIO.parse(fnain, "fasta"):
        cumulative_gc += seq_record.seq.upper().count('G') + seq_record.seq.upper().count('C')
        cumulative_len += len(seq_record.seq)
        contigs += 1

    for seq_record in SeqIO.parse(faain, "fasta"):
        cumulative_len_aa += len(seq_record.seq)
        gene_count += 1

    return [contigs, cumulative_len, round(cumulative_gc / cumulative_len * 100, 2), gene_count, round(cumulative_len_aa * 3 / cumulative_len * 100, 2)]


def check_filename(fnafile):
    """
    Input has to be .fna with no additional . in filename
    """
    if len(fnafile.split(".")) > 2:
        print("fnafile name contains additional '.'")
        sys.exit()
    if fnafile.split(".")[1] != "fna":
        print("fnafile ending needs to be .fna")
        sys.exit()


def rename_header(bestcode, finalfaa, gffout):
    """
    Rename sequence header to ><filename>|<proteinid>
    """
    records = []
    infilebase = os.path.splitext(os.path.basename(bestcode))[0]

    for seq_record in SeqIO.parse(bestcode, "fasta"):
        headerbase = seq_record.description.split("|")[0]
        if headerbase != infilebase:
            seq_record.id = f"{infilebase}|{str(seq_record.id).split()[0]}"
            seq_record.description = ""
            records.append(seq_record)
        elif headerbase == infilebase:
            seq_record.id = str(seq_record.id).split()[0]
            seq_record.description = ""
            records.append(seq_record)

    SeqIO.write(records, finalfaa, "fasta")
    # keep only best gff
    shutil.copy(gffout + "_code" + bestcode.split("_code")[-1], gffout)


def run_genecalling_codes(fnafile, code, gffout):
    """
    Genecalling with prodigal and defined tt
    """
    faaout = gffout.replace(".gff", ".faa")
    if code > 0:
        run_command = f"prodigal -n -q -f gff -i {fnafile} -a {faaout}_code{code} -o {gffout}_code{code} -g {code}"
        call(run_command, shell=True)
    else:  # output for snakemake with -p meta
        run_command = f"prodigal -n -q -f gff -p meta -i {fnafile} -a {faaout}_codemeta -o {gffout}"
        call(run_command, shell=True)
        shutil.copy(gffout, gffout + "_codemeta")


@click.command()
@click.option('--fnafile', '-f', type=click.Path(exists=True), help='Input FNA file')
@click.option('--gffout', '-g', type=click.Path(), help='Output GFF file')
@click.option('--genecalling_statsout', '-gs', '--genecalling-statsout', type=click.Path(), help='Genecalling statistics file')
@click.option('--summary_statsout', '-ss', type=click.Path(), help='Genome stats outfile')
@click.option('--finalfaa', '-fa', type=click.Path(), help='Output FASTA file')
def main(fnafile, gffout, genecalling_statsout, summary_statsout, finalfaa):
    # check input
    check_filename(fnafile)
    # gene calling
    # create output for snakemake
    codes = [0, 1, 4, 6, 11] # pass these from snakemake config instead
    for code in codes:
        run_genecalling_codes(fnafile, code, gffout)
    stats_dict = {}
    faaoutdir = os.path.dirname(gffout)
    for faain in glob.glob(os.path.join(faaoutdir, "*.faa_code*")):
        genome_id = os.path.basename(faain)
        stats_dict[genome_id] = calc_stats(fnafile, faain)

    # gene calling stats to select ttable that yields highest coding density
    stats_df = pd.DataFrame.from_dict(stats_dict, columns=["contigs", "LENbp", "GCperc", "genecount", "CODINGperc"], orient="index")
    stats_df = stats_df.sort_values("CODINGperc", ascending=False)
    stats_df["ttable"] = stats_df.index.map(lambda x: x.split("_")[-1])
    stats_df.insert(0, "query", stats_df.index.map(lambda x: os.path.splitext(x)[0]))
    bestcode = os.path.join(faaoutdir, stats_df.index[0])  # genome id with code appended that had highest coding density
    stats_df.to_csv(genecalling_statsout, sep="\t", index=None)
    stats_df.head(1).to_csv(summary_statsout, sep="\t", index=None)

    # rename headers and cleanup
    rename_header(bestcode, finalfaa, gffout)
    for tempout in glob.glob(f"{gffout}_code*"):
        os.remove(tempout)
        faaout = tempout.replace(".gff", ".faa")
        if os.path.isfile(faaout):
            os.remove(faaout)


if __name__ == '__main__':
    main()

import click
import pandas as pd
import pickle
from Bio import SeqIO
from Bio.Seq import Seq
from sklearn.preprocessing import LabelEncoder


def get_label_mapping(mapping_file):
    """Returns a dictionary mapping encoded labels to their original values"""
    with open(mapping_file) as f:
        return {int(line.split('\t')[0]): line.strip().split('\t')[1] for line in f}


def get_best_kmers(kmers_file):
    """Returns a list of the best kmers"""
    with open(kmers_file) as f:
        return [line.strip() for line in f.readlines()]


def count_kmers(kmers, fna):
    """Counts the occurrences of kmers in a given FASTA file"""
    kmer_counts_genome_dict = {x:0 for x in kmers}
    total_seqs = sum(1 for _ in SeqIO.parse(fna, "fasta"))
    with click.progressbar(SeqIO.parse(fna, "fasta"), label='Counting kmers', length=total_seqs) as records:
        for seq_record in records:
            for kmer in kmers:
                kmer_counts_genome_dict[kmer] += Seq(seq_record.seq).count_overlap(kmer)
    return kmer_counts_genome_dict


def predict(model, kmers, labels, queryseq):
    # Load XGBoost model
    xgbmodel = pickle.load(open(model, 'rb'))

    # Load label mapping
    inv_mapping = get_label_mapping(labels)

    # Load and preprocess query sequence counts
    kmers_list = get_best_kmers(kmers)
    kmer_count_all_dict = {queryseq.split("/")[-1].split(".f")[0]: count_kmers(kmers_list, queryseq)}
    df_qcounts  =  pd.DataFrame.from_dict(kmer_count_all_dict).T.fillna(0)
    df_qcounts = df_qcounts.div(df_qcounts.sum(axis=1), axis=0)
    df_qcounts = df_qcounts.reset_index().drop(['index'], axis=1)

    # Ensure correct order of features
    cols_when_model_builds = xgbmodel.get_booster().feature_names
    df_qcounts = df_qcounts[cols_when_model_builds]

    # Make prediction and print result
    y_pred = xgbmodel.predict(df_qcounts.head(1))
    prediction = inv_mapping[y_pred[0]]
    prob_val = xgbmodel.predict_proba(df_qcounts.head(1))
    print (prob_val)
    proba_dict = {inv_mapping[x]:[prob_val[0,x]] for x,y in inv_mapping.items()}
    return prediction, proba_dict


@click.command()
@click.option('--queryseq', '-q', type=click.Path(exists=True), required=True, help='Path to query sequence')
@click.option('--model', '-m', type=click.Path(exists=True), required=True, help='Path to prebuilt model')
@click.option('--kmers', '-k', type=click.Path(exists=True), required=True, help='Path to kmer file')
@click.option('--labels', '-l', type=click.Path(exists=True), required=True, help='Path to labels')
@click.option('--out', '-o', type=click.Path(exists=False), required=True, help='Path to prediction output')
def main(queryseq, model, kmers, labels, out):
    # get domain level prediction
    prediction_domain, proba_dict_domain = predict(model, kmers, labels, queryseq)
    # if NCDLV then also get order level prediction
    if prediction_domain == "NCLDV":
        model_order = model.replace("domain", "order")
        kmers_order = kmers.replace("domain", "order")
        labels_order = labels.replace("domain", "order")
        prediction_order, proba_dict_order = predict(model_order, kmers_order, labels_order, queryseq)
    else:
        prediction_order = ""
        proba_dict_order = {}
    with open(out, "w") as outfile:
        outfile.write(f"{queryseq.split('/')[-1].split('.f')[0]}\t{prediction_order}___{prediction_domain}")
    proba_dict = {**proba_dict_domain, **proba_dict_order}
    df_proba = pd.DataFrame(proba_dict)
    df_proba["query"] = queryseq.split('/')[-1].split('.f')[0]
    df_proba = df_proba.set_index("query")
    #print (df_proba)
    #df_proba.to_csv("test.tab", sep="\t", index=None)
    df_proba.to_csv(f"{out}.proba", sep="\t", index=True)

if __name__ == "__main__":
    main()

import click
import subprocess
import sys
import pandas as pd


def run_cmd(cmd):
    """Runs a command in the shell and prints standard output and error"""
    sp = subprocess.Popen(cmd, shell=True)
    std_out, std_err = sp.communicate()
    print('std_err: ', std_err)
    print('std_out: ', std_out)


def get_models(modelsin):
    """
    Generate list of all model names
    Some models might have no hits in hmmout, but they should be displayed in count matrix
    """
    models = []
    with open(modelsin, "r") as f:
        models = [line.split()[-1].rstrip() for line in f if line.startswith("NAME")]
    return models


def get_markercompleteness(models, hmmout, query):
    """Returns a dictionary of marker count for each model"""
    count_dict = {x: 0 for x in models}
    seen = []
    with open(hmmout, "r") as f:
        lines = [line.rstrip() for line in f if not line.startswith("#")]
        for line in lines:
            if line.split()[0] not in seen:
                count_dict[line.split()[3]] += 1
                seen.append(line.split()[0])
    count_dict = {query: count_dict}
    return count_dict


@click.command()
@click.option('--queryfaa', '-q', type=click.Path(exists=True), help='Input query FASTA file')
@click.option('--modelscombined', '-m', type=click.Path(exists=True), help='Combined HMM model file')
@click.option('--hmmout', '-h', type=click.Path(), help='Output file for HMM hits')
@click.option('--target', '-t', type=str, help='Target database (UNI56 or UNIREF90)')
@click.option('--countout', '-c', type=click.Path(), help='Output file for marker gene counts')
def main(queryfaa, modelscombined, hmmout, target, countout):
    ### hmmsearch ###
    cutoff = "-E 1e-10"
    if target == "UNI56":
        cutoff = "--cut_ga"
    hmmsearch = ["hmmsearch" +
                 " --noali" +
                 " " + cutoff + " " +
                 " --domtblout " + hmmout +
                 " --cpu 4 " + modelscombined +
                 " " + queryfaa]
    run_cmd(hmmsearch)

    ### get counts ###
    query = queryfaa.split("/")[-1].split(".")[0]
    models = get_models(modelscombined)
    count_dict = get_markercompleteness(models, hmmout, query)
    pd.DataFrame.from_dict(count_dict).T.to_csv(countout, sep="\t")


if __name__ == '__main__':
    main()

import click
import subprocess
from Bio import SeqIO
from collections import defaultdict
import os


def run_cmd(cmd):
    sp = subprocess.Popen(cmd,
        #stderr=subprocess.PIPE,
        #stdout=subprocess.PIPE,
        shell=True)
    std_out, std_err = sp.communicate()
    print('std_err: ', std_err)
    print('std_out: ', std_out)


def run_blastp(queryfaa, refdb, blastpout):
    # no evalue set, what is the default?
    dblastp = ["diamond blastp --threads 8 "
               "--quiet --outfmt 6 "
               "--more-sensitive --query-cover 30 "
               "--subject-cover 30 "
               "-e 1e-10 "
               "-d {} -q {} -o {}".format(refdb, queryfaa, blastpout)]
    run_cmd(dblastp)


def parse_blastp(blastpout):
    queryids_hits_dict = defaultdict(list)
    seen = []
    # there could be multiple queries
    try:
        num_queries = len(set([line.split()[0] for line in open(blastpout)]))
        with open(blastpout, "r") as infile:
            for line in infile:
                queryname = line.split()[0]
                subjectname = line.split()[1]
                pident = float(line.split()[2])
                # Yield equal number of subject ids per query
                if subjectname not in seen and len(queryids_hits_dict[queryname]) <= 100/num_queries and pident!=100:
                    queryids_hits_dict[queryname].append(subjectname)
                    seen.append(subjectname)
    except:
        pass
    # single list with all 100 top subject ids
    besthits = [x for v in queryids_hits_dict.values() for x in v]
    return besthits


def reduce_ref(queryfaa, reffaa, reduced_reffaa_merged, refdb, blastpout):
    run_blastp(queryfaa, refdb, blastpout)
    besthits = parse_blastp(blastpout)
    seenids = []
    outseqs = []
    for seq_record in SeqIO.parse(queryfaa, "fasta"):
        if seq_record.id not in seenids:
            outseqs.append(seq_record)
            seenids.append(seq_record.id)
    for seq_record in SeqIO.parse(reffaa, "fasta"):
        if seq_record.description in besthits and seq_record.id not in seenids:
            outseqs.append(seq_record)
            seenids.append(seq_record.id)
    SeqIO.write(outseqs, reduced_reffaa_merged, "fasta")


@click.command()
@click.option('--queryfaa', '-q', required=True, help='Query sequences from hmmsearch for a single marker')
@click.option('--reffaa', '-r', required=True, help='Ref seqs for a single marker')
@click.option('--refdb', '-d', required=True, help='Diamond formatted db for a single marker')
@click.option('--blastpout', '-b', required=True, help='Output file name for diamond blastp')
@click.option('--reduced_reffaa_merged', '-o', required=True, help='Output file name for reduced reference sequences')
def main(queryfaa, reffaa, refdb, blastpout, reduced_reffaa_merged):
    if os.path.getsize(queryfaa) > 0:
        reduce_ref(queryfaa, reffaa, reduced_reffaa_merged, refdb, blastpout)
    else: 
        print (f"{queryfaa} is empty, no hits, omitting blast")

if __name__ == '__main__':
    main()

import subprocess
import click


def run_cmd(cmd):
    sp = subprocess.Popen(cmd,
        stderr=subprocess.PIPE,
        stdout=subprocess.PIPE,
        shell=True)
    (std_out, std_err) = sp.communicate()
    print('std_err: ', std_err)
    print('std_out: ', std_out)


@click.command()
@click.option("--queryfaamerged", "-q", required=True, type=click.Path(exists=True), help="Input fasta file containing merged fasta sequences")
@click.option("--alnout", "-a", required=True, type=click.Path(), help="Output fasta file for aligned sequences")
@click.option("--trimout", "-t", required=True, type=click.Path(), help="Output fasta file for trimmed sequences")
def main(queryfaamerged, alnout, trimout):
    mafftaln = f"mafft --quiet --thread 4 {queryfaamerged} > {alnout}"
    trimaln = f"trimal -gt 0.1 -in {alnout} -out {trimout}"
    for cmd in [mafftaln, trimaln]:
        run_cmd(cmd)

if __name__ == "__main__":
    main()

import click
import subprocess


def run_cmd(cmd):
    sp = subprocess.Popen(cmd,
        stderr=subprocess.PIPE,
        stdout=subprocess.PIPE,
        shell=True)
    (std_out, std_err) = sp.communicate()
    print('std_err: ', std_err)
    print('std_out: ', std_out)


@click.command()
@click.option("--aln", "-a", type=click.Path(exists=True), help='Alignment file in FASTA format')
@click.option("--treeout", "-t",  type=click.Path(), help='Output file for the tree')
def main(aln, treeout):
    runtree = f"fasttree -lg < {aln} > {treeout}"
    run_cmd(runtree)


if __name__ == "__main__":
    main()

import glob
import click
from ete3 import Tree
from collections import defaultdict
from collections import Counter
import pandas as pd


def flattenkeys_to_list(dict1, labels_dict, taxlevel):
    if taxlevel=="order":
        level = 3
    elif taxlevel=="family":
        level = 2
    elif taxlevel=="genus":
        level = 1
    dictlist = [  dict1[model] for model in dict1.keys() ]
    flattened = list({k: v for d in dictlist for k, v in d.items()}.keys())
    return [labels_dict[tax.split("|")[0]][level] for tax in flattened]


def unpack_family(node, leaves, query, seen):
    """
    get all leaves under a parental node
    """
    children = node.get_children()
    for child in children:
        childname = str(child).replace("--","").replace("/-", "").replace("\-", "").replace("\n","")
        if child.is_leaf() and childname.split("|")[0] != query:
                seen.append(childname)
                leaves.append(childname)
        elif not child.is_leaf() and childname != query and childname not in seen:
            seen.append(childname)
            unpack_family(child, leaves, query, seen) 
        elif child.is_leaf() and childname == query:
            seen.append(childname)
        else:
            pass
    if len(leaves)==0:
        parentnode = node.up
        unpack_family(parentnode, leaves, query, seen) 
    return leaves


@click.command()
@click.option('-q', '--query', required=True, help='Query to search in the tree.')
@click.option('-t', '--tree_outpath', required=True, help='Path to the tree directory.')
@click.option('-o', '--outname', required=True, help='Output file name.')
@click.option('-l', '--ncldv_labels', required=True, help='Path to the NCBI taxonomy labels file.')
def main(query, tree_outpath, outname, ncldv_labels):

    def get_closestrelative(query, leaves, node, tree):
        distance = 10
        closestrelative = "nd"
        for child in leaves:
            if child != query:
                newdist = t.get_distance(child, query, topology_only = False)
                if newdist < distance:
                    distance = newdist
                    closestrelative = child
                else:
                    pass
        return closestrelative, distance


    def get_neighbor(t, modelname, query):
        tree_dict_paralogs = {} # include nearest neighbor for all paralogs
        for node in t.traverse():
            if node.name.split("|")[0] == query:
                queryintree = node.name
                print (queryintree)
                # move up one node in the tree
                # collect all terminal leaves under the parent node
                leaves = []
                seen = []
                leaves = unpack_family(node, leaves, query, seen)
                print (leaves)
                closestrelative, distance = get_closestrelative(queryintree, leaves, node, t)
                print (closestrelative)
                tree_dict_paralogs[closestrelative + ";;;" + queryintree] = distance
        #closestrelativep = min(tree_dict_paralogs, key=tree_dict_paralogs.get) # get only ref that is most similar from all paralogs = shorted distance
        #return {closestrelativep : tree_dict_paralogs[closestrelativep]} # return only ref that is most similar
        return tree_dict_paralogs


    labels_dict = {}
    with open(ncldv_labels, "r") as infile:
        for line in infile:
            taxid = line.split("\t")[0]
            lineage = line.split("\t")[1].replace("\n","").split("|")
            labels_dict[taxid] = lineage

    # all distances including paralogs
    tree_dict_all = {}
    # dict with only single paralog that had shortest distance to reference
    tree_dict_all_best = {}
    for ftree in glob.glob(tree_outpath + "/*.FTWAG"):
        try:
            t = Tree(ftree)
            modelname = ftree.split("/")[-1].split(".")[0]
            print (f"{modelname}")
            tree_dict_all[modelname] = get_neighbor(t, modelname, query)
            print (f"{tree_dict_all}")
            tree_dict_all_best[modelname] = {min(tree_dict_all[modelname], key=tree_dict_all[modelname].get) : tree_dict_all[modelname][min(tree_dict_all[modelname], key=tree_dict_all[modelname].get)]}
        except:
            print (f"Something wrong with {ftree}")
            pass

    # this is the summary result, show number of best hits per tax string
    # cutoff for order: tree dist 2, aln 30 idperc --> needs further evaluation

    print (f"{tree_dict_all_best}@@@@@@@@{query}")
    final_tree_dict = Counter(flattenkeys_to_list(tree_dict_all_best, labels_dict, "order"))
    tree_result = "|".join([":".join([str(x) for x in pair]) for pair in list(final_tree_dict.items())])

    with open(outname, "w") as outfile:
        if len(tree_result) > 0:
            outfile.write(f"{query}\t{tree_result}\ttree\n")
        else:
            outfile.write(f"{query}\tno_hits\ttree\n")
        for model, value in tree_dict_all.items():
            # print lowest distance first
            value = dict(sorted(value.items(), key=lambda item: item[1], reverse=False))
            #outfile.write("\n".join(f"{model}\t{k.split(';;;')[1]}\t{k.split(';;;')[0]}\t{'|'.join(labels_dict[k.split('|')[0]])}\t{v}" for k,v in value.items()) + "\n")
            outfile.write("\n".join([
                f"{model}\t{k.split(';;;')[1]}\t{k.split(';;;')[0]}\t{'|'.join(labels_dict[k.split('|')[0]])}\t{v}"
                for k, v in value.items()]) + "\n")


if __name__ == "__main__":
    main()