import pandas as pd
import ete2
from ete2 import faces, Tree, AttrFace, TreeStyle
import pylab
from matplotlib.colors import hex2color, rgb2hex, hsv_to_rgb, rgb_to_hsv

kelly_colors_hex = [
    0xFFB300, # Vivid Yellow
    0x803E75, # Strong Purple
    0xFF6800, # Vivid Orange
    0xA6BDD7, # Very Light Blue
    0xC10020, # Vivid Red
    0xCEA262, # Grayish Yellow
    0x817066, # Medium Gray

    # The following don't work well for people with defective color vision
    0x007D34, # Vivid Green
    0xF6768E, # Strong Purplish Pink
    0x00538A, # Strong Blue
    0xFF7A5C, # Strong Yellowish Pink
    0x53377A, # Strong Violet
    0xFF8E00, # Vivid Orange Yellow
    0xB32851, # Strong Purplish Red
    0xF4C800, # Vivid Greenish Yellow
    0x7F180D, # Strong Reddish Brown
    0x93AA00, # Vivid Yellowish Green
    0x593315, # Deep Yellowish Brown
    0xF13A13, # Vivid Reddish Orange
    0x232C16, # Dark Olive Green
    ]

def my_layout(node):
    if node.is_leaf():
         # If terminal node, draws its name
         name_face = AttrFace("name")
    else:
         # If internal node, draws label with smaller font size
         name_face = AttrFace("name", fsize=10)
    # Adds the name face to the image at the preferred position
    faces.add_face_to_node(name_face, node, column=0, position="branch-right")
    
def adjust_kelly_brightness(hex_color, val, recon_min, recon_max):
    """set brightness according to change in continuous reconstruction value"""
    h, s, v =  rgb_to_hsv(hex2color('#{0:06X}'.format(hex_color)))
    scale_factor = 1 - (recon_max - val) / (recon_max - recon_min)
    v_new = v - (v * (scale_factor)) 
    return rgb2hex(hsv_to_rgb(pd.np.array([h, s, v_new])))

def get_style():
    ts = TreeStyle()
    # Do not add leaf names automatically
    ts.show_leaf_name = False
    ts.show_scale = True 
    ts.force_topology = False 
    # Use my custom layout
    ts.layout_fn = my_layout
    return ts

def plot_tree(pt_tree, target_node, out):
    #pt_tree, feats, pf2color = get_tree(phenotype = phenotype, feat_list = "top_cor", is_ml_plus_phypat = True, target_node = target_node)
    pt_tree.dist = 0
    target = pt_tree.search_nodes(name = target_node)[0]
    target.render(out + '_tree.pdf',  tree_style = get_style())
    #target.render(out + '_tree.png', tree_style = get_style())
    return target, feats, pf2color

def plot_legend(feats, out, pf2color,  pf_desc = False, pf_acc = True, include_class = False):
    fig = pylab.figure()
    figlegend = pylab.figure(figsize = (9, 6))
    ax = fig.add_subplot(111)
    x = [0,1]
    lines = [ax.plot(x, pd.np.ones(len(x)), 'o', color = "#%06x" % (pf2color[feats.index[i]]))[0] for i in range(len(pf2color))]
    labels= [i for i in feats.index]
    #labels= ["%s" %(feats.loc[:,"Pfam_acc"].iloc[i]) for i in range(feats.shape[0])]
    #if include_class:
    #    labels= ["%s %s" %(labels[i], feats.loc[:, "class"].iloc[i]) for i in range(len(labels))]
    #if pf_desc:
    #    labels = ["%s %s" % (labels[i], pf2short_desc.loc[feats.loc[:,"Pfam_acc"].iloc[i], 1]) for i in range(len(labels))]
    #if pf_acc:
    #    labels = ["%s %s" % (labels[i], pf2acc.loc[feats.loc[:,"Pfam_acc"].iloc[i], 1]) for i in range(len(labels))]

    figlegend.legend(lines, labels, markerscale = 2.5, numpoints = 1, frameon = False)
    #fig.show()
    fig.tight_layout()
    figlegend.savefig(out + "_legend.svg")
    figlegend.savefig(out + "_legend.png")
    return figlegend

def get_tree(phenotype, tree, gain_recon, loss_recon, node_recon, pfam_mapping, feat_list, sample_mapping, threshold = 0.5, target_node = None, are_continuous_features_with_discrete_phenotype = False, max_feats = 10, miscl = None, node_annotation = None):
    #read target feats
    feats = pd.read_csv(feat_list, index_col = 0, sep = "\t")
    pt_tree = ete2.Tree(tree, format = 1)
    pt_tree.ladderize()
    if not node_annotation is None:
        node_table = pd.read_csv(node_annotation, sep = "\t", index_col = 0)
    sample_mapping = pd.read_csv(sample_mapping, index_col = 0, sep = "\t")
    #read node and edge reconstruction matrices
    node_recon = pd.read_csv(node_recon, sep = "\t", index_col = 0)
    gain_recon = pd.read_csv(gain_recon, sep = "\t", index_col = 0)
    gain_recon.index = ["_".join(("_".join(i.split("_")[:-1]), i.split("_")[-1])) for i in gain_recon.index.values]
    loss_recon = pd.read_csv(loss_recon, sep = "\t", index_col = 0)
    loss_recon.index = ["_".join(("_".join(i.split("_")[:-1]), i.split("_")[-1])) for i in loss_recon.index.values]
    #prune to target node
    if target_node is not None:
        pt_tree = pt_tree.search_nodes(name = target_node)[0]
    node2name = dict((i.name, i.name) for i in pt_tree.traverse(strategy = 'preorder'))
    pfams_with_event = set()
    pfam2color = {}
    #set the style of the branches and nodes according to the posterior probability
    top10_feats = feats.iloc[:max_feats,]
    #for visualization of continuous feature get the range of values for each feature
    if are_continuous_features_with_discrete_phenotype:
        recon_min = gain_recon.abs().apply(pd.np.min)
        recon_max = gain_recon.abs().apply(pd.np.max)
    if not miscl is None:
        miscl_m = pd.read_csv(miscl, sep = "\t", index_col = 0)
    for n in pt_tree.traverse():
        #ignore the root
        if n.name == "N1":
            continue
        if not node_annotation is None:
            if n.name in node_table.index:
                for attr,i  in zip(node_table.columns, range(len(node_table.columns))):
                    value = node_table.loc[n.name, attr]
                    if not pd.isnull(value):
                        if value == 0:
                            rf = ete2.CircleFace(radius = 8, style = "circle", color = 'red')
                        elif value == 2:
                            rf = faces.CircleFace(radius = 8, style = "circle", color = 'orange')
                        else:
                            rf = faces.CircleFace(radius = 8, style = "circle", color = 'green')
                    else:
                        rf = faces.CircleFace(radius = 8, style = "circle", color = 'grey')
                    n.add_face(rf, column = i, position = "aligned")

        ns = node_recon.loc[n.name, phenotype] 
        style = ete2.NodeStyle()
        style["shape"] = 'square'
        style['size'] = 10
        if pd.isnull(ns):
            style['fgcolor'] = 'grey'
        elif ns < threshold:
            style['fgcolor'] = 'darkred'
        else:
            style['fgcolor'] = 'green'
        if not n.name == "N1":
            branch_id = n.name + "_" + n.up.name
            if gain_recon.loc[branch_id, phenotype] > threshold:
                style["hz_line_type"] = 1
                style["hz_line_color"] = 'green' 
                style["hz_line_width"] = 3
            elif loss_recon.loc[branch_id, phenotype] > threshold:
                style["hz_line_type"] = 1
                style["hz_line_color"] = 'red'
                style["hz_line_width"] = 3
            else:
                style["hz_line_type"] = 0
                style["hz_line_color"] = 'black'
            n.set_style(style)
            #check if sample was misclassified and add misclassified label
            if not miscl is None:
                if node2name[n.name] in miscl_m.index:
                    tf = faces.TextFace("misclassified")
                    n.add_face(tf, column = 0, position = "branch-right")
            #set species name instead of tax id
            if n.name in sample_mapping.index:
                node2name[n.name] = sample_mapping.loc[n.name,][0]
            #add majority feature gains and losses
            events = []
            for i in range(top10_feats.shape[0]): 
                if not are_continuous_features_with_discrete_phenotype:
                    cf = faces.CircleFace(radius = 8, style = "circle", color = kelly_colors_hex[i])
                    #gain events
                    if gain_recon.loc[branch_id, top10_feats.index[i]] > threshold:
                        pfam2color[top10_feats.index[i]] = kelly_colors_hex[i]
                        tf = faces.TextFace("-")
                        events.append(tf)
                        pfams_with_event.add(node_recon.index[i])
                        events.append(cf)
                    #loss events
                    elif loss_recon.loc[branch_id, top10_feats.index[i]] > threshold:
                        pfam2color[top10_feats.index[i]] = kelly_colors_hex[i]
                        tf = faces.TextFace("-")
                        events.append(tf)
                        pfams_with_event.add(node_recon.index[i])
                        events.append(cf)
                #continuous features
                else: 
                    adjusted_color = adjust_kelly_brightness(kelly_colors_hex[i], abs(loss_recon.loc[branch_id, top10_feats.index[i]]), recon_min.loc[top10_feats.index[i]], recon_max.loc[top10_feats.index[i]])
                    #tf = faces.TextFace(gain_recon.loc[branch_id, top10_feats.index[i]])
                    if loss_recon.loc[branch_id, top10_feats.index[i]] < 0:
                        tf = faces.TextFace("-")
                    else:
                        tf = faces.TextFace("+")
                    cf = faces.CircleFace(radius = 8, style = "circle", color = adjusted_color)
                    pfam2color[top10_feats.index[i]] = kelly_colors_hex[i]
                    pfams_with_event.add(node_recon.index[i])
                    events.append(cf)
                    events.append(tf)
            for i in range(len(events)):
                n.add_face(events[i], column = i, position = "branch-top")
    for n in pt_tree.traverse():
        if n.name in node2name:
            n.name = node2name[n.name]
    #filtered_pfams = filter(lambda i: i in list(pfams_with_event), top10_feats.loc[:,"Pfam_acc"].values)
    #print filtered_pfams
    #filtered_ids = pt_gt2id.loc[filtered_pfams, 0] - 1
    #print filtered_ids
    #top10_feats_with_event = top10_feats.loc[filtered_ids,]
    #process node annotation
    return pt_tree, top10_feats, pfam2color

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser("""visualize target list of features""")
    parser.add_argument("node_recon", help = "node ancestral character state reconstruction") 
    parser.add_argument("gain_recon", help = "gain events ancestral character state reconstruction") 
    parser.add_argument("loss_recon", help = "loss events ancestral character state reconstruction") 
    parser.add_argument("tree", help = "tree with internal nodes labeled") 
    parser.add_argument("pfam_mapping",  help = "feature mapping/list")
    parser.add_argument("feat_list",  help = "list of features")
    parser.add_argument("--target_node", default = "N1", help = "list of features")
    parser.add_argument("phenotype",  help = "target phenotype")
    parser.add_argument("--are_continuous_features_with_discrete_phenotype", action = 'store_true',  help = "set if using continuous features with a discrete phenotype")
    parser.add_argument("threshold", type = float,  help = "threshold to call genotype/phenotype events")
    parser.add_argument("sample_mapping",  help = "mapping between sample ids and names")
    parser.add_argument("out",  help = "output file")
    parser.add_argument("--max_feats", type = int, default = 10, help = "visualize at most max_feats features")
    parser.add_argument("--miscl", help = "table of misclassified samples")
    parser.add_argument("--node_annotation", help = "table of binary features for labeling the nodes")
    a = parser.parse_args()
    pt_tree, feats, pf2color = get_tree(node_recon = a.node_recon, gain_recon = a.gain_recon, loss_recon = a.loss_recon, pfam_mapping = a.pfam_mapping, tree = a.tree, feat_list = a.feat_list, phenotype = a.phenotype, target_node = a.target_node, threshold = a.threshold, sample_mapping = a.sample_mapping, are_continuous_features_with_discrete_phenotype = a.are_continuous_features_with_discrete_phenotype, max_feats = a.max_feats, miscl = a.miscl, node_annotation = a.node_annotation)
    plot_tree(pt_tree, a.target_node, a.out)
    plot_legend(feats, a.out, pf2color) 

[(40, 'ete2.faces.add_face_to_node', 'faces.add_face_to_node', (['name_face', 'node'], {'column': '(0)', 'position': '"""branch-right"""'}), False, 'from ete2 import faces, Tree, AttrFace, TreeStyle\n'), (50, 'ete2.TreeStyle', 'TreeStyle', ([], {}), False, 'from ete2 import faces, Tree, AttrFace, TreeStyle\n'), (68, 'pylab.figure', 'pylab.figure', ([], {}), False, 'import pylab\n'), (69, 'pylab.figure', 'pylab.figure', ([], {'figsize': '(9, 6)'}), False, 'import pylab\n'), (91, 'pandas.read_csv', 'pd.read_csv', (['feat_list'], {'index_col': '(0)', 'sep': '"""\t"""'}), True, 'import pandas as pd\n'), (92, 'ete2.Tree', 'ete2.Tree', (['tree'], {'format': '(1)'}), False, 'import ete2\n'), (96, 'pandas.read_csv', 'pd.read_csv', (['sample_mapping'], {'index_col': '(0)', 'sep': '"""\t"""'}), True, 'import pandas as pd\n'), (98, 'pandas.read_csv', 'pd.read_csv', (['node_recon'], {'sep': '"""\t"""', 'index_col': '(0)'}), True, 'import pandas as pd\n'), (99, 'pandas.read_csv', 'pd.read_csv', (['gain_recon'], {'sep': '"""\t"""', 'index_col': '(0)'}), True, 'import pandas as pd\n'), (101, 'pandas.read_csv', 'pd.read_csv', (['loss_recon'], {'sep': '"""\t"""', 'index_col': '(0)'}), True, 'import pandas as pd\n'), (215, 'argparse.ArgumentParser', 'argparse.ArgumentParser', (['"""visualize target list of features"""'], {}), False, 'import argparse\n'), (35, 'ete2.AttrFace', 'AttrFace', (['"""name"""'], {}), False, 'from ete2 import faces, Tree, AttrFace, TreeStyle\n'), (38, 'ete2.AttrFace', 'AttrFace', (['"""name"""'], {'fsize': '(10)'}), False, 'from ete2 import faces, Tree, AttrFace, TreeStyle\n'), (95, 'pandas.read_csv', 'pd.read_csv', (['node_annotation'], {'sep': '"""\t"""', 'index_col': '(0)'}), True, 'import pandas as pd\n'), (116, 'pandas.read_csv', 'pd.read_csv', (['miscl'], {'sep': '"""\t"""', 'index_col': '(0)'}), True, 'import pandas as pd\n'), (137, 'ete2.NodeStyle', 'ete2.NodeStyle', ([], {}), False, 'import ete2\n'), (140, 'pandas.isnull', 'pd.isnull', (['ns'], {}), True, 'import pandas as pd\n'), (47, 'pandas.np.array', 'pd.np.array', (['[h, s, v_new]'], {}), True, 'import pandas as pd\n'), (163, 'ete2.faces.TextFace', 'faces.TextFace', (['"""misclassified"""'], {}), False, 'from ete2 import faces, Tree, AttrFace, TreeStyle\n'), (172, 'ete2.faces.CircleFace', 'faces.CircleFace', ([], {'radius': '(8)', 'style': '"""circle"""', 'color': 'kelly_colors_hex[i]'}), False, 'from ete2 import faces, Tree, AttrFace, TreeStyle\n'), (195, 'ete2.faces.CircleFace', 'faces.CircleFace', ([], {'radius': '(8)', 'style': '"""circle"""', 'color': 'adjusted_color'}), False, 'from ete2 import faces, Tree, AttrFace, TreeStyle\n'), (125, 'pandas.isnull', 'pd.isnull', (['value'], {}), True, 'import pandas as pd\n'), (133, 'ete2.faces.CircleFace', 'faces.CircleFace', ([], {'radius': '(8)', 'style': '"""circle"""', 'color': '"""grey"""'}), False, 'from ete2 import faces, Tree, AttrFace, TreeStyle\n'), (176, 'ete2.faces.TextFace', 'faces.TextFace', (['"""-"""'], {}), False, 'from ete2 import faces, Tree, AttrFace, TreeStyle\n'), (192, 'ete2.faces.TextFace', 'faces.TextFace', (['"""-"""'], {}), False, 'from ete2 import faces, Tree, AttrFace, TreeStyle\n'), (194, 'ete2.faces.TextFace', 'faces.TextFace', (['"""+"""'], {}), False, 'from ete2 import faces, Tree, AttrFace, TreeStyle\n'), (127, 'ete2.CircleFace', 'ete2.CircleFace', ([], {'radius': '(8)', 'style': '"""circle"""', 'color': '"""red"""'}), False, 'import ete2\n'), (183, 'ete2.faces.TextFace', 'faces.TextFace', (['"""-"""'], {}), False, 'from ete2 import faces, Tree, AttrFace, TreeStyle\n'), (129, 'ete2.faces.CircleFace', 'faces.CircleFace', ([], {'radius': '(8)', 'style': '"""circle"""', 'color': '"""orange"""'}), False, 'from ete2 import faces, Tree, AttrFace, TreeStyle\n'), (131, 'ete2.faces.CircleFace', 'faces.CircleFace', ([], {'radius': '(8)', 'style': '"""circle"""', 'color': '"""green"""'}), False, 'from ete2 import faces, Tree, AttrFace, TreeStyle\n')]