# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

import rdkit
from rdkit import Chem
from rdkit.Chem import Draw
from rdkit.Chem.EState import Fingerprinter
from rdkit.Chem import Descriptors
from rdkit.Chem import rdFMCS
from rdkit.Chem.rdmolops import RDKFingerprint
from rdkit.Chem.Fingerprints import FingerprintMols
from rdkit import DataStructs
from rdkit.Avalon.pyAvalonTools import GetAvalonFP
import numpy as np
import pandas as pd
import pubchempy as pc


# +
def get_similarity(a,b):
    return DataStructs.FingerprintSimilarity(a, b, metric=DataStructs.TanimotoSimilarity)

def user_input_to_smiles(input_cid):
    """Takes a PubChem CID input and outputs the associated SMILES."""
    for compound in pc.get_compounds(input_cid):
        return compound.isomeric_smiles

# import ellie's function 
def input_to_compound(PubChemCID:str, master_df:pd.DataFrame):
    """
    input_to_compound() calculates the distance from input compound to 
        all products and create a new column. 
    Additionally, calculates the average distance between input compound
    and products for each promiscuous enzyme. 
    
    
    Args: 
        PubChemSID (str): contains PubChemSID to be queried
        master_df (pd.Dataframe): pre-calculated model dataframe 
               
    Returns:
        pd.Dataframe: master_df with two new columns 
        1) that contains distance from input compound to all products
        2) that contains avg distance from input compound to products for each promiscuous enzyme
    """
    
    input_SMILES = user_input_to_smiles(PubChemCID) #Ellie's function 
    input_fingerprint = FingerprintMols.FingerprintMol(Chem.rdmolfiles.MolFromSmiles(input_SMILES))
    dist_from_input = []
    avg_dist_from_input = []
    
    #get individual distance between input and all compounds 
    for i in master_df['Fingerprint']:
        dist_from_input.append(get_similarity(i,input_fingerprint))
    
    
    
    #get average distance between input and compounds for each promiscuous enzyme 
    for entry in np.unique(master_df['entry']):
        #filter dataframe with same enzyme 
        enzyme_split = master_df[master_df['entry']==entry]
        number_of_compounds = len(enzyme_split)
        total_distance = 0 
        
        for j in enzyme_split['Fingerprint']:
            total_distance += get_similarity(j,input_fingerprint)
        #get the average distance
        avg_distance = total_distance/number_of_compounds
        #create a list that contains number of compounds 
        avg_list = [avg_distance]*number_of_compounds
        avg_dist_from_input += avg_list 
        
        
    
    master_df['dist_from_input'] = dist_from_input 
    master_df['avg_dist_from_input'] = avg_dist_from_input 
    
    return master_df
    
    
# -

input_to_compound(243, master_df)

master_df = pd.read_csv("../datasets/MASTER_DF.csv")
df.head(10)

# +
# this dataframe does not makes sense! but I just wanted to see my above function works! 
master_avg_dist = []
for entry in np.unique(df.entry):
    enzyme_split = df[df['entry']==entry]
    number_of_compounds = len(enzyme_split)
    distance = 0 
    for j in enzyme_split['PubChem']:
        distance += j 
    avgdist = distance/number_of_compounds
    avgdistlist = [avgdist]*number_of_compounds
    master_avg_dist += avgdistlist

df['master_avg_dist'] = master_avg_dist
    
df
# -



StarcoderdataJupyterScriptsDedupFiltered