src/probe.py

"""Flow probe system: single probe, multi-probe, and branching support.

The probe is advected by the same velocity field as the particles. Its trajectory
is recorded and can be analyzed for brain region transitions.

Features:
  - Single probe: click to place, follows mean flow
  - Multi-probe: initialize N probes in a local neighborhood
  - Branching: when MDN components are highly uncertain (50/50 split), spawn
    a ghost probe that follows the dominant alternative component
  - Live region highlighting in RED (ghost highlights more transparent)
"""

import numpy as np
import vtk


class FlowProbe:
    """A single probe that follows the flow field."""

    def __init__(self, ren: vtk.vtkRenderer, amin: np.ndarray, amax: np.ndarray,
                 color=(0.0, 1.0, 0.3), opacity=0.9, ghost=False, label=""):
        self.ren = ren
        self.amin = amin.astype(np.float32)
        self.amax = amax.astype(np.float32)
        self.active = False
        self.position = None
        self.path: list[np.ndarray] = []
        self.speeds: list[float] = []
        self.raw_field_mags: list[float] = []  # field magnitude (independent of speed scale)
        self._max_path_len = 50000
        self._mesh_overlay = None
        self._highlighted_regions: set[str] = set()
        self._highlight_actors: dict[str, vtk.vtkActor] = {}
        self._check_interval = 30
        self._step_counter = 0
        self.current_regions: set[str] = set()
        self._boundary_check = None
        self._on_region_change = None
        self.ghost = ghost
        self.label = label
        self._color = color
        self.steps_alive = 0
        self._stuck_counter = 0
        self._stuck_threshold = 150  # steps with near-zero movement before warning
        self._stuck_warned = False
        self._stuck_eps = 1e-6  # minimum displacement per step
        # Debounce: require N consecutive detections before entering, N misses before leaving
        self._debounce_enter = 3   # consecutive checks to confirm entry
        self._debounce_leave = 5   # consecutive misses to confirm exit
        self._region_hit_count: dict[str, int] = {}   # key -> consecutive hit count
        self._region_miss_count: dict[str, int] = {}  # key -> consecutive miss count
        self._nearest_max_dist_mm = 5.0  # hard limit for nearest-region snapping (mm)

        diag = float(np.linalg.norm(amax - amin))
        marker_alpha = 0.4 if ghost else 1.0
        trail_alpha = 0.35 if ghost else 0.9

        # Probe marker (sphere)
        self._sphere = vtk.vtkSphereSource()
        self._sphere.SetRadius(0.012 * diag if not ghost else 0.008 * diag)
        self._sphere.SetThetaResolution(16)
        self._sphere.SetPhiResolution(16)
        mapper = vtk.vtkPolyDataMapper()
        mapper.SetInputConnection(self._sphere.GetOutputPort())
        self.marker_actor = vtk.vtkActor()
        self.marker_actor.SetMapper(mapper)
        self.marker_actor.GetProperty().SetColor(*color)
        self.marker_actor.GetProperty().SetOpacity(marker_alpha)
        self.marker_actor.GetProperty().LightingOff()
        self.marker_actor.VisibilityOff()
        ren.AddActor(self.marker_actor)

        # Trail line
        self._trail_points = vtk.vtkPoints()
        self._trail_cells = vtk.vtkCellArray()
        self._trail_pd = vtk.vtkPolyData()
        self._trail_pd.SetPoints(self._trail_points)
        self._trail_pd.SetLines(self._trail_cells)
        trail_mapper = vtk.vtkPolyDataMapper()
        trail_mapper.SetInputData(self._trail_pd)
        self.trail_actor = vtk.vtkActor()
        self.trail_actor.SetMapper(trail_mapper)
        self.trail_actor.GetProperty().SetColor(*color)
        self.trail_actor.GetProperty().SetLineWidth(3.0 if not ghost else 2.0)
        self.trail_actor.GetProperty().SetOpacity(trail_alpha)
        self.trail_actor.GetProperty().LightingOff()
        self.trail_actor.VisibilityOff()
        ren.AddActor(self.trail_actor)

    def set_mesh_overlay(self, mesh_overlay):
        self._mesh_overlay = mesh_overlay

    def set_boundary_check(self, fn):
        self._boundary_check = fn

    def set_on_region_change(self, fn):
        self._on_region_change = fn

    def place(self, position: np.ndarray):
        """Place probe at position and start recording."""
        self.position = np.array(position, dtype=np.float32).ravel()[:3]

        # Boundary validation
        if self._boundary_check is not None and not self._boundary_check(self.position):
            found = False
            diag = float(np.linalg.norm(self.amax - self.amin))
            for scale in [0.01, 0.02, 0.05, 0.1, 0.2]:
                for _ in range(50):
                    offset = np.random.randn(3).astype(np.float32) * scale * diag
                    candidate = np.clip(self.position + offset, self.amin, self.amax)
                    if self._boundary_check(candidate):
                        self.position = candidate
                        found = True
                        break
                if found:
                    break
            if not found:
                self.position = (self.amin + self.amax) / 2.0
                print(f"[probe{self.label}] Could not find valid position, using domain center")

        self.path = [self.position.copy()]
        self.speeds = [0.0]
        self.raw_field_mags = [0.0]
        self.active = True
        self._step_counter = 0
        self.steps_alive = 0
        self._stuck_counter = 0
        self._stuck_warned = False

        # Reset trail
        self._trail_points = vtk.vtkPoints()
        self._trail_cells = vtk.vtkCellArray()
        self._trail_points.InsertNextPoint(*self.position.tolist())
        self._trail_pd.SetPoints(self._trail_points)
        self._trail_pd.SetLines(self._trail_cells)
        self._trail_pd.Modified()

        self.marker_actor.VisibilityOn()
        self.marker_actor.SetPosition(*self.position.tolist())
        self.trail_actor.VisibilityOn()

        self._update_region_highlights()
        tag = " (ghost)" if self.ghost else ""
        print(f"[probe{self.label}{tag}] placed at "
              f"({self.position[0]:.2f}, {self.position[1]:.2f}, {self.position[2]:.2f})")

    def step(self, sampler, dt_step: float):
        """Advect one step using sampler."""
        if not self.active or self.position is None:
            return
        V = sampler.sample_vec(self.position[None, :])
        velocity = V[0]
        raw_mag = float(np.linalg.norm(velocity))

        new_pos = self.position + velocity * dt_step
        new_pos = np.clip(new_pos, self.amin, self.amax)

        # Boundary constraint
        if self._boundary_check is not None and not self._boundary_check(new_pos):
            half_pos = self.position + velocity * dt_step * 0.5
            half_pos = np.clip(half_pos, self.amin, self.amax)
            if self._boundary_check(half_pos):
                new_pos = half_pos
            else:
                return  # hit boundary

        self.position = new_pos.astype(np.float32)
        speed = float(np.linalg.norm(self.position - self.path[-1])) if self.path else 0.0

        # Stuck / weak flow detection
        if speed < self._stuck_eps and raw_mag < self._stuck_eps:
            self._stuck_counter += 1
            if self._stuck_counter >= self._stuck_threshold and not self._stuck_warned:
                self._stuck_warned = True
                tag = f" (ghost)" if self.ghost else ""
                print(f"\n[probe{self.label}{tag}] Flow is very weak or has ended here. "
                      f"The probe is stuck at ({self.position[0]:.1f}, {self.position[1]:.1f}, {self.position[2]:.1f}).")
                print(f"[probe{self.label}{tag}] Try placing the probe deeper in the brain "
                      f"where flow is stronger (press 'c' to clear, then 'g' + click).\n")
        else:
            self._stuck_counter = 0

        if len(self.path) < self._max_path_len:
            self.path.append(self.position.copy())
            self.speeds.append(speed)
            self.raw_field_mags.append(raw_mag)

        self.marker_actor.SetPosition(*self.position.tolist())

        # Append to trail
        idx = self._trail_points.InsertNextPoint(*self.position.tolist())
        if idx > 0:
            self._trail_cells.InsertNextCell(2)
            self._trail_cells.InsertCellPoint(idx - 1)
            self._trail_cells.InsertCellPoint(idx)
        self._trail_points.Modified()
        self._trail_cells.Modified()
        self._trail_pd.Modified()

        # Periodic region check
        self._step_counter += 1
        self.steps_alive += 1
        if self._step_counter % self._check_interval == 0:
            self._update_region_highlights()

    def _point_in_bbox(self, point, bounds):
        """Fast bounding-box containment check. bounds is (xmin,xmax,ymin,ymax,zmin,zmax)."""
        return (bounds[0] <= point[0] <= bounds[1] and
                bounds[2] <= point[1] <= bounds[3] and
                bounds[4] <= point[2] <= bounds[5])

    def _update_region_highlights(self):
        if self._mesh_overlay is None or self.position is None:
            return

        # --- Raw detection (what region key is at the probe right now?) ---
        raw_key = None
        if hasattr(self._mesh_overlay, 'get_region_at_point'):
            raw_key = self._mesh_overlay.get_region_at_point(self.position)
            if raw_key is None and hasattr(self._mesh_overlay, 'find_nearest_region'):
                raw_key = self._mesh_overlay.find_nearest_region(
                    self.position, search_radius=2,
                    max_distance_mm=self._nearest_max_dist_mm)
        else:
            for key in self._mesh_overlay.get_all_region_keys():
                if hasattr(self._mesh_overlay, 'fast_point_in_mesh'):
                    if self._mesh_overlay.fast_point_in_mesh(self.position, key):
                        raw_key = key
                        break
                elif self._mesh_overlay.point_in_mesh(self.position, key):
                    raw_key = key
                    break

        raw_detected = {raw_key} if raw_key else set()

        # --- Debounce: require consecutive detections to enter, consecutive misses to leave ---
        # Update hit/miss counters for detected key
        for key in raw_detected:
            self._region_hit_count[key] = self._region_hit_count.get(key, 0) + 1
            self._region_miss_count.pop(key, None)

        # Update miss counters for keys that were NOT detected this tick
        for key in list(self._region_hit_count.keys()):
            if key not in raw_detected:
                self._region_miss_count[key] = self._region_miss_count.get(key, 0) + 1
                self._region_hit_count[key] = 0

        # Determine stable set: regions that passed the entry threshold
        # and have not yet exceeded the leave threshold
        new_regions = set()
        for key in set(list(self._region_hit_count.keys()) +
                       list(self._highlighted_regions)):
            hits = self._region_hit_count.get(key, 0)
            misses = self._region_miss_count.get(key, 0)
            if key in self._highlighted_regions:
                # Already highlighted — keep it unless misses exceed threshold
                if misses < self._debounce_leave:
                    new_regions.add(key)
            else:
                # Not yet highlighted — add if hits exceed entry threshold
                if hits >= self._debounce_enter:
                    new_regions.add(key)

        # Clean up stale counters
        for key in list(self._region_miss_count.keys()):
            if self._region_miss_count[key] > self._debounce_leave + 2:
                self._region_miss_count.pop(key, None)
                self._region_hit_count.pop(key, None)

        # --- Apply changes (enter/leave) ---
        left = self._highlighted_regions - new_regions
        left_names = []
        for key in left:
            if key in self._highlight_actors:
                try:
                    self.ren.RemoveActor(self._highlight_actors[key])
                except Exception:
                    pass
                del self._highlight_actors[key]
            name = self._mesh_overlay.get_region_name(key)
            if hasattr(self._mesh_overlay, 'get_hemisphere_label'):
                hemi = self._mesh_overlay.get_hemisphere_label(key, self.position)
                if hemi:
                    name = f"{name} ({hemi})"
            left_names.append(name)
            tag = " [branch]" if self.ghost else ""
            print(f"[probe{self.label}] LEFT: {name}{tag}")

        entered = new_regions - self._highlighted_regions
        entered_names = []
        highlight_opacity = 0.12 if self.ghost else 0.25
        for key in entered:
            poly = None
            if hasattr(self._mesh_overlay, 'get_hemisphere_polydata'):
                poly = self._mesh_overlay.get_hemisphere_polydata(key, self.position)
            if poly is None:
                poly = self._mesh_overlay.get_polydata(key)
            if poly is not None:
                mapper = vtk.vtkPolyDataMapper()
                mapper.SetInputData(poly)
                actor = vtk.vtkActor()
                actor.SetMapper(mapper)
                actor.GetProperty().SetColor(1.0, 0.6, 0.0)  # orange (all regions)
                actor.GetProperty().SetOpacity(highlight_opacity)
                actor.GetProperty().LightingOff()
                self.ren.AddActor(actor)
                self._highlight_actors[key] = actor
            name = self._mesh_overlay.get_region_name(key)
            if hasattr(self._mesh_overlay, 'get_hemisphere_label'):
                hemi = self._mesh_overlay.get_hemisphere_label(key, self.position)
                if hemi:
                    name = f"{name} ({hemi})"
            entered_names.append(name)
            tag = " [branch]" if self.ghost else ""
            # Detailed entry log with position context
            pos_detail = ""
            if self.position is not None and self._mesh_overlay is not None:
                try:
                    center = self._mesh_overlay.get_mesh_center(key)
                    bounds = self._mesh_overlay.get_mesh_bounds(key)
                    if center is not None and bounds is not None:
                        extent = [bounds[1]-bounds[0], bounds[3]-bounds[2],
                                  bounds[5]-bounds[4]]
                        char_size = sum(extent) / 3.0
                        dist = float(np.linalg.norm(self.position - center))
                        depth = max(0.0, 1.0 - min(dist / (char_size * 0.5), 1.0))
                        pos_parts = []
                        diff = self.position - center
                        if abs(diff[2]) > extent[2] * 0.15:
                            pos_parts.append("dorsal" if diff[2] > 0 else "ventral")
                        if abs(diff[0]) > extent[0] * 0.15:
                            pos_parts.append("lateral-R" if diff[0] > 0 else "lateral-L")
                        if abs(diff[1]) > extent[1] * 0.15:
                            pos_parts.append("anterior" if diff[1] > 0 else "posterior")
                        pos_str = "-".join(pos_parts) if pos_parts else "central"
                        pos_detail = f"  [{pos_str}, depth={depth:.0%}]"
                except Exception:
                    pass
            print(f"[probe{self.label}{tag}] ENTERED: {name}{pos_detail}")

        self._highlighted_regions = new_regions
        self.current_regions = new_regions

        # --- Extra parcellation subregion highlighting ---
        # Extra parcellation subregions also get orange, same as main regions.
        if (self._mesh_overlay is not None and
                hasattr(self._mesh_overlay, '_extra') and
                self._mesh_overlay._extra is not None and
                self.position is not None and new_regions):
            try:
                hier = self._mesh_overlay.get_hierarchical_regions_at_point(self.position)
                sub = hier.get("subregion")
                sub_mesh = hier.get("subregion_mesh")
                sub_key = f"_extra_{sub['label_id']}" if sub else None

                # Remove old subregion highlight if changed
                old_sub_key = getattr(self, '_current_subregion_key', None)
                if old_sub_key and old_sub_key != sub_key:
                    if old_sub_key in self._highlight_actors:
                        try:
                            self.ren.RemoveActor(self._highlight_actors[old_sub_key])
                        except Exception:
                            pass
                        del self._highlight_actors[old_sub_key]

                # Add new subregion highlight (orange, same as all regions)
                if sub_key and sub_mesh and sub_key not in self._highlight_actors:
                    display_mesh = sub_mesh
                    try:
                        bounds = [0.0] * 6
                        sub_mesh.GetBounds(bounds)
                        x_extent = bounds[1] - bounds[0]
                        if x_extent > 10.0:
                            x_mid = (bounds[0] + bounds[1]) / 2.0
                            plane = vtk.vtkPlane()
                            plane.SetOrigin(x_mid, 0, 0)
                            if self.position[0] >= x_mid:
                                plane.SetNormal(1, 0, 0)
                            else:
                                plane.SetNormal(-1, 0, 0)
                            clipper = vtk.vtkClipPolyData()
                            clipper.SetInputData(sub_mesh)
                            clipper.SetClipFunction(plane)
                            clipper.SetInsideOut(False)
                            clipper.Update()
                            clipped = clipper.GetOutput()
                            if clipped and clipped.GetNumberOfPoints() > 0:
                                display_mesh = clipped
                    except Exception:
                        pass

                    mapper = vtk.vtkPolyDataMapper()
                    mapper.SetInputData(display_mesh)
                    actor = vtk.vtkActor()
                    actor.SetMapper(mapper)
                    actor.GetProperty().SetColor(1.0, 0.6, 0.0)  # orange
                    actor.GetProperty().SetOpacity(0.3)
                    actor.GetProperty().LightingOff()
                    self.ren.AddActor(actor)
                    self._highlight_actors[sub_key] = actor
                    tag = " [branch]" if self.ghost else ""
                    hemi_str = ""
                    if self.position is not None and self.position[0] >= 0:
                        hemi_str = " (right hemisphere)"
                    elif self.position is not None:
                        hemi_str = " (left hemisphere)"
                    print(f"[probe{self.label}{tag}] SUBREGION: {sub['name']}{hemi_str}")

                # Hide coarser main-atlas regions when we have a finer subregion
                if sub_key and sub_key in self._highlight_actors:
                    hidden = set()
                    for rkey in new_regions:
                        if rkey in self._highlight_actors and not rkey.startswith("_extra_"):
                            self._highlight_actors[rkey].VisibilityOff()
                            hidden.add(rkey)
                    self._hidden_orange_keys = hidden
                elif not sub_key:
                    for rkey in getattr(self, '_hidden_orange_keys', set()):
                        if rkey in self._highlight_actors:
                            self._highlight_actors[rkey].VisibilityOn()
                    self._hidden_orange_keys = set()

                self._current_subregion_key = sub_key
            except Exception:
                pass

        # --- Red hotspot: clip mesh near probe position ---
        # Instead of coloring an entire region red, clip the mesh surface
        # within a sphere around the probe and show that patch in red.
        self._update_hotspot()

        if self._on_region_change and (entered_names or left_names):
            self._on_region_change(entered_names, left_names,
                                   is_ghost=self.ghost, label=self.label)

    def _update_hotspot(self):
        """Color highlighted region meshes with a distance-based heatmap.

        Vertices near the probe are red, fading smoothly to orange further away.
        Uses per-vertex RGBA scalars — no clipping, no extra actors.
        """
        from vtkmodules.util.numpy_support import vtk_to_numpy, numpy_to_vtk

        if self.position is None or not self._highlight_actors:
            return

        probe_pos = self.position.astype(np.float64)
        fade_radius = 15.0  # mm — distance over which red fades to orange

        for key, actor in self._highlight_actors.items():
            mapper = actor.GetMapper()
            if mapper is None:
                continue
            poly = mapper.GetInput()
            if poly is None or poly.GetNumberOfPoints() < 3:
                continue

            pts_vtk = poly.GetPoints()
            if pts_vtk is None:
                continue
            verts = vtk_to_numpy(pts_vtk.GetData()).astype(np.float64)

            # Distance from each vertex to probe
            dists = np.linalg.norm(verts - probe_pos, axis=1)
            t = np.clip(dists / fade_radius, 0.0, 1.0)  # 0=at probe, 1=far

            # Color gradient: red (1,0,0) at probe → orange (1,0.6,0) far
            r = np.full(len(t), 255, np.uint8)
            g = (t * 0.6 * 255).astype(np.uint8)
            b = np.zeros(len(t), np.uint8)

            # Alpha: brighter near probe, dimmer far away
            base_opacity = 0.45 if not self.ghost else 0.2
            far_opacity = 0.2 if not self.ghost else 0.08
            alpha_f = far_opacity + (base_opacity - far_opacity) * (1.0 - t)
            a = (np.clip(alpha_f, 0.0, 1.0) * 255).astype(np.uint8)

            rgba = np.column_stack([r, g, b, a])
            scalars = numpy_to_vtk(rgba, deep=True)
            scalars.SetNumberOfComponents(4)
            scalars.SetName("HeatmapColors")
            poly.GetPointData().SetScalars(scalars)

            mapper.SetColorModeToDirectScalars()
            mapper.SetScalarModeToUsePointData()
            mapper.ScalarVisibilityOn()
            actor.GetProperty().LightingOff()
            # Override the flat color — let scalars drive everything
            actor.GetProperty().SetOpacity(1.0)
            poly.Modified()

    def clear(self):
        self.active = False
        self.position = None
        self.path = []
        self.speeds = []
        self.raw_field_mags = []
        self._step_counter = 0
        self.steps_alive = 0
        self.marker_actor.VisibilityOff()
        self.trail_actor.VisibilityOff()
        self._trail_points = vtk.vtkPoints()
        self._trail_cells = vtk.vtkCellArray()
        self._trail_pd.SetPoints(self._trail_points)
        self._trail_pd.SetLines(self._trail_cells)
        self._trail_pd.Modified()
        for key, actor in self._highlight_actors.items():
            try:
                self.ren.RemoveActor(actor)
            except Exception:
                pass
        self._highlight_actors.clear()
        self._highlighted_regions.clear()
        self.current_regions.clear()

    def destroy(self):
        """Remove all actors from renderer."""
        self.clear()
        try:
            self.ren.RemoveActor(self.marker_actor)
        except Exception:
            pass
        try:
            self.ren.RemoveActor(self.trail_actor)
        except Exception:
            pass

    def get_path_array(self) -> np.ndarray:
        if not self.path:
            return np.zeros((0, 3), np.float32)
        return np.array(self.path, dtype=np.float32)

    def get_speeds_array(self) -> np.ndarray:
        if not self.speeds:
            return np.zeros(0, np.float32)
        return np.array(self.speeds, dtype=np.float32)

    def get_field_mags_array(self) -> np.ndarray:
        if not self.raw_field_mags:
            return np.zeros(0, np.float32)
        return np.array(self.raw_field_mags, dtype=np.float32)


class ProbeSystem:
    """Manages single/multi probes and branching behavior.

    Modes:
      - single: one probe following mean flow
      - multi: N probes in local neighborhood, all following mean flow
      - branching: when PI uncertainty is high, spawn ghost probes

    For state-propagation mode, defaults to multi(4) + branching.
    """

    def __init__(self, ren: vtk.vtkRenderer, win: vtk.vtkRenderWindow,
                 amin: np.ndarray, amax: np.ndarray):
        self.ren = ren
        self.win = win
        self.amin = amin.astype(np.float32)
        self.amax = amax.astype(np.float32)
        self.probes: list[FlowProbe] = []
        self.ghost_probes: list[FlowProbe] = []
        self._mesh_overlay = None
        self._boundary_check = None
        self._on_region_change = None
        self._branching_enabled = False
        self._multi_count = 1
        self._branch_threshold = 0.35  # max ratio between top 2 PI components
        self._branch_min_pi = 0.25     # min weight of 2nd component
        self._branch_check_interval = 50
        self._branch_step_counter = 0
        self._ghost_color = (0.4, 0.7, 1.0)  # pale blue for ghosts
        self._pi_field = None
        self._mus_samplers = None

    def set_mesh_overlay(self, overlay):
        self._mesh_overlay = overlay

    def set_boundary_check(self, fn):
        self._boundary_check = fn

    def set_on_region_change(self, fn):
        self._on_region_change = fn

    def set_branching(self, enabled: bool, pi_field=None, mus_samplers=None):
        """Enable/disable branching.

        Args:
            enabled: toggle branching
            pi_field: the PI weight grid (G,G,G,K) numpy array
            mus_samplers: list of TriLinearSampler for each component
        """
        self._branching_enabled = enabled
        self._pi_field = pi_field
        self._mus_samplers = mus_samplers
        print(f"[probe-system] branching {'ON' if enabled else 'OFF'}")

    def set_multi_count(self, n: int):
        self._multi_count = max(1, n)
        print(f"[probe-system] multi-probe count: {self._multi_count}")

    @property
    def active(self) -> bool:
        return any(p.active for p in self.probes)

    @property
    def path(self):
        """Return path of first probe (for backward compat)."""
        return self.probes[0].path if self.probes else []

    def place(self, position: np.ndarray):
        """Place probe(s) at position."""
        self.clear()
        diag = float(np.linalg.norm(self.amax - self.amin))
        jitter_scale = 0.04 * diag

        for i in range(self._multi_count):
            if i == 0:
                pos = position.copy()
                label = "" if self._multi_count == 1 else f"#{i+1}"
            else:
                offset = np.random.randn(3).astype(np.float32) * jitter_scale
                pos = np.clip(position + offset, self.amin, self.amax)
                label = f"#{i+1}"

            probe = FlowProbe(self.ren, self.amin, self.amax,
                              color=(0.0, 1.0, 0.3), ghost=False, label=label)
            if self._mesh_overlay:
                probe.set_mesh_overlay(self._mesh_overlay)
            if self._boundary_check:
                probe.set_boundary_check(self._boundary_check)
            if self._on_region_change:
                probe.set_on_region_change(self._on_region_change)
            probe.place(pos)
            self.probes.append(probe)

    def step(self, sampler, dt_step: float, pi_sampler=None):
        """Step all probes."""
        for p in self.probes:
            if p.active:
                p.step(sampler, dt_step)

        for g in self.ghost_probes:
            if g.active:
                # Ghost probes follow their specific component sampler
                comp_sampler = getattr(g, '_comp_sampler', sampler)
                g.step(comp_sampler, dt_step)

        # Prune ghost probes that have been alive > 200 steps but did not diverge
        diag = float(np.linalg.norm(self.amax - self.amin))
        prune_dist = 0.03 * diag
        to_remove = []
        for g in self.ghost_probes:
            if not g.active or g.position is None:
                continue
            if g.steps_alive > 200:
                for p in self.probes:
                    if not p.active or p.position is None:
                        continue
                    dist = float(np.linalg.norm(g.position - p.position))
                    if dist < prune_dist:
                        print("[branch] pruned ghost - did not diverge")
                        to_remove.append(g)
                        break
        for g in to_remove:
            g.destroy()
            self.ghost_probes.remove(g)

        # Check for branching
        if self._branching_enabled:
            self._branch_step_counter += 1
            if self._branch_step_counter % self._branch_check_interval == 0:
                self._check_branching(sampler, dt_step)

    def _check_branching(self, mean_sampler, dt_step):
        """Check if any active probe is at a high-uncertainty location and branch."""
        if self._pi_field is None or self._mus_samplers is None:
            return
        if not self.probes:
            return

        from .field_loader import TriLinearSampler

        for p in self.probes:
            if not p.active or p.position is None:
                continue
            # Sample PI weights at probe position
            pos = p.position[None, :]
            pi_vals = mean_sampler.sample_vec(pos)  # dummy, we need PI
            # Actually sample PI field directly
            try:
                K = self._pi_field.shape[-1]
                weights = np.zeros(K, np.float32)
                for k in range(K):
                    pi_k = self._pi_field[..., k:k+1]
                    # Use sampler to interpolate
                    w = mean_sampler.sample_scalar(
                        self._pi_field[..., k],
                        p.position[None, :]
                    )
                    weights[k] = float(w[0])
            except Exception:
                continue

            # Normalize
            ws = weights.sum()
            if ws <= 0:
                continue
            weights /= ws

            # Sort to find top 2
            sorted_idx = np.argsort(weights)[::-1]
            w1 = weights[sorted_idx[0]]
            w2 = weights[sorted_idx[1]] if len(sorted_idx) > 1 else 0.0

            # Check if uncertain enough
            if w2 < self._branch_min_pi:
                continue
            ratio = w2 / max(w1, 1e-9)
            if ratio < self._branch_threshold:
                continue

            # Check we haven't already branched from near this position
            already_branched = False
            for g in self.ghost_probes:
                if g.active and g.path:
                    d = np.linalg.norm(g.path[0] - p.position)
                    if d < 0.02 * float(np.linalg.norm(self.amax - self.amin)):
                        already_branched = True
                        break
            if already_branched:
                continue

            # Branch! Create ghost probe following 2nd component
            comp_idx = int(sorted_idx[1])
            if comp_idx >= len(self._mus_samplers):
                continue

            ghost = FlowProbe(self.ren, self.amin, self.amax,
                              color=self._ghost_color, ghost=True,
                              label=f"~branch(comp{comp_idx+1})")
            if self._mesh_overlay:
                ghost.set_mesh_overlay(self._mesh_overlay)
            if self._boundary_check:
                ghost.set_boundary_check(self._boundary_check)
            if self._on_region_change:
                ghost.set_on_region_change(self._on_region_change)
            ghost._comp_sampler = self._mus_samplers[comp_idx]
            ghost.place(p.position.copy())
            self.ghost_probes.append(ghost)

            print(f"[branch] SPLIT at ({p.position[0]:.2f}, {p.position[1]:.2f}, "
                  f"{p.position[2]:.2f}): comp{sorted_idx[0]+1}={w1:.2f} vs "
                  f"comp{comp_idx+1}={w2:.2f}")
            if self._on_region_change:
                self._on_region_change(
                    [f"BRANCH: comp{comp_idx+1} (weight={w2:.2f})"], [],
                    is_ghost=True, label="branch"
                )

    def clear(self):
        """Clear all probes."""
        for p in self.probes:
            p.destroy()
        for g in self.ghost_probes:
            g.destroy()
        self.probes.clear()
        self.ghost_probes.clear()
        self._branch_step_counter = 0

    def freeze(self):
        """Stop all probes from moving (keep path and highlights intact)."""
        for p in self.probes:
            p.active = False
        for g in self.ghost_probes:
            g.active = False
        print("[probe-system] probes frozen")

    def get_all_probes(self) -> list[FlowProbe]:
        """Return all probes (main + ghost) for analysis."""
        return self.probes + self.ghost_probes

    def get_primary_probe(self) -> FlowProbe | None:
        """Return the first active probe."""
        for p in self.probes:
            if p.active:
                return p
        return None

    def get_all_paths(self) -> list[tuple[np.ndarray, bool, str]]:
        """Return [(path_array, is_ghost, label), ...] for all probes."""
        result = []
        for p in self.probes:
            if p.path:
                result.append((p.get_path_array(), False, p.label))
        for g in self.ghost_probes:
            if g.path:
                result.append((g.get_path_array(), True, g.label))
        return result