{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import torch\n",
    "from rdkit.Chem import MolFromSmiles\n",
    "from rdkit import Chem\n",
    "from rdkit.Chem import AllChem\n",
    "import pickle as pkl\n",
    "degrees = [0, 1, 2, 3, 4, 5]\n",
    "\n",
    "def bond_features(bond, use_chirality=True):\n",
    "    bt = bond.GetBondType()\n",
    "    bond_feats = [\n",
    "        bt == Chem.rdchem.BondType.SINGLE, bt == Chem.rdchem.BondType.DOUBLE,\n",
    "        bt == Chem.rdchem.BondType.TRIPLE, bt == Chem.rdchem.BondType.AROMATIC,\n",
    "        bond.GetIsConjugated(),\n",
    "        bond.IsInRing()\n",
    "    ]\n",
    "    if use_chirality:\n",
    "        bond_feats = bond_feats + one_of_k_encoding_unk(\n",
    "            str(bond.GetStereo()),\n",
    "            [\"STEREONONE\", \"STEREOANY\", \"STEREOZ\", \"STEREOE\"])\n",
    "    return np.array(bond_feats)\n",
    "\n",
    "def one_of_k_encoding(x, allowable_set):\n",
    "    if x not in allowable_set:\n",
    "        raise Exception(\"input {0} not in allowable set{1}:\".format(\n",
    "            x, allowable_set))\n",
    "    return [x == s for s in allowable_set]\n",
    "\n",
    "\n",
    "def one_of_k_encoding_unk(x, allowable_set):\n",
    "    \"\"\"Maps inputs not in the allowable set to the last element.\"\"\"\n",
    "    if x not in allowable_set:\n",
    "        x = allowable_set[-1]\n",
    "    return [x == s for s in allowable_set]\n",
    "\n",
    "\n",
    "def atom_features(atom,\n",
    "                  bool_id_feat=False,\n",
    "                  explicit_H=False,\n",
    "                  use_chirality=True):\n",
    "    results = one_of_k_encoding_unk(\n",
    "        atom.GetSymbol(),\n",
    "        [\n",
    "        'B',\n",
    "        'C',\n",
    "        'N',\n",
    "        'O',\n",
    "        'F',\n",
    "        'Si',\n",
    "        'P',\n",
    "        'S',\n",
    "        'Cl',\n",
    "        'As',\n",
    "        'Se',\n",
    "        'Br',\n",
    "        'Te',\n",
    "        'I',\n",
    "        'At',\n",
    "        'other'\n",
    "        ]) + one_of_k_encoding(atom.GetDegree(),\n",
    "                                [0, 1, 2, 3, 4, 5]) + \\\n",
    "                [atom.GetFormalCharge(), atom.GetNumRadicalElectrons()] + \\\n",
    "                one_of_k_encoding_unk(atom.GetHybridization(), [\n",
    "                Chem.rdchem.HybridizationType.SP, Chem.rdchem.HybridizationType.SP2,\n",
    "                Chem.rdchem.HybridizationType.SP3, Chem.rdchem.HybridizationType.\n",
    "                                    SP3D, Chem.rdchem.HybridizationType.SP3D2,'other'\n",
    "                ]) + [atom.GetIsAromatic()]\n",
    "    # In case of explicit hydrogen(QM8, QM9), avoid calling `GetTotalNumHs`\n",
    "    if not explicit_H:\n",
    "        results = results + one_of_k_encoding_unk(atom.GetTotalNumHs(),\n",
    "                                                    [0, 1, 2, 3, 4])\n",
    "    if use_chirality:\n",
    "        try:\n",
    "            results = results + one_of_k_encoding_unk(\n",
    "                atom.GetProp('_CIPCode'),\n",
    "                ['R', 'S']) + [atom.HasProp('_ChiralityPossible')]\n",
    "        except:\n",
    "            results = results + [False, False\n",
    "                                    ] + [atom.HasProp('_ChiralityPossible')]\n",
    "\n",
    "    return np.array(results)\n",
    "\n",
    "class Node(object):\n",
    "    __slots__ = ['ntype', 'features', '_neighbors', 'rdkit_ix']\n",
    "    def __init__(self, ntype, features, rdkit_ix):\n",
    "        self.ntype = ntype\n",
    "        self.features = features\n",
    "        self._neighbors = []\n",
    "        self.rdkit_ix = rdkit_ix\n",
    "\n",
    "    def add_neighbors(self, neighbor_list):\n",
    "        for neighbor in neighbor_list:\n",
    "            self._neighbors.append(neighbor)\n",
    "            neighbor._neighbors.append(self)\n",
    "\n",
    "    def get_neighbors(self, ntype):\n",
    "        return [n for n in self._neighbors if n.ntype == ntype]\n",
    "    \n",
    "class MolGraph(object):\n",
    "    def __init__(self):\n",
    "        self.nodes = {} # dict of lists of nodes, keyed by node type\n",
    "\n",
    "    def new_node(self, ntype, features=None, rdkit_ix=None):\n",
    "        new_node = Node(ntype, features, rdkit_ix)\n",
    "        self.nodes.setdefault(ntype, []).append(new_node)\n",
    "        return new_node\n",
    "\n",
    "    def add_subgraph(self, subgraph):\n",
    "        old_nodes = self.nodes\n",
    "        new_nodes = subgraph.nodes\n",
    "        for ntype in set(old_nodes.keys()) | set(new_nodes.keys()):\n",
    "            old_nodes.setdefault(ntype, []).extend(new_nodes.get(ntype, []))\n",
    "\n",
    "    def sort_nodes_by_degree(self, ntype):\n",
    "        nodes_by_degree = {i : [] for i in degrees}\n",
    "        for node in self.nodes[ntype]:\n",
    "            nodes_by_degree[len(node.get_neighbors(ntype))].append(node)\n",
    "\n",
    "        new_nodes = []\n",
    "        for degree in degrees:\n",
    "            cur_nodes = nodes_by_degree[degree]\n",
    "            self.nodes[(ntype, degree)] = cur_nodes\n",
    "            new_nodes.extend(cur_nodes)\n",
    "\n",
    "        self.nodes[ntype] = new_nodes\n",
    "\n",
    "    def feature_array(self, ntype):\n",
    "        assert ntype in self.nodes\n",
    "        return np.array([node.features for node in self.nodes[ntype]])\n",
    "\n",
    "    def rdkit_ix_array(self):\n",
    "        return np.array([node.rdkit_ix for node in self.nodes['atom']])\n",
    "\n",
    "    def neighbor_list(self, self_ntype, neighbor_ntype):\n",
    "        assert self_ntype in self.nodes and neighbor_ntype in self.nodes\n",
    "        neighbor_idxs = {n : i for i, n in enumerate(self.nodes[neighbor_ntype])}\n",
    "        return [[neighbor_idxs[neighbor]\n",
    "                 for neighbor in self_node.get_neighbors(neighbor_ntype)]\n",
    "                for self_node in self.nodes[self_ntype]]\n",
    "def array_rep_from_smiles(molgraph):\n",
    "    \"\"\"Precompute everything we need from MolGraph so that we can free the memory asap.\"\"\"\n",
    "    #molgraph = graph_from_smiles_tuple(tuple(smiles))\n",
    "    degrees = [0,1,2,3,4,5]\n",
    "    arrayrep = {'atom_features' : molgraph.feature_array('atom'),\n",
    "                'bond_features' : molgraph.feature_array('bond'),\n",
    "                'atom_list'     : molgraph.neighbor_list('molecule', 'atom'),\n",
    "                'rdkit_ix'      : molgraph.rdkit_ix_array()}\n",
    "\n",
    "    for degree in degrees:\n",
    "        arrayrep[('atom_neighbors', degree)] = \\\n",
    "            np.array(molgraph.neighbor_list(('atom', degree), 'atom'), dtype=int)\n",
    "        arrayrep[('bond_neighbors', degree)] = \\\n",
    "            np.array(molgraph.neighbor_list('bond', 'atom'), dtype=int)\n",
    "        # print(arrayrep[('bond_neighbors', degree)])\n",
    "        # print(len(arrayrep[('bond_neighbors', degree)]))\n",
    "    return arrayrep\n",
    "def graph_from_smiles(smiles):\n",
    "    graph = MolGraph()\n",
    "    mol = MolFromSmiles(smiles)\n",
    "    if not mol:\n",
    "        raise ValueError(\"Could not parse SMILES string:\", smiles)\n",
    "    atoms_by_rd_idx = {}\n",
    "    for atom in mol.GetAtoms():\n",
    "        new_atom_node = graph.new_node('atom', features=atom_features(atom), rdkit_ix=atom.GetIdx())\n",
    "        atoms_by_rd_idx[atom.GetIdx()] = new_atom_node\n",
    "\n",
    "    for bond in mol.GetBonds():\n",
    "        atom1_node = atoms_by_rd_idx[bond.GetBeginAtom().GetIdx()]\n",
    "        atom2_node = atoms_by_rd_idx[bond.GetEndAtom().GetIdx()]\n",
    "        new_bond_node = graph.new_node('bond', features=bond_features(bond))\n",
    "        new_bond_node.add_neighbors((atom1_node, atom2_node))\n",
    "        atom1_node.add_neighbors((atom2_node,))\n",
    "\n",
    "    mol_node = graph.new_node('molecule')\n",
    "    mol_node.add_neighbors(graph.nodes['atom'])\n",
    "    return graph\n",
    "def gen_descriptor_data(smilesList):\n",
    "\n",
    "    smiles_to_fingerprint_array = {}\n",
    "\n",
    "    for i,smiles in enumerate(smilesList):\n",
    "#         if i > 5:\n",
    "#             print(\"Due to the limited computational resource, submission with more than 5 molecules will not be processed\")\n",
    "#             break\n",
    "        smiles = Chem.MolToSmiles(Chem.MolFromSmiles(smiles), isomericSmiles=True)\n",
    "        molgraph = graph_from_smiles(smiles)\n",
    "        molgraph.sort_nodes_by_degree('atom')\n",
    "        arrayrep = array_rep_from_smiles(molgraph)\n",
    "\n",
    "        smiles_to_fingerprint_array[smiles] = arrayrep\n",
    "    return smiles_to_fingerprint_array\n",
    "def save_smiles_dicts(smilesList,filename):\n",
    "    #first need to get the max atom length\n",
    "    max_atom_len = 0\n",
    "    max_bond_len = 0\n",
    "    num_atom_features = 0\n",
    "    num_bond_features = 0\n",
    "    smiles_to_rdkit_list = {}\n",
    "\n",
    "    smiles_to_fingerprint_features = gen_descriptor_data(smilesList)\n",
    "\n",
    "    for smiles,arrayrep in smiles_to_fingerprint_features.items():\n",
    "        atom_features = arrayrep['atom_features']\n",
    "        bond_features = arrayrep['bond_features']\n",
    "\n",
    "        rdkit_list = arrayrep['rdkit_ix']\n",
    "        smiles_to_rdkit_list[smiles] = rdkit_list\n",
    "\n",
    "        atom_len,num_atom_features = atom_features.shape\n",
    "        bond_len,num_bond_features = bond_features.shape\n",
    "\n",
    "        if atom_len > max_atom_len:\n",
    "            max_atom_len = atom_len\n",
    "        if bond_len > max_bond_len:\n",
    "            max_bond_len = bond_len\n",
    "\n",
    "    #then add 1 so I can zero pad everything\n",
    "    \n",
    "    max_atom_len = 115\n",
    "    max_bond_len = 118\n",
    "    \n",
    "    max_atom_index_num = max_atom_len\n",
    "    max_bond_index_num = max_bond_len\n",
    "\n",
    "    max_atom_len += 1\n",
    "    max_bond_len += 1\n",
    "\n",
    "    smiles_to_atom_info = {}\n",
    "    smiles_to_bond_info = {}\n",
    "\n",
    "    smiles_to_atom_neighbors = {}\n",
    "    smiles_to_bond_neighbors = {}\n",
    "\n",
    "    smiles_to_atom_mask = {}\n",
    "\n",
    "    degrees = [0,1,2,3,4,5]\n",
    "    #then run through our numpy array again\n",
    "    for smiles,arrayrep in smiles_to_fingerprint_features.items():\n",
    "        mask = np.zeros((max_atom_len))\n",
    "\n",
    "        #get the basic info of what\n",
    "        #    my atoms and bonds are initialized\n",
    "        atoms = np.zeros((max_atom_len,num_atom_features))\n",
    "        bonds = np.zeros((max_bond_len,num_bond_features))\n",
    "\n",
    "        #then get the arrays initlialized for the neighbors\n",
    "        atom_neighbors = np.zeros((max_atom_len,len(degrees)))\n",
    "        bond_neighbors = np.zeros((max_bond_len,2))\n",
    "\n",
    "        #now set these all to the last element of the list, which is zero padded\n",
    "        atom_neighbors.fill(max_atom_index_num)\n",
    "        bond_neighbors.fill(max_bond_index_num)\n",
    "\n",
    "        atom_features = arrayrep['atom_features']\n",
    "        bond_features = arrayrep['bond_features']\n",
    "\n",
    "        for i,feature in enumerate(atom_features):\n",
    "            mask[i] = 1.0\n",
    "            atoms[i] = feature\n",
    "\n",
    "        for j,feature in enumerate(bond_features):\n",
    "            bonds[j] = feature\n",
    "\n",
    "        atom_neighbor_count = 0\n",
    "        bond_neighbor_count = 0\n",
    "        working_atom_list = []\n",
    "        working_bond_list = []\n",
    "        for degree in degrees:\n",
    "            atom_neighbors_list = arrayrep[('atom_neighbors', degree)]\n",
    "            bond_neighbors_list = arrayrep[('bond_neighbors', degree)]\n",
    "\n",
    "            if len(atom_neighbors_list) > 0:\n",
    "\n",
    "                for i,degree_array in enumerate(atom_neighbors_list):\n",
    "                    for j,value in enumerate(degree_array):\n",
    "                        atom_neighbors[atom_neighbor_count,j] = value\n",
    "                    atom_neighbor_count += 1\n",
    "    \n",
    "        \n",
    "        if len(bond_neighbors_list) > 0:\n",
    "            for i,degree_array in enumerate(bond_neighbors_list):\n",
    "                for j,value in enumerate(degree_array):\n",
    "                    bond_neighbors[bond_neighbor_count,j] = value\n",
    "                bond_neighbor_count += 1\n",
    "        # print(bond_neighbors)\n",
    "\n",
    "        #then add everything to my arrays\n",
    "        smiles_to_atom_info[smiles] = atoms\n",
    "        smiles_to_bond_info[smiles] = bonds\n",
    "\n",
    "        smiles_to_atom_neighbors[smiles] = atom_neighbors\n",
    "        smiles_to_bond_neighbors[smiles] = bond_neighbors\n",
    "        \n",
    "        smiles_to_atom_mask[smiles] = mask\n",
    "\n",
    "    del smiles_to_fingerprint_features\n",
    "    feature_dicts = {}\n",
    "#     feature_dicts['smiles_to_atom_mask'] = smiles_to_atom_mask\n",
    "#     feature_dicts['smiles_to_atom_info']= smiles_to_atom_info\n",
    "    feature_dicts = {\n",
    "        'smiles_to_atom_mask': smiles_to_atom_mask,\n",
    "        'smiles_to_atom_info': smiles_to_atom_info,\n",
    "        'smiles_to_bond_info': smiles_to_bond_info,\n",
    "        'smiles_to_atom_neighbors': smiles_to_atom_neighbors,\n",
    "        'smiles_to_bond_neighbors': smiles_to_bond_neighbors,\n",
    "        'smiles_to_rdkit_list': smiles_to_rdkit_list\n",
    "    }\n",
    "    # if os.path.isdir(filename):\n",
    "    # os.makedirs(os.path.dirname(filename),exist_ok=True)\n",
    "    pkl.dump(feature_dicts,open(filename,\"wb\"))\n",
    "    print('feature dicts file saved as '+ filename)\n",
    "    return feature_dicts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Index(['ACH', 'S_max', 'S_mean', 'S_sum', 'cell_line', 'cell_line_name',\n",
       "       'chembl_id_col', 'chembl_id_row', 'css_col', 'css_ri', 'css_row',\n",
       "       'drug_col', 'drug_row', 'ic50_col', 'ic50_row', 'id',\n",
       "       'molecule_chembl_id', 'molecule_structures_col',\n",
       "       'molecule_structures_row', 'ri_col', 'ri_row', 'schembl_id_col',\n",
       "       'schembl_id_row', 'study_name', 'synergy_bliss', 'synergy_hsa',\n",
       "       'synergy_loewe', 'synergy_zip', 'tissue_name'],\n",
       "      dtype='object')"
      ]
     },
     "execution_count": 60,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "\n",
    "dataframe = pd.read_pickle(f'{EXPORT_DATA_PATH}/data_drugcomb.pkl')\n",
    "\n",
    "dataframe.columns\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "feature dicts file saved as ../data/smiles_graph_data.pkl\n"
     ]
    }
   ],
   "source": [
    "smiles = dataframe['molecule_structures_col'].dropna().values\n",
    "smiles = np.append(smiles,dataframe['molecule_structures_row'].dropna().values)\n",
    "smiles_unique = np.unique(smiles)\n",
    "_ = save_smiles_dicts(smiles_unique, f'{EXPORT_DATA_PATH}/smiles_graph_data.pkl')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
