# More than 80% of the code here was copied from coretracker/coreutils/Faces.py

# which was written for displaying a tree with a list of sequences

import argparse

import math

from PyQt5.QtCore import Qt, QPointF

from PyQt5.QtWidgets import (

QGraphicsRectItem, QGraphicsLineItem, QGraphicsSimpleTextItem, QGraphicsTextItem)

from PyQt5.QtGui import (QPen, QLinearGradient,

QColor, QBrush, QFont)

from ete3 import faces, Tree, TreeStyle, NodeStyle

from ete3.treeview.main import COLOR_SCHEMES

from collections import OrderedDict, defaultdict as ddict

import os

import sys

import re

from Bio import AlignIO, SeqIO

from Bio.Data import CodonTable

import logging

from ete3 import TextFace, AttrFace

# change this to another color if you want to set how node with at least one mark are displayed

MARKED_NODE_COLOR = "red"

RES_COL_WIDTH = 14

_aafgcolors = {

}

_aabgcolors = {

}

_ntfgcolors = {

}

_ntbgcolors = {

}

def _get_codon_fgcolors(codontable):

"""Get colon foreground color"""

return dict((k, _aafgcolors[v]) for (k, v) in codontable.items())

def _get_codon_bgcolors(codontable, spec_codon_col):

"""Get colon background color"""

return dict((k, spec_codon_col.get(k, _aabgcolors.get(v, "#FFFFFF"))) for (k, v) in codontable.items())

class _LineItem(QGraphicsLineItem):

self.setPen(QPen(QColor('#000000')))

class LineFace(faces.Face):

return self.height

class _LineItem(QGraphicsLineItem):

self.setPen(QPen(QColor('#000000')))

class SequenceFace(faces.StaticItemFace):

return '<span style="color:%s;">%s</span>'%(fgcolor, html)

class List90Face(faces.StaticItemFace):

self.width = seq_width

def format_tree(tree, alignment, al_len_dict, edpos, codontable={}, colors=None, codon_col={}, text="C-to-U RNA editing", ic_contents=[]):

return t, ts

def display_tree(tree, aligndict, al_len_dict, eddict, outfile, gcode, colormap, dpi=1200, width=500, scale=100):

t.render(outfile, dpi=dpi, tree_style=ts, w=width)

def get_ed_pos(edpath):

return ed

if __name__ == '__main__':

im_choices = ("svg", "pdf", "png", "html")

parser = argparse.ArgumentParser(

"EdAlignViz", description='Display a codon-by-codon alignment with a phylogenetic tree, and show the list of edited position')

parser.add_argument('-t', '--tree', required=True,

dest='tree', help="Path to the phylogenetic tree")

parser.add_argument('-a', '--align', nargs="+", required=True,

dest='align', help="List of sequence alignments")

parser.add_argument('-m', '--mark_ext', nargs="?", const="truth", dest='mark_ext',

help="The extension of the file in which position to mark are saved. If it is provided without argument, a '.truth' extension is used. Note that the basename of the file should match the one of the respective alignment.")

parser.add_argument('-o', '--output', default="output", dest='output',

help="Path to the output file where to save the plot")

parser.add_argument('--fmt', choices=im_choices,

default="svg", help="Image output format")

parser.add_argument('--dpi', type=int, default=300, help="DPI for image quality")

parser.add_argument('-w','--width', type=float, default=10000, help="Width of the image")

parser.add_argument('--scale', type=int, default=100, help="Scale of the phylogenetic tree (zoom on X axis to increase branch length representation.)")

parser.add_argument('--gcode', default=1, type=int,

help="Genetic code table to use. Only accept Standard table (NCBI) as for now")

parser.add_argument('-v', '--verbose', dest="verbose",

action="store_true", help="Whether debugging text should be printed")

parser.add_argument(

'--colormap', help="Name of a file that map a codon to a specific hex color used as background. Space separated values")

args = parser.parse_args()

colormap = {}

if args.colormap:

with open(args.colormap) as COLOR:

for line in COLOR:

line = line.strip().split()

if line:

colormap[line[0]] = line[1]

out = args.output.rsplit(".", 1)

if len(out) == 1:

out = out[0] + "." + args.fmt

elif len(out) == 2 and out[-1].lower() not in im_choices:

out = out[0] + ".svg"

else:

out = out[0] + "."+out[-1].lower()

s = args.tree

try:

tree = Tree(s)

except Exception as e:

import unicodedata

with open(args.tree) as T:

s = T.read().strip().replace("'", '')

s = re.sub(r'[^\x00-\x7F]+',' ', s)

tree = Tree(s)

al_per_gene = {}

eddict = {}

al_len = {}

gcode = CodonTable.unambiguous_dna_by_id[args.gcode]

gcode = gcode.forward_table

spec_and_descr = {}

for al in args.align:

if os.path.isfile(al):

try:

gname, ext = al.split('.', 1)

gene = os.path.basename(gname)

cur_msa = AlignIO.read(al, "fasta")

# add alignment length

al_len[gene] = cur_msa.get_alignment_length()

spec_and_descr.update(dict((x.id, x.description)

for x in cur_msa))

# add sequences for each gene

al_per_gene[gene] = dict((x.id, str(x.seq)) for x in cur_msa)

# get edited positions in each gene

if args.mark_ext:

ed = get_ed_pos(gname+"."+args.mark_ext.lstrip(","))

if ed:

eddict[gene] = ed

except Exception as e:

print(e)

logging.warning(

"File %s cannot be read as fasta alignment. Skipping the file" % al)

# now try to reset the leaves name in the tree

common_name = []

for leaf in tree:

new_id = [x for x, y in spec_and_descr.items() if leaf.name in y]

if new_id:

common_name.append(new_id[0])

leaf.add_feature("fullname", leaf.name)

leaf.name = new_id[0]

tree.prune(common_name)

#print(tree.get_ascii(attributes=["name", "fullname"]))

if len(tree) < len(spec_and_descr):

logging.warning("The following genomes are missing in the tree:")

for s in set.symmetric_difference(set(tree.get_leaf_names()), spec_and_descr.keys()):

logging.warning(spec_and_descr[s])

display_tree(tree, al_per_gene, al_len, eddict, out, gcode, colormap, dpi=args.dpi, width=args.width, scale=args.scale)