From 18bf5a130bb4e90c35a998e3ac3127ac4f23b2b7 Mon Sep 17 00:00:00 2001
From: Thomas Holder <thomas@thomas-holder.de>
Date: Tue, 19 Dec 2023 17:44:37 +0100
Subject: [PATCH] WIP: preserve_classified, mmCIF classify_entities

Use mmCIF polypeptide information for atom classification if
`preserve_classified` setting is on.

Limitations: Bypasses auto_show_classified

Related to https://github.com/schrodinger/pymol-open-source/issues/317
---
 layer1/SettingInfo.h         |  1 +
 layer2/CifMoleculeReader.cpp | 18 ++++++++++++++++++
 layer3/Selector.cpp          | 16 +++++++++++++++-
 3 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/layer1/SettingInfo.h b/layer1/SettingInfo.h
index 8ced50d08..a6eabf41b 100644
--- a/layer1/SettingInfo.h
+++ b/layer1/SettingInfo.h
@@ -896,6 +896,7 @@ enum {
   REC_i( 786, cartoon_smooth_cylinder_window          , global    , 2 ),
   REC_i( 787, isosurface_algorithm                    , global    , 0, 0, 2 ),
   REC_b( 788, cell_centered                           , global    , false ),
+  REC_b( 789, preserve_classified                     , global    , 0 ),
 
 
 #ifdef SETTINGINFO_IMPLEMENTATION
diff --git a/layer2/CifMoleculeReader.cpp b/layer2/CifMoleculeReader.cpp
index c60e8e844..2ca6b39e2 100644
--- a/layer2/CifMoleculeReader.cpp
+++ b/layer2/CifMoleculeReader.cpp
@@ -1413,6 +1413,22 @@ static void add_missing_ca_sub(PyMOLGlobals * G,
   }
 }
 
+/**
+ * Set the atom classification according to the entity annotation.
+ */
+static void classify_entities(
+    PyMOLGlobals* G, pymol::vla<AtomInfoType>& atInfo, CifContentInfo& info)
+{
+  for (int i = 0, n = atInfo.size(); i < n; ++i) {
+    auto& atom = atInfo[i];
+    const char* entity_id = LexStr(G, atom.custom);
+    if (info.is_polypeptide(entity_id)) {
+      assert(!(atom.flags & cAtomFlag_class));
+      atom.flags |= cAtomFlag_polymer | cAtomFlag_protein;
+    }
+  }
+}
+
 /**
  * Read missing residues / full sequence
  *
@@ -2125,6 +2141,8 @@ static ObjectMolecule *ObjectMoleculeReadCifData(PyMOLGlobals * G,
     if (!I->DiscreteFlag && !SettingGetGlobal_i(G, cSetting_retain_order)) {
       add_missing_ca(G, I->AtomInfo, info);
     }
+
+    classify_entities(G, I->AtomInfo, info);
   } else if ((csets = read_chem_comp_atom_model(G, datablock, &I->AtomInfo))) {
     info.type = CIF_CHEM_COMP;
   } else {
diff --git a/layer3/Selector.cpp b/layer3/Selector.cpp
index 95e0f2de6..0d05ebbdb 100644
--- a/layer3/Selector.cpp
+++ b/layer3/Selector.cpp
@@ -1082,6 +1082,9 @@ int SelectorClassifyAtoms(PyMOLGlobals * G, int sele, int preserve,
     n_dummies = cNDummyAtoms;
   }
   a = 0;
+
+  int n_preserved = 0;
+
   while(a < I->Table.size()) {
     obj = I->Obj[I->Table[a].model];
     at = I->Table[a].atom;
@@ -1119,7 +1122,11 @@ int SelectorClassifyAtoms(PyMOLGlobals * G, int sele, int preserve,
       a0++;
       a1--;
 
-      mask = 0;
+      mask = (ai->flags & cAtomFlag_class);
+
+      if (mask && SettingGet<bool>(G, cSetting_preserve_classified)) {
+        ++n_preserved;
+      } else
       if(!ai->hetatm && AtomInfoKnownProteinResName(LexStr(G, ai->resn)))
         mask = cAtomFlag_polymer | cAtomFlag_protein;
       else if(!ai->hetatm && AtomInfoKnownNucleicResName(LexStr(G, ai->resn)))
@@ -1352,6 +1359,13 @@ int SelectorClassifyAtoms(PyMOLGlobals * G, int sele, int preserve,
     }
     a++;
   }
+
+  if (n_preserved) {
+    PRINTFB(G, FB_Selector, FB_Details)
+    " %s-Detail: Preserved class flags on %d atoms\n", __func__,
+        n_preserved ENDFB(G);
+  }
+
   return true;
 }