Fix the coloring schems to make sure it is consistent and refactored _break_up_by_substrings function to improve understandability

Author: Nikhil172913832

manim/mobject/text/tex_mobject.py

@@ -363,74 +363,113 @@ def _break_up_by_substrings(self) -> Self:
                 curr_index = new_index
                 i += 1
             elif tex_string.strip().startswith(("^", "_")):
-                # Handle consecutive scripts as a group
-                script_group = [tex_string]
-                j = i + 1
-                while j < len(self.tex_strings) and self.tex_strings[
-                    j
-                ].strip().startswith(("^", "_")):
-                    script_group.append(self.tex_strings[j])
-                    j += 1
-
-                # Calculate total submobjects needed for all scripts
-                total_script_submobs = sum(
-                    len(
-                        SingleStringMathTex(
-                            s,
-                            tex_environment=self.tex_environment,
-                            tex_template=self.tex_template,
-                        ).submobjects
-                    )
-                    for s in script_group
-                )
+                # Handle consecutive scripts as a group, matching by Y-position
+                script_group, j = self._group_consecutive_scripts(i)
+                total_script_submobs = self._total_submobs_for_scripts(script_group)
+                script_pool = self.submobjects[curr_index:curr_index + total_script_submobs]
+
+                self._assign_script_group(script_group, script_pool, new_submobjects)
 
-                # Get the pool of available submobjects for all scripts
-                script_pool = self.submobjects[
-                    curr_index : curr_index + total_script_submobs
-                ]
-
-                # Process each script in the group
-                for script_tex in script_group:
-                    script_mob = SingleStringMathTex(
-                        script_tex,
-                        tex_environment=self.tex_environment,
-                        tex_template=self.tex_template,
-                    )
-                    script_num_submobs = len(script_mob.submobjects)
-
-                    if script_num_submobs > 0 and len(script_pool) > 0:
-                        # Select submobjects by Y position
-                        is_superscript = script_tex.strip().startswith("^")
-                        sorted_pool = sorted(
-                            script_pool,
-                            key=lambda mob: mob.get_center()[1],
-                            reverse=is_superscript,  # highest first for ^, lowest first for _
-                        )
-
-                        # Take the first script_num_submobs from sorted pool
-                        selected = sorted_pool[:script_num_submobs]
-                        script_mob.submobjects = selected
-
-                        # Remove selected submobjects from pool
-                        for sel in selected:
-                            if sel in script_pool:
-                                script_pool.remove(sel)
-
-                    new_submobjects.append(script_mob)
-
-                # Update indices
                 curr_index += total_script_submobs
-                i = j  # Skip past all processed scripts
+                i = j
             else:
-                # Normal (non-script) processing
-                sub_tex_mob.submobjects = self.submobjects[curr_index:new_index]
-                new_submobjects.append(sub_tex_mob)
-                curr_index = new_index
-                i += 1
+                # Base element processing: check if followed by scripts
+                next_is_script = (i + 1 < len(self.tex_strings) and
+                                self.tex_strings[i + 1].strip().startswith(("^", "_")))
+
+                if next_is_script and num_submobs > 0:
+                    script_group, j = self._group_consecutive_scripts(i + 1)
+                    total_script_submobs = self._total_submobs_for_scripts(script_group)
+                    total_needed = num_submobs + total_script_submobs
+
+                    all_submobs = self.submobjects[curr_index:curr_index + total_needed]
+
+                    # Only use special handling if scripts have content (non-empty)
+                    if total_script_submobs > 0 and len(all_submobs) == total_needed:
+                        # LaTeX may render base+scripts in unexpected order
+                        # Find base by Y-position: closest to baseline (Y=0)
+                        base_submob = min(all_submobs, key=lambda m: abs(m.get_center()[1]))
+                        sub_tex_mob.submobjects = [base_submob]
+
+                        script_pool = [m for m in all_submobs if m != base_submob]
+                        new_submobjects.append(sub_tex_mob)
+
+                        self._assign_script_group(script_group, script_pool, new_submobjects)
+
+                        curr_index += total_needed
+                        i = j
+                    else:
+                        # Fallback if counts don't match
+                        sub_tex_mob.submobjects = self.submobjects[curr_index:new_index]
+                        new_submobjects.append(sub_tex_mob)
+                        curr_index = new_index
+                        i += 1
+                else:
+                    sub_tex_mob.submobjects = self.submobjects[curr_index:new_index]
+                    new_submobjects.append(sub_tex_mob)
+                    curr_index = new_index
+                    i += 1
 
         self.submobjects = new_submobjects
         return self
 
+    def _group_consecutive_scripts(self, start_index: int) -> tuple[list[str], int]:
+        """Collect consecutive script tex_strings starting at ``start_index``.
+
+        Returns the list of scripts and the index just after the group.
+        Scripts are tex strings starting with '^' or '_'.
+        """
+        script_group = [self.tex_strings[start_index]]
+        j = start_index + 1
+        while j < len(self.tex_strings) and self.tex_strings[j].strip().startswith(("^", "_")):
+            script_group.append(self.tex_strings[j])
+            j += 1
+        return script_group, j
+
+    def _total_submobs_for_scripts(self, script_group: list[str]) -> int:
+        """Calculate total submobject count for a group of script strings.
+
+        Creates temporary SingleStringMathTex instances to inspect counts.
+        """
+        total = 0
+        for s in script_group:
+            total += len(
+                SingleStringMathTex(
+                    s, tex_environment=self.tex_environment, tex_template=self.tex_template
+                ).submobjects
+            )
+        return total
+
+    def _assign_script_group(self, script_group: list[str], script_pool: list[VMobject], new_submobjects: list[VMobject]) -> None:
+        """Assign submobjects from ``script_pool`` to scripts in ``script_group``.
+
+        Selection strategy:
+        - Superscripts ('^'): Pick highest Y-position items (above baseline)
+        - Subscripts ('_'): Pick lowest Y-position items (below baseline)
+
+        Selected submobjects are removed from pool to prevent reuse.
+        """
+        for script_tex in script_group:
+            script_mob = SingleStringMathTex(
+                script_tex, tex_environment=self.tex_environment, tex_template=self.tex_template
+            )
+            script_num_submobs = len(script_mob.submobjects)
+
+            if script_num_submobs > 0 and len(script_pool) > 0:
+                is_superscript = script_tex.strip().startswith("^")
+                # Sort by Y-position: reverse=True for superscripts (highest first)
+                sorted_pool = sorted(script_pool, key=lambda mob: mob.get_center()[1], reverse=is_superscript)
+
+                selected = sorted_pool[:script_num_submobs]
+                script_mob.submobjects = selected
+
+                # Remove selected items from pool
+                for sel in selected:
+                    if sel in script_pool:
+                        script_pool.remove(sel)
+
+            new_submobjects.append(script_mob)
+
     def get_parts_by_tex(
         self, tex: str, substring: bool = True, case_sensitive: bool = True
     ) -> VGroup: