from rdkit import Chem
from rdkit.Chem import AllChem, DataStructs, MolToSmiles
import numpy as np
from igraph import Graph
import numpy as np
import logging
logging.basicConfig(level=logging.DEBUG)

threshold = 0.5

suppl = Chem.SDMolSupplier('syk.sdf')

fps = [AllChem.GetMorganFingerprintAsBitVect(mol, 2) for mol in suppl]
smiles = [MolToSmiles(mol) for mol in suppl]

mat = []
for i in range(len(fps)):
    mat.append([DataStructs.FingerprintSimilarity(fps[i],fps[j]) for j in range(len(fps))])

sm = np.array(mat)

def search_root(sm, remain):
    cand = None
    max_c = 0
    for l in remain:
        num_c = len(np.where(sm[l] > threshold)[0].tolist())
        if num_c > max_c:
            max_c = num_c
            cand = l
    return cand

def check_edge(sm, current, root):
    cand = None
    max_c = 0
    candidates = set(np.where(sm[root] > threshold)[0].tolist())
    for c in candidates:
        if c in current:
            if sm[root, c] > max_c:
                max_c = sm[root, c]
                cand = c
    return cand

edges = []
sim_edges = []
np.fill_diagonal(sm, 0)
remain = set(range(len(sm[0])))
current = set()
new = set()
ith = 1

while remain:
    logging.debug("### {} th ###".format(ith))
    ith += 1
    root = search_root(sm, remain)
    if root is None:
        break
    logging.debug("root: {}".format(root))

    cand = check_edge(sm, current, root)
    if cand:
        logging.debug("connect {} and {}".format(cand, root))
        edges.append((cand, root))
        current.remove(cand)
        sm[:, cand] = 0

    remain.remove(root)
    sm[:, root] = 0

    for i in np.where(sm[root] > threshold)[0].tolist():
        if i in remain:
            edges.append((root, i))
            sim_edges.append(sm[root, i])
            new.add(i)

    logging.debug("new nodes: {}".format(new))
    current |= new
    remain -= current
    new = set()

g = Graph(edges)
g.vs["smiles"] = smiles
g.es["similarity"] = sim_edges
g.save("test.gml")

from pychembldb import *
from copy import copy
from random import shuffle, random
from math import log10
#Inhibition of recombinant Syk
#Bioorg. Med. Chem. Lett. (2009) 19:1944-1949
assay = chembldb.query(Assay).filter_by(chembl_id="CHEMBL1022010").one()

with open("node.csv","w") as f:
    f.write('"ID","pIC50","ALOGP","MWT","SMILES"\n')

    for act in assay.activities:
        if act.standard_relation == "=":
            f.write('"{}",{},{},{},{}\n'.format(act.compound.molecule.chembl_id,
                                            9 - log10(act.standard_value),
                                            act.compound.molecule.property.alogp, 
                                            act.compound.molecule.property.mw_freebase,
                                            act.compound.molecule.structure.canonical_smiles))

with open("network.sif", "w") as f:
    smiles_list = [ act.compound.molecule.structure.canonical_smiles for act in assay.activities if act.standard_relation == "="]
    for i in range(len(smiles_list)):
        tl = copy(smiles_list)
        tl[i:i+1] = []
        shuffle(tl)
        f.write('{} randomnetwork {}\n'.format(smiles_list[i], " ".join(tl[:int(random() * 5)])))

from rdkit import Chem
from rdkit.Chem import AllChem, DataStructs
from sklearn.decomposition import PCA
from sklearn.cluster import AffinityPropagation
from ggplot import *
import numpy as np
import pandas as pd

suppl = Chem.SDMolSupplier('syk.sdf')

fps = []
for mol in suppl:
    fp = AllChem.GetMorganFingerprintAsBitVect(mol, 2)
    arr = np.zeros((1,))
    DataStructs.ConvertToNumpyArray(fp, arr)
    fps.append(arr)

print len(fps)
pca = PCA(n_components=2)
x = pca.fit(fps).transform(fps)
af = AffinityPropagation().fit(x)

d = pd.DataFrame(x)
d.columns = ["PCA1", "PCA2"]
d["labels"] = pd.Series(af.labels_)
g = ggplot(aes(x="PCA1", y="PCA2", color="labels"), data=d) + geom_point() + xlab("PCA1") + ylab("PCA2")
ggsave("ap.png", g)