#! /usr/bin/env python

"""
# ---------------------------------------------------------------------------
#
# funciones_network_science.py
#
# Archivo con todas las funciones que he utilizado para aplicar métricas de 
# ciencia de redes a redes bipartitas y proyectadas.
#
# TFG: Ciencia de redes y reposicionamiento de fármacos: potencial a través 
# de la medicina de redes
#
# María Marín Tercero
# ----------------------------------------------------------------------------
"""

import pandas as pd
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
import networkx as nx
from networkx.algorithms import bipartite
import nxpd
    
# =================================================================================

def asociaciones(archivo, var1, var2, nombre_var1, nombre_var2, ax):
    """
    Esta función toma un archivo que representa asociaciones entre dos grupos disjuntos de nodos. Finalmente, representa en un histograma la distribución del número de asociaciones de un grupo de nodos con nodos de la otra clase. 
    Datos de entrada: 
    1. archivo: Archivo con asociaciones entre nodos organizado en dos columnas, una para cada grupo de nodos 
    2. var1: Nombre de una columna
    3. var2: Nombre de la otra columna
    4. nombre_var1: Nombre de los nodos representados en la columna 1
    5 nombre_var2: Nombre de los nodos representados en la columna 2
    6. ax: Variable empleada para la representación organizada en un subplot
    
    """
    conteo = archivo.groupby(var1)[var2].count().reset_index() #cuento el número de asociaciones del grupo var1 con el grupo var2
    
    #Represento los resultados
    ax.hist(conteo[var2], bins=20, color='#79C4FF', edgecolor='black')
    ax.set_xlabel('Nº de ' + nombre_var2, fontsize=12)
    ax.set_ylabel(nombre_var1, fontsize=12)
    ax.set_title('Distribución de '+ r'$\mathbf{' + nombre_var2 + '}$' + " por " + r'$\mathbf{' + nombre_var1 + '}$' ,  fontsize=12)

# =================================================================================

def matriz_spl_enfermedades_bip(archivo, graf):
    """
    Esta función genera una matriz bidimensional que contiene los Shortest Path 
    Lenghts de todos los nodos enfermedades de una red bipartita compuesta por dos 
    grupos disjuntos de nodos.
    Datos de entrada:
    1. archivo: Archivo con los datos sobre los enlaces entre los nodos de la red
    2. graf: Red bipartita 
    """
    
    enfermedades = list(set(archivo['dis'])) #lista de enfermedades (columna "dis" del archivo)

    matriz_spl = np.zeros((len(enfermedades), len(enfermedades))) #Inicializo una matriz de ceros para almacenar los SPL

    #Calculo los SPL entre todas las parejas de enfermedades
    for i in range(len(enfermedades)): #enfermedad_i
        for j in range(i + 1, len(enfermedades)): #enfermedad_i+1
            enfermedad_i = enfermedades[i] #guardo el nodo enfermedad_i en la fila i de la matriz
            enfermedad_j = enfermedades[j] #guardo el nodo enfermedad_i+1 en la columna j de la matriz
            try: #me aseguro de que exista un camino entre ambas enfermedades
                spl = nx.shortest_path_length(graf, source=enfermedad_i, target=enfermedad_j) #calculo su spl
                matriz_spl[i, j] = spl #añado el valor a la matriz en su posición correspondiente
                matriz_spl[j, i] = spl #añado el valor a la matriz en su posición correspondiente
            except nx.NetworkXNoPath: #si no hay camino, continúo
                pass
    #Creo un DataFrame con la matriz de SPLs
    df_matriz_spl = pd.DataFrame(matriz_spl, index=enfermedades, columns=enfermedades)

    return df_matriz_spl

# =================================================================================

def matriz_spl_var_bip(archivo, graf, nombre_var):
    """
    Esta función genera una matriz bidimensional que contiene los Shortest Path 
    Lenghts para el grupo de nodos distinto a enfermedades en una red bipartita 
    que contiene dos grupos diferenets de nodos, siendo uno enfermedades. El otro
    grupo varía en función de la red. 
    Datos de entrada:
    1. archivo: Archivo con los datos sobre los enlaces entre los nodos de la red
    2.  graf: Red bipartita
    3. nombre_var: Identificador del grupo de nodos distinto a enfermedades en el archivo
    """

    var = list(set(archivo[str(nombre_var)])) #lista de nodos del grupo disjunto a enfermedades (columna con el identificador añadido en los datos de entrada)

    matriz_spl = np.zeros((len(var), len(var))) #Inicializo una matriz de ceros para almacenar los SPL

    #Calculo los SPL entre todas las parejas de los nodos seleccionados
    for i in range(len(var)): #nodos_i
        for j in range(i + 1, len(var)): #nodo_i+1
            var_i = var[i] #guardo el nodo_i en la fila i de la matriz
            var_j =  var[j] #guardo el nodo_i+1 en la columna j de la matriz
            try: #me aseguro de que exista un camino entre ambos nodos
                spl = nx.shortest_path_length(graf, source=var_i, target=var_j) #calculo su spl
                matriz_spl[i, j] = spl #añado el valor a la matriz en su posición correspondiente
                matriz_spl[j, i] = spl #añado el valor a la matriz en su posición correspondiente
            except nx.NetworkXNoPath: #si no hay camino, continúo
                pass

    #Creo un DataFrame con la matriz de SPLs
    df_matriz_spl = pd.DataFrame(matriz_spl, index=var, columns=var)

    return df_matriz_spl

# =================================================================================

def matriz_spl_enfermedades_proj(archivo, graf):
    """
    Esta función genera una matriz bidimensional que contiene los Shortest Path 
    Lenghts para las enfermedades en una red proyectada. 
    Datos de entrada:
    1. archivo: Archivo con los datos sobre los enlaces entre los nodos de la red
    2. graf: Red bipartita
    """
    all_shortest_paths = dict(nx.all_pairs_shortest_path_length(graf)) #creo un diccionario con los spls entre todas las parejas de nodos de la red 
    
    enfermedades = list(graf.nodes) # Obtengo la lista de nodos enfermedades
    
    matriz_spl = np.zeros((len(enfermedades), len(enfermedades))) #Inicializo una matriz de ceros para almacenar los SPL
   
    for i, enfermedad_i in enumerate(enfermedades): #enfermedad_i
        for j, enfermedad_j in enumerate(enfermedades): #enfermedad_j
            matriz_spl[i, j] = all_shortest_paths[enfermedad_i].get(enfermedad_j, 0) #añado a la matriz el spl entre la enfermedad_i y la enfermedad_j según el diccionario con el total de spls
    
    #Creo un DataFrame con la matriz de SPLs
    df_matriz_spl = pd.DataFrame(matriz_spl, index=enfermedades, columns=enfermedades)
    return df_matriz_spl

# =================================================================================
def degree_bip(graf, archivo, nombre_graf, nombre_var, ax, ax1):
    """
    Esta función calcula y representa la distribución de grados de dos 
    grupos disjuntos de nodos en una red bipartita, siendo uno de esos grupos
    el correspondiente a enfermedades
    Datos de entrada:
    1. graf: Red bipartita
    2. archivo: Archivo con los datos sobre los enlaces entre los nodos de la red
    3. nombre_graf: Nombre de la red
    4. nombre_var: Nombre del grupo disjunto de nodos distinto a enfermedades
    5. ax: Variable empleada para la representación organizada de la primera distribución en un subplot
    6. ax1: Variable empleada para la representación organizada  de la segunda distribución en un subplot
    """
    degrees_dict = dict(graf.degree())  # diccionario con los grados de los nodos del grafo
    
    nodos_enf = set(archivo.iloc[:, 0])  # selecciono el set de nodos enfermedades en el archivo (que se encuentran en la primera columna)
    nodos_var = set(archivo.iloc[:, 1])  # selecciono el set de nodos del grupo variable en el archivo (que se encuentran en la segunda columna)

    # Filtro los degrees para nodos de enfermedades y variables
    degree_enf = [degrees_dict[clave] for clave in nodos_enf if clave in degrees_dict]
    degree_var = [degrees_dict[clave] for clave in nodos_var if clave in degrees_dict]

    # Represento la distribución de grados de las enfermedades
    G_lista_enf = pd.DataFrame(list(zip(nodos_enf, degree_enf)), columns=['node', 'degree'])
    G_plot_enf = G_lista_enf.groupby('degree').count()
    ax.scatter(G_plot_enf.index, G_plot_enf['node'], marker='o', color='#79C4FF', s=10)
    ax.set_yscale('log')
    ax.set_xscale('log')
    ax.set_xlabel('Grado', fontsize=10)
    ax.set_ylabel('Nº nodos', fontsize=10)
    ax.set_title('Distribución de grados de ' + r'$\mathbf{' + "enfermedades" + '}$' +' en\n' + r'$\mathbf{' + "red" + '}$' + " " + r'$\mathbf{' + "bipartita" + '}$'+ " "+ r'$\mathbf{' + nombre_graf + '}$' , fontsize=12)
    
    # Represento la distribución de grados del grupo disjunto de nodos distinto a enfermedades
    G_lista_var = pd.DataFrame(list(zip(nodos_var, degree_var)), columns=['node', 'degree'])
    G_plot_var = G_lista_var.groupby('degree').count()
    ax1.scatter(G_plot_var.index, G_plot_var['node'], marker='o', color='#79C4FF', s=10)
    ax1.set_yscale('log')
    ax1.set_xscale('log')
    ax1.set_xlabel('Grado', fontsize=10)
    ax1.set_ylabel('Nº nodos', fontsize=10)
    ax1.set_title('Distribución de grados de ' +  r'$\mathbf{' + nombre_var + '}$' + ' en\n' + r'$\mathbf{' + " red" + '}$' + " "+ r'$\mathbf{' + "bipartita" + '}$'+ " " + r'$\mathbf{' + nombre_graf + '}$', fontsize=12)


# =================================================================================

def degree_proj(graf, nombre_graf, ax):
    """
    Esta función calcula y representa la distribución de grados de las enfermedades
    en una red proyectada.
    Datos de entrada:
    1. graf: Red bipartita
    2. nombre_graf: Nombre de la red
    3. ax: Variable empleada para la representación organizada de la distribución en un subplot
    """
    degrees_dict = dict(graf.degree())  # diccionario con los grados de los nodos del grafo
    
    nodos = list(degrees_dict.keys())  # lista con los nodos de la red (claves del diccionario anterior)
    degrees = list(degrees_dict.values())  # lista con los grados de los nodos de la red (valores del diccionario anterior)
    
    # Represento la distribución de grados
    G_lista = pd.DataFrame(list(zip(nodos, degrees)), columns=['node', 'degree'])
    G_plot = G_lista.groupby('degree').count()
    ax.scatter(G_plot.index, G_plot['node'], marker='o', color='#79C4FF', s=10)
    ax.set_yscale('log')
    ax.set_xscale('log')
    ax.set_xlabel('Grado', fontsize=10)
    ax.set_ylabel('Nº nodos', fontsize=10)
    ax.set_title('Distribución de grados de ' + r'$\mathbf{' + "enfermedades" + '}$'+  ' en\n' + r'$\mathbf{' + " red "+ '}$'+ " "+ r'$\mathbf{' + "proyectada" + '}$' + " " + r'$\mathbf{' + nombre_graf + '}$',  fontsize=12)

# =================================================================================

def average_spl(matriz_spl, nombre_graf, nombre_var, ax, bipartita = True):
    """
    Esta función realiza una distribución de los Average SPL d los nodos de un grafo.
    Datos de entrada: 
    1. matriz_spl: Matriz con la información sobre los SPLs entre todas las parejas de nodos
    2. nombre_graf: Nombre de la red
    3. nombre_var: Nombre del tipo de nodos representado
    4. ax: Variable empleada para la representación organizada de la distribución en un subplot
    5. Bipartita: dependiendo de si es igualue o False, añade el título correspondiente a una red bipartita o una proyectada, respectivamente.
    """
    average_spl_por_nodo = matriz_spl.mean(axis=0) #Calculo el Average_SPL de cada nodo (media de cada columna de la matriz de SPLs)
    
    df_average_spl = pd.DataFrame({'Nodo': matriz_spl.index, 'AverageSPL': average_spl_por_nodo.values}) #creo un DataFrame con los resultados
    
     # Represento la distribución de los Average Shortest Path Lengths
    ax.hist(df_average_spl['AverageSPL'], bins=20, color='#79C4FF', edgecolor='black')
    ax.set_xlabel('Average Shortest Path Lengths', style='italic', fontsize=12)
    ax.set_ylabel('Nº nodos', fontsize=12)
    if bipartita == True:
        ax.set_title('Distribución de Average SPLs de ' + r'$\mathbf{' + nombre_var + '}$'+  '  en\n' + r'$\mathbf{' + " red "+ '}$'+ " "+ r'$\mathbf{' + "bipartita" + '}$' + " " + r'$\mathbf{' + nombre_graf + '}$',  fontsize=12)
    else:
         ax.set_title('Distribución de Average SPLs de ' + r'$\mathbf{' + "enfermedades" + '}$'+  ' en\n' + r'$\mathbf{' + " red "+ '}$'+ " "+ r'$\mathbf{' + "proyectada" + '}$' + " " + r'$\mathbf{' + nombre_graf + '}$',  fontsize=12)

# =================================================================================

def centrality_b_bip(graf, archivo, nombre_graf, nombre_var, ax, ax1):
    """
    Esta función calcula y representa la distribución de centralidad de intermediación 
    de dos grupos disjuntos de nodos en una red bipartita, siendo uno de esos grupos
    el correspondiente a enfermedades
    Datos de entrada:
    1. graf: Red bipartita
    2. archivo: Archivo con los datos sobre los enlaces entre los nodos de la red
    3. nombre_graf: Nombre de la red
    4. nombre_var: Nombre del grupo disjunto de nodos distinto a enfermedades
    5. ax: Variable empleada para la representación organizada de la primera distribución en un subplot
    6. ax1: Variable empleada para la representación organizada  de la segunda distribución en un subplot
    
    """
    #Calculo la centralidad de intermediación de todos los nodos del grafo
    bc = nx.betweenness_centrality(graf, k=None, normalized=True, weight=None, endpoints=False, seed=None)

    nodos_enf = set(archivo.iloc[:, 0]) #selecciono el set de nodos enfermedades en el archivo (que se encuentran en la primera columna)
    nodos_var = set(archivo.iloc[:, 1]) #selecciono el set de nodos del grupo variable en el archivo (que se encuentran en la segunda columna)

    #Filtro la centralidad de intermediación para ambos grupos de nodos
    bc_enf = {clave: bc[clave] for clave in nodos_enf if clave in bc}
    bc_var = {clave: bc[clave] for clave in nodos_var if clave in bc}

    #Creo DataFrames para la centralidad de intermediación de ambos grupos de nodos
    df_bc_enf = pd.DataFrame({'Nodo': list(bc_enf.keys()), 'Centralidad de intermediación': list(bc_enf.values())})
    df_bc_var = pd.DataFrame({'Nodo': list(bc_var.keys()), 'Centralidad de intermediación': list(bc_var.values())})

    #Represento la centralidad de intermediación de las enfermedades
    ax.hist(df_bc_enf['Centralidad de intermediación'], bins=20, color='#79C4FF', edgecolor='black')
    ax.set_xlabel('Centralidad de intermediación', style='italic', fontsize=10)
    ax.set_ylabel('Nº nodos', fontsize=10)
    ax.set_title('Distribución de centralidad de intermediación de ' + r'$\mathbf{' + "enfermedades" + '}$' +' en\n' + r'$\mathbf{' + "red" + '}$' + " " + r'$\mathbf{' + "bipartita" + '}$'+ " "+ r'$\mathbf{' + nombre_graf + '}$' , fontsize=10)

    #Represento la centralidad de intermediación del grupo disjunto de nodos distinto a enfermedades
    ax1.hist(df_bc_var['Centralidad de intermediación'], bins=20, color='#79C4FF', edgecolor='black')
    ax1.set_xlabel('Centralidad de intermediación', style='italic', fontsize=10)
    ax1.set_ylabel('Nº nodos', fontsize=10)
    ax1.set_title('Distribución de centralidad de intermediación de ' +  r'$\mathbf{' + nombre_var + '}$' + ' en\n' + r'$\mathbf{' + " red" + '}$' + " "+ r'$\mathbf{' + "bipartita" + '}$'+ " " + r'$\mathbf{' + nombre_graf + '}$', fontsize=10)

# =================================================================================

def centrality_b_proj(graf, nombre_graf, ax):
    """
    Esta función calcula y representa la distribución de centralidad de intermediación 
    de enfermedades en una red proyectada.
    Datos de entrada:
    1. graf: Red bipartita
    2. nombre_graf: Nombre de la red
    3. ax: Variable empleada para la representación organizada de la distribución en un subplot
    """
    #Calculo la centralidad de intermediación de todos los nodos del grafo
    bc= nx.betweenness_centrality(graf, k=None, normalized=True, weight=None, endpoints=False, seed=None)
    
    #Creo un DataFrame para la centralidad de intermediación de las enfermedades
    df_bc = pd.DataFrame({'Nodo': list(bc.keys()), 'Centralidad de intermediación': list(bc.values())})
    
    #Represento la centralidad de intermediación de las enfermedades
    ax.hist(df_bc['Centralidad de intermediación'], bins=20, color='#79C4FF', edgecolor='black')
    ax.set_xlabel('Centralidad de intermediación',style='italic', fontsize=10)
    ax.set_ylabel('Nº nodos', fontsize=10)
    ax.set_title('Distribución de centralidad de intermediación de ' + r'$\mathbf{' + "enfermedades" + '}$'+  ' en\n' + r'$\mathbf{' + " red "+ '}$'+ " "+ r'$\mathbf{' + "proyectada" + '}$' + " " + r'$\mathbf{' + nombre_graf + '}$',  fontsize=10)

# =================================================================================

def centrality_c_bip(graf, archivo, nombre_graf, nombre_var, ax, ax1):
    """
    Esta función calcula y representa la distribución de centralidad de cercanía 
    de dos grupos disjuntos de nodos en una red bipartita, siendo uno de esos grupos
    el correspondiente a enfermedades
    Datos de entrada:
    1. graf: Red bipartita
    2. archivo: Archivo con los datos sobre los enlaces entre los nodos de la red
    3. nombre_graf: Nombre de la red
    4. nombre_var: Nombre del grupo disjunto de nodos distinto a enfermedades
    5. ax: Variable empleada para la representación organizada de la primera distribución en un subplot
    6. ax1: Variable empleada para la representación organizada  de la segunda distribución en un subplot
    
    """
    #Calculo la centralidad de intermediación de todos los nodos del grafo
    cc = nx.closeness_centrality(graf, u=None, distance=None, wf_improved=True)

    nodos_enf = set(archivo.iloc[:, 0]) #selecciono el set de nodos enfermedades en el archivo (que se encuentran en la primera columna)
    nodos_var = set(archivo.iloc[:, 1]) #selecciono el set de nodos del grupo variable en el archivo (que se encuentran en la segunda columna)

    #Filtro la centralidad de cercanía para ambos grupos de nodos
    cc_enf = {clave: cc[clave] for clave in nodos_enf if clave in cc}
    cc_var = {clave: cc[clave] for clave in nodos_var if clave in cc}

    #Creo DataFrames para la centralidad de cercanía de ambos grupos de nodos
    df_cc_enf = pd.DataFrame({'Nodo': list(cc_enf.keys()), 'Centralidad de cercanía': list(cc_enf.values())})
    df_cc_var = pd.DataFrame({'Nodo': list(cc_var.keys()), 'Centralidad de cercanía': list(cc_var.values())})

    #Represento la centralidad de intermediación de las enfermedades
    ax.hist(df_cc_enf['Centralidad de cercanía'], bins=20, color='#79C4FF', edgecolor='black')
    ax.set_xlabel('Centralidad de cercanía', style='italic', fontsize=10)
    ax.set_ylabel('Nº nodos', fontsize=10)
    ax.set_title('Distribución de centralidad de cercanía de ' + r'$\mathbf{' + "enfermedades" + '}$' +' en\n' + r'$\mathbf{' + "red" + '}$' + " " + r'$\mathbf{' + "bipartita" + '}$'+ " "+ r'$\mathbf{' + nombre_graf + '}$' , fontsize=10)

    #Represento la centralidad de cercanía del grupo disjunto de nodos distinto a enfermedades
    ax1.hist(df_cc_var['Centralidad de cercanía'], bins=20, color='#79C4FF', edgecolor='black')
    ax1.set_xlabel('Centralidad de cercanía', style='italic', fontsize=10)
    ax1.set_ylabel('Nº nodos', fontsize=10)
    ax1.set_title('Distribución de centralidad de cercanía de ' +  r'$\mathbf{' + nombre_var + '}$' + ' en\n' + r'$\mathbf{' + " red" + '}$' + " "+ r'$\mathbf{' + "bipartita" + '}$'+ " " + r'$\mathbf{' + nombre_graf + '}$', fontsize=10)
    
# =================================================================================

def centrality_c_proj(graf, nombre_graf,ax):
    """
    Esta función calcula y representa la distribución de centralidad de cercanía 
    de enfermedades en una red proyectada.
    Datos de entrada:
    1. graf: Red bipartita
    2. nombre_graf: Nombre de la red
    3. ax: Variable empleada para la representación organizada de la distribución en un subplot
    """
    #Calculo la centralidad de cercanía de todos los nodos del grafo
    cc= nx.closeness_centrality(graf, u=None, distance=None, wf_improved=True)
    #Creo un DataFrame para la centralidad de cercanía de las enfermedades
    df_cc = pd.DataFrame({'Nodo': list(cc.keys()), 'Centralidad de cercanía': list(cc.values())})
    #Represento la centralidad de cercanía de alas enfermedades
    ax.hist(df_cc['Centralidad de cercanía'], bins=20, color='#79C4FF', edgecolor='black')
    ax.set_xlabel('Centralidad de cercanía',style='italic', fontsize=10)
    ax.set_ylabel('Nº nodos', fontsize=10)
    ax.set_title('Distribución de centralidad de cercanía de ' + r'$\mathbf{' + "enfermedades" + '}$'+  ' en\n' + r'$\mathbf{' + " red "+ '}$'+ " "+ r'$\mathbf{' + "proyectada" + '}$' + " " + r'$\mathbf{' + nombre_graf + '}$',  fontsize=10)

# =================================================================================