#Libraries----------------------------------------------------------------------

"""
Dependencies and modules necessary for analytical functions to work
"""

#Cheminformatics
import rdkit
from rdkit import Chem
from rdkit.Chem import AllChem
from rdkit.Chem import Draw
from rdkit.Chem.Draw import IPythonConsole
from rdkit.Chem import Descriptors
from rdkit.Chem import PandasTools
from rdkit import DataStructs

#Data processing
import pandas as pd

import matplotlib.pyplot as plt
from matplotlib import gridspec

# Clustering
from scipy.cluster.hierarchy import dendrogram, linkage


#Similarity analysis----------------------------------------------------------------------


def similarity_dendogram(data):

    """
    Function takes the data file provided by the data_prep function and plots a dendogram based on compound similarity.
    Fingerprinting is based on Morgan fingerprints and similarity search is based on Tanimoto Similarity.
    #Input: data frame generated by data_prep function.
    #Output: dendogram and a data frame with similarity values.
    """

    radius=2
    nBits=1024

    #generate fingerprints for database of compounds
    ECFP6_data = [AllChem.GetMorganFingerprintAsBitVect(mol,radius=radius, nBits=nBits) for mol in data['Rdkit_mol']]
    data["Morgan_fpt"]=ECFP6_data


    #build array with similarity scores
    length=len(ECFP6_data)
    array=pd.DataFrame(index=range(length),columns=range(length))
    array.columns=list(data.CID)
    array.index=list(data.CID)

    linkage_array=np.empty(shape=(length,length))
    for i in range(length):
        var1=list(data.CID)[i]
        mol1=list(data.Morgan_fpt)[i]
        for j in range(length):
            #calculate similarity
            var2=list(data.CID)[j]
            mol2=list(data.Morgan_fpt)[j]
            similarity=DataStructs.FingerprintSimilarity(mol1,mol2)
            array.iloc[i,j]=similarity
            linkage_array[i,j]=similarity

    linked = linkage(linkage_array,'single')
    labelList = list(data.CID)


   
    plt.figure(figsize=(8,15))

    #Draw dendogram
    ax1=plt.subplot()
    plot=dendrogram(linked,  
                orientation='left',
                labels=labelList,
                distance_sort='descending',
                show_leaf_counts=True)

    ax1.spines['left'].set_visible(False)
    ax1.spines['top'].set_visible(False)
    ax1.spines['right'].set_visible(False)
    plt.title('Similarity clustering',fontsize=20,weight='bold')
    plt.tick_params ('both',width=2,labelsize=8)
    plt.tight_layout()
    plt.show() 
   
    return (array)