from rdkit import Chem
from rdkit.Chem.AllChem import MolFromSmiles
from rdkit.Chem.Fingerprints.FingerprintMols import FingerprintMol
from rdkit.DataStructs import FingerprintSimilarity
import pandas as pd
import matplotlib.pyplot as plt
    
plot_colours = ['#F7977A', '#F9AD81', '#FDC68A', '#FFF79A', '#C4DF9B', 
    '#A2D39C', '#82CA9D', '#7BCDC8', '#6ECFF6', '#7EA7D8', '#8493CA',
    '#8882BE', '#A187BE', '#BC8DBF', '#F49AC2', '#F6989D']
    
def calculate_fps(df):
    # first convert SMILES to molecules
    df['mols'] = [MolFromSmiles(smiles) for smiles in df['SMILES']]
    # then convert molecules to fingerprints
    df['fps'] = [FingerprintMol(mol) for mol in df['mols']]
    
def generate_dissimilarities(leads, xenos):
    # generate dissimilarities
    leads['dissim'] = [{xeno['Name']: FingerprintSimilarity(xeno['fps'], lead['fps'])
                for _, xeno in xenos.iterrows()}
                for _, lead in leads.iterrows()]
    # find max similarity, sort by this
    leads['max_dissim'] = [max(ds.values()) for ds in leads['dissim']]
    leads.sort('max_dissim', ascending=False, inplace=True)

def plot_dissimilarities(row, label_x_distance = 0.025, label_y_distance = 0.5):
    # new plot
    fig = plt.figure(figsize=(8,4), dpi=96)
    ax = fig.add_subplot(111)

    # sort similarities
    dissim = pd.Series(row['dissim'])
    dissim.sort(ascending=False)
    

    # plot the axis
    plt.plot([0,1], [0,0], 'k|-', zorder=1)
    plt.text(0, -0.05, "0 = no similarity", 
        size=10, va='top', ha='left')
    plt.text(1, -0.05, "identical = 1", 
        size=10, va='top', ha='right')
    plt.text(dissim[0], -0.05, "best = %.2f" % dissim[0],
        size=10, va='top', ha='center')

    # plot labels
    label_last = 2
    n = len(dissim)
    for name, pt in dissim.iteritems():
        # calculate where to plot the label
        label_pt = pt
        if label_last - label_pt < label_x_distance:
            label_pt = label_last - label_x_distance
        label_last = label_pt

        if label_pt > 0:
            # plot the label line
            line_x = [pt,pt,label_pt,label_pt]
            line_y = [0,label_y_distance*n/len(dissim),
                        label_y_distance*n/len(dissim),label_y_distance]

            # plot the line, label
            plt.plot(line_x, line_y, plot_colours[n % len(plot_colours)],
                 zorder=1)
            plt.text(label_pt, line_y[-1], " "+name, rotation=90, 
                 size=8, va='bottom', ha='center')
        plt.plot(pt, 0, plot_colours[n % len(plot_colours)], marker='o', 
            zorder=1, ms=3)
        n-=1

    # finalise plot
    plt.xlim([-0.02,1.02])
    plt.ylim([-0.02,2])
    plt.axis('off')
    plt.tight_layout()
    
