From bf48ba703589d85379ac2aa2514a3118fb78d951 Mon Sep 17 00:00:00 2001
From: valasatava <yana.valasatava@gmail.com>
Date: Wed, 6 Nov 2024 17:53:56 -0800
Subject: [PATCH] allow retrieving all similarity cutoffs

---
 .../data/repository/GroupRepository.java      | 54 ++++++++++---------
 1 file changed, 29 insertions(+), 25 deletions(-)

diff --git a/src/main/java/org/rcsb/idmapper/backend/data/repository/GroupRepository.java b/src/main/java/org/rcsb/idmapper/backend/data/repository/GroupRepository.java
index a10dfeb..2cc1bbb 100644
--- a/src/main/java/org/rcsb/idmapper/backend/data/repository/GroupRepository.java
+++ b/src/main/java/org/rcsb/idmapper/backend/data/repository/GroupRepository.java
@@ -17,24 +17,16 @@ public class GroupRepository extends AnyRepository {
     private final Multimap<String, String> groupToProvenance = HashMultimap.create(EXPECTED_KEYS, EXPECTED_VALUES_PER_KEY);
 
     private final Map<Input.AggregationMethod,
-            Multimap<String, String>> identity = new HashMap<>(); // members ID -> group IDs
-
-    private final Map<Input.AggregationMethod,
-            Map<Integer, // similarity cutoff
-            Multimap<String, String>>> similarity = new HashMap<>(); // members ID -> group IDs
+            Map<Integer, // similarity cutoff (can be NULL for identity-based groups)
+            Multimap<String, String>>> memberToGroup = new HashMap<>(); // members ID -> group IDs
 
     public void addGroupProvenance(String groupId, String provenanceId) {
         groupToProvenance.put(groupId, provenanceId);
     }
 
     public void addGroupMembers(Input.AggregationMethod method, Integer cutoff, String gId, List<String> mIds) {
-        Multimap<String, String> map;
-        if (cutoff != null) {
-            map = similarity.computeIfAbsent(method, k -> new HashMap<>())
-                    .computeIfAbsent(cutoff, k -> HashMultimap.create());
-        } else {
-            map = identity.computeIfAbsent(method, k -> HashMultimap.create());
-        }
+        Multimap<String, String> map = memberToGroup.computeIfAbsent(method, k -> new HashMap<>())
+                .computeIfAbsent(cutoff, k -> HashMultimap.create());
         mIds.forEach(id -> map.put(id, gId));
     }
 
@@ -42,22 +34,34 @@ public Collection<String> getGroupToProvenance(String groupId) {
         return groupToProvenance.get(groupId);
     }
 
-    public Collection<String> getMemberToGroup(Input.AggregationMethod method, Integer cutoff, String mId) {
-        if (cutoff != null)
-            return similarity.get(method).get(cutoff).get(mId);
-        else
-            return identity.get(method).asMap().getOrDefault(mId, new ArrayList<>());
+    private boolean isSimilarityBasedGrouping(Input.AggregationMethod method) {
+        return method == Input.AggregationMethod.sequence_identity;
     }
 
-    public Long countGroups(Input.AggregationMethod method) {
-        if (method == Input.AggregationMethod.sequence_identity) {
-            AtomicReference<Long> count = new AtomicReference<>(0L);
-            similarity.get(method).keySet()
-                    .forEach(cutoff -> count.updateAndGet(v -> v + similarity.get(method).get(cutoff)
-                            .values().stream().distinct().count()));
-            return count.get();
+    public Collection<String> getMemberToGroup(Input.AggregationMethod method, Integer cutoff, String mId) {
+        if (isSimilarityBasedGrouping(method)) {
+            if (cutoff == null) {
+                List<String> all = new ArrayList<>();
+                memberToGroup.get(method).keySet()
+                        .forEach(c -> {
+                            Multimap<String, String> map = memberToGroup.get(method).get(c);
+                            if (map.containsKey(mId)) {
+                                all.addAll(map.get(mId));
+                            }
+                        });
+                return all;
+            }
+            else return memberToGroup.get(method).get(cutoff).get(mId);
         } else {
-            return identity.get(method).values().stream().distinct().count();
+            return memberToGroup.get(method).get(null).get(mId);
         }
     }
+
+    public Long countGroups(Input.AggregationMethod method) {
+        AtomicReference<Long> count = new AtomicReference<>(0L);
+        memberToGroup.get(method).keySet()
+                .forEach(cutoff -> count.updateAndGet(v -> v + memberToGroup.get(method).get(cutoff)
+                        .values().stream().distinct().count()));
+        return count.get();
+    }
 }