From 4017a468652e03317207c98f0c75ddc0a1261bf2 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Wed, 28 Apr 2021 15:46:27 +0100
Subject: [PATCH] Preserve molecule ordering when adding molecules. [ref #193]

---
 python/BioSimSpace/_SireWrappers/_molecule.py | 29 ++++++++++++++++++-
 .../BioSimSpace/_SireWrappers/_molecules.py   | 10 +++++--
 2 files changed, 35 insertions(+), 4 deletions(-)

diff --git a/python/BioSimSpace/_SireWrappers/_molecule.py b/python/BioSimSpace/_SireWrappers/_molecule.py
index 7186d45b3..2db4ba2ac 100644
--- a/python/BioSimSpace/_SireWrappers/_molecule.py
+++ b/python/BioSimSpace/_SireWrappers/_molecule.py
@@ -131,6 +131,11 @@ def __add__(self, other):
 
         # A single Molecule object.
         elif type(other) is Molecule:
+            mol = self.copy()
+            mol._renumber(_SireMol.MolNum.getUniqueNumber())
+            molecules[0] = mol
+            other = other.copy()
+            other._renumber(_SireMol.MolNum.getUniqueNumber())
             molecules.append(other)
 
         # A Molecules object.
@@ -139,7 +144,13 @@ def __add__(self, other):
 
         # A list of Molecule objects.
         elif type(other) is list and all(isinstance(x, Molecule) for x in other):
-            molecules.extend(other)
+            mol = self.copy()
+            mol._renumber(_SireMol.MolNum.getUniqueNumber())
+            molecules[0] = mol
+            for mol in other:
+                mol = mol.copy()
+                mol._renumber(_SireMol.MolNum.getUniqueNumber())
+                molecules.append(mol)
 
         # Unsupported.
         else:
@@ -1111,6 +1122,22 @@ def makeCompatibleWith(self, molecule, property_map={}, overwrite=True,
             # Finally, commit the changes to the internal object.
             self._sire_object = edit_mol.commit()
 
+    def _renumber(self, mol_num):
+        """Renumber the molecule with a unique MolNum.
+
+           Parameters
+           ----------
+
+           mol_num : Sire.Mol.MolNum
+               The molecule number.
+        """
+        if type(mol_num) is not _SireMol.MolNum:
+            raise TypeError("'mol_num' must be of type 'Sire.Mol.MolNum'")
+
+        edit_mol = self._sire_object.edit()
+        edit_mol.renumber(mol_num)
+        self._sire_object = edit_mol.commit()
+
     def _getPropertyMap0(self):
         """Generate a property map for the lambda = 0 state of the merged molecule."""
 
diff --git a/python/BioSimSpace/_SireWrappers/_molecules.py b/python/BioSimSpace/_SireWrappers/_molecules.py
index 85ea3d95f..448899f09 100644
--- a/python/BioSimSpace/_SireWrappers/_molecules.py
+++ b/python/BioSimSpace/_SireWrappers/_molecules.py
@@ -124,16 +124,20 @@ def __add__(self, other):
 
         # A Molecule object.
         elif type(other) is _Molecule:
-            molecules.add(other._sire_object)
+            mol = other.copy()
+            mol._renumber(_SireMol.MolNum.getUniqueNumber())
+            molecules.add(mol._sire_object)
 
         # A Molecules object.
         elif type(other) is Molecules:
             molecules.add(other._sire_object)
 
         # A list of Molecule objects.
-        elif type(other) is list and all(isinstance(x, Molecule) for x in other):
+        elif type(other) is list and all(isinstance(x, _Molecule) for x in other):
             for molecule in other:
-                molecules.add(molecule._sire_object)
+                mol = molecule.copy()
+                mol._renumber(_SireMol.MolNum.getUniqueNumber())
+                molecules.add(mol._sire_object)
 
         # Unsupported.
         else: