# -*- coding: utf-8 -*-

"""

Created on Fri Aug 27 14:06:17 2021

@author: Shen Wanxiang

"""

import csv

from collections import defaultdict

from pprint import pprint

import pandas as pd

from scipy.cluster.hierarchy import dendrogram, linkage, to_tree

from scipy.spatial.distance import squareform

itol_header = '''TREE_COLORS

SEPARATOR TAB

#First 3 fields define the node id, type and color

#Possible types are:

#'range': defines a colored range (colored background for labels/clade)

#'clade': defines color/style for all branches in a clade

#'branch': defines color/style for a single branch

#'label': defines font color/style for the leaf label

#'label_background': defines the leaf label background color

#The following additional fields are required:

#for 'range', field 4 defines the colored range label (used in the legend)

#The following additional fields are optional:

#for 'label', field 4 defines the font style ('normal',''bold', 'italic' or 'bold-italic') and field 5 defines the numeric scale factor for the font size (eg. with value 2, font size for that label will be 2x the standard size)

#for 'clade' and 'branch', field 4 defines the branch style ('normal' or 'dashed') and field 5 defines the branch width scale factor (eg. with value 0.5, branch width for that clade will be 0.5 the standard width)

DATA

#NODE_ID TYPE COLOR LABEL_OR_STYLE SIZE_FACTOR

'''

def _getNewick(node, newick, parentdist, leaf_names):

    if node.is_leaf():

        return "%s:%.2f%s" % (leaf_names[node.id], parentdist - node.dist, newick)

    else:

        if len(newick) > 0:

            newick = "):%.2f%s" % (parentdist - node.dist, newick)

        else:

            newick = ");"

        newick = _getNewick(node.get_left(), newick, node.dist, leaf_names)

        newick = _getNewick(node.get_right(), ",%s" % (newick), node.dist, leaf_names)

        newick = "(%s" % (newick)

        return newick

def mp2newick(mp, treefile = 'phenotype_tree', leaf_names = None):

    df = mp.df_embedding[['colors','Subtypes']]

    if leaf_names == None:

        leaf_names = mp.flist #use final list only

    else:

        assert len(leaf_names) == len(df), 'leaf names should be a list with a length of %s' % len(df)

        df.index = leaf_names

    linkage_matrix = mp.Z

    tree = to_tree(linkage_matrix, rd=False)

    newick = _getNewick(tree, "", tree.dist, leaf_names = leaf_names)

    # write newick file for itol

    with open(treefile + '.nwk', 'w') as f:

        f.write(newick)

    # write dataset file for itol

    df['TYPE'] = 'clade'

    df['STYLE'] = 'normal'

    df = df[['TYPE', 'colors', 'STYLE']]

    with open(treefile + '.txt', 'w') as f:

        f.write(itol_header)

    df.to_csv(treefile + '.txt', mode = 'a', header=None, sep='\t')

    return df

def tree(): 

    return defaultdict(tree)

def tree_add(t, path):

    for node in path:

        t = t[node]

def pprint_tree(tree_instance):

    def dicts(t): return {k: dicts(t[k]) for k in t}

    pprint(dicts(tree_instance))

def dfs_to_tree(dfs):

    t = tree()

    for i in range(len(dfs)):

        row = dfs.iloc[i].dropna().tolist()

        tree_add(t, row)

    return t

def tree_to_newick(root):

    items = []

    for k in root.keys():

        s = ''

        if len(root[k].keys()) > 0:

            sub_tree = tree_to_newick(root[k])

            if sub_tree != '':

                s += '(' + sub_tree + ')'

        s += k

        items.append(s)

    return ','.join(items)

def dfs_to_weightless_newick(dfs):

    t = dfs_to_tree(dfs)

    newick_tree = tree_to_newick(t)

    return newick_tree

if __name__ == '__main__':

    species_list = pd.read_csv('./species.list.csv', header=None,index_col=0)[1].to_list()

    dfs = pd.Series(species_list).apply(lambda x: dict([i.split('__') for i in x.split('|')])).apply(pd.Series)

    level_dict = {'k':'kingdom', 'p':'phylum', 'c':'class' ,'o':'order' ,'f':'family' ,'g': 'genus','s': 'species'}

    dfs = dfs.rename(columns=level_dict)

    nwk_string = dfs_to_weightless_newick(dfs)

    with open("1.nwk", "w") as f:

        f.write(nwk_string)