From 93a0d3324948e019b0ca819e8dac275c335a76e6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?M=C3=A1t=C3=A9=20Balajti?= <mate.balajti@unibas.ch>
Date: Fri, 19 Jan 2024 13:48:28 +0100
Subject: [PATCH] feat: add assumed mates #160

---
 htsinfer/get_library_type.py   | 18 +++++++++++++++---
 htsinfer/models.py             |  2 ++
 tests/test_get_library_type.py |  6 +++---
 3 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/htsinfer/get_library_type.py b/htsinfer/get_library_type.py
index e5fd2da..e4aea81 100644
--- a/htsinfer/get_library_type.py
+++ b/htsinfer/get_library_type.py
@@ -224,15 +224,15 @@ def _align_mates(self):
         LOGGER.debug(f"Number of aligned reads file 2: {aligned_mate2}")
         LOGGER.debug(f"Number of concordant reads: {concordant}")
 
-        self._update_relationship(
+        self._update_relationship_type(
             concordant, min(aligned_mate1, aligned_mate2)
         )
 
         samfile1.close()
         samfile2.close()
 
-    def _update_relationship(self, concordant, aligned_reads):
-        """Helper function to update relationship based on alignment."""
+    def _update_relationship_type(self, concordant, aligned_reads):
+        """Helper function to update relationship and type."""
         try:
             ratio = concordant / aligned_reads
         except ZeroDivisionError:
@@ -253,6 +253,18 @@ def _update_relationship(self, concordant, aligned_reads):
                 self.results.relationship = (
                     StatesTypeRelationship.not_mates
                 )
+        if self.results.relationship == (
+            StatesTypeRelationship.split_mates
+        ) and (
+            self.results.file_1 == StatesType.single and
+            self.results.file_2 == StatesType.single
+        ) or (
+            self.results.file_1 == StatesType.not_available and
+            self.results.file_2 == StatesType.not_available
+        ):
+            # Update first and second relationship
+            self.results.file_1 = StatesType.first_mate_assumed
+            self.results.file_2 = StatesType.second_mate_assumed
 
     class AlignedSegment:
         """Placeholder class for mypy "Missing attribute"
diff --git a/htsinfer/models.py b/htsinfer/models.py
index 6b8c2f9..eceaab1 100644
--- a/htsinfer/models.py
+++ b/htsinfer/models.py
@@ -155,9 +155,11 @@ class StatesType(Enum):
             that the library represents a single-end library.
     """
     first_mate = "first_mate"
+    first_mate_assumed = "first_mate_assumed"
     mixed_mates = "mixed_mates"
     not_available = None
     second_mate = "second_mate"
+    second_mate_assumed = "second_mate_assumed"
     single = "single"
 
 
diff --git a/tests/test_get_library_type.py b/tests/test_get_library_type.py
index e1a0dce..ef1e299 100644
--- a/tests/test_get_library_type.py
+++ b/tests/test_get_library_type.py
@@ -167,7 +167,7 @@ def test_evaluate_mate_relationship_not_available(self, tmpdir):
             StatesTypeRelationship.not_available
         )
 
-    def test_update_relationship_not_mates(self, tmpdir):
+    def test_update_relationship_type_not_mates(self, tmpdir):
         """Test update_relationship logic."""
         CONFIG.args.path_1_processed = FILE_IDS_NOT_MATCH_1
         CONFIG.args.path_2_processed = FILE_MATE_2
@@ -185,8 +185,8 @@ def test_update_relationship_not_mates(self, tmpdir):
         concordant = 0
         read_counter = 20
 
-        # Call the _update_relationship method
-        test_instance._update_relationship(concordant, read_counter)
+        # Call the _update_relationship_type method
+        test_instance._update_relationship_type(concordant, read_counter)
 
         assert (
             test_instance.results.relationship ==