# A virtual X-server XVFB is used in case you do not have X-server
from xvfbwrapper import Xvfb

vdisplay = Xvfb()
vdisplay.start()

# launch stuff inside virtual display here

#vdisplay.stop()


from ete3 import Tree, faces, TreeStyle, NodeStyle
from ete3 import ClusterTree, RectFace, AttrFace, ProfileFace, TextFace
from ete3.treeview.faces import add_face_to_node
import pandas as pd
import numpy as np
import colorsys

t = Tree()
t.populate(7,names_library=['A','B','C','D','E','F','G'])
## %%liline is used for showing plots in ipythonnotebook.
## t.render(file_name="tree.pdf") # will save tree into pdf file
t.render(file_name="%%inline")

t_str = t.write(outfile=None, format=0)
t_str

'(((B:1,A:1)1:1,(G:1,F:1)1:1)1:1,(E:1,(D:1,C:1)1:1)1:1);'

t = Tree(t_str)
ts = TreeStyle()
ts.show_leaf_name = True
ts.show_branch_length = True
ts.show_branch_support = True
t.render(file_name="%%inline", tree_style=ts)

t.write(outfile="tree.nw", format=0)
t = Tree("tree.nw")
ts.mode = "c"
ts.arc_start = -180 # 0 degrees = 3 o'clock
ts.arc_span = 180
t.render(file_name="%%inline", w=500, tree_style=ts)

ts = TreeStyle()
ts.show_leaf_name = True
ts.show_branch_length = True
ts.show_branch_support = True

# Draws nodes as small red spheres of diameter equal to 10 pixels
for n in t.traverse():  # Traverse each node and set attribute for each type of nodes
	if n.is_leaf(): # Decide if leaf node
		nstyle = NodeStyle()
		nstyle["shape"] = "sphere"
		nstyle["size"] = 10
		nstyle["fgcolor"] = "darkred"
		n.set_style(nstyle)
	else:
		nstyle = NodeStyle()
		nstyle["shape"] = "square"
		nstyle["size"] = 15
		nstyle["fgcolor"] = "orange"
		n.set_style(nstyle)        

t.img_style["size"] = 30
t.img_style["fgcolor"] = "blue"
	
t.render(file_name="%%inline", w=500, tree_style=ts)

t = Tree(t_str)

nameMap = {'A': 'American', 'B': 'Britain', 'C':'China', 
		   'D':'Dutch', 'E':'Egypt','F':'France','G':'German'}

ts = TreeStyle()
ts.show_leaf_name = False
ts.show_branch_length = True
ts.show_branch_support = True

# Creates my own layout function. I will use all previously created
# faces and will set different node styles depending on the type of
# node.
def mylayout(node):   
	# If node is a leaf, add the nodes name and a its scientific
	# name
	if node.is_leaf():
		# We can also create faces on the fly
		newName = nameMap.get(node.name, node.name)
		newNameFace = faces.TextFace(newName)
		faces.add_face_to_node(newNameFace, node, column=1, aligned=True)

		# Sets the style of leaf nodes
		node.img_style["size"] = 12
		node.img_style["shape"] = "sphere"
		node.img_style["fgcolor"] = "blue"
	#If node is an internal node
	else:
		# Sets the style of internal nodes
		node.img_style["size"] = 1
		node.img_style["shape"] = "circle"
		node.img_style["fgcolor"] = "darkred"

ts.layout_fn = mylayout

t.img_style["size"] = 30
t.img_style["fgcolor"] = "black"
	
t.render(file_name="%%inline", w=600, tree_style=ts)

t = Tree(t_str)

nameMap = {'A': 'American', 'B': 'Britain', 'C':'China', 
		   'D':'Dutch', 'E':'Egypt','F':'France','G':'German'}

colorMap = { 'American': '#ACFFFF',
			 'Britain': '#ACACFF',
			 'China': '#ACACAC',
			 'Dutch': '#59ACAC',
			 'Egypt': '#5959AC',
			 'France': '#595959',
			 'German': '#065959'}

ts = TreeStyle()
ts.show_leaf_name = False
ts.show_branch_length = True
ts.show_branch_support = True

# Creates my own layout function. I will use all previously created
# faces and will set different node styles depending on the type of
# node.
def mylayout(node):   
	#Change branch color
	node.img_style["hz_line_color"] = 'orange'  # change horizontal branch color 
	node.img_style["vt_line_color"] = 'red' # Change vertical branch color
	
	# If node is a leaf, add the nodes name and a its scientific
	# name
	if node.is_leaf():
		# We can also create faces on the fly
		newName = nameMap.get(node.name, node.name)
		newNameFace = faces.TextFace(newName)
		faces.add_face_to_node(newNameFace, node, column=1, aligned=True)

		# Sets the style of leaf nodes
		node.img_style["size"] = 12
		node.img_style["shape"] = "sphere"
		node.img_style["fgcolor"] = "blue"
		node.img_style["bgcolor"] = colorMap[newName]
		
		node.img_style["hz_line_color"] = 'blue'  # change branch color 
	#If node is an internal node
	else:
		# Sets the style of internal nodes
		node.img_style["size"] = 1
		node.img_style["shape"] = "circle"
		node.img_style["fgcolor"] = "darkred"

ts.layout_fn = mylayout

t.img_style["size"] = 30
t.img_style["fgcolor"] = "black"
	
t.render(file_name="%%inline",tree_style=ts)

nameFace = AttrFace("name", fsize=12) #Set leaf node attribute

def setup_heatmap(tree, tree_style, header, center_value=0.0, nameMap ={}, nameLabel = '',
				  color_up=0.7, color_down=0.2, color_center="white"):
	DEFAULT_COLOR_SATURATION = 0.5
	BASE_LIGHTNESS = 0.7
	def gradient_color(value, max_value, saturation=0.5, hue=0.1):    
		def rgb2hex(rgb):
			return '#%02x%02x%02x' % rgb
		def hls2hex(h, l, s):
			return rgb2hex( tuple(map(lambda x: int(x*255), 
						  colorsys.hls_to_rgb(h, l, s))))
	
		lightness = 1 - (value * BASE_LIGHTNESS) / max_value
		return hls2hex(hue, lightness, DEFAULT_COLOR_SATURATION)


	# Calculate max gradient value from the ClusterTree matrix
	maxv = abs(center_value - tree.arraytable._matrix_max)
	minv = abs(center_value - tree.arraytable._matrix_min)
	if center_value <= tree.arraytable._matrix_min:
		MAX_VALUE = minv + maxv
	else:
		MAX_VALUE = max(maxv, minv)
		
	# Add heatmap colors to tree
	cols_add_before_heat = 0
	if nameMap:
		cols_add_before_heat = 1
	for lf in tree:
		if nameMap:
			longNameFace = faces.TextFace(nameMap.get(lf.name, lf.name))
			lf.add_face(longNameFace, column=0, position="aligned")
			
		for i, value in enumerate(getattr(lf, "profile", [])):
			if value > center_value:
				color = gradient_color(abs(center_value - value), MAX_VALUE, 
									   hue=color_up)
			elif value < center_value:
				color = gradient_color(abs(center_value - value), MAX_VALUE, 
									   hue=color_down)
			else:
				color = color_center
			lf.add_face(RectFace(20, 20, color, color), position="aligned", 
						column=i+cols_add_before_heat)
			# Uncomment to add numeric values to the matrix
			#lf.add_face(TextFace("%0.2f "%value, fsize=5), position="aligned", column=i)
		lf.add_face(nameFace, column=i+cols_add_before_heat+1, position="aligned")
		
	if nameMap and nameLabel:
		nameF = TextFace(nameLabel, fsize=7)
		#nameF.rotation = -90
		tree_style.aligned_header.add_face(nameF, column=0)
	# Add header 
	for i, name in enumerate(header):
		nameF = TextFace(name, fsize=7)
		nameF.rotation = -90
		tree_style.aligned_header.add_face(nameF, column=i+cols_add_before_heat)
 #-------------END setup_heatmap----------------------------------------------        

data = pd.read_table("matrix", header=0, index_col=0)
data.index.name = "#Names"  #修改第一行的名字使其符合ETE的要求
data_mat = data.to_csv(None, sep="\t", float_format="%.2f")
header = list(data.columns.values)  #获取列的名字用于标记

data

data_mat

t = ClusterTree(t_str, data_mat)

ts = TreeStyle()
ts.show_leaf_name = False
ts.show_branch_length = True
ts.show_branch_support = True

setup_heatmap(t, ts, header, center_value=0, color_up=0.9, color_down=0.3, color_center="white")

t.render(file_name="%%inline", tree_style=ts)

def mylayout_only(node):   
	#Change branch color
	node.img_style["hz_line_color"] = 'orange'  # change horizontal branch color 
	node.img_style["vt_line_color"] = 'red' # Change vertical branch color
	
	# If node is a leaf, add the nodes name and a its scientific
	# name
	if node.is_leaf():
		newName = nameMap.get(node.name)
		node.img_style["size"] = 12
		node.img_style["shape"] = "sphere"
		node.img_style["fgcolor"] = "blue"
		node.img_style["bgcolor"] = colorMap[newName]        
		node.img_style["hz_line_color"] = 'blue'  # change branch color 
	#If node is an internal node
	else:
		# Sets the style of internal nodes
		node.img_style["size"] = 1
		node.img_style["shape"] = "circle"
		node.img_style["fgcolor"] = "darkred"



t = ClusterTree(t_str, data_mat)

ts = TreeStyle()
ts.show_leaf_name = False
ts.show_branch_length = True
ts.show_branch_support = True
ts.layout_fn = mylayout_only

setup_heatmap(t, ts, header, center_value=0, color_up=0.9, color_down=0.3, 
			  color_center="white", nameMap=nameMap, nameLabel="Full")

t.render(file_name="%%inline", tree_style=ts)

nw = """
(((Dre:0.008339,Dme:0.300613)1.000000:0.596401,
(Cfa:0.640858,Hsa:0.753230)1.000000:0.182035)1.000000:0.106234,
((Dre:0.271621,Cfa:0.046042)1.000000:0.953250,
(Hsa:0.061813,Mms:0.110769)1.000000:0.204419)1.000000:0.973467);

vdisplay.stop()