import networkx as nx
import networkx.algorithms.isomorphism as iso
import numpy as np
import pandas as pd
from scipy.sparse.csgraph import laplacian
from scipy.linalg import eigh
from scipy.integrate import quad
from sklearn.metrics import pairwise_distances
import warnings

import numba
from numba import jit, float32

def IM_dist(G1, G2):
    adj1 = nx.to_numpy_array(G1)
    adj2 = nx.to_numpy_array(G2)
    hwhm = 0.08
    
    N = len(adj1)
    # get laplacian matrix
    L1 = laplacian(adj1, normed=False)
    L2 = laplacian(adj2, normed=False)

    # get the modes for the positive-semidefinite laplacian
    w1 = np.sqrt(np.abs(eigh(L1)[0][1:]))
    w2 = np.sqrt(np.abs(eigh(L2)[0][1:]))

    # we calculate the norm for both spectrum
    norm1 = (N - 1) * np.pi / 2 - np.sum(np.arctan(-w1 / hwhm))
    norm2 = (N - 1) * np.pi / 2 - np.sum(np.arctan(-w2 / hwhm))

    # define both spectral densities
    density1 = lambda w: np.sum(hwhm / ((w - w1) ** 2 + hwhm ** 2)) / norm1
    density2 = lambda w: np.sum(hwhm / ((w - w2) ** 2 + hwhm ** 2)) / norm2

    func = lambda w: (density1(w) - density2(w)) ** 2
    return np.sqrt(quad(func, 0, np.inf, limit=100)[0])

def build_milestone_net(subgraph, init_node):
    '''
    Args:
        subgraph     - a connected component of the graph, csr_matrix
        init_node    - root node
    Returns:
        df_subgraph  - dataframe of milestone network
    '''

    if len(subgraph)==1:
        warnings.warn('Singular node.')
        return []
    else:
        # Dijkstra's Algorithm
        unvisited = {node: {'parent':None,
                            'score':np.inf,
                            'distance':np.inf} for node in subgraph.nodes}
        current = init_node
        currentScore = 0
        currentDistance = 0
        unvisited[current]['score'] = currentScore

        milestone_net = []
        while True:
            for neighbour in subgraph.neighbors(current):
                if neighbour not in unvisited: continue
                newScore = currentScore + subgraph[current][neighbour]['weight']
                if unvisited[neighbour]['score'] > newScore:
                    unvisited[neighbour]['score'] = newScore
                    unvisited[neighbour]['parent'] = current
                    unvisited[neighbour]['distance'] = currentDistance+1

            if len(unvisited)<len(subgraph):
                milestone_net.append([unvisited[current]['parent'],
                                      current,
                                      unvisited[current]['distance']])
            del unvisited[current]
            if not unvisited: break
            current, currentScore, currentDistance = \
                sorted([(i[0],i[1]['score'],i[1]['distance']) for i in unvisited.items()],
                        key = lambda x: x[1])[0]
        return np.array(milestone_net)

def comp_pseudotime(G, node, df):
    connected_comps = nx.node_connected_component(G, node)
    subG = G.subgraph(connected_comps)
    milestone_net = build_milestone_net(subG,node)

    # compute pseudotime
    pseudotime = - np.ones(len(df))
    for i in range(len(milestone_net)):
        _from, _to = milestone_net[i,:2]
        _from, _to = int(_from), int(_to)

        idc = (df['from']==_from)&(df['to']==_to)
        if np.sum(idc)>0:
            pseudotime[idc] = df['percentage'].values[idc] + milestone_net[i,-1] - 1
            
        idc = (df['from']==_to)&(df['to']==_from)
        if np.sum(idc)>0:
            pseudotime[idc] = 1-df['percentage'].values[idc] + milestone_net[i,-1] - 1        

        if np.any(df['from']==_from):
            idc = (df['from']==_from)&(df['to']==_from)
            pseudotime[idc] =  milestone_net[i,-1] - 1
    if len(milestone_net)>0 and np.any((df['from']==_to)&(df['to']==_to)):
        idc = (df['from']==_to)&(df['to']==_to)
        pseudotime[idc] =  milestone_net[i,-1]
    return pseudotime

def topology(G_true, G_pred, is_GED=True):   
    res = {}
    
    # 1. Isomorphism with same initial node
    def comparison(N1, N2):    
        if N1['is_init'] != N2['is_init']:
            return False
        else:
            return True
    score_isomorphism = int(nx.is_isomorphic(G_true, G_pred, node_match=comparison))
    res['ISO score'] = score_isomorphism
    
    # 2. GED (graph edit distance)
    if len(G_true)>10 or len(G_pred)>10:
        warnings.warn("Didn't calculate graph edit distances for large graphs.")
        res['GED score'] = np.nan  
    else:
        max_num_oper = len(G_true)
        score_GED = 1 - np.min([nx.graph_edit_distance(G_pred, G_true, node_match=comparison),
                            max_num_oper]) / max_num_oper
        res['GED score'] = score_GED
        
    # 3. Hamming-Ipsen-Mikhailov distance
    if len(G_true)==len(G_pred):
        score_IM = 1-IM_dist(G_true, G_pred)
        score_IM = np.maximum(0, score_IM)
    else:
        score_IM = 0
    res['score_IM'] = score_IM
    return res

@jit((float32[:,:], float32[:,:]), nopython=True, nogil=True)
def _rand_index(true, pred):
    n = true.shape[0]
    m_true = true.shape[1]
    m_pred = pred.shape[1]
    RI = 0.0
    for i in range(1, n-1):
        for j in range(i, n):
            RI_ij = 0.0
            for k in range(m_true):
                RI_ij += true[i,k]*true[j,k]
            for k in range(m_pred):
                RI_ij -= pred[i,k]*pred[j,k]     
            RI += 1-np.abs(RI_ij)
    return RI / (n*(n-1)/2.0)


def get_GRI(true, pred):
    '''
    Params:
        ture - [n_samples, n_cluster_1] for proportions or [n_samples, ] for grouping
        pred - [n_samples, n_cluster_2] for estimated proportions
    '''    
    if len(true)!=len(pred):
        raise ValueError('Inputs should have same lengths!')
        
    if len(true.shape)==1:
        true = pd.get_dummies(true).values
    if len(pred.shape)==1:
        pred = pd.get_dummies(pred).values
        
    true = np.sqrt(true).astype(np.float32)
    pred = np.sqrt(pred).astype(np.float32)

    return _rand_index(true, pred)

import numpy as np
import pandas as pd
from sklearn.metrics.cluster import adjusted_rand_score
from sklearn.preprocessing import LabelEncoder
import sys#; sys.path.insert(0, '../..')
sys.path.append('../..')
from VITAE import load_data

type_dict = {
    # dyno
    'dentate':'UMI', 
    'immune':'UMI', 
    'planaria_full':'UMI', 
    'planaria_muscle':'UMI',
    'aging':'non-UMI', 
    'cell_cycle':'non-UMI',
    'fibroblast':'non-UMI', 
    'germline':'non-UMI',    
    'human_embryos':'non-UMI', 
    'mesoderm':'non-UMI',
    
    # dyngen    
    "cycle_1":'non-UMI', 
    "cycle_2":'non-UMI', 
    "cycle_3":'non-UMI',
    "linear_1":'non-UMI', 
    "linear_2":'non-UMI', 
    "linear_3":'non-UMI', 
    "trifurcating_1":'non-UMI', 
    "trifurcating_2":'non-UMI', 
    "bifurcating_1":'non-UMI', 
    'bifurcating_2':'non-UMI',
    "bifurcating_3":'non-UMI', 
    "converging_1":'non-UMI',
    
    # our model
    'linear':'UMI',
    'bifurcation':'UMI',
    'multifurcating':'UMI',
    'tree':'UMI',
}
source_dict = {
    'dentate':'real', 
    'immune':'real', 
    'planaria_muscle':'real',
    'planaria_full':'real',
    'aging':'real', 
    'cell_cycle':'real',
    'fibroblast':'real', 
    'germline':'real',    
    'human_embryos':'real', 
    'mesoderm':'real',
        
    "cycle_1":'dyngen', 
    "cycle_2":'dyngen', 
    "cycle_3":'dyngen',
    "linear_1":'dyngen', 
    "linear_2":'dyngen', 
    "linear_3":'dyngen', 
    "trifurcating_1":'dyngen', 
    "trifurcating_2":'dyngen', 
    "bifurcating_1":'dyngen', 
    'bifurcating_2':'dyngen',
    "bifurcating_3":'dyngen', 
    "converging_1":'dyngen',
    
    'linear':'our model',
    'bifurcation':'our model',
    'multifurcating':'our model',
    'tree':'our model',
}
df = pd.DataFrame()
for data_name in type_dict.keys():
    data = load_data('data/', data_name)
    milestone_net = data['milestone_network']    
    le = LabelEncoder()
    grouping = le.fit_transform(data['grouping'])
    begin_node_true = le.transform([data['root_milestone_id']])[0]
    if 'w' in milestone_net.columns:
        grouping = None
    if milestone_net is not None:
        milestone_net['from'] = le.transform(milestone_net['from'])
        milestone_net['to'] = le.transform(milestone_net['to'])  
    G_true = nx.Graph()
    G_true.add_nodes_from(np.unique(milestone_net[['from','to']].values.flatten()))
    if grouping is None:
        G_true.add_edges_from(list(
            milestone_net[~pd.isna(milestone_net['w'])].groupby(['from', 'to']).count().index))
    # otherwise, 'milestone_net' indicates edges
    else:
        if milestone_net is not None:             
            G_true.add_edges_from(list(
                milestone_net.groupby(['from', 'to']).count().index))
    G_true.remove_edges_from(nx.selfloop_edges(G_true))
    nx.set_node_attributes(G_true, False, 'is_init')
    G_true.nodes[begin_node_true]['is_init'] = True

        
    for method in ['monocle','paga','slingshot']:
                
        _df = pd.read_csv('result/other methods/%s/%s_progressions.csv'%(data_name,method), 
                          index_col=0).astype({'from':str,'to':str})   
        pred_milestone_net = pd.read_csv('result/other methods/%s/%s_milestone_network.csv'%(data_name,method), 
                                         index_col=0).sort_index().astype({'from':str,'to':str})
        if method=='paga':
            _df['from'] = _df['from'].apply(lambda x:int(x.split('-')[0]) if type(x)==str else x)    
            _df['to'] = _df['to'].apply(lambda x:int(x.split('-')[0]) if type(x)==str else x)  
            pred_milestone_net['from'] = pred_milestone_net['from'].apply(
                lambda x:int(x.split('-')[0]) if type(x)==str else x)  
            pred_milestone_net['to'] = pred_milestone_net['to'].apply(
                lambda x:int(x.split('-')[0]) if type(x)==str else x)              

        le = LabelEncoder()
        le.fit(np.unique(pred_milestone_net[['from','to']].values.flatten()))
        pred_milestone_net['from'] = le.transform(pred_milestone_net['from'])
        pred_milestone_net['to'] = le.transform(pred_milestone_net['to'])
        _df['from'] = le.transform(_df['from'])
        _df['to'] = le.transform(_df['to'])
        
        G_pred = nx.Graph()
        G_pred.add_nodes_from(np.unique(_df[['from', 'to']].values.flatten()))
        G_pred.add_edges_from(pred_milestone_net[['from', 'to']].values)
        G_pred.remove_edges_from(nx.selfloop_edges(G_pred))  
        nx.set_node_attributes(G_pred, False, 'is_init')
        G_pred.nodes[pred_milestone_net.iloc[0]['from']]['is_init'] = True
        
        pseudotime = comp_pseudotime(nx.from_numpy_array(nx.to_numpy_array(G_pred)), 
                                     pred_milestone_net.iloc[0]['from'], 
                                     _df.copy())
        
        # 1. topology
        res = topology(G_true, G_pred)

        # 2. Milestones assignment
        if grouping is None:
            milestones_true = milestone_net['from'].values.copy()
            milestones_true[(milestone_net['from']!=milestone_net['to'])
                           &(milestone_net['w']<0.5)] = milestone_net[(milestone_net['from']!=milestone_net['to'])
                                                                      &(milestone_net['w']<0.5)]['to'].values
        else:
            milestones_true = grouping
        milestones_true = milestones_true[pseudotime!=-1]
        milestones_pred = _df['from'].values.copy()
        milestones_pred[_df['percentage']>0.5] = _df[_df['percentage']>0.5]['to'].values.copy()
        milestones_pred = milestones_pred[pseudotime!=-1]
        res['ARI'] = (adjusted_rand_score(milestones_true, milestones_pred) + 1)/2

        n_samples = len(pseudotime)
        w = np.zeros((n_samples,n_samples))
        w[np.arange(n_samples), _df['to']] = _df['percentage']
        w[np.arange(n_samples), _df['from']] = 1 - _df['percentage']
        if grouping is None:
            n_samples = len(milestone_net)
            prop = np.zeros((n_samples,n_samples))
            prop[np.arange(n_samples), milestone_net['to']] = 1-milestone_net['w']
            prop[np.arange(n_samples), milestone_net['from']] = np.where(np.isnan(milestone_net['w']), 1, milestone_net['w'])
            res['GRI'] = get_GRI(prop, w)
        else:
            res['GRI'] = get_GRI(grouping, w)
        
        # 3. Correlation between geodesic distances / Pseudotime   
        if 'cycle' in data_name or 'immune' in data_name:
            res['PDT score'] = np.nan
        else:
            if grouping is None:
                pseudotime_true = milestone_net['from'].values + 1 - milestone_net['w'].values
                pseudotime_true[np.isnan(pseudotime_true)] = milestone_net[pd.isna(milestone_net['w'])]['from'].values            
            else:
                pseudotime_true = - np.ones(len(grouping))
                nx.set_edge_attributes(G_true, values = 1, name = 'weight')
                connected_comps = nx.node_connected_component(G_true, begin_node_true)
                subG = G_true.subgraph(connected_comps)
                milestone_net_true = build_milestone_net(subG, begin_node_true)
                if len(milestone_net_true)>0:
                    pseudotime_true[grouping==int(milestone_net_true[0,0])] = 0
                    for i in range(len(milestone_net_true)):
                        pseudotime_true[grouping==int(milestone_net_true[i,1])] = milestone_net_true[i,-1]
            pseudotime_true = pseudotime_true[pseudotime>-1]
            pseudotime_pred = pseudotime[pseudotime>-1]
            if np.std(pseudotime_true)==0 or np.std(pseudotime_pred)==0:
                res['PDT score'] = 0
            else:
                res['PDT score'] = (np.corrcoef(pseudotime_true,pseudotime_pred)[0,1]+1)/2
    
        res['method'] = method
        res['data'] = data_name            
        res['source'] = source_dict[data_name]       
        df = df.append(pd.DataFrame(res, index=[0]),ignore_index=True)        

df.to_csv('result/result_other_methods.csv')