From 231f0bccd8b0d8964b9de7ada434db17f6e0c2e7 Mon Sep 17 00:00:00 2001
From: Yannik Hampe <yannik.hampe@knecon.com>
Date: Wed, 17 Apr 2024 12:50:27 +0200
Subject: [PATCH] RED-8402: hotfix for fixing exception caused by layout parser
 changes

---
 .../service/document/NerEntitiesAdapter.java    | 17 +++++++++--------
 1 file changed, 9 insertions(+), 8 deletions(-)

diff --git a/redaction-service-v1/redaction-service-server-v1/src/main/java/com/iqser/red/service/redaction/v1/server/service/document/NerEntitiesAdapter.java b/redaction-service-v1/redaction-service-server-v1/src/main/java/com/iqser/red/service/redaction/v1/server/service/document/NerEntitiesAdapter.java
index 5fa03602..24e46ccf 100644
--- a/redaction-service-v1/redaction-service-server-v1/src/main/java/com/iqser/red/service/redaction/v1/server/service/document/NerEntitiesAdapter.java
+++ b/redaction-service-v1/redaction-service-server-v1/src/main/java/com/iqser/red/service/redaction/v1/server/service/document/NerEntitiesAdapter.java
@@ -12,7 +12,8 @@ import com.iqser.red.service.redaction.v1.server.client.model.NerEntitiesModel;
 import com.iqser.red.service.redaction.v1.server.model.NerEntities;
 import com.iqser.red.service.redaction.v1.server.model.document.TextRange;
 import com.iqser.red.service.redaction.v1.server.model.document.nodes.Document;
-import com.iqser.red.service.redaction.v1.server.model.document.nodes.Section;
+import com.iqser.red.service.redaction.v1.server.model.document.nodes.NodeType;
+import com.iqser.red.service.redaction.v1.server.model.document.nodes.SemanticNode;
 import com.iqser.red.service.redaction.v1.server.model.document.textblock.TextBlock;
 
 import lombok.AccessLevel;
@@ -43,7 +44,7 @@ public class NerEntitiesAdapter {
      */
     public NerEntities toNerEntities(NerEntitiesModel nerEntitiesModel, Document document) {
 
-        return new NerEntities(addOffsetsAndFlatten(getStringStartOffsetsForMainSections(document),
+        return new NerEntities(addOffsetsAndFlatten(getStringStartOffsetsForMainSectionsHeadersFooters(document),
                                                     nerEntitiesModel).map(nerEntityModel -> new NerEntities.NerEntity(nerEntityModel.getValue(),
                                                                                                                       new TextRange(nerEntityModel.getStartOffset(),
                                                                                                                                     nerEntityModel.getEndOffset()),
@@ -160,11 +161,11 @@ public class NerEntitiesAdapter {
     }
 
 
-    private static Stream<EntityRecognitionEntity> addOffsetsAndFlatten(List<Integer> stringOffsetsForMainSections, NerEntitiesModel nerEntitiesModel) {
+    private static Stream<EntityRecognitionEntity> addOffsetsAndFlatten(List<Integer> stringOffsetsForMainSectionsHeadersFooters, NerEntitiesModel nerEntitiesModel) {
 
         nerEntitiesModel.getData()
                 .forEach((sectionNumber, listOfNerEntities) -> listOfNerEntities.forEach(entityRecognitionEntity -> {
-                    int newStartOffset = entityRecognitionEntity.getStartOffset() + stringOffsetsForMainSections.get(sectionNumber);
+                    int newStartOffset = entityRecognitionEntity.getStartOffset() + stringOffsetsForMainSectionsHeadersFooters.get(sectionNumber);
                     entityRecognitionEntity.setStartOffset(newStartOffset);
                     entityRecognitionEntity.setEndOffset(newStartOffset + entityRecognitionEntity.getValue().length());
                 }));
@@ -174,11 +175,11 @@ public class NerEntitiesAdapter {
     }
 
 
-    private static List<Integer> getStringStartOffsetsForMainSections(Document document) {
+    private static List<Integer> getStringStartOffsetsForMainSectionsHeadersFooters(Document document) {
 
-        return document.getMainSections()
-                .stream()
-                .map(Section::getTextBlock)
+        return document.streamChildren()
+                .filter(child -> (child.getType().equals(NodeType.FOOTER) ||child.getType().equals(NodeType.HEADER) ||child.getType().equals(NodeType.SECTION)))
+                .map(SemanticNode::getTextBlock)
                 .map(TextBlock::getTextRange)
                 .map(TextRange::start)
                 .toList();