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

from scipy.spatial.distance import squareform

import seaborn as sns

from highcharts import Highchart

import pandas as pd

import numpy as np

import os

from aggmap.utils.logtools import print_info

def plot_scatter(mp, htmlpath = './', htmlname = None, radius = 2, enabled_data_labels = False):

    '''

    mp: the object of mp

    htmlpath: the figure path, not include the prefix of 'html'

    htmlname: the name 

    radius: int, defaut:3, the radius of scatter dot

    '''

    title = '2D emmbedding of %s based on %s method' % (mp.ftype, mp.emb_method)

    subtitle = 'number of %s: %s, metric method: %s' % (mp.ftype, len(mp.flist), mp.metric)

    name = '%s_%s_%s_%s_%s' % (mp.ftype,len(mp.flist), mp.metric, mp.emb_method, 'scatter')

    if not os.path.exists(htmlpath):

        os.makedirs(htmlpath)

    if htmlname:

        name = htmlname + '_' + name 

    filename = os.path.join(htmlpath, name)

    print_info('generate file: %s' % filename)

    colormaps = mp.colormaps

    df = mp.df_scatter

    H = Highchart(width=1000, height=850)

    H.set_options('chart', {'type': 'scatter', 'zoomType': 'xy'})    

    H.set_options('title', {'text': title})

    H.set_options('subtitle', {'text': subtitle})

    H.set_options('xAxis', {'title': {'enabled': True,'text': 'X', 'style':{'fontSize':20}},

                           'labels':{'style':{'fontSize':20}}, 

                           'gridLineWidth': 1,

                           'startOnTick': True,

                           'endOnTick': True,

                           'showLastLabel': True})

    H.set_options('yAxis', {'title': {'text': 'Y', 'style':{'fontSize':20}},

                            'labels':{'style':{'fontSize':20}}, 

                            'gridLineWidth': 1,})

#     H.set_options('legend', {'layout': 'horizontal','verticalAlign': 'top','align':'right','floating': False,

#                              'backgroundColor': "(Highcharts.theme && Highcharts.theme.legendBackgroundColor) || '#FFFFFF'",

#                              'borderWidth': 1})

    H.set_options('legend', {'align': 'right', 'layout': 'vertical',

                             'margin': 1, 'verticalAlign': 'top', 'y':40,

                              'symbolHeight': 12, 'floating': False,})

    H.set_options('plotOptions', {'scatter': {'marker': {'radius': radius,

                                                         'states': {'hover': {'enabled': True,

                                                                              'lineColor': 'rgb(100,100,100)'}}},

                                              'states': {'hover': {'marker': {'enabled': False} }},

                                              'tooltip': {'headerFormat': '<b>{series.name}</b><br>',

                                                          'pointFormat': '{point.IDs}'}},

                                  'series': {'turboThreshold': 5000, 

                                             'dataLabels': {'enabled': enabled_data_labels, 'format': '{point.IDs}'}}

                                 })

    for subtype, color in colormaps.items():

        dfi = df[df['Subtypes'] == subtype]

        if len(dfi) == 0:

            continue

        data = dfi.to_dict('records')

        H.add_data_set(data, 'scatter', subtype, color=color)

    H.save_file(filename)

    print_info('save html file to %s' % filename)

    return H

def plot_grid(mp, htmlpath = './', htmlname = None, enabled_data_labels = False):

    '''

    mp: the object of mp

    htmlpath: the figure path

    '''    

    if not os.path.exists(htmlpath):

        os.makedirs(htmlpath)    

    title = 'Assignment of %s by %s emmbedding result' % (mp.ftype, mp.emb_method)

    subtitle = 'number of %s: %s, metric method: %s' % (mp.ftype, len(mp.flist), mp.metric)    

    name = '%s_%s_%s_%s_%s' % (mp.ftype,len(mp.flist), mp.metric, mp.emb_method, 'mp')

    if htmlname:

        name = name = htmlname + '_' + name   

    filename = os.path.join(htmlpath, name)

    print_info('generate file: %s' % filename)

    colormaps = mp.colormaps

    df = mp.df_grid

    H = Highchart(width=1000, height=850)

    H.set_options('chart', {'type': 'heatmap', 'zoomType': 'xy'})

    H.set_options('title', {'text': title})

    H.set_options('subtitle', {'text': subtitle})

#     H.set_options('xAxis', {'title': '', 

#                             'min': 0, 'max': mp._S.fmap_shape[1]-1,

#                             'allowDecimals':False,

#                             'labels':{'style':{'fontSize':20}}})

#     H.set_options('yAxis', {'title': '', 'tickPosition': 'inside', 

#                             'min': 0, 'max': mp._S.fmap_shape[0]-1,

#                             'reversed': True,

#                             'allowDecimals':False,

#                             'labels':{'style':{'fontSize':20}}})

    H.set_options('xAxis', {'title': None,                         

                            'min': 0, 'max': mp.fmap_shape[1]-1,

                            'startOnTick': False,

                            'endOnTick': False,    

                            'allowDecimals':False,

                            'labels':{'style':{'fontSize':20}}})

    H.set_options('yAxis', {'title': {'text': ' ', 'style':{'fontSize':20}}, 

                            'startOnTick': False,

                            'endOnTick': False,

                            'gridLineWidth': 0,

                            'reversed': True,

                            'min': 0, 'max': mp.fmap_shape[0]-1,

                            'allowDecimals':False,

                            'labels':{'style':{'fontSize':20}}})

    H.set_options('legend', {'align': 'right', 'layout': 'vertical',

                             'margin': 1, 'verticalAlign': 'top', 

                             'y': 60, 'symbolHeight': 12, 'floating': False,})

    H.set_options('tooltip', {'headerFormat': '<b>{series.name}</b><br>',

                              'pointFormat': '{point.v}'})

    H.set_options('plotOptions', {'series': {'turboThreshold': 5000,

                                             'dataLabels': {'enabled': enabled_data_labels, 

                                                            'format': '{point.v}',

                                                            'style': {'textOutline':False, 'color': 'black'},

                                                           }

                                            }

                                 }

                 )

    for subtype, color in colormaps.items():

        dfi = df[df['Subtypes'] == subtype]

        if len(dfi) == 0:

            continue

        H.add_data_set(dfi.to_dict('records'), 'heatmap', 

                       name = subtype,

                       #dataLabels = {'enabled': True, 'color': color}

                       color = color,

                      )

    H.save_file(filename)

    print_info('save html file to %s' % filename)

    return H

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 = 'mytree'):

    dist_matrix = mp.dist_matrix

    leaf_names = mp.flist

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

    dists = squareform(dist_matrix)

    linkage_matrix = linkage(dists, 'complete')

    tree = to_tree(linkage_matrix, rd=False)

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

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

        f.write(newick)

    df.to_excel(treefile + '.xlsx')

def plot_tree(mp, htmlpath = './', htmlname = None):

    pass