#! /usr/bin/env python

"""
# ---------------------------------------------------------------------------
#
# funciones_reposicionamiento.py
#
# Archivo con todas las funciones que he utilizado para identificar fármacos
candidatos de reposicionamiento a un subconjunto de enfermedades neurológicas:
demencia, bipolaridad, epilepsia y esquizofrenia.

# 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
import funciones_reposicionamiento
import re
    
# =================================================================================


def farmacos_candidatos(df_combinado, nombre_enf, arch):
    """
    Esta función permite obtener un DataFrame con las siguientes columnas:
    1. Enfermedades de estudio: enfermedades para las que buscamos fármacos de
    reposicionamiento.
    2. Fármacos candidatos: medicamentos candidatos de reposicionamiento para la
    patología situada en su misma fila, pero en la columna 1.
    3. Enfermedades candidatas: patologías asociadas a los fármacos situados en 
    su misma fila, pero en la columna 2. 
    
    Datos de entrada: 
    1. df_combinado: DataFrame con las siguientes columnas: enfermedades, fármacos, 
    tratamiento, dc, dc_mean, dc_std, dc_zscore, targets en módulo (el DataFrame 
    que se obtiene como resultado en la función targets_modulo.
    2. nombre_enf: nombre de la enfermedad.
    3. arch: archivo que contiene las relaciones entre fármacos y las patologías que
    tratan
    """
    #Filtro los datos para generar dos nuevos dataframes, uno con los datos de los fármacos que se empleen como tratamiento de la enfermedad (identificada por su nombre según nombre_enf), y otro con los datos para el resto de fármacos.
    farmacos_tratamiento = df_combinado[(df_combinado['Enfermedades'] == nombre_enf) & (df_combinado['Tratamiento'] == 'yes')]
    farmacos_no_tratamiento = df_combinado[(df_combinado['Enfermedades'] == nombre_enf) & (df_combinado['Tratamiento'] == 'unknown')]
    
    #Calculo los valores correspondiente al primer cuartil y al mínimo de la separación de los fármacos empleados como tratamiento al módulo de la enfermedad
    primer_cuartil_tratamiento = farmacos_tratamiento['dc'].quantile(0.25)
    minimo_tratamiento = farmacos_tratamiento['dc'].min()
    
    if nombre_enf == "Demencia": #si la enfermedad es Demencia
        
        #filtro los fármacos tratamiento con una cercanía mayor o igual a la separación mínima de los fármacos tratamiento, que además no presenten ninguna de sus proteínas diana en el módulo de la enfermedad, y con un z-score de la separación al módulo de la enfermedad menor o igual a -0.15 (lo que indica qye son fármacos significativamente próximos al módulo).
        farmacos_seleccionados = farmacos_no_tratamiento[(farmacos_no_tratamiento["dc"] >= minimo_tratamiento) & 
                                           (farmacos_no_tratamiento["Target en módulo"] == "No") &
                                           (farmacos_no_tratamiento["dc_zscore"] <= -0.15)]
    else: #si la enfermedad no es Demencia
        #filtro los fármacos tratamiento con una cercanía mayor o igual al primer cuartil de la separación de los fármacos tratamiento, que además no presenten ninguna de sus proteínas diana en el módulo de la enfermedad, y con un z-score de la separación al módulo de la enfermedad menor o igual a -0.15 (lo que indica qye son fármacos significativamente próximos al módulo).
        farmacos_seleccionados = farmacos_no_tratamiento[(farmacos_no_tratamiento["dc"] >= primer_cuartil_tratamiento) & 
                                           (farmacos_no_tratamiento["Target en módulo"] == "No") &
                                           (farmacos_no_tratamiento["dc_zscore"] <= -0.15)]
        
    #Ordeno los fármacos de forma ascendente según su z-score y selecciono los 5 primeros
    farmacos_seleccionados = farmacos_seleccionados.sort_values(by="dc_zscore").head(5)
    
    # Filtro el DataFrame arch para obtener solo las filas correspondientes a los fármacos seleccionados
    arch_filtrado = arch[arch['dru'].isin(farmacos_seleccionados['Fármacos'])]

    # Creo el DataFrame con los resultados obtenidos
    df_candidatos = pd.DataFrame({
        "Enfermedades de estudio": [nombre_enf] * len(arch_filtrado),
        "Fármacos candidatos": arch_filtrado['dru'].tolist(),
        "Enfermedades candidatas": arch_filtrado['dis'].tolist()
    })
    
    return df_candidatos

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


def spl_candidatos(spl_files, id_enfermedades, arch_candidatos):
    """
    Esta función genera un DataFrame compuesto por seis columnas:
    1. Enfermedades de estudio: enfermedades para las que buscamos fármacos de
    reposicionamiento.
    2. Fármacos candidatos: medicamentos candidatos de reposicionamiento para la
    patología situada en su misma fila, pero en la columna 1.
    3. Enfermedades candidatas: patologías asociadas a los fármacos situados en 
    su misma fila, pero en la columna 2. 
    4. SPL proyectada enf-gen: Shortest Path Length en la red proyectada de
    enfermedades relacionadas en función de genes patológicos entre las enfermedades 
    de la columna 1 y 3 situadas en la misma fila.
    5. SPL proyectada enf-dru: Shortest Path Length en la red proyectada de
    enfermedades relacionadas en función de fármacos entre las enfermedades 
    de la columna 1 y 3 situadas en la misma fila.
    6. SPL proyectada enf-sym: Shortest Path Length en la red proyectada de
    enfermedades relacionadas en función de síntomas entre las enfermedades 
    de la columna 1 y 3 situadas en la misma fila.
    
    Datos de entrada:
    1. spl_files: lista con el nombre de los archivos que contienen el SPL entre 
    todos los nodos de cada red (values). 
    2. id_enfermedades: código identificador de las enfermedades. 
    3. arch_candidatos: archivo que relaciona en cada fila la enfermedad neurológica de
    estudio, un fármaco candidato de reposicionamiento y una de las enfermedades que tratan
    dicho fármaco.
    """
    spls_gen = [] #lista en la que añadiré el SPL en la red proyectada de enfermedades - genes
    spls_dru = [] #lista en la que añadiré el SPL en la red proyectada de enfermedades - fármacos
    spls_sym =[] #lista en la que añadiré el SPL en la red proyectada de enfermedades - síntomas
    
    spl_enfermedad_gen, spl_enfermedad_dru, spl_enfermedad_sym = spl_files #asocio cada elemento de la lista spl_files a una variable
    
    df_resultados = arch_candidatos.copy() #copio el DataFrame con la información sobre los fármacos candidatos de reposicionamiento (arch_candidatos)
    
    for i, enfermedad in enumerate(arch_candidatos["Enfermedades de estudio"]): #para cada enfermedad neurológica de estudio
        id_enfermedad = id_enfermedades[enfermedad] #obtengo su id en el diccionario de identificadores de enfermedades
        
        fila_gen = spl_enfermedad_gen.columns.get_loc(id_enfermedad) #localizo el id en el archivo de fármacos candidatos
        columna_gen = arch_candidatos["Enfermedades candidatas"][i] #obtengo la enfermedad tratada por dicho fármaco en su misma fila
        
        if columna_gen in spl_enfermedad_gen.columns.tolist(): #si dicha enfermedad forma parte d ela red proyectada de enfermedades y genes
            spls_gen.append(spl_enfermedad_gen.loc[fila_gen, columna_gen]) #aádo a la lista spl_gen el SPL entre la enfermedad neurológica estudiada y la enfermedad tratada por el fármaco candidato de reposicionamiento. Para ello, localizo el SPL en la fila de la enfermedad estudiada y la columna de la enfermedad tratada por el fármaco.
        else:
            spls_gen.append("Unknown") #si la enfermedad tratada no está en la red pruyectada, añado "unknown" a la lista spl_gen
            
        #repito el proceso anterior para la red proyectada de enfermedades - fármacos
        fila_dru = spl_enfermedad_dru.columns.get_loc(id_enfermedad)
        columna_dru = arch_candidatos["Enfermedades candidatas"][i]
        
        if columna_dru in spl_enfermedad_dru.columns.tolist():
            spls_dru.append(spl_enfermedad_dru.loc[fila_dru, columna_dru])
        else:
            spls_dru.append("Unknown")
        
        #repito el proceso anterior para la red proyectada de enfermedades - síntomas
        fila_sym = spl_enfermedad_sym.columns.get_loc(id_enfermedad)
        columna_sym = arch_candidatos["Enfermedades candidatas"][i]
        
        if columna_sym in spl_enfermedad_sym.columns.tolist():
            spls_sym.append(spl_enfermedad_sym.loc[fila_sym, columna_sym])
        else:
            spls_sym.append("Unknown")
    
    #añado los spls entre las enfermedades en cada red al DataFrame de resultados
    df_resultados["SPL proyectada enf - gen"] = spls_gen
    df_resultados["SPL proyectada enf - dru"] = spls_dru
    df_resultados["SPL proyectada enf - sym"] = spls_sym
    
    return df_resultados

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

def distribucion_enf_candidatas(arch_candidatos, nombre_enf):
    """
    Esta función crea un DataFrame que contiene los fármacos candidatos de 
    reposcionamiento para una enfermedad neurológica, el total de enfermedades tratadas
    por dicho fármaco, y el número de estas enfermedades conectadas a la enfermedad
    neurológica en las redes proyectadas de enfermedades - genes, enfermedades - fármacos
    y enfermedades síntomas.
    
    Datos de entrada:
    1. arch_candidatos: DataFrame con las enfermedades neurológicas estudiadas, 
    sus fármacos candidatos de reposcionamiento, las enfermedades tratadas por cada uno de 
    los fármacos candidatos, y el SPL entre las enfermedades neurológicas y cada una de las 
    enfermedades tratadas por los fármacos candidatos en las tres redes proyectadas.
    2. nombre_enf: nombre de la enfermedad neurológica cuyos candidatos de reposicionamiento
    se encuentran en la lista anterior. 
    """
    
   # Filtrar las filas que incluyen enfermedades tratadas por fármacos candidatos de reposicionamiento para la enfermedad neurológica nombre_enf
    filt_enf = arch_candidatos[arch_candidatos["Enfermedades de estudio"] == nombre_enf]

    #Obtengo la lista de fármacos candidatos para reposicionamiento de nombre_enf
    lista_candidatos = set(filt_enf["Fármacos candidatos"])
    
    # Filtrar las filas que incluyen enfermedades tratadas por fármacos candidatos de reposicionamiento para la enfermedad neurológica nombre_enf que se conecten a nombre_enf en la red proyectada enf-genes.
    filt_gen = filt_enf[filt_enf["SPL proyectada enf - gen"] != "Unknown"]
                               
    # Filtrar las filas que incluyen enfermedades tratadas por fármacos candidatos de reposicionamiento para la enfermedad neurológica nombre_enf que se conecten a nombre_enf en la red proyectada enf-fármacos.
    filt_dru = filt_enf[filt_enf["SPL proyectada enf - dru"] != "Unknown"]
    
    # Filtrar las filas que incluyen enfermedades tratadas por fármacos candidatos de reposicionamiento para la enfermedad neurológica nombre_enf que se conecten a nombre_enf en la red proyectada enf-síntomas.
    filt_sym = filt_enf[filt_enf["SPL proyectada enf - sym"] != "Unknown"]
   
    #listas en la que añadiré el número de enfermedades tratadas por cada fármaco candidato que pueden conectarse a nombre_enf en cada red proyectada.
    enf_gen_lista = []  #conexión con nombre_enf en red proyectada enfermedades-genes.
    enf_dru_lista = []  #conexión con nombre_enf en red proyectada enfermedades-fármacos.
    enf_sym_lista = []  #conexión con nombre_enf en red proyectada enfermedades-síntomas.
    
    #listas en la que añadiré el número total de enfermedades tratadas por cada fármaco candidato que pueden conectarse a nombre_enf al menos en una de las redes proyectadas
    total_enf = []
                            
    for farmaco in list(lista_candidatos):  # para cada fármaco de la lista de candidatos de reposicionamiento para nombre_enf
        #cuento el número de enfermedades tratadas por uno de los fármacos candidatos que pueden conectarse a nombre_enf en cada red proyectada
        #el número de enfermedades para cada fármaco en cada caso se corresponde con el número de líneas en el que aparece dicho fármaco en los DataFrames filtrados
        enf_gen = len(filt_gen[filt_gen["Fármacos candidatos"] == farmaco])  #conexión con nombre_enf en red proyectada enfermedades-genes.
        enf_dru = len(filt_dru[filt_dru["Fármacos candidatos"] == farmaco])  #conexión con nombre_enf en red proyectada enfermedades-fármacos.
        enf_sym = len(filt_sym[filt_sym["Fármacos candidatos"] == farmaco])  #conexión con nombre_enf en red proyectada enfermedades-síntomas.
        enf_total = len(filt_enf[filt_enf["Fármacos candidatos"] == farmaco]) #total de enfermedades asociadas al fármaco
        # añado los resultados a las listas correspondientes para cada fármaco
        total_enf.append(enf_total)
        enf_gen_lista.append(enf_gen) 
        enf_dru_lista.append(enf_dru)
        enf_sym_lista.append(enf_sym)

    # creo el DataFrame de resultados
    df_candidatos = pd.DataFrame({
        "Fármacos candidatos " + str(nombre_enf): list(lista_candidatos),
        "Total enfermedades tratadas": total_enf,
        "Enfermedades en red proyectada enf-genes": enf_gen_lista,
        "Enfermedades en red proyectada enf-dru": enf_dru_lista,
        "Enfermedades en red proyectada enf-sym": enf_sym_lista,
     })

    return df_candidatos

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