From fd2089073275234c81dc88cec8d36a635d9f4172 Mon Sep 17 00:00:00 2001
From: PedroBarbosa <psbpedrobarbosa@gmail.com>
Date: Mon, 22 Apr 2024 21:12:58 +0100
Subject: [PATCH] Fix bug with individual genotype correction

---
 dress/datasetgeneration/config_evolution.py   |  2 +-
 .../grammars/pwm_perturbation_grammar.py      | 28 +++----------------
 2 files changed, 5 insertions(+), 25 deletions(-)

diff --git a/dress/datasetgeneration/config_evolution.py b/dress/datasetgeneration/config_evolution.py
index 3b64b98..a4eab2e 100644
--- a/dress/datasetgeneration/config_evolution.py
+++ b/dress/datasetgeneration/config_evolution.py
@@ -114,7 +114,7 @@ def correct_phenotypes(
                     genotype_to_phenotype=rep.genotype_to_phenotype,  # type: ignore
                 )
                 is_valid = (
-                    True if _is_valid_individual(ind, original_seq, excluded_regions, rs) else False
+                    True if _is_valid_individual(_ind, original_seq, excluded_regions, rs) else False
                 )
 
             new_pop.append(_ind)  # type: ignore
diff --git a/dress/datasetgeneration/grammars/pwm_perturbation_grammar.py b/dress/datasetgeneration/grammars/pwm_perturbation_grammar.py
index 76c903e..20d6b0f 100644
--- a/dress/datasetgeneration/grammars/pwm_perturbation_grammar.py
+++ b/dress/datasetgeneration/grammars/pwm_perturbation_grammar.py
@@ -152,7 +152,7 @@ def remove_diffunit_overlaps(
                 if isinstance(d, MotifSNV) and d.is_redundant(seq):
                     d.sample_new_nucleotide(rs)
 
-                if isinstance(d, MotifDeletion):
+                if isinstance(d, (MotifDeletion, MotifAblation)):
                     current_r = [
                         1,
                         d.position[0],
@@ -162,8 +162,8 @@ def remove_diffunit_overlaps(
                     ]
 
                 elif isinstance(
-                    d, (MotifInsertion, MotifSubstitution, MotifAblation, MotifSNV)
-                ):
+                    d, (MotifInsertion, MotifSubstitution, MotifSNV)
+                ):  
                     current_r = [1, d.position, d.position + 1, d.get_size(), str(d)]
 
                 else:
@@ -360,17 +360,7 @@ def perturb(self, seq: str, position_tracker: int) -> str:
             )
 
         def adjust_index(self, ss_indexes: List[list]) -> List[list]:
-            _ss_idx = list(itertools.chain(*ss_indexes))
-
-            adj = [
-                (
-                    ss - self.get_size()
-                    if isinstance(ss, int) and self.position[0] < ss
-                    else ss
-                )
-                for ss in _ss_idx
-            ]
-            return [adj[0:2], adj[2:4], adj[4:6]]
+            return ss_indexes
 
         def get_size(self) -> int:
             return self.position[1] - self.position[0]
@@ -417,16 +407,6 @@ def get_distance_to_cassette(self) -> int:
         def __str__(self):
             return f"MotifSubstitution[{self.position},{self.nucleotides[0]},{self.get_size()},{self.get_location()},{self.get_distance_to_cassette()}]"
 
-    #     grammar_nodes = []
-    # grammar_map = {'snv_weight': MotifSNV,
-    #                'deletion_weight': MotifDeletion,
-    #                'insertion_weight': MotifInsertion,
-    #                'motif_ablation_weight': MotifAblation,
-    #                'motif_substitution_weight': MotifSubstitution}
-
-    # for _arg, nodename in  grammar_map.items():
-    #     if kwargs[_arg] > 0:
-    #         grammar_nodes.append(weight(kwargs[_arg])(nodename))
     g = extract_grammar(
         [
             weight(kwargs["snv_weight"])(MotifSNV),