From 35f7af4e817caebd05cdda1b13ef9cc3a0ad5d70 Mon Sep 17 00:00:00 2001
From: LindeSchoenmaker <99398231+LindeSchoenmaker@users.noreply.github.com>
Date: Fri, 28 Jun 2024 17:21:31 +0200
Subject: [PATCH] Add other intermediate generator to wrapper

---
 .../intermediate_generators/intermediator.py  | 38 +++++++++++++++++--
 1 file changed, 35 insertions(+), 3 deletions(-)

diff --git a/src/konnektor/network_tools/intermediate_generators/intermediator.py b/src/konnektor/network_tools/intermediate_generators/intermediator.py
index 82242c9..c76e122 100644
--- a/src/konnektor/network_tools/intermediate_generators/intermediator.py
+++ b/src/konnektor/network_tools/intermediate_generators/intermediator.py
@@ -1,7 +1,11 @@
 
 from rdkit import Chem
 
-# Intermediate generator
+# R-group enumeration intermediate generator
+from rgroupinterm import rgroupenumeration
+from rgroupinterm import pruners
+
+# STONED SELFIES intermediate generator
 from generator import generation
 from generator import scoring
 #notes
@@ -12,8 +16,36 @@
 
 from gufe import SmallMoleculeComponent
 
-
 class Intermediator():
+    def __init__(self, enumerate_kekule=False, permutate=False, insert_small=False):
+        self.enumerate_kekule = enumerate_kekule
+        self.permutate = permutate
+        self.insert_small = insert_small 
+
+    def generate_intermediate(self, molA: SmallMoleculeComponent, molB: SmallMoleculeComponent):
+
+        rdmolA = Chem.MolFromSmiles(Chem.MolToSmiles(molA.to_rdkit()))
+        rdmolB = Chem.MolFromSmiles(Chem.MolToSmiles(molB.to_rdkit()))
+
+        generator = rgroupenumeration.EnumRGroups(self.enumerate_kekule, self.permutate, self.insert_small)
+        df_interm, _ = generator.generate_intermediates([rdmolA, rdmolB])
+
+        pruner = pruners.BasePruner(
+            [pruners.TanimotoScorer(transformer=pruners.HarmonicMeanTransformer(exponent=4))],
+            topn=1)
+        
+        df_interm['Parent_1'] = rdmolA
+        df_interm['Parent_2'] = rdmolB
+        df_interm['Pair'] = 0
+        pruned_df = pruner(df_interm)
+        rd_mol_intermediate = pruned_df['Intermediate'].values[0]
+        rd_mol_intermediate = Chem.AddHs(rd_mol_intermediate)
+        Chem.AllChem.EmbedMolecule(rd_mol_intermediate)
+
+        return SmallMoleculeComponent.from_rdkit(rdkit=rd_mol_intermediate,
+                                                 name=molA.name+"_"+molB.name+"_intermediate")
+        
+class Intermediator_STONEDSELFIES():
     def __init__(self, scoring_method=None, fp_type:str="ECFP4", num_tries:int=5,
                  num_random_smiles: int= 10, collect_bidirectional:bool=True,
                  exponent_path:int=4, n_rounds:int=1, num_random_samples:int=500,
@@ -87,4 +119,4 @@ def generate_intermediate(self, molA: SmallMoleculeComponent, molB: SmallMolecul
         selected_intermediate = next(iter(sorted_smiles_dict))
         rd_mol_intermediate = Chem.MolFromSmiles(selected_intermediate)
         return SmallMoleculeComponent.from_rdkit(rdkit=rd_mol_intermediate,
-                                                 name=molA.name+"_"+molB.name+"_intermediate")
\ No newline at end of file
+                                                 name=molA.name+"_"+molB.name+"_intermediate")