Merge pull request #15 from Brain-Modulation-Lab/feat/segments-amplitude

enh: divide amplitute bet segments

Author: Lucia-Poma

src/clinical_dbs_annotator/ui/amplitude_split_widget.py

@@ -54,7 +54,7 @@ def __init__(self, amp_edit: QLineEdit, parent=None):
     # Public API
     # ------------------------------------------------------------------
 
-    def update_cathodes(self, cathode_labels: list[str]) -> None:
+    def update_cathodes(self, cathode_labels: list[str], is_single_grouped_directional: bool = False) -> None:
         """Rebuild the rows to match the current set of active cathodes.
 
         Called by the parent view whenever the electrode canvas changes.
@@ -68,8 +68,10 @@ def update_cathodes(self, cathode_labels: list[str]) -> None:
         # Filter out "case" — CASE as cathode does not get a percentage row
         cathode_labels = [lbl for lbl in cathode_labels if lbl.lower() != "case"]
 
-        if len(cathode_labels) <= 1:
-            # Single or no cathode → hide everything
+        # Use the passed parameter to determine if this is a single grouped directional contact
+
+        # Hide widget if no cathodes, or single non-grouped cathode
+        if len(cathode_labels) == 0 or (len(cathode_labels) == 1 and not is_single_grouped_directional):
             self._clear_rows()
             self._cathode_labels = []
             self._percentages.clear()
@@ -86,7 +88,11 @@ def update_cathodes(self, cathode_labels: list[str]) -> None:
 
         # Compute new default percentages
         self._cathode_labels = list(cathode_labels)
-        self._redistribute_percentages()
+        if len(cathode_labels) == 1 and is_single_grouped_directional:
+            # For single grouped directional contact, set 100% to enable segment splitting
+            self._percentages[cathode_labels[0]] = 100.0
+        else:
+            self._redistribute_percentages()
 
         # Rebuild UI rows
         self._rebuild_rows()
@@ -104,9 +110,12 @@ def get_amplitude_text(self) -> str:
           e.g. ``"1.5_1.0"``.
         """
         total_text = self._amp_edit.text().strip()
+
+        # For single or no cathode, return total amplitude
         if len(self._cathode_labels) <= 1:
             return total_text
 
+        # Multiple cathodes - use regular logic
         try:
             total_amp = float(total_text)
         except (ValueError, TypeError):
@@ -151,29 +160,102 @@ def _rebuild_rows(self) -> None:
         """Tear down and recreate all percentage rows."""
         self._clear_rows()
 
-        for i, lbl in enumerate(self._cathode_labels):
+        # Check if we have a single grouped contact (E1, E2, etc.)
+        if (len(self._cathode_labels) == 1 and
+            len(self._cathode_labels[0]) >= 2 and
+            self._cathode_labels[0][0] == 'E' and
+            self._cathode_labels[0][1:].isdigit()):
+            # Create segment rows for grouped contact
+            self._create_segment_rows_only(self._cathode_labels[0])
+        else:
+            # Create regular rows for individual contacts/segments
+            for i, lbl in enumerate(self._cathode_labels):
+                self._create_regular_row(lbl, i)
+
+        self._refresh_ma_values()
+
+    def _create_regular_row(self, lbl: str, index: int) -> None:
+        """Create a regular percentage row for a single contact."""
+        row_w = QWidget()
+        row_layout = QHBoxLayout(row_w)
+        row_layout.setContentsMargins(20, 0, 0, 0)
+        row_layout.setSpacing(4)
+
+        # Contact label
+        name_lbl = QLabel(f"<b>{lbl}</b>")
+        name_lbl.setMinimumWidth(40)
+        row_layout.addWidget(name_lbl)
+
+        # Percentage edit
+        pct_edit = QLineEdit()
+        pct_edit.setMaximumWidth(55)
+        pct_edit.setAlignment(Qt.AlignRight)
+        pct_edit.setValidator(QDoubleValidator(0.0, 100.0, 1))
+        pct_edit.setText(f"{self._percentages.get(lbl, 0.0):g}")
+        row_layout.addWidget(pct_edit)
+
+        pct_label = QLabel("%")
+        row_layout.addWidget(pct_label)
+
+        arrow_lbl = QLabel("->")
+        row_layout.addWidget(arrow_lbl)
+
+        # Computed mA label (read-only)
+        ma_label = QLabel("")
+        ma_label.setMinimumWidth(60)
+        ma_label.setStyleSheet("color: #64748b;")
+        row_layout.addWidget(ma_label)
+
+        row_layout.addStretch(1)
+
+        is_last = (index == len(self._cathode_labels) - 1)
+
+        # Connect editing
+        if is_last:
+            pct_edit.setReadOnly(True)
+            pct_edit.setStyleSheet("background: transparent; border: none; color: #64748b;")
+        else:
+            pct_edit.editingFinished.connect(
+                lambda lbl_=lbl: self._on_pct_edited(lbl_)
+            )
+            pct_edit.textChanged.connect(
+                lambda _text, lbl_=lbl: self._on_pct_text_changed(lbl_)
+            )
+
+        self._rows[lbl] = (row_w, pct_edit, ma_label)
+        self._layout.addWidget(row_w)
+
+    def _create_segment_rows_only(self, parent_lbl: str) -> None:
+        """Create only segment rows for a grouped directional contact (no main row)."""
+        seg_labels = ["a", "b", "c"]
+        parent_percentage = self._percentages.get(parent_lbl, 0.0)
+        segment_percentage = round(parent_percentage / 3.0, 1)  # Equal split by default, rounded to 1 decimal
+
+        for seg_idx, seg_label in enumerate(seg_labels):
+            seg_lbl = f"{parent_lbl}{seg_label}"
+
             row_w = QWidget()
             row_layout = QHBoxLayout(row_w)
-            row_layout.setContentsMargins(20, 0, 0, 0)
+            row_layout.setContentsMargins(20, 0, 0, 0)  # Regular indentation (not extra)
             row_layout.setSpacing(4)
 
-            # Contact label
-            name_lbl = QLabel(f"<b>{lbl}</b>")
+            # Segment label
+            name_lbl = QLabel(f"<b>{seg_lbl}</b>")
             name_lbl.setMinimumWidth(40)
             row_layout.addWidget(name_lbl)
 
-            # Percentage edit
+            # Percentage edit (editable for segments)
             pct_edit = QLineEdit()
             pct_edit.setMaximumWidth(55)
             pct_edit.setAlignment(Qt.AlignRight)
             pct_edit.setValidator(QDoubleValidator(0.0, 100.0, 1))
-            pct_edit.setText(f"{self._percentages.get(lbl, 0.0):g}")
+            pct_edit.setText(f"{segment_percentage:g}")
             row_layout.addWidget(pct_edit)
 
             pct_label = QLabel("%")
             row_layout.addWidget(pct_label)
 
-            arrow_lbl = QLabel("→")
+            arrow_lbl = QLabel("->")
             row_layout.addWidget(arrow_lbl)
 
             # Computed mA label (read-only)
@@ -184,24 +266,25 @@ def _rebuild_rows(self) -> None:
 
             row_layout.addStretch(1)
 
-            is_last = (i == len(self._cathode_labels) - 1)
+            # Store segment percentages
+            self._percentages[seg_lbl] = segment_percentage
 
-            # Connect editing
+            # Connect editing for segment rows
+            is_last = (seg_idx == len(seg_labels) - 1)
             if is_last:
                 pct_edit.setReadOnly(True)
                 pct_edit.setStyleSheet("background: transparent; border: none; color: #64748b;")
             else:
                 pct_edit.editingFinished.connect(
-                    lambda lbl_=lbl: self._on_pct_edited(lbl_)
+                    lambda lbl_=seg_lbl: self._on_pct_edited(lbl_)
                 )
                 pct_edit.textChanged.connect(
-                    lambda _text, lbl_=lbl: self._on_pct_text_changed(lbl_)
+                    lambda _text, lbl_=seg_lbl: self._on_pct_text_changed(lbl_)
                 )
 
-            self._rows[lbl] = (row_w, pct_edit, ma_label)
+            self._rows[seg_lbl] = (row_w, pct_edit, ma_label)
             self._layout.addWidget(row_w)
 
-        self._refresh_ma_values()
 
     def _on_pct_edited(self, edited_label: str) -> None:
         """Called when user finishes editing a percentage field."""
@@ -266,7 +349,8 @@ def _refresh_ma_values(self) -> None:
         except (ValueError, TypeError):
             total_amp = 0.0
 
-        for lbl in self._cathode_labels:
+        # Update all rows (both main contacts and segments)
+        for lbl in list(self._rows.keys()):
             row_data = self._rows.get(lbl)
             if not row_data:
                 continue
@@ -351,18 +435,30 @@ def get_cathode_labels(canvas) -> list[str]:
 
     if model.is_directional:
         for contact_idx in range(model.num_contacts):
-            seg_states = [
-                canvas.contact_states.get((contact_idx, seg), ContactState.OFF)
-                for seg in range(3)
-            ]
-            if all(s == ContactState.CATHODIC for s in seg_states):
-                labels.append(f"E{contact_idx}")
-                continue
-            seg_labels = ["a", "b", "c"]
-            for seg, seg_state in enumerate(seg_states):
-                if seg_state == ContactState.CATHODIC:
-                    labels.append(f"E{contact_idx}{seg_labels[seg]}")
+            # Check if this contact level is directional (has segments)
+            is_contact_directional = model.is_level_directional(contact_idx)
+
+            if is_contact_directional:
+                # This contact has segments - check individual segments
+                seg_states = [
+                    canvas.contact_states.get((contact_idx, seg), ContactState.OFF)
+                    for seg in range(3)
+                ]
+                if all(s == ContactState.CATHODIC for s in seg_states):
+                    # All segments are cathodic - show as ring contact (no segment suffix)
+                    labels.append(f"E{contact_idx}")
+                    continue
+                seg_labels = ["a", "b", "c"]
+                for seg, seg_state in enumerate(seg_states):
+                    if seg_state == ContactState.CATHODIC:
+                        labels.append(f"E{contact_idx}{seg_labels[seg]}")
+            else:
+                # This contact is a ring contact - no segments
+                state = canvas.contact_states.get((contact_idx, 0), ContactState.OFF)
+                if state == ContactState.CATHODIC:
+                    labels.append(f"E{contact_idx}")
     else:
+        # Non-directional model - all contacts are ring contacts
         for contact_idx in range(model.num_contacts):
             state = canvas.contact_states.get((contact_idx, 0), ContactState.OFF)
             if state == ContactState.CATHODIC:

src/clinical_dbs_annotator/views/step1_view.py

@@ -121,19 +121,39 @@ def refresh_theme_icons(self) -> None:
         if btn:
             btn.setIcon(self._create_settings_icon())
 
+    def _is_single_grouped_directional(self, cathode_labels: list[str], canvas) -> bool:
+        """Check if a single cathode label represents a grouped directional contact."""
+        if len(cathode_labels) != 1 or not canvas or not canvas.model:
+            return False
+
+        lbl = cathode_labels[0]
+        # Check if this is a grouped contact (no segment suffix) that could be directional
+        if len(lbl) >= 2 and lbl[0] == 'E' and lbl[1:].isdigit():
+            try:
+                contact_idx = int(lbl[1:])
+                return (canvas.model.is_directional and
+                       canvas.model.is_level_directional(contact_idx))
+            except (ValueError, IndexError):
+                pass
+        return False
+
     def _on_left_canvas_validation(self, is_valid: bool, error_msg: str) -> None:
         """Callback when left electrode canvas validation state changes."""
         self._left_selection_valid = is_valid
         self.update_configuration_display()
         if hasattr(self, "left_amp_split"):
-            self.left_amp_split.update_cathodes(get_cathode_labels(self.left_canvas))
+            left_labels = get_cathode_labels(self.left_canvas)
+            left_is_grouped = self._is_single_grouped_directional(left_labels, self.left_canvas)
+            self.left_amp_split.update_cathodes(left_labels, left_is_grouped)
 
     def _on_right_canvas_validation(self, is_valid: bool, error_msg: str) -> None:
         """Callback when right electrode canvas validation state changes."""
         self._right_selection_valid = is_valid
         self.update_configuration_display()
         if hasattr(self, "right_amp_split"):
-            self.right_amp_split.update_cathodes(get_cathode_labels(self.right_canvas))
+            right_labels = get_cathode_labels(self.right_canvas)
+            right_is_grouped = self._is_single_grouped_directional(right_labels, self.right_canvas)
+            self.right_amp_split.update_cathodes(right_labels, right_is_grouped)
 
     def _setup_ui(self) -> None:
         """Set up the UI layout."""
@@ -567,13 +587,15 @@ def _format_configuration_html(self, canvas: ElectrodeCanvas) -> str:
                 state = canvas.contact_states[contact_id]
                 state_str = "Anodic (+)" if state == ContactState.ANODIC else "Cathodic (-)"
 
-                if model.is_directional:
+                if model.is_directional and model.is_level_directional(contact_idx):
+                    # This contact has segments - show segment label
                     segment_labels = ['a', 'b', 'c']
                     contact_name = f"C{contact_idx}{segment_labels[seg_idx]}"
                 else:
+                    # This contact is a ring contact - no segment label
                     contact_name = f"C{contact_idx}"
 
-                lines.append(f"  • {contact_name}: {state_str}")
+                lines.append(f"  • {contact_name}: {state_str}")
         else:
             lines.append("<i>No active contacts</i>")
 
@@ -611,21 +633,34 @@ def _get_anode_cathode_texts(self, canvas: ElectrodeCanvas) -> tuple[str, str]:
 
         if model.is_directional:
             for contact_idx in range(model.num_contacts):
-                seg_states = [canvas.contact_states.get((contact_idx, seg), ContactState.OFF) for seg in range(3)]
-                if all(s == ContactState.ANODIC for s in seg_states):
-                    anode_items.append(f"E{contact_idx}")
-                    continue
-                if all(s == ContactState.CATHODIC for s in seg_states):
-                    cathode_items.append(f"E{contact_idx}")
-                    continue
+                # Check if this contact level is directional (has segments)
+                is_contact_directional = model.is_level_directional(contact_idx)
+
+                if is_contact_directional:
+                    # This contact has segments - check individual segments
+                    seg_states = [canvas.contact_states.get((contact_idx, seg), ContactState.OFF) for seg in range(3)]
+                    if all(s == ContactState.ANODIC for s in seg_states):
+                        anode_items.append(f"E{contact_idx}")
+                        continue
+                    if all(s == ContactState.CATHODIC for s in seg_states):
+                        cathode_items.append(f"E{contact_idx}")
+                        continue
 
-                seg_labels = ["a", "b", "c"]
-                for seg, seg_state in enumerate(seg_states):
-                    if seg_state == ContactState.ANODIC:
-                        anode_items.append(f"E{contact_idx}{seg_labels[seg]}")
-                    elif seg_state == ContactState.CATHODIC:
-                        cathode_items.append(f"E{contact_idx}{seg_labels[seg]}")
+                    seg_labels = ["a", "b", "c"]
+                    for seg, seg_state in enumerate(seg_states):
+                        if seg_state == ContactState.ANODIC:
+                            anode_items.append(f"E{contact_idx}{seg_labels[seg]}")
+                        elif seg_state == ContactState.CATHODIC:
+                            cathode_items.append(f"E{contact_idx}{seg_labels[seg]}")
+                else:
+                    # This contact is a ring contact - no segments
+                    state = canvas.contact_states.get((contact_idx, 0), ContactState.OFF)
+                    if state == ContactState.ANODIC:
+                        anode_items.append(f"E{contact_idx}")
+                    elif state == ContactState.CATHODIC:
+                        cathode_items.append(f"E{contact_idx}")
         else:
+            # Non-directional model - all contacts are ring contacts
             for contact_idx in range(model.num_contacts):
                 state = canvas.contact_states.get((contact_idx, 0), ContactState.OFF)
                 if state == ContactState.ANODIC:
@@ -709,9 +744,13 @@ def apply_tokens(text: str, state: ContactState) -> None:
         canvas.update()
         # Refresh amplitude split widget for this canvas
         if canvas is self.left_canvas and hasattr(self, "left_amp_split"):
-            self.left_amp_split.update_cathodes(get_cathode_labels(self.left_canvas))
+            left_labels = get_cathode_labels(self.left_canvas)
+            left_is_grouped = self._is_single_grouped_directional(left_labels, self.left_canvas)
+            self.left_amp_split.update_cathodes(left_labels, left_is_grouped)
         elif canvas is self.right_canvas and hasattr(self, "right_amp_split"):
-            self.right_amp_split.update_cathodes(get_cathode_labels(self.right_canvas))
+            right_labels = get_cathode_labels(self.right_canvas)
+            right_is_grouped = self._is_single_grouped_directional(right_labels, self.right_canvas)
+            self.right_amp_split.update_cathodes(right_labels, right_is_grouped)
 
     def reset_all(self):
         """Reset all contacts and case"""