From b4f079c3c2a5b26828a81b1c950d0f7e40c6314a Mon Sep 17 00:00:00 2001
From: Philipp Schramm <Philipp.Schramm@iqser.com>
Date: Thu, 27 Oct 2022 14:15:42 +0200
Subject: [PATCH] RSS-118: Refactored method in Section for adding AI entries

---
 .../v1/server/redaction/model/Section.java    | 87 +++++++++++++------
 1 file changed, 61 insertions(+), 26 deletions(-)

diff --git a/redaction-service-v1/redaction-service-server-v1/src/main/java/com/iqser/red/service/redaction/v1/server/redaction/model/Section.java b/redaction-service-v1/redaction-service-server-v1/src/main/java/com/iqser/red/service/redaction/v1/server/redaction/model/Section.java
index ba4d3f94..0a84f380 100644
--- a/redaction-service-v1/redaction-service-server-v1/src/main/java/com/iqser/red/service/redaction/v1/server/redaction/model/Section.java
+++ b/redaction-service-v1/redaction-service-server-v1/src/main/java/com/iqser/red/service/redaction/v1/server/redaction/model/Section.java
@@ -80,39 +80,29 @@ public class Section {
     private ManualRedactions manualRedactions;
 
 
+    @Deprecated
     @SuppressWarnings("unused")
     @ThenAction
     public void addAiEntities(@Argument(ArgumentType.TYPE) String type, @Argument(ArgumentType.TYPE) String asType) {
+        redactOrRecommendAiEntities(type, asType, false, 0, null, null);
+    }
 
-        Set<Entity> entitiesOfType = nerEntities.stream().filter(nerEntity -> nerEntity.getType().equals(type)).collect(Collectors.toSet());
-        List<String> values = entitiesOfType.stream().map(Entity::getWord).collect(Collectors.toList());
-        Set<Entity> found = EntitySearchUtils.findEntities(searchText,
-                new SearchImplementation(values, dictionary.isCaseInsensitiveDictionary(asType)),
-                dictionary.getType(asType),
-                new FindEntityDetails(asType, headline, sectionNumber, false, false, Engine.NER, EntityType.RECOMMENDATION));
-        EntitySearchUtils.clearAndFindPositions(found, searchableText, dictionary, manualRedactions);
-        found = found.stream().filter(f -> !f.isFalsePositive()).collect(Collectors.toSet());
 
-        Set<Entity> finalResult = new HashSet<>();
+    @SuppressWarnings("unused")
+    @ThenAction
+    public void recommendAiEntities(@Argument(ArgumentType.TYPE) String type, @Argument(ArgumentType.TYPE) String asType) {
+        redactOrRecommendAiEntities(type, asType, false, 0, null, null);
+    }
 
-        // Only keep Entities with correct offsets from AI Service.
-        for (Entity current : found) {
-            boolean foundSameOffsets = false;
-            for (Entity entity : entitiesOfType) {
-                if (entity.getStart().equals(current.getStart()) && entity.getEnd().equals(current.getEnd())) {
-                    foundSameOffsets = true;
-                    break;
-                }
-            }
-            if (foundSameOffsets) {
-                finalResult.add(current);
-            }
-        }
+    @SuppressWarnings("unused")
+    @ThenAction
+    public void redactAiEntities(@Argument(ArgumentType.TYPE) String type,
+                                 @Argument(ArgumentType.TYPE) String asType,
+                                 @Argument(ArgumentType.RULE_NUMBER) int ruleNumber,
+                                 @Argument(ArgumentType.STRING) String reason,
+                                 @Argument(ArgumentType.LEGAL_BASIS) String legalBasis) {
 
-        var nonOverlapping = EntitySearchUtils.findNonOverlappingMatchEntities(entities, finalResult);
-        EntitySearchUtils.addEntitiesWithHigherRank(entities, nonOverlapping, dictionary);
-        EntitySearchUtils.removeEntitiesContainedInLarger(entities);
-        nerEntities.removeAll(entitiesOfType);
+        redactOrRecommendAiEntities(type, asType, true, ruleNumber, reason, legalBasis);
     }
 
 
@@ -1395,6 +1385,51 @@ public class Section {
     }
 
 
+    private void redactOrRecommendAiEntities(String type, String asType, boolean redact, int ruleNumber, String reason, String legalBasis) {
+
+        Set<Entity> entitiesOfType = nerEntities.stream().filter(nerEntity -> nerEntity.getType().equals(type)).collect(Collectors.toSet());
+        List<String> values = entitiesOfType.stream().map(Entity::getWord).collect(Collectors.toList());
+        Set<Entity> found = EntitySearchUtils.findEntities(searchText,
+                                                           new SearchImplementation(values, dictionary.isCaseInsensitiveDictionary(asType)),
+                                                           dictionary.getType(asType),
+                                                           new FindEntityDetails(asType, headline, sectionNumber, false, false, Engine.NER, EntityType.RECOMMENDATION));
+        EntitySearchUtils.clearAndFindPositions(found, searchableText, dictionary, manualRedactions);
+        found = found.stream().filter(f -> !f.isFalsePositive()).collect(Collectors.toSet());
+
+        Set<Entity> finalResult = new HashSet<>();
+
+        // Only keep Entities with correct offsets from AI Service.
+        for (Entity current : found) {
+            boolean foundSameOffsets = false;
+            for (Entity entity : entitiesOfType) {
+                if (entity.getStart().equals(current.getStart()) && entity.getEnd().equals(current.getEnd())) {
+                    foundSameOffsets = true;
+                    break;
+                }
+            }
+            if (foundSameOffsets) {
+                finalResult.add(current);
+            }
+        }
+
+        var nonOverlapping = EntitySearchUtils.findNonOverlappingMatchEntities(entities, finalResult);
+
+        if (redact) {
+            nonOverlapping.forEach(entity -> {
+                entity.setRedaction(true);
+                entity.setEntityType(EntityType.ENTITY);
+                entity.setMatchedRule(ruleNumber);
+                entity.setRedactionReason(reason);
+                entity.setLegalBasis(legalBasis);
+            });
+        }
+
+        EntitySearchUtils.addEntitiesWithHigherRank(entities, nonOverlapping, dictionary);
+        EntitySearchUtils.removeEntitiesContainedInLarger(entities);
+        nerEntities.removeAll(entitiesOfType);
+    }
+
+
     @Retention(RetentionPolicy.RUNTIME)
     @Target(ElementType.METHOD)
     public @interface WhenCondition {