import glob
import sys
from Bio import SeqIO
from collections import defaultdict

alndir = sys.argv[1] # dir with aligned multifa
aln_concat = sys.argv[2] # name of concat aln

def get_aln_len(aln):
    for seq_record in SeqIO.parse(aln, "fasta"):
        return len(seq_record.seq)

def get_taxa(alndir):
    taxa = set()
    for aln in glob.glob(alndir + "/*.*"):
        for seq_record in SeqIO.parse(aln, "fasta"):
            taxa.add(seq_record.id.split("|")[0])
    return list(taxa)

def get_aln(aln, taxa, final_dict, aln_len):
    for taxon in taxa:
        seq = "?" * aln_len
        for seq_record in SeqIO.parse(aln, "fasta"):
            if taxon == seq_record.id.split("|")[0]:
                seq = str(seq_record.seq)
                break
        final_dict[taxon].append(seq)
    return final_dict

final_dict = defaultdict(list)
for aln in glob.glob(alndir + "/*"):
    try:
        aln_len = get_aln_len(aln)
        final_dict = get_aln(aln, get_taxa(alndir), final_dict, aln_len)
    except:
        if aln.split(".")[-1] not in ["faa", "fa", "fasta", "fna", "mafft", "mafft01", "aln", "afa"]:
            print("Check content of dir, all need to be aligned multifasta with typical suffix: fna, fasta, fa, faa...")
        else:
            print("Something else is wrong")

# create concatenated aln
concat_dict = {x:"".join(y) for x,y in final_dict.items()}
with open(aln_concat, "w") as outfile:
    for seqid, seq in concat_dict.items():
        outfile.write(">" + seqid + "\n" + seq + "\n")

import sys
from ete3 import Tree
import ete3 as ete
import os.path
import warnings
warnings.filterwarnings("ignore", ".*is with a literal.*", SyntaxWarning)

treein = sys.argv[1] # tree in newick format
queryids_or_pattern = sys.argv[2] # file with leaf names or pattern to look for leaves
hexcolor =  sys.argv[3] # color for identified branches
itolout = sys.argv[4] # name of outfile

def get_alltaxa(t, queryids_or_pattern):
    target = []
    if os.path.isfile(queryids_or_pattern):
        with open(queryids_or_pattern, "r") as infile:
            target = [line.strip() for line in infile]
    else:
        for node in t.traverse():
            if node.is_leaf() and node.name.startswith(queryids_or_pattern):
                target.append(node.name)
    return target


def unpack_family(parentnode, leaves, target):
    """
    get all leaves under a parental node
    """
    children = parentnode.get_children()
    for child in children:
        if child.is_leaf() and child.name not in leaves:
            leaves.append(child.name)
        else:
            unpack_family(child, leaves, target) 
    return leaves


def check_monophyly(node, target, mono_leaves):
    # move up one node in the tree
    parentnode = node.up
    # collect all terminal leaves under the parent node
    leaves = unpack_family(parentnode, [], target)
    if len(leaves) == len([x for x in leaves if x in target]):
        mono_leaves.extend(leaves)
        check_monophyly(parentnode, target, mono_leaves)
    return mono_leaves

def main():
    t = Tree(treein)
    target = get_alltaxa(t, queryids_or_pattern)
    mono_leaves_all = []
    for node in t.traverse():
        if node.name in target and node.name not in [x for sublist in mono_leaves_all for x in sublist]:
            try:
                mono_leaves_all.append(check_monophyly(node, target, []))
            except:
                mono_leaves_all.append([node.name])
    #write to output
    queriesinclades = [x for sublist in mono_leaves_all for x in sublist]
    singletons = [x for x in target if x not in queriesinclades]
    with open(itolout, "w") as outfile:
        for monotclade in mono_leaves_all:
            if len(monotclade) > 1:
                outfile.write("|".join(list(set(monotclade))) + ",clade,#"+hexcolor+",normal,2\n")
        for singleton in singletons:
            outfile.write(singleton + ",branch,#"+hexcolor+",normal,2\n")

if __name__ == '__main__':
    main()

import glob
import sys
from ete3 import Tree
from collections import defaultdict
from collections import Counter
import pandas as pd
import warnings
warnings.filterwarnings("ignore", ".*is with a literal.*", SyntaxWarning)

def main():
    tree_in = sys.argv[1]
    queries_in = sys.argv[2]
    outname = sys.argv[3]

    def unpack_family(node, leaves, query, seen, queries):
        """
        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 not in queries:
                    seen.append(childname)
                    leaves.append(childname)
            elif not child.is_leaf() and childname not in seen:
                seen.append(childname)
                unpack_family(child, leaves, query, seen, queries)
            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, queries)
        return leaves


    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, query, queries):
        for node in t.traverse():
            if node.name == query:
                child = str(node.name).replace("--","").replace("/-", "").replace("\-", "").replace("\n","")
                queryintree = node.name
                # move up one node in the tree
                # collect all terminal leaves under the parent node
                leaves = []
                seen = []
                leaves = unpack_family(node, leaves, query, seen, queries)
                closestrelative, distance = get_closestrelative(queryintree, leaves, node, t)
                return [closestrelative, distance]

    # get queries
    with open(queries_in) as infile:
        queries = infile.read().splitlines()

    # all distances
    tree_dict_all = {}
    t = Tree(tree_in)
    for query in queries:
        tree_dict_all[query] =  get_neighbor(t, query, queries)
    df = pd.DataFrame.from_dict(tree_dict_all).fillna("nd").T
    with open(outname + ".pairs", "w") as outfile:
        for query, bestref in tree_dict_all.items():
            if bestref:
                outfile.write(f"{query}\t{bestref[0]}\t{bestref[1]}\n")
            else:
                outfile.write(f"{query}\tnd\tnd\n")


if __name__ == "__main__":
    main()

from scipy.cluster import hierarchy
import numpy as np
import subprocess
import sys
import pandas as pd
import glob
from collections import defaultdict
import os

os.environ['QT_QPA_PLATFORM'] = 'offscreen'

queryfaadir = sys.argv[1]
modelscombined = sys.argv[2] # e.g. pathtofiles + "models/GVOGuni9.hmm"
hmmout = sys.argv[3] # e.g. pathtofiles + "Fadolivirus_GVOGuni9.out"
countout = sys.argv[4] # summarized hit counts per model from hmmout
removedtaxa = sys.argv[5] # list of removed taxa with info on why removed these were removed
minmarker = float(sys.argv[6]) # percentage of markers to be present for a genome to be included in downstream analysis$
maxsdup = int(sys.argv[7]) # max copy number of any marker, if exceeded respective genome will be excluded
maxdupl = float(sys.argv[8]) # percentage of markers that can be present in multiple copies
outitol = sys.argv[9]


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
    """
    with open(modelsin, "r") as f:
        models = [line.split()[-1].rstrip() for line in f if line.startswith("NAME")]
    return models


def get_taxa (queryfaadir):
    taxa = [taxon.split("/")[-1].split(".")[0] for taxon in glob.glob(queryfaadir + "/*.faa")]
    return taxa


def get_markercompleteness (models, hmmout, query):
    """
    Get copy numbers for each marker
    """
    # add 0s to include models that are not in hmmout
    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 and line.split("|")[0]==query:
                count_dict[line.split()[3]] += 1
                seen.append(line.split()[0])
    return count_dict


def filter_taxa(df, models, minmarker, maxsdup, maxdupl):
    """ Quality filtering"""
    filteredtaxa_dict = defaultdict(list)
    df['max'] = df.max(axis=1)
    filtered_taxa = df[df['max'] > maxsdup].index
    filteredtaxa_dict.update({taxon: ["maxsdup:{}".format(df.at[taxon, 'max'])] for taxon in filtered_taxa})
    df.drop(['max'], inplace=True, axis=1)
    df['duplications'] = (df > 1).sum(axis=1)
    filtered_taxa = df[df['duplications'] / len(models) > maxdupl].index
    filteredtaxa_dict.update({taxon: ["maxdupl:{:.4f}".format(df.at[taxon, 'duplications'] / len(models))] for taxon in filtered_taxa})
    df.drop(['duplications'], inplace=True, axis=1)
    df[df > 1] = 1
    df['sum'] = df.sum(axis=1)
    filtered_taxa = df[df['sum'] < len(models) * minmarker].index
    filteredtaxa_dict.update({taxon: ["completeness:{:.4f}".format(df.at[taxon, 'sum'] / len(models))] for taxon in filtered_taxa})
    return filteredtaxa_dict


def convert2itol(flabels, outfile):
    with open(outfile, "w") as outfile:
        outfile.write("DATASET_HEATMAP\n"
                        "SEPARATOR COMMA\n"
                        "DATASET_LABEL,Count\n"
                        "COLOR,#ff0000\n"
                        "COLOR_MIN,#ff0000\n"
                        "COLOR_MAX,#0000ff\n"
                        "FIELD_LABELS," + ",".join(flabels) + "\n"
                        "DATA\n"
                       )


def reorder_cols(df):
    """Reorder columns in countmatrix using hierarchical clustering"""
    Z = hierarchy.linkage(df.T, method='ward')
    reordered_ind = hierarchy.leaves_list(Z)
    df = df[df.columns[reordered_ind]]
    return df


def main():
    models = get_models(modelscombined)
    taxa = get_taxa(queryfaadir)

    count_dict = {query: get_markercompleteness(models, hmmout, query) for query in taxa}
    df = pd.DataFrame.from_dict(count_dict).T
    df.to_csv(countout, sep='\t', index=True)

    filteredtaxa_dict = filter_taxa(df, models, minmarker, maxsdup, maxdupl)
    with open(removedtaxa, 'w') as outfile:
        for taxon, reasons in filteredtaxa_dict.items():
            outfile.write(f'{taxon}\t{";".join(reasons)}\n')

    df.drop('sum', axis=1, inplace=True)
    reordered_df = reorder_cols(df)
    convert2itol(list(df.columns), outitol)
    with open(outitol, "a") as outitolA:
        df.to_csv(outitolA, header = False, sep = ",")

if __name__ == '__main__':
    main()

shell:
  """
  (echo "hmmsearch --noali {params.hmmsearch_cutoff} --domtblout  {output.hmmsearchout} --cpu {params.hmmsearch_cpu} {input.modelshmm} {input.faacombined}"
  hmmsearch --noali {params.hmmsearch_cutoff} --domtblout  {output.hmmsearchout} --cpu {params.hmmsearch_cpu} {input.modelshmm} {input.faacombined}) &> {log}
  """

shell:
  """
  (echo "mafft{params.mafftv} {params.mafft} --quiet {params.mafft_thread} {input.hitsfaa} > {output.aln}"
  mafft{params.mafftv} {params.mafft} --quiet {params.mafft_thread} {input.hitsfaa} > {output.aln} || touch {output.aln}
  echo "trimal {params.trimal_gt} -in {output.aln} -out {output.trimmedaln}"
  trimal {params.trimal_gt} -in {output.aln} -out {output.trimmedaln} || touch {output.trimmedaln} ) >& {log}
  """

shell:
  """
  (echo "fasttree {params.proteintrees} {input.trimmedaln} > {output.tree}"
  mkdir -p {params.outdirt}
  fasttree {params.proteintrees} {input.trimmedaln} > {params.outpath}.treefile || touch {output.tree}
  cp {params.outpath}.treefile {output.tree}) >& {log}
  """

shell:
  """
  (echo "iqtree --prefix {params.outpath} -m {params.proteintrees} -T 1 -fast -s {input.trimmedaln}"
  mkdir -p {params.outdirt}
  iqtree --prefix {params.outpath} -m {params.proteintrees} -T 1 -fast -s {input.trimmedaln} || touch {params.outpath}.treefile
  cp {params.outpath}.treefile {output.tree}
  touch {output.tree}) >& {log}
  """