examples/roi_flow_mode.py

#!/usr/bin/env python
"""ROI Flow Mode — manifold flow visualization with ROI activation panel.

Single-window visualization with two viewports:
  Left:  MDN particle flow through a learned neural manifold
  Right: ROI activation spheres showing corresponding brain region activity

Place a probe in the manifold (follows flow or manual path), freeze it,
and ask the LLM to interpret what the resulting ROI contribution shift means.
After LLM analysis, animated particles flow between donor/receiver ROIs.

Usage:
    python examples/roi_flow_mode.py [--hq] [--debug]

Controls:
    G then click — place probe in manifold (follows flow)
    M then click — manual path mode (click to extend path)
    Shift+G     — freeze probe, compute ROI delta, ask LLM
    V           — toggle ROI anim mode (path ↔ particles)
    1 / 2       — save current path as Path A / B
    D           — compare Path A vs Path B (cyan=A, magenta=B)
    C           — clear all
    +/-         — speed scale
    Q/Esc       — quit

Data:
    Run `python scripts/download_roi_flow_data.py` first to download all
    required data from HuggingFace.
"""

import argparse
import json
import sys
import threading
from pathlib import Path

import numpy as np
import vtk
from vtkmodules.util.numpy_support import numpy_to_vtk
from vtkmodules.vtkInteractionStyle import vtkInteractorStyleTrackballCamera
from vtkmodules.vtkFiltersGeneral import vtkSplineFilter
from vtkmodules.vtkFiltersCore import vtkTubeFilter

# Add project root to path
PROJECT_ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(PROJECT_ROOT))

from src.field_loader import load_field, TriLinearSampler
from src.colormaps import turbo_rgb01
from src.roi_flow import ManifoldToROIKNN, ROIFlowAnalyzer, ROIFlowLLM
# Reuse visualization utilities from src/main.py
from src.main import build_cloud, add_window_legend, VtkTextOverlay, wrap_inside

# Default paths — all ROI flow data lives in data/roi_flow/
DATA_DIR = PROJECT_ROOT / "data"
ROI_FLOW_DIR = DATA_DIR / "roi_flow"

DEFAULT_META = ROI_FLOW_DIR / "mdn_universal_raw_grid64_meta.json"
DEFAULT_OOS = ROI_FLOW_DIR / "universal_soul_2sdm_rest_points.ply"
DEFAULT_PROBE_EMBED = ROI_FLOW_DIR / "probe_embed.npy"
DEFAULT_PROBE_ROI = ROI_FLOW_DIR / "probe_roi.npy"
DEFAULT_PROBE_ROI_CENTERS = ROI_FLOW_DIR / "probe_roi_centers.npy"


# ---------------------------------------------------------------------------
# Utilities
# ---------------------------------------------------------------------------

def _load_points_any(path: Path) -> np.ndarray:
    """Load points from .npy, .ply, .obj, .stl files using VTK."""
    if path.suffix == ".npy":
        return np.load(str(path)).astype(np.float32)
    readers = {
        ".ply": vtk.vtkPLYReader, ".obj": vtk.vtkOBJReader,
        ".stl": vtk.vtkSTLReader, ".vtk": vtk.vtkPolyDataReader,
    }
    reader_cls = readers.get(path.suffix)
    if reader_cls is None:
        raise ValueError(f"Unsupported format: {path.suffix}")
    reader = reader_cls()
    reader.SetFileName(str(path))
    reader.Update()
    from vtkmodules.util.numpy_support import vtk_to_numpy
    return vtk_to_numpy(reader.GetOutput().GetPoints().GetData()).astype(np.float32)


def _map_roi_to_brain_regions(roi_centers: np.ndarray) -> list[str]:
    """Map ROI centroids to Allen brain regions."""
    try:
        from src.mesh_overlay import FlowMeshOverlay
        alignment_file = DATA_DIR / "brain_alignment.json"
        mesh_dir = DATA_DIR / "meshes"
        if not alignment_file.exists() or not mesh_dir.exists():
            raise FileNotFoundError("Brain meshes not available")

        ren = vtk.vtkRenderer()
        win = vtk.vtkRenderWindow()
        win.SetOffScreenRendering(1)
        win.AddRenderer(ren)

        overlay = FlowMeshOverlay(ren=ren, win=win, mesh_dir=mesh_dir,
                                   alignment_file=alignment_file)
        grid_cache = DATA_DIR / "label_grid_cache.npz"
        if not overlay.load_label_grid(grid_cache):
            overlay.build_label_grid()
            overlay.save_label_grid(grid_cache)

        names = []
        for i, pos in enumerate(roi_centers):
            key = overlay.get_region_at_point(pos)
            if key is None:
                key = overlay.find_nearest_region(pos, search_radius=6)
            if key is not None:
                region_name = overlay.get_region_name(key)
                hemi = "left" if pos[0] < 0 else "right"
                ap = "anterior" if pos[1] > 0 else "posterior"
                names.append(f"ROI_{i} ({region_name}, {hemi}, {ap})")
            else:
                hemi = "left" if pos[0] < 0 else "right"
                names.append(f"ROI_{i} ({hemi})")
        mapped = sum(1 for n in names if ", " in n)
        print(f"[roi-map] Mapped {mapped}/{len(names)} ROIs to brain regions")
        return names
    except Exception as e:
        print(f"[roi-map] Could not map ROIs to brain regions: {e}")
        names = []
        for i, pos in enumerate(roi_centers):
            hemi = "left" if pos[0] < 0 else "right"
            ap = "anterior" if pos[1] > 0 else "posterior"
            si = "superior" if pos[2] > 0 else "inferior"
            names.append(f"ROI_{i} ({hemi}, {ap}, {si})")
        return names


def _build_spline_trail(pts_array: np.ndarray, diag: float):
    """Build a smooth yellow spline+tube actor from a sequence of points."""
    if pts_array is None or len(pts_array) < 2:
        return None
    vpts = vtk.vtkPoints()
    vpts.SetData(numpy_to_vtk(pts_array.astype(np.float32), deep=True))
    pl = vtk.vtkPolyLine()
    pl.GetPointIds().SetNumberOfIds(len(pts_array))
    for i in range(len(pts_array)):
        pl.GetPointIds().SetId(i, i)
    lines = vtk.vtkCellArray()
    lines.InsertNextCell(pl)
    poly = vtk.vtkPolyData()
    poly.SetPoints(vpts)
    poly.SetLines(lines)

    spl = vtkSplineFilter()
    spl.SetInputData(poly)
    spl.SetSubdivideToLength()
    spl.SetLength(max(diag * 0.01, 1e-6))
    spl.Update()

    tube = vtkTubeFilter()
    tube.SetInputConnection(spl.GetOutputPort())
    tube.SetNumberOfSides(12)
    tube.SetRadius(diag * 0.004)
    tube.CappingOn()
    tube.Update()

    mapper = vtk.vtkPolyDataMapper()
    mapper.SetInputConnection(tube.GetOutputPort())
    actor = vtk.vtkActor()
    actor.SetMapper(mapper)
    actor.GetProperty().SetColor(1.0, 1.0, 0.0)
    actor.GetProperty().SetOpacity(0.85)
    actor.GetProperty().LightingOff()
    return actor


# ---------------------------------------------------------------------------
# ROI Panel (right side of split window)
# ---------------------------------------------------------------------------

class ROIPanel:
    """Renders ROI spheres colored/sized by activation (white-center bicolor)."""

    # Original bicolor scheme: white center, orange positive, blue negative
    POS_COLOR = np.array([255, 120, 0], np.float64)   # orange
    NEG_COLOR = np.array([0, 100, 255], np.float64)    # blue
    WHITE = np.array([255, 255, 255], np.float64)
    GRAY = np.array([200, 200, 200], np.float64) / 255.0

    def __init__(self, ren: vtk.vtkRenderer, centers: np.ndarray,
                 names: list[str]):
        self.ren = ren
        self.centers = centers.astype(np.float64)
        self.names = names
        self.n_rois = len(names)
        self.base_radius = 0.6
        self.radius_scale = 1.0
        self._spheres: list[vtk.vtkActor] = []
        self._setup()

    def _setup(self):
        for i in range(self.n_rois):
            sphere = vtk.vtkSphereSource()
            sphere.SetCenter(*self.centers[i])
            sphere.SetRadius(self.base_radius)
            sphere.SetPhiResolution(16)
            sphere.SetThetaResolution(16)
            sphere.Update()
            mapper = vtk.vtkPolyDataMapper()
            mapper.SetInputConnection(sphere.GetOutputPort())
            actor = vtk.vtkActor()
            actor.SetMapper(mapper)
            actor.GetProperty().SetColor(*self.GRAY)
            actor.GetProperty().SetOpacity(0.7)
            self.ren.AddActor(actor)
            self._spheres.append(actor)

    @staticmethod
    def _bicolor(value: float, abs_max: float):
        """White-center bicolor: white->orange (positive), white->blue (negative)."""
        if abs_max < 1e-12:
            return ROIPanel.GRAY
        t = np.clip(abs(value) / abs_max, 0.0, 1.0)
        if value >= 0:
            rgb = ROIPanel.WHITE * (1.0 - t) + ROIPanel.POS_COLOR * t
        else:
            rgb = ROIPanel.WHITE * (1.0 - t) + ROIPanel.NEG_COLOR * t
        return rgb / 255.0

    def update_values(self, values: np.ndarray):
        if values is None or len(values) != self.n_rois:
            return
        av = np.abs(values)
        abs_max = float(np.percentile(av, 97)) + 1e-8
        for i in range(self.n_rois):
            v = float(values[i])
            color = self._bicolor(v, abs_max)
            radius = self.base_radius + self.radius_scale * min(abs(v) / abs_max, 1.0)
            self._spheres[i].GetProperty().SetColor(*color)
            mapper = self._spheres[i].GetMapper()
            src = vtk.vtkSphereSource()
            src.SetCenter(*self.centers[i])
            src.SetRadius(radius)
            src.SetPhiResolution(16)
            src.SetThetaResolution(16)
            src.Update()
            mapper.SetInputData(src.GetOutput())
            mapper.Update()

    def dim_spheres(self, opacity: float = 0.15):
        for s in self._spheres:
            s.GetProperty().SetOpacity(opacity)

    def restore_spheres(self, opacity: float = 0.7):
        for s in self._spheres:
            s.GetProperty().SetOpacity(opacity)

    def reset_colors(self):
        for i in range(self.n_rois):
            self._spheres[i].GetProperty().SetColor(*self.GRAY)
            self._spheres[i].GetProperty().SetOpacity(0.7)


# ---------------------------------------------------------------------------
# ROI Flow Dots — animated particles flowing between donor/receiver ROIs
# ---------------------------------------------------------------------------

class ROIFlowDots:
    """Particles flowing from donor (negative delta) to receiver (positive delta) ROIs.

    Matches the original simulate_mdn_flow.py behavior: 120k particles,
    turbo colormap, endpoint fade, continuous emission from donor→receiver pairs.
    """

    def __init__(self, ren: vtk.vtkRenderer, roi_centers: np.ndarray,
                 max_particles: int = 120000):
        self.ren = ren
        self.C = roi_centers.astype(np.float32)
        self.Nmax = max_particles
        self.rng = np.random.default_rng(42)

        self.pos = np.zeros((self.Nmax, 3), np.float32)
        self.p0 = np.zeros((self.Nmax, 3), np.float32)
        self.dest = np.zeros((self.Nmax, 3), np.float32)
        self.alive = np.zeros(self.Nmax, bool)
        self.age = np.zeros(self.Nmax, np.float32)
        self.life = np.ones(self.Nmax, np.float32)
        self.speed = np.zeros(self.Nmax, np.float32)  # for colormap

        self.emitter_on = False
        self._neg_idx = np.array([], np.int32)
        self._pos_idx = np.array([], np.int32)
        self._pairs_probs = None
        self._emit_rate = 2000.0
        self._emit_accum = 0.0
        self._dt = 1.0 / 60.0
        # Endpoint fade parameters (matching original)
        self._fade_start = 0.08   # fade in over first 8%
        self._fade_end = 0.08     # fade out over last 8%
        # Burst mode (on by default — short sharp bursts with visible gap)
        self.burst_mode = True
        self.burst_period = 0.6       # seconds per cycle (fast)
        self.burst_emit_frac = 0.25   # emit during first 25% — short burst, long gap
        self._burst_timer = 0.0
        self.capture_accel = 2.5      # fast capture so particles clear out during gap
        # Mid-segment fade: make particles transparent in the middle of their journey
        # so you can clearly see source and destination endpoints
        self._mid_fade = True

        # Use build_cloud from src/main.py
        init_rgba = np.zeros((self.Nmax, 4), np.uint8)
        self._pts, self._colors, self._pd, self.actor = build_cloud(
            np.zeros((self.Nmax, 3), np.float32), rgba=init_rgba)
        self.actor.GetProperty().SetPointSize(1.2)
        self.actor.VisibilityOff()
        self.ren.AddActor(self.actor)

    def start_from_delta(self, delta: np.ndarray, top_frac: float = 0.15,
                         emit_rate: float = 2000.0):
        v = delta.astype(np.float64)
        k = max(1, int(len(v) * top_frac))
        idx_strong = np.argsort(-np.abs(v))[:k]
        v_str = v[idx_strong]
        eps = 1e-8
        self._pos_idx = idx_strong[v_str > eps]
        self._neg_idx = idx_strong[v_str < -eps]

        if len(self._pos_idx) == 0 or len(self._neg_idx) == 0:
            return

        p_don = np.abs(v[self._neg_idx])
        p_don /= p_don.sum()
        p_recv = np.abs(v[self._pos_idx])
        p_recv /= p_recv.sum()
        pairs = np.outer(p_don, p_recv).ravel()
        pairs /= pairs.sum()
        self._pairs_probs = pairs

        # Scale emit rate by delta magnitude (like original)
        l1 = float(np.sum(np.abs(v)))
        self._emit_rate = max(100.0, min(emit_rate * max(l1, 0.1), 8000.0))

        self.alive[:] = False
        self._emit_accum = 0.0
        self.emitter_on = True
        self.actor.VisibilityOn()
        print(f"[roi-flow] Started: {len(self._neg_idx)} donors -> "
              f"{len(self._pos_idx)} receivers, rate={self._emit_rate:.0f}/s")

    def tick(self):
        if not self.emitter_on:
            return

        # Burst phase: emit only during first fraction of each cycle
        if self.burst_mode:
            self._burst_timer = (self._burst_timer + self._dt) % max(self.burst_period, 1e-6)
            emit_phase = (self._burst_timer / max(self.burst_period, 1e-6)) < self.burst_emit_frac
        else:
            emit_phase = True

        # Advance alive particles
        idx = np.nonzero(self.alive)[0]
        if len(idx) > 0:
            # Accelerate during capture phase of burst
            speed_mul = 1.0
            if self.burst_mode and not emit_phase:
                speed_mul = self.capture_accel
            self.age[idx] += self._dt * speed_mul
            t = np.clip(self.age[idx] / self.life[idx], 0.0, 1.0)
            te = t * t * (3.0 - 2.0 * t)  # smoothstep easing
            seg = self.dest[idx] - self.p0[idx]
            self.pos[idx] = self.p0[idx] + seg * te[:, None]
            self.alive[idx[t >= 1.0]] = False

        # Spawn new particles (only during emit phase)
        Nd, Nr = len(self._neg_idx), len(self._pos_idx)
        if Nd > 0 and Nr > 0:
            self._emit_accum += self._emit_rate * self._dt
            n_new = int(self._emit_accum)
            if emit_phase and n_new > 0:
                self._emit_accum -= n_new
                free = np.nonzero(~self.alive)[0]
                use = free[:n_new]
                if len(use) > 0:
                    choice = self.rng.choice(Nd * Nr, size=len(use),
                                             replace=True, p=self._pairs_probs)
                    src_roi = self._neg_idx[choice // Nr]
                    dst_roi = self._pos_idx[choice % Nr]
                    self.p0[use] = self.C[src_roi]
                    self.dest[use] = self.C[dst_roi]
                    self.life[use] = self.rng.uniform(0.8, 1.8, size=len(use)).astype(np.float32)
                    self.age[use] = 0.0
                    self.alive[use] = True
                    seg_len = np.linalg.norm(self.dest[use] - self.p0[use], axis=1)
                    self.speed[use] = seg_len / (self.life[use] + 1e-8)
            elif not emit_phase:
                # Still accumulate but don't spend — creates burst on next emit phase
                pass

        self._update_display()

    def _update_display(self):
        idx = np.nonzero(self.alive)[0]
        P = np.zeros((self.Nmax, 3), np.float32)
        rgba = np.zeros((self.Nmax, 4), np.uint8)

        if len(idx) > 0:
            P[idx] = self.pos[idx]
            # Color from precomputed speed via turbo colormap
            spd = self.speed[idx]
            s_max = float(np.percentile(spd, 97)) if len(spd) > 0 else 1.0
            t = np.clip(spd / max(s_max, 1e-8), 0.0, 1.0)
            rgb = turbo_rgb01(t)

            # Endpoint fade + mid-segment dip
            # Particles are opaque near source and destination but fade in the
            # middle of their journey so you can clearly see where they come
            # from and where they go.
            prog = np.clip(self.age[idx] / self.life[idx], 0.0, 1.0)
            fs, fe = self._fade_start, self._fade_end

            # Base endpoint fade
            alpha = np.where(prog < fs, prog / max(fs, 1e-6),
                             np.where(prog > (1.0 - fe),
                                      (1.0 - prog) / max(fe, 1e-6), 1.0))

            # Mid-segment transparency dip (U-shaped alpha along journey)
            if getattr(self, '_mid_fade', False):
                # mid_alpha: 1.0 at endpoints, dips to 0.3 at center
                mid = 0.5
                spread = 0.35  # width of the dip
                dist_from_mid = np.abs(prog - mid) / spread
                mid_alpha = np.clip(0.3 + 0.7 * dist_from_mid, 0.3, 1.0)
                alpha = alpha * mid_alpha

            alpha = (np.clip(alpha, 0.0, 1.0) * 220).astype(np.uint8)

            rgba[idx, :3] = rgb
            rgba[idx, 3] = alpha

        self._pts.SetData(numpy_to_vtk(P, deep=True))
        self._colors.DeepCopy(numpy_to_vtk(rgba, deep=True))
        self._colors.Modified()
        self._pd.Modified()

    def stop(self):
        self.emitter_on = False
        self.alive[:] = False
        self.actor.VisibilityOff()

    def is_active(self):
        return self.emitter_on


# ---------------------------------------------------------------------------
# ROI Path Animation — animate ROI activation along the probe path
# ---------------------------------------------------------------------------

class ROIPathAnimation:
    """Snap between START and END ROI activation states — no intermediates.

    Fast toggle between the brain state at the beginning vs end of the path.
    200 most dynamic ROIs are animated; the rest stay dim gray.
    Active ROIs are vivid turbo-colored and blow up to 3.5× base radius.
    Quiet ROIs shrink to near-invisible gray dots.
    """

    N_TOP = 200             # number of ROIs to animate
    TOGGLE_TICKS = 40       # ticks to hold each state before flipping
    RADIUS_MIN_FRAC = 0.3   # quiet ROI radius multiplier
    RADIUS_MAX_FRAC = 3.5   # loudest ROI radius multiplier

    def __init__(self, roi_panel: ROIPanel, knn: 'ManifoldToROIKNN'):
        self.panel = roi_panel
        self.knn = knn
        self.active = False
        self._start_vals: np.ndarray | None = None   # (n_rois,)
        self._end_vals: np.ndarray | None = None     # (n_rois,)
        self._top_mask: np.ndarray | None = None     # (n_rois,) bool
        self._showing_end = False
        self._tick_counter = 0
        self._abs_max_start = 1.0
        self._abs_max_end = 1.0

    # ----- resampling helper (also used by ROIPathCompare) -----
    @staticmethod
    def _resample_path(pts: np.ndarray, n: int) -> np.ndarray:
        diffs = np.linalg.norm(np.diff(pts, axis=0), axis=1)
        cum = np.concatenate([[0.0], np.cumsum(diffs)])
        total_len = cum[-1]
        if total_len < 1e-9:
            return pts[:1]
        sample_dists = np.linspace(0, total_len, n)
        out = np.zeros((n, 3), np.float32)
        for i, sd in enumerate(sample_dists):
            idx = np.clip(np.searchsorted(cum, sd, side='right') - 1,
                          0, len(pts) - 2)
            seg_len = cum[idx + 1] - cum[idx]
            t = (sd - cum[idx]) / max(seg_len, 1e-9)
            out[i] = pts[idx] * (1 - t) + pts[idx + 1] * t
        return out

    def build_from_path(self, path_points: list[np.ndarray],
                        n_samples: int = 60):
        """Query ROI values at path start and end."""
        if len(path_points) < 2:
            return

        pts = np.array(path_points, dtype=np.float32)
        start_pt = pts[0]
        end_pt = pts[-1]

        start_vals = self.knn.query(start_pt)
        end_vals = self.knn.query(end_pt)

        # Select top N_TOP ROIs by absolute difference start→end
        diff = np.abs(end_vals - start_vals)
        n_top = min(self.N_TOP, self.panel.n_rois)
        top_indices = np.argsort(-diff)[:n_top]
        mask = np.zeros(self.panel.n_rois, bool)
        mask[top_indices] = True

        self._start_vals = start_vals
        self._end_vals = end_vals
        self._top_mask = mask
        self._abs_max_start = float(np.percentile(np.abs(start_vals[mask]), 97)) + 1e-8
        self._abs_max_end = float(np.percentile(np.abs(end_vals[mask]), 97)) + 1e-8
        self._showing_end = False
        self._tick_counter = 0
        self.active = True

        # Dim non-selected ROIs immediately
        for i in range(self.panel.n_rois):
            if not mask[i]:
                self.panel._spheres[i].GetProperty().SetColor(0.3, 0.3, 0.32)
                self.panel._spheres[i].GetProperty().SetOpacity(0.06)
                mapper = self.panel._spheres[i].GetMapper()
                src = vtk.vtkSphereSource()
                src.SetCenter(*self.panel.centers[i])
                src.SetRadius(self.panel.base_radius * 0.25)
                src.SetPhiResolution(6)
                src.SetThetaResolution(6)
                src.Update()
                mapper.SetInputData(src.GetOutput())
                mapper.Update()

        # Apply start state immediately
        self._apply_state(False)

        print(f"[roi-anim] Start/end toggle: {n_top} active ROIs, "
              f"flip every {self.TOGGLE_TICKS} ticks")

    def _apply_state(self, show_end: bool):
        """Snap all selected ROIs to start or end state."""
        vals = self._end_vals if show_end else self._start_vals
        abs_max = self._abs_max_end if show_end else self._abs_max_start

        for i in range(self.panel.n_rois):
            if not self._top_mask[i]:
                continue

            v = float(vals[i])
            intensity = min(abs(v) / abs_max, 1.0)

            # Pure turbo color — no gray blending for strong signals
            rgb_t = turbo_rgb01(np.array([intensity]))[0]
            tr, tg, tb = rgb_t[0] / 255.0, rgb_t[1] / 255.0, rgb_t[2] / 255.0

            # Blend: very low intensity stays slightly gray, rest is full turbo
            if intensity < 0.08:
                r, g, b = 0.45, 0.45, 0.47
            else:
                sat = min(intensity * 1.5, 1.0)
                r = 0.45 * (1 - sat) + tr * sat
                g = 0.45 * (1 - sat) + tg * sat
                b = 0.47 * (1 - sat) + tb * sat

            # Radius: massive difference between quiet and active
            radius = self.panel.base_radius * (
                self.RADIUS_MIN_FRAC
                + (self.RADIUS_MAX_FRAC - self.RADIUS_MIN_FRAC) * intensity
            )

            # Opacity
            opacity = 0.12 + 0.88 * intensity

            self.panel._spheres[i].GetProperty().SetColor(r, g, b)
            self.panel._spheres[i].GetProperty().SetOpacity(opacity)

            mapper = self.panel._spheres[i].GetMapper()
            src = vtk.vtkSphereSource()
            src.SetCenter(*self.panel.centers[i])
            src.SetRadius(radius)
            src.SetPhiResolution(16)
            src.SetThetaResolution(16)
            src.Update()
            mapper.SetInputData(src.GetOutput())
            mapper.Update()

        self._showing_end = show_end

    def tick(self):
        """Toggle between start and end states on a timer."""
        if not self.active:
            return

        self._tick_counter += 1
        if self._tick_counter >= self.TOGGLE_TICKS:
            self._tick_counter = 0
            self._apply_state(not self._showing_end)

    def stop(self):
        self.active = False
        self._start_vals = None
        self._end_vals = None
        self.panel.restore_spheres()
        self.panel.reset_colors()

    def is_active(self):
        return self.active


# ---------------------------------------------------------------------------
# ROI Path Comparison — animate difference between two paths
# ---------------------------------------------------------------------------

class ROIPathCompare:
    """Animate the difference between two saved paths.

    Shows how Path A vs Path B differ in their ROI activation trajectories.
    Color encodes which path activates each ROI more:
      - Cyan  = Path A dominant
      - Magenta = Path B dominant
      - Gray  = similar
    Sphere size pulses with magnitude of the difference.
    """

    N_TOP = 100
    RADIUS_MIN_FRAC = 0.6
    RADIUS_MAX_FRAC = 2.5

    def __init__(self, roi_panel: ROIPanel, knn: 'ManifoldToROIKNN'):
        self.panel = roi_panel
        self.knn = knn
        self.active = False
        self._diff_frames: np.ndarray | None = None  # (n, n_rois)
        self._top_mask: np.ndarray | None = None
        self._global_max: float = 1.0
        self._n_frames = 0
        self._cursor = 0.0
        self._direction = 1
        self._tick_accum = 0.0

    def build(self, path_a: list[np.ndarray], path_b: list[np.ndarray],
              n_samples: int = 60):
        """Build comparison animation from two paths."""
        if len(path_a) < 2 or len(path_b) < 2:
            print("[compare] Both paths need at least 2 points.")
            return

        pts_a = ROIPathAnimation._resample_path(
            np.array(path_a, np.float32), n_samples)
        pts_b = ROIPathAnimation._resample_path(
            np.array(path_b, np.float32), n_samples)

        frames_a = np.zeros((n_samples, self.panel.n_rois), np.float32)
        frames_b = np.zeros((n_samples, self.panel.n_rois), np.float32)
        for i in range(n_samples):
            frames_a[i] = self.knn.query(pts_a[i])
            frames_b[i] = self.knn.query(pts_b[i])

        # Difference: positive = A stronger, negative = B stronger
        diff = frames_a - frames_b

        # Top N most different ROIs
        max_diff = np.max(np.abs(diff), axis=0)
        n_top = min(self.N_TOP, self.panel.n_rois)
        top_idx = np.argsort(-max_diff)[:n_top]
        mask = np.zeros(self.panel.n_rois, bool)
        mask[top_idx] = True

        self._diff_frames = diff
        self._top_mask = mask
        self._global_max = float(np.percentile(max_diff[mask], 97)) + 1e-8
        self._n_frames = n_samples
        self._cursor = 0.0
        self._direction = 1
        self._tick_accum = 0.0
        self.active = True

        # Dim non-selected
        for i in range(self.panel.n_rois):
            if not mask[i]:
                self.panel._spheres[i].GetProperty().SetColor(0.3, 0.3, 0.3)
                self.panel._spheres[i].GetProperty().SetOpacity(0.06)
                mapper = self.panel._spheres[i].GetMapper()
                src = vtk.vtkSphereSource()
                src.SetCenter(*self.panel.centers[i])
                src.SetRadius(self.panel.base_radius * 0.35)
                src.SetPhiResolution(8)
                src.SetThetaResolution(8)
                src.Update()
                mapper.SetInputData(src.GetOutput())

        print(f"[compare] Built comparison: {n_samples} frames, "
              f"{n_top} active ROIs")

    def tick(self):
        if not self.active or self._diff_frames is None:
            return

        self._tick_accum += 0.5
        if self._tick_accum < 1.0:
            return
        self._tick_accum -= 1.0

        self._cursor += self._direction
        if self._cursor >= self._n_frames - 1:
            self._cursor = float(self._n_frames - 1)
            self._direction = -1
        elif self._cursor <= 0:
            self._cursor = 0.0
            self._direction = 1

        idx_lo = int(self._cursor)
        idx_hi = min(idx_lo + 1, self._n_frames - 1)
        frac = self._cursor - idx_lo
        cur = (self._diff_frames[idx_lo] * (1 - frac)
               + self._diff_frames[idx_hi] * frac)

        gm = self._global_max

        for i in range(self.panel.n_rois):
            if not self._top_mask[i]:
                continue
            d = float(cur[i])
            mag = min(abs(d) / gm, 1.0)

            # Cyan (A dominant) ↔ gray ↔ Magenta (B dominant)
            if d > 0:
                # Path A stronger → cyan
                r = 0.5 * (1 - mag)
                g = 0.5 + 0.5 * mag
                b = 0.5 + 0.5 * mag
            else:
                # Path B stronger → magenta
                r = 0.5 + 0.5 * mag
                g = 0.5 * (1 - mag)
                b = 0.5 + 0.5 * mag

            radius = self.panel.base_radius * (
                self.RADIUS_MIN_FRAC
                + (self.RADIUS_MAX_FRAC - self.RADIUS_MIN_FRAC) * mag)
            opacity = 0.2 + 0.8 * mag

            self.panel._spheres[i].GetProperty().SetColor(r, g, b)
            self.panel._spheres[i].GetProperty().SetOpacity(opacity)

            mapper = self.panel._spheres[i].GetMapper()
            src = vtk.vtkSphereSource()
            src.SetCenter(*self.panel.centers[i])
            src.SetRadius(radius)
            src.SetPhiResolution(16)
            src.SetThetaResolution(16)
            src.Update()
            mapper.SetInputData(src.GetOutput())
            mapper.Update()

    def stop(self):
        self.active = False
        self._diff_frames = None
        self.panel.restore_spheres()
        self.panel.reset_colors()

    def is_active(self):
        return self.active


# ---------------------------------------------------------------------------
# Main application
# ---------------------------------------------------------------------------

def main():
    ap = argparse.ArgumentParser(description="ROI Flow Mode — manifold + ROI visualization")

    ap.add_argument("--meta", type=Path, default=DEFAULT_META)
    ap.add_argument("--oos", type=Path, default=DEFAULT_OOS)
    ap.add_argument("--probe-embed", type=Path, default=DEFAULT_PROBE_EMBED)
    ap.add_argument("--probe-roi", type=Path, default=DEFAULT_PROBE_ROI)
    ap.add_argument("--probe-roi-centers", type=Path, default=DEFAULT_PROBE_ROI_CENTERS)
    ap.add_argument("--roi-names", type=Path, default=None,
                    help="JSON list of R ROI names (auto-mapped to brain regions if not given)")
    ap.add_argument("--probe-k", type=int, default=256)
    ap.add_argument("--probe-sigma", type=float, default=0.0)
    ap.add_argument("--dt", type=float, default=1.0)
    ap.add_argument("--speed-scale", type=float, default=1.0)
    ap.add_argument("--max-step-frac", type=float, default=0.01)
    ap.add_argument("--fps", type=int, default=60)
    ap.add_argument("--stride", type=int, default=3)
    ap.add_argument("--respawn-jitter", type=float, default=0.015)
    ap.add_argument("--window-size", type=int, nargs=2, default=[1600, 800])
    ap.add_argument("--hq", action="store_true", help="Use gpt-5.4 instead of gpt-5.4-mini")
    ap.add_argument("--debug", action="store_true")

    args = ap.parse_args()
    model = "gpt-5.4" if args.hq else "gpt-5.4-mini"
    print(f"[config] ROI Flow Mode | LLM model: {model}")

    # Check data exists
    for name, path in [("probe-embed", args.probe_embed),
                       ("probe-roi", args.probe_roi),
                       ("probe-roi-centers", args.probe_roi_centers)]:
        if not path.exists():
            print(f"\n[error] Missing: {path}")
            print(f"  Run: python scripts/download_roi_flow_data.py")
            sys.exit(1)

    # ---------- Load MDN field ----------
    print(f"[field] loading {args.meta} ...")
    if not args.meta.exists():
        print(f"\n[error] Missing: {args.meta}")
        print(f"  Run: python scripts/download_roi_flow_data.py")
        sys.exit(1)
    fld = load_field(args.meta)
    G = fld["G"]
    amin, amax = fld["amin"], fld["amax"]
    diag = float(np.linalg.norm(amax - amin))
    sampler = TriLinearSampler(fld["mean"], amin, amax)
    V_all = fld["mean"].reshape(-1, 3)
    vmax_mean = float(np.percentile(np.linalg.norm(V_all, axis=1), 99.5))
    target_step = args.max_step_frac * max(diag, 1e-9)
    print(f"[field] grid={G}, diag={diag:.4f}, vmax={vmax_mean:.6f}")

    # ---------- Load OOS points ----------
    oos_pts = _load_points_any(args.oos)
    print(f"[oos] {oos_pts.shape[0]} points")

    # Seed particles from OOS∩TRAIN overlap (like original script)
    from scipy.spatial import cKDTree
    train_pts = fld.get("TRAIN")
    if train_pts is not None and len(train_pts) > 0:
        overlap_radius = 0.01 * max(diag, 1e-9)
        tree = cKDTree(train_pts.astype(np.float32))
        dists, _ = tree.query(oos_pts.astype(np.float32), k=1)
        cand = oos_pts[dists <= overlap_radius]
        if len(cand) > 100:
            seed_pts = cand.astype(np.float32)
            print(f"[seed] {len(seed_pts)} OOS∩TRAIN overlap points")
        else:
            in_bounds = np.all((oos_pts >= amin) & (oos_pts <= amax), axis=1)
            seed_pts = oos_pts[in_bounds] if in_bounds.sum() > 100 else oos_pts
            print(f"[seed] {len(seed_pts)} OOS points (overlap too sparse)")
    else:
        in_bounds = np.all((oos_pts >= amin) & (oos_pts <= amax), axis=1)
        seed_pts = oos_pts[in_bounds] if in_bounds.sum() > 100 else oos_pts
        print(f"[seed] {len(seed_pts)} OOS points (no TRAIN)")

    # ---------- Load probe data ----------
    X_embed = np.load(str(args.probe_embed)).astype(np.float32)
    Y_roi = np.load(str(args.probe_roi)).astype(np.float32)
    C_centers = np.load(str(args.probe_roi_centers)).astype(np.float32)
    print(f"[probe] embed: {X_embed.shape}, roi: {Y_roi.shape}, centers: {C_centers.shape}")
    assert X_embed.shape[0] == Y_roi.shape[0]
    assert C_centers.shape[0] == Y_roi.shape[1]
    n_rois = Y_roi.shape[1]

    # ROI names
    if args.roi_names and args.roi_names.exists():
        roi_names = json.loads(args.roi_names.read_text(encoding="utf-8"))
    else:
        print("[roi-map] Mapping ROI centroids to brain regions...")
        roi_names = _map_roi_to_brain_regions(C_centers)
    assert len(roi_names) == n_rois

    knn = ManifoldToROIKNN(X_embed, Y_roi, k=args.probe_k, sigma=args.probe_sigma)
    analyzer = ROIFlowAnalyzer(roi_names, C_centers)
    llm = ROIFlowLLM(model=model, debug=args.debug)
    _sigma_str = "auto" if knn.sigma == 0 else f"{knn.sigma:.4f}"
    print(f"[knn] built with k={knn.k}, sigma={_sigma_str}")

    # ---------- VTK setup ----------
    win = vtk.vtkRenderWindow()
    win.SetSize(*args.window_size)
    win.SetWindowName("mindVisualizer — ROI Flow Mode")

    ren_manifold = vtk.vtkRenderer()
    ren_manifold.SetViewport(0.0, 0.0, 0.55, 1.0)
    ren_manifold.SetBackground(0.0, 0.0, 0.0)
    win.AddRenderer(ren_manifold)

    ren_roi = vtk.vtkRenderer()
    ren_roi.SetViewport(0.55, 0.0, 1.0, 1.0)
    ren_roi.SetBackground(0.02, 0.02, 0.04)
    win.AddRenderer(ren_roi)

    iren = vtk.vtkRenderWindowInteractor()
    iren.SetRenderWindow(win)
    style = vtkInteractorStyleTrackballCamera()
    iren.SetInteractorStyle(style)

    # ---------- OOS overlay (reuse build_cloud from src/main.py) ----------
    n_oos = len(oos_pts)
    oos_rgba = np.zeros((n_oos, 4), np.uint8)
    oos_rgba[:, 0] = 100; oos_rgba[:, 1] = 100; oos_rgba[:, 2] = 120; oos_rgba[:, 3] = 72
    _, oos_colors, oos_pd, oos_actor = build_cloud(oos_pts, point_size=1.5, rgba=oos_rgba)
    ren_manifold.AddActor(oos_actor)

    # ---------- MDN Particles ----------
    rng = np.random.default_rng(0)
    n_particles = min(len(seed_pts), 12000)
    sel = rng.choice(len(seed_pts), size=n_particles,
                     replace=len(seed_pts) < n_particles)
    P = seed_pts[sel].copy()
    overlap_sigma = args.respawn_jitter * max(diag, 1e-9)
    P += rng.standard_normal(P.shape).astype(np.float32) * overlap_sigma
    np.clip(P, amin, amax, out=P)

    ttl_lo, ttl_hi = 30, 120
    ttl = rng.integers(ttl_lo, ttl_hi + 1, size=n_particles, dtype=np.int32)
    ages = rng.integers(0, ttl_hi, size=n_particles, dtype=np.int32)

    # Initial color: white with low alpha (like original)
    init_rgba = np.tile(np.array([[255, 255, 255, 48]], np.uint8), (n_particles, 1))
    pts_vtk, colors_arr, p_pd, p_actor = build_cloud(P, point_size=2.0, rgba=init_rgba)
    ren_manifold.AddActor(p_actor)

    dt = [args.dt]
    speed_scale = [args.speed_scale]

    # ---------- ROI Panel ----------
    roi_panel = ROIPanel(ren_roi, C_centers, roi_names)

    # ---------- ROI Flow Dots (particle mode — alternative) ----------
    roi_flow_dots = ROIFlowDots(ren_roi, C_centers)

    # ---------- ROI Path Animation (default mode) ----------
    roi_path_anim = ROIPathAnimation(roi_panel, knn)
    roi_anim_mode = ["path"]  # "path" (default) or "particles"

    # ---------- ROI Path Comparison ----------
    roi_path_compare = ROIPathCompare(roi_panel, knn)
    saved_paths = {"A": None, "B": None}  # saved path point lists
    saved_trail_actors = {"A": None, "B": None}  # trail actors for saved paths

    # ---------- Probe state ----------
    probe_mode = [None]  # None, "flow", "manual"
    probe_pos = [None]
    probe_path = []
    probe_start_roi = [None]
    probe_frozen = [False]
    trail_actor_ref = [None]  # reference to current spline trail actor
    placement_mode = [None]  # None, "flow", "manual" — set by G/M key, consumed by click

    # Probe marker — proportional to manifold size
    probe_radius = diag * 0.008
    probe_sphere = vtk.vtkSphereSource()
    probe_sphere.SetRadius(probe_radius)
    probe_sphere.SetPhiResolution(16)
    probe_sphere.SetThetaResolution(16)
    probe_mapper = vtk.vtkPolyDataMapper()
    probe_mapper.SetInputConnection(probe_sphere.GetOutputPort())
    probe_actor = vtk.vtkActor()
    probe_actor.SetMapper(probe_mapper)
    probe_actor.GetProperty().SetColor(1.0, 1.0, 0.0)
    probe_actor.VisibilityOff()
    ren_manifold.AddActor(probe_actor)

    # ---------- Text overlays (reuse from src/main.py) ----------
    overlay = VtkTextOverlay(ren_manifold, max_lines=6)

    # Legend (top-left, white Courier)
    add_window_legend(ren_manifold, [
        "G then click  place probe (flow)",
        "M then click  manual path mode",
        "Shift+G       freeze & interpret",
        "V             toggle ROI anim mode",
        "1 / 2         save path A / B",
        "D             compare A vs B",
        "C             clear all",
        "+/-           speed  |  Q quit",
    ], font_px=12)

    gpt_pending = {"result": None}

    # ROI mode legend (top-left of ROI viewport)
    roi_legend_actor = vtk.vtkTextActor()
    roi_legend_actor.SetInput("")
    roi_leg_tp = roi_legend_actor.GetTextProperty()
    roi_leg_tp.SetFontFamilyToCourier()
    roi_leg_tp.SetFontSize(11)
    roi_leg_tp.SetColor(0.8, 0.8, 0.8)
    roi_leg_tp.SetOpacity(0.7)
    roi_leg_tp.SetJustificationToLeft()
    roi_leg_tp.SetVerticalJustificationToTop()
    roi_legend_actor.GetPositionCoordinate().SetCoordinateSystemToNormalizedViewport()
    roi_legend_actor.GetPositionCoordinate().SetValue(0.02, 0.97)
    ren_roi.AddActor(roi_legend_actor)

    # Static legend for ROI viewport (bottom-left)
    add_window_legend(ren_roi, [
        "V  switch view mode",
    ], font_px=11)

    def _update_roi_legend():
        if roi_flow_dots.is_active():
            roi_legend_actor.SetInput(
                "PARTICLE FLOW\n"
                "Particles: donor -> receiver\n"
                "Burst mode | V to switch")
        elif roi_path_anim.is_active():
            roi_legend_actor.SetInput(
                "PATH ANIMATION\n"
                "Flipping START <-> END\n"
                "V to switch view")
        elif roi_path_compare.is_active():
            roi_legend_actor.SetInput(
                "PATH COMPARISON\n"
                "Cyan=A  Magenta=B")
        else:
            roi_legend_actor.SetInput("")

    # LLM result text panel in the ROI viewport (bottom, green text)
    gpt_text_actor = vtk.vtkTextActor()
    gpt_text_actor.SetInput("")
    gpt_tp = gpt_text_actor.GetTextProperty()
    gpt_tp.SetFontFamilyToCourier()
    gpt_tp.SetFontSize(11)
    gpt_tp.SetColor(0.3, 0.9, 0.3)  # green
    gpt_tp.SetOpacity(0.9)
    gpt_tp.SetJustificationToLeft()
    gpt_tp.SetVerticalJustificationToBottom()
    gpt_text_actor.GetPositionCoordinate().SetCoordinateSystemToNormalizedViewport()
    gpt_text_actor.GetPositionCoordinate().SetValue(0.02, 0.02)
    gpt_text_actor.VisibilityOff()
    ren_roi.AddActor(gpt_text_actor)

    def _word_wrap(text: str, width: int = 50, max_lines: int = 18) -> str:
        """Word-wrap text to fit the ROI viewport."""
        import textwrap
        lines = []
        for paragraph in text.split('\n'):
            wrapped = textwrap.wrap(paragraph, width=width) or ['']
            lines.extend(wrapped)
        if len(lines) > max_lines:
            lines = lines[:max_lines - 1] + ['...']
        return '\n'.join(lines)

    # ---------- Helper: remove trail ----------
    def _remove_trail():
        if trail_actor_ref[0] is not None:
            ren_manifold.RemoveActor(trail_actor_ref[0])
            trail_actor_ref[0] = None

    # ---------- Helper: rebuild trail from path ----------
    def _rebuild_trail():
        _remove_trail()
        if len(probe_path) >= 2:
            pts_arr = np.array(probe_path, dtype=np.float32)
            actor = _build_spline_trail(pts_arr, diag)
            if actor is not None:
                ren_manifold.AddActor(actor)
                trail_actor_ref[0] = actor

    # ---------- Helper: place probe ----------
    def _place_probe(pos, mode):
        pos = np.clip(pos, amin, amax)
        probe_mode[0] = mode
        probe_frozen[0] = False
        probe_pos[0] = pos.copy()
        probe_path.clear()
        probe_path.append(pos.copy())
        probe_start_roi[0] = knn.query(pos)

        roi_panel.update_values(probe_start_roi[0])
        roi_flow_dots.stop()
        roi_path_anim.stop()
        roi_path_compare.stop()
        _remove_trail()

        probe_sphere.SetCenter(*pos)
        probe_actor.VisibilityOn()

        overlay.clear_log()
        overlay.hide_gpt()
        gpt_pending["result"] = None

        mode_label = "following flow" if mode == "flow" else "manual (click to extend)"
        overlay.add_log(f"Probe placed [{mode_label}]")
        print(f"[probe] placed at ({pos[0]:.4f}, {pos[1]:.4f}, {pos[2]:.4f}) [{mode}]")

    # ---------- Helper: freeze and analyze ----------
    def _freeze_and_analyze():
        if probe_mode[0] is None or probe_frozen[0] or probe_pos[0] is None:
            return
        probe_frozen[0] = True

        end_roi = knn.query(probe_pos[0])
        delta = analyzer.compute_delta(probe_start_roi[0], end_roi)

        overlay.add_log("Probe frozen — asking LLM...")
        print("[probe] frozen, computing ROI delta...")

        context = analyzer.build_llm_context(delta)
        print(f"\n{context}\n")

        # Start ROI animation (path mode or particle mode)
        if roi_anim_mode[0] == "path":
            roi_flow_dots.stop()
            roi_path_anim.build_from_path(probe_path, n_samples=60)
        else:
            roi_path_anim.stop()
            roi_panel.update_values(delta)
            roi_flow_dots.start_from_delta(delta)
            roi_panel.dim_spheres(0.15)
        _update_roi_legend()

        # Rebuild trail as smooth spline now that path is complete
        _rebuild_trail()

        def _gpt_worker():
            gpt_pending["result"] = llm.interpret_roi_flow(context)

        t = threading.Thread(target=_gpt_worker, daemon=True)
        t.start()

    # ---------- Helper: reset probe (keep saved paths) ----------
    def _reset_probe():
        probe_mode[0] = None
        probe_frozen[0] = False
        probe_pos[0] = None
        probe_path.clear()
        probe_start_roi[0] = None
        probe_actor.VisibilityOff()
        _remove_trail()
        roi_panel.reset_colors()
        roi_panel.restore_spheres()
        roi_flow_dots.stop()
        roi_path_anim.stop()
        roi_path_compare.stop()
        roi_legend_actor.SetInput("")
        gpt_text_actor.SetInput("")
        gpt_text_actor.VisibilityOff()
        overlay.clear_log()
        overlay.hide_gpt()
        gpt_pending["result"] = None

    # ---------- Helper: clear ----------
    def _clear_all():
        probe_mode[0] = None
        probe_frozen[0] = False
        probe_pos[0] = None
        probe_path.clear()
        probe_start_roi[0] = None
        probe_actor.VisibilityOff()
        _remove_trail()
        roi_panel.reset_colors()
        roi_panel.restore_spheres()
        roi_flow_dots.stop()
        roi_path_anim.stop()
        roi_path_compare.stop()
        # Remove saved path trails
        for k in ("A", "B"):
            if saved_trail_actors[k] is not None:
                ren_manifold.RemoveActor(saved_trail_actors[k])
                saved_trail_actors[k] = None
        overlay.clear_log()
        overlay.hide_gpt()
        gpt_text_actor.SetInput("")
        gpt_text_actor.VisibilityOff()
        roi_legend_actor.SetInput("")
        gpt_pending["result"] = None

    # ---------- Callbacks ----------
    trail_rebuild_counter = [0]

    def on_key(obj, event):
        key = iren.GetKeySym()
        shift = bool(iren.GetShiftKey())

        if key == "g" and not shift:
            # Enter flow placement mode — next click places probe
            placement_mode[0] = "flow"
            overlay.add_log("Click to place probe (flow mode)")

        elif key == "m" and not shift:
            if probe_mode[0] == "manual" and not probe_frozen[0] and probe_pos[0] is not None:
                # Already in manual mode with active probe — extend path
                placement_mode[0] = "manual_extend"
                overlay.add_log("Click to add path point")
            else:
                # Start new manual path
                placement_mode[0] = "manual"
                overlay.add_log("Click to place probe (manual mode)")

        elif key == "G" or (key == "g" and shift):
            _freeze_and_analyze()

        elif key in ("c", "C"):
            placement_mode[0] = None
            _clear_all()

        elif key in ("plus", "equal"):
            speed_scale[0] *= 1.25
            overlay.add_log(f"Speed: {speed_scale[0]:.2f}")
        elif key == "minus":
            speed_scale[0] /= 1.25
            overlay.add_log(f"Speed: {speed_scale[0]:.2f}")

        elif key == "v" and not shift:
            # Toggle between path animation and particle flow mode
            if roi_anim_mode[0] == "path":
                roi_anim_mode[0] = "particles"
                overlay.add_log("ROI mode: particles")
            else:
                roi_anim_mode[0] = "path"
                overlay.add_log("ROI mode: path animation")
            # If frozen, restart with new mode
            if probe_frozen[0] and probe_pos[0] is not None:
                end_roi = knn.query(probe_pos[0])
                delta = analyzer.compute_delta(probe_start_roi[0], end_roi)
                if roi_anim_mode[0] == "path":
                    roi_flow_dots.stop()
                    roi_panel.restore_spheres()
                    roi_path_anim.build_from_path(probe_path, n_samples=60)
                else:
                    roi_path_anim.stop()
                    roi_panel.update_values(delta)
                    roi_flow_dots.start_from_delta(delta)
                    roi_panel.dim_spheres(0.15)
                _update_roi_legend()

        elif key == "1" and not shift:
            # Save current path as Path A, then reset probe for next path
            if len(probe_path) >= 2:
                saved_paths["A"] = [p.copy() for p in probe_path]
                overlay.add_log(f"Path A saved ({len(probe_path)} pts)")
                print(f"[compare] Path A saved: {len(probe_path)} points")
                if saved_trail_actors["A"] is not None:
                    ren_manifold.RemoveActor(saved_trail_actors["A"])
                ta = _build_spline_trail(np.array(probe_path, np.float32), diag)
                if ta is not None:
                    ta.GetProperty().SetColor(0.0, 0.9, 0.9)  # cyan
                    ta.GetProperty().SetOpacity(0.6)
                    ren_manifold.AddActor(ta)
                    saved_trail_actors["A"] = ta
                # Reset probe so next path starts fresh
                _reset_probe()
            else:
                overlay.add_log("Need a path first (place+freeze probe)")

        elif key == "2" and not shift:
            # Save current path as Path B, then reset probe for next path
            if len(probe_path) >= 2:
                saved_paths["B"] = [p.copy() for p in probe_path]
                overlay.add_log(f"Path B saved ({len(probe_path)} pts)")
                print(f"[compare] Path B saved: {len(probe_path)} points")
                if saved_trail_actors["B"] is not None:
                    ren_manifold.RemoveActor(saved_trail_actors["B"])
                tb = _build_spline_trail(np.array(probe_path, np.float32), diag)
                if tb is not None:
                    tb.GetProperty().SetColor(0.9, 0.0, 0.9)  # magenta
                    tb.GetProperty().SetOpacity(0.6)
                    ren_manifold.AddActor(tb)
                    saved_trail_actors["B"] = tb
                _reset_probe()
            else:
                overlay.add_log("Need a path first (place+freeze probe)")

        elif key == "d" and not shift:
            # Compare Path A vs Path B with LLM analysis
            if saved_paths["A"] is None or saved_paths["B"] is None:
                overlay.add_log("Save Path A (1) and Path B (2) first")
            else:
                roi_flow_dots.stop()
                roi_path_anim.stop()
                roi_panel.restore_spheres()
                roi_path_compare.build(saved_paths["A"], saved_paths["B"])
                overlay.add_log("Comparing Path A (cyan) vs B (magenta)...")
                print("[compare] Started path comparison animation")

                # Compute deltas and LLM context for both paths
                pa, pb = saved_paths["A"], saved_paths["B"]
                roi_start_a = knn.query(np.array(pa[0], np.float32))
                roi_end_a = knn.query(np.array(pa[-1], np.float32))
                delta_a = analyzer.compute_delta(roi_start_a, roi_end_a)
                ctx_a = analyzer.build_llm_context(delta_a)

                roi_start_b = knn.query(np.array(pb[0], np.float32))
                roi_end_b = knn.query(np.array(pb[-1], np.float32))
                delta_b = analyzer.compute_delta(roi_start_b, roi_end_b)
                ctx_b = analyzer.build_llm_context(delta_b)

                print(f"\n--- PATH A ---\n{ctx_a}\n--- PATH B ---\n{ctx_b}\n")

                # LLM comparison in background
                def _compare_worker():
                    gpt_pending["result"] = llm.compare_two_paths(ctx_a, ctx_b)

                t = threading.Thread(target=_compare_worker, daemon=True)
                t.start()
                overlay.add_log("Asking LLM to compare paths...")
                _update_roi_legend()

        elif key in ("q", "Escape"):
            iren.TerminateApp()

    iren.AddObserver("KeyPressEvent", on_key)

    def on_click(obj, event):
        if placement_mode[0] is None:
            return  # No placement pending — normal click behavior

        x, y = iren.GetEventPosition()
        picker = vtk.vtkWorldPointPicker()
        picker.Pick(x, y, 0, ren_manifold)
        pos = np.array(picker.GetPickPosition(), dtype=np.float32)
        pos = np.clip(pos, amin, amax)

        mode = placement_mode[0]
        placement_mode[0] = None  # Consume the placement

        if mode == "flow":
            _place_probe(pos, "flow")
        elif mode == "manual":
            _place_probe(pos, "manual")
        elif mode == "manual_extend":
            probe_pos[0] = pos.copy()
            probe_path.append(pos.copy())
            probe_sphere.SetCenter(*pos)
            _rebuild_trail()
            roi_vec = knn.query(pos)
            roi_panel.update_values(roi_vec)
            overlay.add_log(f"Path point {len(probe_path)}")

    iren.AddObserver("LeftButtonPressEvent", on_click)

    # ---------- Timer loop ----------
    def on_timer(obj, event):
        nonlocal P, ages, ttl

        # Check for GPT result
        if gpt_pending["result"] is not None:
            text = gpt_pending["result"]
            gpt_pending["result"] = None
            overlay.show_gpt(text)
            overlay.add_log("LLM interpretation ready.")
            # Show in ROI viewport as well
            gpt_text_actor.SetInput(_word_wrap(text, width=45, max_lines=16))
            gpt_text_actor.VisibilityOn()
            print(f"\n--- ROI FLOW INTERPRETATION ---\n{text}\n--- END ---\n")

        # Advect particles
        V = sampler.sample_vec(P)
        step = dt[0] * speed_scale[0] * (target_step / max(vmax_mean, 1e-9))
        P[:] += V * step
        np.clip(P, amin, amax, out=P)

        # Death/respawn
        ages += 1
        dead = ages >= ttl
        if np.any(dead):
            n_dead = dead.sum()
            sel2 = rng.integers(0, len(seed_pts), size=n_dead)
            base = seed_pts[sel2]
            J = rng.standard_normal((n_dead, 3)).astype(np.float32) * overlap_sigma
            P[dead] = np.clip(base + J, amin, amax)
            ages[dead] = 0
            ttl[dead] = rng.integers(ttl_lo, ttl_hi + 1, size=n_dead, dtype=np.int32)

        # Update colors (vectorized turbo colormap)
        speeds = np.linalg.norm(V, axis=1)
        s_max = float(np.percentile(speeds, 97)) + 1e-8
        t_vals = np.clip(speeds / s_max, 0.0, 1.0)
        rgb = turbo_rgb01(t_vals)
        alpha = np.full(n_particles, 200, np.uint8)
        rgba = np.concatenate([rgb, alpha[:, None]], axis=1)

        pts_vtk.SetData(numpy_to_vtk(P, deep=True))
        colors_arr.DeepCopy(numpy_to_vtk(rgba, deep=True))
        colors_arr.Modified()
        p_pd.Modified()

        # Advect probe (flow mode)
        if probe_mode[0] == "flow" and not probe_frozen[0] and probe_pos[0] is not None:
            pos = probe_pos[0]
            v = sampler.sample_vec(pos.reshape(1, 3))[0]
            pos += v * step
            np.clip(pos, amin, amax, out=pos)
            probe_pos[0] = pos
            probe_sphere.SetCenter(*pos)
            probe_path.append(pos.copy())

            # Rebuild trail periodically (every 20 steps) for smooth spline
            trail_rebuild_counter[0] += 1
            if trail_rebuild_counter[0] >= 20:
                trail_rebuild_counter[0] = 0
                _rebuild_trail()

            if len(probe_path) % 5 == 0:
                roi_vec = knn.query(pos)
                roi_panel.update_values(roi_vec)

        # Tick ROI animation (path, particles, or compare — whichever is active)
        if roi_path_anim.is_active():
            roi_path_anim.tick()
        if roi_flow_dots.is_active():
            roi_flow_dots.tick()
        if roi_path_compare.is_active():
            roi_path_compare.tick()

        win.Render()

    iren.Initialize()
    iren.CreateRepeatingTimer(int(1000 / args.fps))
    iren.AddObserver("TimerEvent", on_timer)

    ren_manifold.ResetCamera()
    ren_roi.ResetCamera()

    print("\n=== ROI Flow Mode ===")
    print("  G then click  — place probe (follows flow)")
    print("  M then click  — manual path mode (click to extend)")
    print("  Shift+G       — freeze & interpret ROI flow")
    print("  V             — toggle ROI anim: path ↔ particles")
    print("  1 / 2         — save current path as A / B")
    print("  D             — compare Path A vs Path B")
    print("  C             — clear all")
    print("  +/-           — speed scale")
    print("  Q             — quit")
    print("=" * 40)

    win.Render()
    iren.Start()


if __name__ == "__main__":
    main()