import os, sys, re
from Bio import SeqIO


path = './'
aa_files = ['YR_tree/all_YR.fas',
            'comparable_trees/RT.fas',
            'comparable_trees/RT_outgrouped.fas',
            'comparable_trees/YR.fas',
            'comparable_trees/YR_outgrouped.fas']


for aa_file in aa_files:
    corrected_filename = aa_file.split('.')[0] + '_corrected.fas'
    corrected_file = open(path + corrected_filename,'wt')
    records = SeqIO.parse(path + aa_file, 'fasta')
    for record in records:
        record.description = re.sub('\:','_',record.description)
        record.description = re.sub('\;','_',record.description)
        record.description = re.sub('\|','_',record.description)
        record.description = re.sub('\*','_',record.description)
        record.description = re.sub('\-','_',record.description)
        record.description = re.sub('\,','_',record.description)
        record.description = re.sub('\/','_',record.description)
        record.description = re.sub(' ','_',record.description)
        record.description = re.sub('\.','_',record.description)
        corrected_file.write('>' + record.description.split(' ')[0] + '\n')
        sequence = re.sub('\*','X',str(record.seq))
        sequence = re.sub('J','X',sequence)
        sequence = re.sub('-','',sequence)
        corrected_file.write(sequence + '\n')
    corrected_file.close()

from Bio import SeqIO
minimum_length = 90
path = './'
trimmed_records = []

records = list(SeqIO.parse(path + 'YR_tree/all_YR_corrected.fas', 'fasta'))

for record in records:
    res_count = 0
    sequence = str(record.seq)
    for res in sequence:
        if not res == '-':
            res_count += 1
            
    if res_count >= minimum_length:
        trimmed_records.append(record)

SeqIO.write(trimmed_records, path + 'YR_tree/all_YR_corrected.fas', 'fasta')

from ete2 import Tree, NodeStyle, TreeStyle, TextFace, PhyloTree, CircleFace, PieChartFace
import sys, numpy
path = './'

def blen_cutoff(node,brlns_cutoff_factor):
    blens = []
    for br in node.traverse():
        if not br.name == node.name and not br.dist == 0:
            blens.append(float(br.dist))
    return numpy.median(blens)*brlns_cutoff_factor


tree_string_files = [
                     'YR_tree/all_YR_corrected_reviewed.tre',
                     #'comparable_trees/RT_corrected_reviewed.tre',
                     'comparable_trees/RT_outgrouped_corrected_reviewed.tre',
                     #'comparable_trees/RTYR_corrected_reviewed.tre',
                     #'comparable_trees/YR_corrected_reviewed.tre',
                     'comparable_trees/YR_outgrouped_corrected_reviewed.tre'
                    ]
   


for tree_string_file in tree_string_files:
    support_cutoff = 0.01
    if tree_string_file ==  'YR_tree/all_YR_corrected_reviewed.tre':
        support_cutoff = 0.9 #higher cutoff for phylo-placing than for just phylo presenting
    elif tree_string_file ==  'comparable_trees/YR_outgrouped_corrected_reviewed.tre':
        support_cutoff = 0.7 #higher cutoff for phylo-placing than for just phylo presenting
    
        
    tree_string = open(path + tree_string_file, 'r').read()
    t = Tree(tree_string, format=1)
    print len(t)
    node_number = 1 #We are going to give running numbers to all the nodes for easier access later on
    
    for n in t.traverse():
        if n.is_leaf and 'rypton' in n.name:
            n.name = n.name + '_crypton' # This will make the annotation easier
        if n.support < support_cutoff:
            n.delete()                   # multifrucate unsupported bifrucarions
            
    print_collapsed_nodes = 1            # will make a table of the collaped nodes, named by running numbers, and their 
                                         # leaf content
    
    for n in t.traverse():
        if not n.is_leaf():
            organisms_in_node = {       # will be used to count taxa in collapsed nodes     
                    'PAT': 0,
                    'BAS': 0,
                    'RET': 0,
                    'Mflo': 0,
                    'Dviv': 0,
                    'Ebre': 0,
                    'Rcul': 0,
                    'Avit': 0,
                    'Asuu': 0,
                    'Bxyl': 0,
                    'Hbac': 0,
                    'Haor': 0,
                    'Hcon': 0,
                    'Srat': 0,
                    'briC': 0,
                    'Cele': 0,
                    'Cjap': 0,
                    'Bmal': 0,
                    'Cbre': 0,
                    'Crem': 0,
                    'Gpal': 0,
                    'Lsig': 0,
                    'Ppac': 0,
                    'Lloa': 0,
                    'Mhap': 0,
                    'Minc': 0,
                    'C5sp': 0,
                    'Dimm': 0,
                    'Ooch': 0,
                    'Ovol': 0,
                    'Otip': 0,
                    'Pred': 0,
                    'C11sp': 0,
                    'Cjap': 0,
                    'Cang': 0,
                    'Tmur': 0,
                    'Tspi': 0,
                    'Wban': 0,
                    'Acas': 0,
                    'Agam': 0,
                    'Acal': 0,
                    'Dpul': 0,
                    'Dsim': 0,
                    'Lgi': 0,
                    'Nvi': 0,
                    'Nve': 0,
                    'Ath': 0,
                    'Vcar': 0,
                    'Alyr': 0,
                    'Hduj': 0
                    }
            elements_in_node = {     #will be used to count element types in collapsed nodes
                    'dirs1_like': 0,
                    'pat1_like': 0,
                    'PAT': 0,
                    'viper_or_Ngaro': 0,
                    'kangaroo': 0,
                    'DIRS': 0,
                    'domains_inverted': 0,
                    'two_domains': 0,
                    'YR_only': 0,
                    'crypton': 0,
                    'DERIVED': 0,    
                    'GYPSY': 0,
                    'COPIA': 0,
                    'BEL': 0
                    }
            # node will be collapsed only if all the leafs can be placed in only one of these arrays
            dirs1_hirarchy = ('dirs1_like', 'DIRS', 'two_domains', 'YR_only','DERIVED')
            pat_like_hirarchy = ('pat1_like', 'DIRS', 'two_domains', 'YR_only','DERIVED')
            toc_hirarchy = ('PAT', 'DIRS', 'two_domains', 'YR_only','DERIVED')
            viper_or_Ngaro_hirarchy = ('viper_or_Ngaro', 'two_domains', 'YR_only','DERIVED')
            kangaroo_hirarchy = ('kangaroo', 'domains_inverted','two_domains', 'YR_only','DERIVED', 'DIRS')
            DIRS_hirarchy = ('DIRS', 'two_domains', 'YR_only','DERIVED')
            domains_inverted_hirarchy = ('domains_inverted','two_domains', 'YR_only','DERIVED')
            two_domains_hirarchy = ('two_domains','domains_inverted', 'YR_only','DERIVED')
            crypton_hirarchy = ('crypton', 'YR_only')
            YR_hirarchy = crypton_hirarchy 
            GYPSY_hirarchy = ('GYPSY')
            COPIA_hirarchy = ('COPIA')
            BEL_hirarchy = ('BEL')
        
        
            n.name = 'node' + str(node_number) #naming the node
            node_number += 1
        
            # prepare species counts and element counts to be added to all nodes
            for leaf in n:
                    for orgn in organisms_in_node.keys(): # orgn is the species key
                        search_length = len(orgn)
                        if leaf.name[:search_length] == orgn: # if begining of the leaf name same as orgn
                            organisms_in_node[orgn] += 1      # count it
                    for element in elements_in_node.keys():   
                        search_length = len(element)                    
                        if leaf.name[-1*search_length:] == element:  # if end of the leaf name same as element
                            elements_in_node[element] += 1           # count it 
            for orgn in organisms_in_node.keys(): # get read of species codes with zero count 
                if organisms_in_node[orgn] == 0:
                    organisms_in_node.pop(orgn, None) 
            for element in elements_in_node.keys(): # get read of elements with zero count
                if elements_in_node[element] == 0:
                    elements_in_node.pop(element, None)
            n.add_features(species=organisms_in_node, elements=elements_in_node) # add speceis count and element counts
                                                                                 # as node features
    # write tree with node names for future reference        
    t.write(outfile=path + tree_string_file.split('.')[0] + '_internal_nodes.tre', format=1, features=["support"])
    collapsed_nodes_file = open(path + tree_string_file.split('.')[0] + '_collapsed_nodes.txt','wt')
    collapsed_nodes = []
    

    # decside which node to collapse based on the presence absence of conflicting element forms
    for n in t.traverse():
        if not n.is_leaf():
            collapse = 1

            if 'dirs1_like' in n.elements.keys():
                for name in n.elements.keys():
                    if not name in dirs1_hirarchy:
                        collapse = 0
            elif 'pat1_like' in n.elements.keys():
                for name in n.elements.keys():
                    if not name in pat_like_hirarchy:
                        collapse = 0 
            elif 'PAT' in n.elements.keys():
                for name in n.elements.keys():
                    if not name in toc_hirarchy:
                        collapse = 0
            elif 'kangaroo' in n.elements.keys():
                for name in n.elements.keys():
                    if not name in kangaroo_hirarchy:
                        collapse = 0
            elif 'viper_or_Ngaro' in n.elements.keys():
                for name in n.elements.keys():
                    if not name in viper_or_Ngaro_hirarchy:
                        collapse = 0
            elif 'DIRS' in n.elements.keys():
                for name in n.elements.keys():
                    if not name in DIRS_hirarchy:
                        collapse = 0
            elif 'domains_inverted' in n.elements.keys():
                for name in n.elements.keys():
                    if not name in domains_inverted_hirarchy:
                        collapse = 0
            elif 'two_domains' in n.elements.keys():
                for name in n.elements.keys():
                    if not name in two_domains_hirarchy:
                        collapse = 0
            elif 'COPIA' in n.elements.keys():
                for name in n.elements.keys():
                    if not name in COPIA_hirarchy:
                        collapse = 0
            elif 'BEL' in n.elements.keys():
                for name in n.elements.keys():
                    if not name in BEL_hirarchy:
                        collapse = 0
            elif 'GYPSY' in n.elements.keys():
                for name in n.elements.keys():
                    if not name in GYPSY_hirarchy:
                        collapse = 0
            elif 'crypton' in n.elements.keys():
                for name in n.elements.keys():
                    if not name in crypton_hirarchy:
                        collapse = 0
            elif 'YR_only' in n.elements.keys():
                for name in n.elements.keys():
                    if not name in YR_hirarchy:
                        collapse = 0
            if collapse == 1:
                if tree_string_file == 'YR_tree/all_YR_corrected_reviewed.tre' and n.support < 0.95:
                    n_descendants = n.get_descendants()
                    cutoff = blen_cutoff(n,4)
                    for d in n_descendants:
                         if d.dist > cutoff:
                             d.detach()
                if print_collapsed_nodes == 1: #print all leaves in a file. It will be used later 
                    for leaf in n:
                        if not 'node' in leaf.name:
                            collapsed_nodes.append(leaf.name + '\t' + n.name + '\t')
                for C in n.get_descendants(): # and remove their leaves
                        C.detach() 
        
    
    
    
    ts = TreeStyle()
    ts.show_leaf_name = False
    #ts.scale = 150
    ts.show_branch_support = False
    t.ladderize(1)
    R = t.get_midpoint_outgroup()

    ns = NodeStyle()
    ns['hz_line_width'] = 5
    ns['vt_line_width'] = 5
    ns['fgcolor'] = 'black'
    ns['size'] = 0
    
    
    t.set_outgroup(R)

    for node in t.traverse(): 
        node.set_style(ns)
        if node.is_leaf():           
            if not 'node' in node.name:    #add faces to leaves
                color = 'gray'
                if 'rypton' in node.name:
                    color = 'sandybrown'
                    collapsed_nodes_file.write(node.name + '\t' + node.name + '\t' + 'Crypton' + '\n')
                elif 'kangaroo' in node.name:
                    color = 'red'
                    collapsed_nodes_file.write(node.name + '\t' + node.name + '\t' + 'kangaroo' + '\n')
                elif 'dirs1' in node.name:
                    color = 'darkviolet'
                elif 'PAT' in node.name:
                    color = 'green'
                elif 'Ngaro' in node.name:
                    color = 'black'
                    if ('BAS' in node.name or 'RET' in node.name) and not 'DIRS' in node.name:
                        collapsed_nodes_file.write(node.name + '\t' + node.name + '\t' + 'Ngaro' + '\n')
                elif 'ypsy' in node.name:
                    node.name = 'GYPSY'
                    color = 'black'
                name_parts = node.name.split('_')
                name = name_parts[0]
                get_name_part = 0
                for part in name_parts:
                    if get_name_part == 1:
                        name = name + '_' + part
                    elif part == '1':
                        get_name_part = 1
                hola = TextFace(name, fsize=15, fgcolor=color)    
                #node.add_face(hola, column=0, position = "branch-right")
                if node.name == 'GYPSY':
                    node.add_face(hola, column=0, position = "branch-right")
                elif color == 'gray':
                    node.add_face(CircleFace(radius = 10, color = 'gray', style = 'sphere'), column=0, position = "branch-right")
                else:
                    node.add_face(CircleFace(radius = 10, color = color), column=0, position = "branch-right")
                
                
                if 'RET_' in node.name or 'PAT_' in node.name or 'BAS_' in node.name:
                    if 'dirs1_like' in node.name:
                        collapsed_nodes_file.write(node.name + '\t' + node.name + '\t' + 'dirs1_like' + '\n')
                    elif 'pat1_like' in node.name:
                        collapsed_nodes_file.write(node.name + '\t' + node.name + '\t' + 'pat1_like' + '\n')
                    elif '_PAT' in node.name:
                        collapsed_nodes_file.write(node.name + '\t' + node.name + '\t' + 'PAT' + '\n')
                    else:
                        collapsed_nodes_file.write(node.name + '\t' + node.name + '\t' + 'None' + '\n')
            elif 'node' in node.name:  # add faces to collapsed nodes
                bcolor = 'white'
                species = ''
                count = 0
                for key in node.species.keys():
                    species =species + key + '_' + str(node.species[key]) + '_'
                    count += node.species[key]
                elements = ''
                for key in node.elements.keys():
                    elements = elements + key + '_' + str(node.elements[key]) + ' '
                verdict = 'None'
                verdict_color = 'black'
                if 'dirs1_like' in elements:
                    verdict = 'dirs1_like'
                    verdict_color = 'darkviolet'
                elif 'pat1_like' in elements:
                    verdict = 'pat1_like'
                    verdict_color = 'blue'
                elif 'PAT' in elements:
                    verdict = 'PAT'
                    verdict_color = 'green'
                elif 'kangaroo' in elements:
                    verdict = 'kangaroo'
                    verdict_color = 'red'
                elif 'Ngaro' in elements:
                    verdict = 'Ngaro'
                    verdict_color = 'black'
                elif 'crypton' in elements:
                    verdict = 'Crypton'
                    verdict_color = 'sandybrown'
                elif 'BEL' in elements:
                    verdict = 'BEL'
                    verdict_color = 'black'
                elif 'COPIA' in elements:
                    verdict = 'COPIA'
                    verdict_color = 'black'
                elif 'GYPSY' in elements:
                    verdict = 'GYPSY'
                    verdict_color = 'black'
                for line in collapsed_nodes:
                    if line.split('\t')[1] == node.name:
                        collapsed_nodes_file.write(line + verdict + '\n')
                el = TextFace(elements, fsize=10, fgcolor = 'dimgray')
                if bcolor == 'black':
                    el.fgcolor = 'white'
                sp = TextFace(species, fsize=10, fgcolor = 'dimgray')
                if count < 15:
                    count = 15
                #if len(species) > 0:
                    #species_line_to_pie_chart_face(node, species,species_groupings,float(count)*1+5)
                #sp.border.width = 1
                
                vr = TextFace(verdict, fsize=count, fgcolor=verdict_color)
                #if verdict == 'pat1_like':
                    #vr.background.color = 'lightgray'
                #node.add_face(sp, column=2, position = "branch-right")
                
                #node.add_face(el, column=2, position = "branch-right")
                if not verdict == 'None' and not verdict == 'COPIA' and not verdict == 'BEL' and not verdict == 'GYPSY':
                    node.add_face(CircleFace(radius = 10, color = verdict_color), column=0, position = "branch-right")
                    #node.add_face(vr, column=1, position = "branch-right")
                   
                elif not verdict == 'COPIA' and not verdict == 'BEL' and not verdict == 'GYPSY':
                    node.add_face(CircleFace(radius = 10, color = 'gray', style = 'sphere'), column=0, position = "branch-right")
                else:
                    node.add_face(vr, column=1, position = "branch-right")
                hola = TextFace(node.name, fsize=10, fgcolor='gray')    
                #node.add_face(hola, column=0, position = "branch-right")
                score = TextFace(str(node.support),fsize=6, fgcolor='black')
                #node.add_face(score,column=0, position = "branch-right")
                if node.support == 1:
                    node.add_face(CircleFace(radius = 5, color = 'black'),column=0, position = "float")
                elif node.support < 1 and node.support >= 0.95:
                    node.add_face(CircleFace(radius = 5, color = 'dimgray'),column=0, position = "float")
                elif node.support < 0.95 and node.support >= 0.9:
                    node.add_face(CircleFace(radius = 5, color = 'gray'),column=0, position = "float")
                elif node.support < 0.9 and node.support >= 0.8:
                    node.add_face(CircleFace(radius = 5, color = 'gainsboro'),column=0, position = "float")
        else: #add faces to internal nodes
            hola = TextFace(node.name, fsize=10, fgcolor='gray')    
            #node.add_face(hola, column=0, position = "branch-right")
            if node.support == 1:
                    node.add_face(CircleFace(radius = 5, color = 'black'),column=0, position = "float")
            elif node.support < 1 and node.support >= 0.95:
                    node.add_face(CircleFace(radius = 5, color = 'dimgray'),column=0, position = "float")
            elif node.support < 0.95 and node.support >= 0.9:
                    node.add_face(CircleFace(radius = 5, color = 'gray'),column=0, position = "float")
            elif node.support < 0.9 and node.support >= 0.8:
                    node.add_face(CircleFace(radius = 5, color = 'gainsboro'),column=0, position = "float")
            
    
    t.dist = 0
    t.support = 0
    t.show(tree_style = ts)
    #t.render(path + tree_string_file.split('.')[0] + '_Figure.tif', w = 1000, tree_style = ts)
    collapsed_nodes_file.close()