import os
import numpy as np
import pandas as pd
import kmapper as km
from sklearn.manifold import TSNE
import networkx as nx
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.cluster import AgglomerativeClustering

##clean tha data concataneted by matlab
def clean_matlab(data,label):
    '''
    data and label are outputs of ws.m, wm.m, or sm.m files
    This function removes nans.
    '''
    y=data.shape[1]
    mask=np.isnan(data)
    clean_data=data[~mask].reshape(-1,y)
    clean_label=label[~mask[:,0]]
    return clean_data,clean_label

##SNE Transformation&Mapper Graph
def tsne_mapper(cdata,interval,overlap,clust):
    '''
    cdata is output of clean_matlab function.
    interval, overlap, clust are parameters of mapper.
    This function applies the tsne transformation and returns the mapper graph.
    '''
    #start SNE
    tsne = TSNE()
    data_tsne = pd.DataFrame(tsne.fit_transform(cdata),columns=['C0','C1'])
    #start Mapper
    mapper = km.KeplerMapper(verbose=1) 
    graph = mapper.map(lens=data_tsne, X=cdata, cover=km.Cover(interval, overlap), clusterer=cluster.AgglomerativeClustering(n_clusters=clust))
    return graph

##Modularity
def modularity(clabel,graph):
    '''
    clabel is output of clean_matlab function. The labesl are either 1 representing task 1 or 2 representing task 2. The order of tasks can be found at the end of each matlab file.
    graph is the mapper graph returned by tsne_mapper function.
    This function returns modularity score.
    '''
    task1 = sum(clabel==1)
    task1_members = []
    task2_members = []
    #task membership vectors
    for node in graph['nodes']:
        if sum(np.array(graph['nodes'][node])<=task1)<len(graph['nodes'][node])/2:
            task2_members.append(node)
        else:
            task1_members.append(node)
    #mapper graph to networkx graph
    graph_nx = km.adapter.to_nx(graph)
    #modularity calculation
    return nx.community.modularity(graph_nx,[set(task2_members),set(task1_members)])


#Centrality
def centrality(clabel,graph):
    '''
    clabel is output of clean_matlab function. The labesl are either 1 representing task 1 or 2 representing task 2. The order of tasks can be found at the end of each matlab file.
    graph is the mapper graph returned by tsne_mapper function.
    This function returns the degree, eigenvector, betweenness, closeness centrality scores of a mapper graph.
    '''
    task1 = sum(clabel==1)
    graph_nx = km.adapter.to_nx(graph)
    #task membership series
    task_membership = pd.Series(index=graph['nodes'])
    for node in graph['nodes']:
        if sum(np.array(graph['nodes'][node])<task1)<len(graph['nodes'][node])/2:
            task_membership[node] = 2
        else:
            task_membership[node] = 1
    #centrality types
    centrality_functions = [nx.degree_centrality,
                            nx.eigenvector_centrality,
                            nx.betweenness_centrality,
                            nx.closeness_centrality]
    centrality_type = ['degree','eigenvector','betweenness','closeness']
   # centrality dataframe
    deg_df = pd.DataFrame(index=graph['nodes'])
    for func,centtype in zip(centrality_functions,centrality_type):
        if func!=nx.eigenvector_centrality:
            cent = func(graph_nx)
            deg_df[centtype] = cent.values()
        else:
            cent = func(graph_nx, max_iter=1000)
            deg_df[centtype] = cent.values()
    deg_df['Task'] = task_membership
    return deg_df