tracescope/query.py

"""
TraceQuery — Programmatic query API for the TraceScope semantic space.

After running the full pipeline once, wrap the AnalysisResult in a TraceQuery
to get fast, LLM-agent-friendly methods for querying the computed space:

    query = TraceQuery(result, embedding_provider)
    lookup = query.get_lookup()
    path_info = query.explain_path(["hello", "how are you", "goodbye"])
    flow_info = query.query_flow_at("tell me about work")
    dir_info  = query.query_direction_at(["start text", "end text"])
    sim       = query.path_similarity(["a","b"], ["x","y"])
"""

from __future__ import annotations

import logging
from typing import Any, Dict, List, Optional

import numpy as np

from tracescope.models.analysis import AnalysisResult
from tracescope.providers.embedding import EmbeddingProvider
from tracescope.analysis.explainer import compute_cosine_similarity

logger = logging.getLogger(__name__)


class TraceQuery:
    """Fast programmatic query interface over a computed TraceScope space.

    Args:
        result: AnalysisResult from AnalysisPipeline.analyze().
        embedding_provider: Same provider used during pipeline run.
        explainer: Optional SemanticExplainer for LLM-powered explanations.
    """

    def __init__(
        self,
        result: AnalysisResult,
        embedding_provider: EmbeddingProvider,
        explainer=None,
    ):
        self._result = result
        self._emb = embedding_provider
        self._explainer = explainer

        # Pre-compute lookup data
        self._lookup = self._build_lookup()

    # ─── Lookup table ─────────────────────────────────────────────────

    def _build_lookup(self) -> dict:
        r = self._result
        pts = r.projected_3d

        # Axis info
        axis_labels = list(r.axis_info.labels)
        axis_mins = pts.min(axis=0).tolist()
        axis_maxs = pts.max(axis=0).tolist()

        # Cluster info
        cluster_info = []
        for c in range(r.clusters.n_clusters):
            indices = [i for i, l in enumerate(r.clusters.labels) if l == c]
            label = r.cluster_labels[c] if c < len(r.cluster_labels) else f"Cluster {c}"
            centroid_3d = r.cluster_centroids_3d[c].tolist() if r.cluster_centroids_3d is not None else None
            sample_texts = [r.session.entries[i].text for i in indices[:5]]
            cluster_info.append({
                "id": c,
                "label": label,
                "centroid_3d": centroid_3d,
                "size": len(indices),
                "sample_texts": sample_texts,
            })

        lookup = {
            "axis_labels": axis_labels,
            "axis_ranges": [
                {"axis": axis_labels[i], "min": axis_mins[i], "max": axis_maxs[i]}
                for i in range(3)
            ],
            "clusters": cluster_info,
            "n_points": len(r.session),
            "embedding_model": r.embedding_model,
            "embedding_dim": r.embeddings.shape[1],
            "has_flow": r.velocity_grid is not None,
            "flow_bounds": {
                "axis_min": r.axis_min.tolist() if r.axis_min is not None else None,
                "axis_max": r.axis_max.tolist() if r.axis_max is not None else None,
            },
        }

        # Score channels (only if any exist)
        score_channels = r.score_channels
        if score_channels:
            score_info = {}
            for ch in score_channels:
                entry_vals = [v for v in r.get_entry_scores(ch) if v is not None]
                path_vals = list(r.get_path_scores(ch).values())
                all_vals = entry_vals + path_vals
                score_info[ch] = {
                    "entry_count": len(entry_vals),
                    "path_count": len(path_vals),
                    "min": round(min(all_vals), 4) if all_vals else None,
                    "max": round(max(all_vals), 4) if all_vals else None,
                    "mean": round(sum(all_vals) / len(all_vals), 4) if all_vals else None,
                }
            lookup["score_channels"] = score_info

        return lookup

    def get_lookup(self) -> dict:
        """Return the pre-computed lookup table with space metadata."""
        return self._lookup

    # ─── Internal helpers ─────────────────────────────────────────────

    def _embed_texts(self, texts: List[str]) -> np.ndarray:
        """Embed texts using the same provider as the pipeline."""
        return self._emb.embed_batch(texts)

    def _project_to_3d(self, embeddings: np.ndarray) -> np.ndarray:
        """Project high-dim embeddings to 3D using the stored fitted reducer.

        Falls back to nearest-neighbor interpolation if no reducer is stored.
        """
        reducer = self._result.fitted_reducer
        if reducer is not None and hasattr(reducer, "transform"):
            # L2-normalize to match pipeline preprocessing
            norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
            norms[norms == 0] = 1.0
            normed = embeddings / norms
            return reducer.transform(normed).astype(np.float32)

        # Fallback: weighted k-NN interpolation in high-dim space
        logger.warning("No fitted reducer available, using k-NN projection fallback")
        return self._knn_project(embeddings)

    def _knn_project(self, new_embeddings: np.ndarray, k: int = 5) -> np.ndarray:
        """Project new points via weighted k-NN in embedding space."""
        orig_embs = self._result.embeddings
        orig_3d = self._result.projected_3d

        # Normalize both
        orig_norms = np.linalg.norm(orig_embs, axis=1, keepdims=True)
        orig_norms[orig_norms == 0] = 1.0
        orig_normed = orig_embs / orig_norms

        new_norms = np.linalg.norm(new_embeddings, axis=1, keepdims=True)
        new_norms[new_norms == 0] = 1.0
        new_normed = new_embeddings / new_norms

        # Cosine similarity matrix: (M, N)
        sims = new_normed @ orig_normed.T

        projected = np.zeros((len(new_embeddings), 3), dtype=np.float32)
        for i in range(len(new_embeddings)):
            top_k = np.argsort(-sims[i])[:k]
            weights = np.maximum(sims[i, top_k], 0.0)
            w_sum = weights.sum()
            if w_sum > 0:
                weights /= w_sum
            else:
                weights = np.ones(k) / k
            projected[i] = (weights[:, None] * orig_3d[top_k]).sum(axis=0)

        return projected

    def _compute_axis_pcts(self, point_3d: np.ndarray) -> dict:
        """Compute axis percentages for a 3D point."""
        pts = self._result.projected_3d
        mins = pts.min(axis=0)
        maxs = pts.max(axis=0)
        ranges = maxs - mins
        ranges[ranges == 0] = 1.0
        pct = np.clip(((point_3d - mins) / ranges) * 100, 0, 100)
        labels = self._result.axis_info.labels
        return {labels[i]: round(float(pct[i]), 1) for i in range(3)}

    def _compute_cluster_distances(self, point_3d: np.ndarray) -> dict:
        """Compute closeness % to each cluster centroid."""
        r = self._result
        pts = r.projected_3d
        max_dist = float(np.linalg.norm(pts.max(axis=0) - pts.min(axis=0)))
        if max_dist == 0:
            max_dist = 1.0

        distances = {}
        for c in range(r.clusters.n_clusters):
            centroid = r.cluster_centroids_3d[c]
            dist = float(np.linalg.norm(point_3d - centroid))
            closeness = max(0, (1 - dist / max_dist) * 100)
            label = r.cluster_labels[c] if c < len(r.cluster_labels) else f"Cluster {c}"
            distances[label] = round(closeness, 1)
        return distances

    def _find_nearest_points(self, point_3d: np.ndarray, k: int = 5) -> List[dict]:
        """Find k nearest original points to a 3D position."""
        pts = self._result.projected_3d
        dists = np.linalg.norm(pts - point_3d, axis=1)
        nearest = np.argsort(dists)[:k]
        results = []
        for idx in nearest:
            idx = int(idx)
            entry = self._result.session.entries[idx]
            d = {
                "index": idx,
                "text": entry.text,
                "distance": round(float(dists[idx]), 4),
                "cluster": self._result.clusters.labels[idx],
            }
            if entry.scores:
                d["scores"] = entry.scores
            return_path_id = entry.path_id
            if return_path_id is not None and return_path_id in self._result.session.path_scores:
                d["path_scores"] = self._result.session.path_scores[return_path_id]
            results.append(d)
        return results

    def _sample_velocity(self, point_3d: np.ndarray) -> Optional[np.ndarray]:
        """Sample velocity from the velocity grid at a 3D point."""
        r = self._result
        if r.velocity_grid is None or r.axis_min is None or r.axis_max is None:
            return None

        from tracescope.visualization.flow_field import FlowFieldSystem
        # Use the static sample_velocity method logic directly
        grid = r.velocity_grid
        G = grid.shape[0]
        span = r.axis_max - r.axis_min

        result_v = np.zeros(3, dtype=np.float32)
        nx = (point_3d[0] - r.axis_min[0]) / span[0] * (G - 1) if span[0] > 0 else 0
        ny = (point_3d[1] - r.axis_min[1]) / span[1] * (G - 1) if span[1] > 0 else 0
        nz = (point_3d[2] - r.axis_min[2]) / span[2] * (G - 1) if span[2] > 0 else 0

        ix = int(np.floor(nx))
        iy = int(np.floor(ny))
        iz = int(np.floor(nz))
        tx = nx - ix
        ty = ny - iy
        tz = nz - iz

        ix = np.clip(ix, 0, G - 2)
        iy = np.clip(iy, 0, G - 2)
        iz = np.clip(iz, 0, G - 2)

        for dx in range(2):
            for dy in range(2):
                for dz in range(2):
                    w = ((1 - tx) if dx == 0 else tx) * \
                        ((1 - ty) if dy == 0 else ty) * \
                        ((1 - tz) if dz == 0 else tz)
                    result_v += w * grid[ix + dx, iy + dy, iz + dz]

        return result_v

    def _estimate_settling_times(
        self,
        attractors: List[dict],
        n_samples: int = 500,
        max_steps: int = 400,
        eps: float = 1e-3,
        seed: int = 42,
    ) -> List[Optional[float]]:
        """For ``n_samples`` random points inside the bounding box, run
        a fixed-step forward integration and record how many steps it
        takes to either drop below ``eps`` velocity or exit the box.

        Returns a list of length ``n_samples`` where each entry is the
        settled step count (as a float), or ``None`` if the trajectory
        did not converge inside ``max_steps``.

        Used internally by ``topology_summary``; safe to call without a
        flow field (returns an empty list)."""
        r = self._result
        if r.velocity_grid is None or r.axis_min is None or r.axis_max is None:
            return []
        rng = np.random.default_rng(int(seed))
        axis_min = np.asarray(r.axis_min, dtype=np.float64)
        axis_max = np.asarray(r.axis_max, dtype=np.float64)
        span = axis_max - axis_min
        if not np.all(span > 0):
            return []

        G = r.velocity_grid.shape[0]
        # Coarse "in-basin" mask if attractors expose basin_mask;
        # used to early-exit a trajectory the moment it enters one.
        basin_union = None
        if attractors:
            for att in attractors:
                bm = att.get("basin_mask")
                if bm is None or bm.shape != (G, G, G):
                    continue
                basin_union = bm if basin_union is None else (basin_union | bm)

        dt = 0.05  # flow-time per step — matches scale of velocity grid
        out: List[Optional[float]] = []
        for _ in range(int(n_samples)):
            x = axis_min + rng.random(3) * span
            settled: Optional[float] = None
            for step in range(int(max_steps)):
                v = self._sample_velocity(x.astype(np.float32))
                if v is None:
                    break
                speed = float(np.linalg.norm(v))
                if speed < eps:
                    settled = float(step)
                    break
                if basin_union is not None:
                    gi = (
                        np.clip(((x - axis_min) / span * (G - 1)).astype(int),
                                0, G - 1)
                    )
                    if bool(basin_union[gi[0], gi[1], gi[2]]):
                        settled = float(step)
                        break
                x = x + np.asarray(v, dtype=np.float64) * dt
                if np.any(x < axis_min) or np.any(x > axis_max):
                    break
            out.append(settled)
        return out

    def _decompose_vector(
        self,
        point_3d: np.ndarray,
        direction: np.ndarray,
        include_nearby: bool = True,
    ) -> dict:
        """Decompose a direction vector into axis, cluster, and nearby-point components.

        Used by both query_flow_at (MDN velocity) and query_direction_at (path estimate).
        """
        r = self._result
        labels = r.axis_info.labels
        speed = float(np.linalg.norm(direction))

        # (a) Axis-relative decomposition
        axis_decomposition = []
        for i in range(3):
            component = float(direction[i])
            axis_decomposition.append({
                "axis_label": labels[i],
                "component": round(component, 6),
                "magnitude": round(abs(component), 6),
                "direction": "+" if component >= 0 else "-",
            })

        # (b) Cluster-relative decomposition
        cluster_pull = []
        for c in range(r.clusters.n_clusters):
            centroid = r.cluster_centroids_3d[c]
            to_cluster = centroid - point_3d
            dist = float(np.linalg.norm(to_cluster))
            if dist > 0 and speed > 0:
                to_cluster_normed = to_cluster / dist
                alignment = float(np.dot(direction, to_cluster_normed)) / speed
            else:
                alignment = 0.0

            label = r.cluster_labels[c] if c < len(r.cluster_labels) else f"Cluster {c}"
            interpretation = "toward" if alignment > 0.1 else ("away from" if alignment < -0.1 else "neutral")
            cluster_pull.append({
                "cluster_label": label,
                "alignment": round(alignment, 4),
                "distance": round(dist, 4),
                "interpretation": interpretation,
            })

        result = {
            "axis_decomposition": axis_decomposition,
            "cluster_pull": sorted(cluster_pull, key=lambda x: -abs(x["alignment"])),
        }

        # (c) Nearby points
        if include_nearby:
            pts = r.projected_3d
            dists = np.linalg.norm(pts - point_3d, axis=1)
            nearest = np.argsort(dists)[:5]
            nearby = []
            for idx in nearest:
                idx = int(idx)
                to_point = pts[idx] - point_3d
                dist = float(dists[idx])
                if dist > 0 and speed > 0:
                    to_normed = to_point / dist
                    vel_alignment = float(np.dot(direction, to_normed)) / speed
                else:
                    vel_alignment = 0.0

                # "would pass through" if alignment > 0.8 and point is close
                median_dist = float(np.median(dists))
                would_pass = vel_alignment > 0.8 and dist < median_dist * 0.5

                nearby.append({
                    "text": r.session.entries[idx].text,
                    "distance": round(dist, 4),
                    "velocity_alignment": round(vel_alignment, 4),
                    "would_pass_through": would_pass,
                })
            result["nearby_points"] = nearby

        return result

    # ─── Method 1: explain_path ───────────────────────────────────────

    def explain_path(self, texts: List[str]) -> dict:
        """Project a list of texts into the semantic space and explain the path.

        Args:
            texts: Ordered list of texts forming a semantic path.

        Returns:
            dict with path_3d, points (per-point info), and explanation.
        """
        if len(texts) < 2:
            raise ValueError("Need at least 2 texts to form a path")

        embeddings = self._embed_texts(texts)
        projected = self._project_to_3d(embeddings)

        points = []
        control_points = []
        for i, (text, pt) in enumerate(zip(texts, projected)):
            axis_pcts = self._compute_axis_pcts(pt)
            cluster_dists = self._compute_cluster_distances(pt)
            nearest = self._find_nearest_points(pt, k=3)

            points.append({
                "text": text,
                "position_3d": pt.tolist(),
                "axis_percentages": axis_pcts,
                "cluster_distances": cluster_dists,
                "nearest_texts": nearest,
            })

            # Build control_points for the explainer prompt
            control_points.append({
                "axis_pcts": [int(axis_pcts[l]) for l in self._result.axis_info.labels],
                "cluster_distances": [
                    (label, int(pct)) for label, pct in cluster_dists.items()
                ],
            })

        result = {
            "path_3d": projected.tolist(),
            "points": points,
        }

        # Generate LLM explanation if explainer available
        if self._explainer is not None:
            try:
                explanation = self._explainer.explain_probe_multi(
                    axis_labels=list(self._result.axis_info.labels),
                    control_points=control_points,
                )
                result["explanation"] = explanation
            except Exception as e:
                logger.warning(f"Path explanation failed: {e}")
                result["explanation"] = None
        else:
            result["explanation"] = None

        return result

    # ─── Method 2: query_flow_at ──────────────────────────────────────

    def query_flow_at(self, text: str) -> dict:
        """Query the MDN flow field at the position of a given text.

        Args:
            text: Input text to embed and query.

        Returns:
            dict with position, velocity, speed, axis/cluster/nearby decomposition.
        """
        embeddings = self._embed_texts([text])
        projected = self._project_to_3d(embeddings)
        pt = projected[0]

        velocity = self._sample_velocity(pt)
        if velocity is None:
            raise RuntimeError(
                "No velocity grid available. This can happen if: "
                "(1) train_flow=False was passed to pipeline.analyze(), "
                "(2) the flow model training failed silently (check for warnings about "
                "unsupported kernels or missing PyTorch), or "
                "(3) the result was loaded from a save that didn't include a velocity grid."
            )

        speed = float(np.linalg.norm(velocity))
        decomposition = self._decompose_vector(pt, velocity, include_nearby=True)

        return {
            "text": text,
            "position_3d": pt.tolist(),
            "velocity": velocity.tolist(),
            "speed": round(speed, 6),
            "source": "flow_field",
            "axis_percentages": self._compute_axis_pcts(pt),
            "cluster_distances": self._compute_cluster_distances(pt),
            **decomposition,
        }

    # ─── Method 3: query_direction_at ─────────────────────────────────

    def query_direction_at(self, texts: List[str]) -> dict:
        """Estimate movement direction from a sequence of texts (no flow field needed).

        Args:
            texts: Two or more texts defining a path.

        Returns:
            dict with estimated direction, decomposition (same format as query_flow_at).
        """
        if len(texts) < 2:
            raise ValueError("Need at least 2 texts to estimate direction")

        embeddings = self._embed_texts(texts)
        projected = self._project_to_3d(embeddings)

        # Estimate direction: average of consecutive differences
        diffs = np.diff(projected, axis=0)  # (N-1, 3)
        avg_direction = diffs.mean(axis=0)

        # Use last point as the reference position
        last_pt = projected[-1]
        magnitude = float(np.linalg.norm(avg_direction))
        decomposition = self._decompose_vector(last_pt, avg_direction, include_nearby=True)

        return {
            "texts": texts,
            "path_3d": projected.tolist(),
            "position_3d": last_pt.tolist(),
            "estimated_direction": avg_direction.tolist(),
            "estimated_magnitude": round(magnitude, 6),
            "source": "path_estimate",
            "axis_percentages": self._compute_axis_pcts(last_pt),
            "cluster_distances": self._compute_cluster_distances(last_pt),
            **decomposition,
        }

    # ─── Method 4: score_summary ─────────────────────────────────────

    def score_summary(self, channel: str) -> dict:
        """Get a statistical summary of a score channel across the space.

        Args:
            channel: Score channel name (e.g. "success", "error_rate").

        Returns:
            dict with per-cluster and per-path score breakdowns.
        """
        r = self._result
        if channel not in r.score_channels:
            raise ValueError(f"Score channel '{channel}' not found. "
                             f"Available: {r.score_channels}")

        entry_scores = r.get_entry_scores(channel)
        path_score_map = r.get_path_scores(channel)

        # Per-cluster breakdown
        cluster_stats = []
        for c in range(r.clusters.n_clusters):
            indices = [i for i, l in enumerate(r.clusters.labels) if l == c]
            vals = [entry_scores[i] for i in indices if entry_scores[i] is not None]
            label = r.cluster_labels[c] if c < len(r.cluster_labels) else f"Cluster {c}"
            stat = {
                "cluster": label,
                "count": len(vals),
                "mean": round(sum(vals) / len(vals), 4) if vals else None,
                "min": round(min(vals), 4) if vals else None,
                "max": round(max(vals), 4) if vals else None,
            }
            cluster_stats.append(stat)

        # Per-path breakdown
        path_stats = []
        path_ids = set(e.path_id for e in r.session.entries if e.path_id is not None)
        for pid in sorted(path_ids):
            indices = [i for i, e in enumerate(r.session.entries) if e.path_id == pid]
            entry_vals = [entry_scores[i] for i in indices if entry_scores[i] is not None]
            path_val = path_score_map.get(pid)
            path_label = None
            for i in indices:
                pl = r.session.entries[i].metadata.get("path_label")
                if pl:
                    path_label = pl
                    break
            stat = {
                "path_id": pid,
                "path_label": path_label or f"Path {pid}",
                "path_score": path_val,
                "entry_scores_count": len(entry_vals),
                "entry_mean": round(sum(entry_vals) / len(entry_vals), 4) if entry_vals else None,
            }
            path_stats.append(stat)

        all_vals = [v for v in entry_scores if v is not None]
        return {
            "channel": channel,
            "total_entries_with_score": len(all_vals),
            "overall_mean": round(sum(all_vals) / len(all_vals), 4) if all_vals else None,
            "total_paths_with_score": len(path_score_map),
            "cluster_breakdown": cluster_stats,
            "path_breakdown": path_stats,
        }

    # ─── Method 5: path_similarity ────────────────────────────────────

    def path_similarity(self, path_a: List[str], path_b: List[str]) -> dict:
        """Compare two semantic paths using high-dimensional embedding vectors.

        Pure vector computation — no 3D projection or flow field involved.

        Args:
            path_a: First ordered list of texts.
            path_b: Second ordered list of texts.

        Returns:
            dict with frechet_distance, mean_cosine_similarity,
            direction_similarity, start/end_similarity, overall_score.
        """
        if len(path_a) < 2 or len(path_b) < 2:
            raise ValueError("Both paths need at least 2 texts")

        emb_a = self._embed_texts(path_a)
        emb_b = self._embed_texts(path_b)

        # Normalize
        emb_a = emb_a / (np.linalg.norm(emb_a, axis=1, keepdims=True) + 1e-12)
        emb_b = emb_b / (np.linalg.norm(emb_b, axis=1, keepdims=True) + 1e-12)

        # 1. Discrete Frechet distance
        frechet = self._discrete_frechet(emb_a, emb_b)

        # 2. Mean pairwise cosine similarity (DTW-aligned)
        alignment = self._dtw_alignment(emb_a, emb_b)
        aligned_sims = [
            float(np.dot(emb_a[i], emb_b[j]))
            for i, j in alignment
        ]
        mean_cos = float(np.mean(aligned_sims)) if aligned_sims else 0.0

        # 3. Direction similarity (average direction vectors)
        dir_a = (emb_a[-1] - emb_a[0])
        dir_b = (emb_b[-1] - emb_b[0])
        norm_a = np.linalg.norm(dir_a)
        norm_b = np.linalg.norm(dir_b)
        if norm_a > 0 and norm_b > 0:
            direction_sim = float(np.dot(dir_a, dir_b) / (norm_a * norm_b))
        else:
            direction_sim = 0.0

        # 4. Endpoint similarities
        start_sim = float(np.dot(emb_a[0], emb_b[0]))
        end_sim = float(np.dot(emb_a[-1], emb_b[-1]))

        # 5. Overall score (weighted combination, normalized to [0, 1])
        # Frechet is a distance, convert to similarity: 1/(1+d)
        frechet_sim = 1.0 / (1.0 + frechet)
        overall = (
            0.3 * mean_cos +
            0.3 * direction_sim +
            0.15 * start_sim +
            0.15 * end_sim +
            0.1 * frechet_sim
        )
        # Clamp to [0, 1]
        overall = max(0.0, min(1.0, (overall + 1) / 2))  # shift from [-1,1] to [0,1]

        return {
            "frechet_distance": round(frechet, 6),
            "mean_cosine_similarity": round(mean_cos, 6),
            "direction_similarity": round(direction_sim, 6),
            "start_similarity": round(start_sim, 6),
            "end_similarity": round(end_sim, 6),
            "overall_score": round(overall, 6),
        }

    @staticmethod
    def _discrete_frechet(P: np.ndarray, Q: np.ndarray) -> float:
        """Compute discrete Frechet distance between two paths in embedding space.

        Uses cosine distance (1 - cosine_similarity) as the point distance.
        """
        n, m = len(P), len(Q)
        ca = np.full((n, m), -1.0)

        def _dist(i, j):
            return 1.0 - float(np.dot(P[i], Q[j]))

        def _c(i, j):
            if ca[i, j] > -0.5:
                return ca[i, j]
            d = _dist(i, j)
            if i == 0 and j == 0:
                ca[i, j] = d
            elif i == 0:
                ca[i, j] = max(_c(0, j - 1), d)
            elif j == 0:
                ca[i, j] = max(_c(i - 1, 0), d)
            else:
                ca[i, j] = max(min(_c(i - 1, j), _c(i - 1, j - 1), _c(i, j - 1)), d)
            return ca[i, j]

        return _c(n - 1, m - 1)

    @staticmethod
    def _dtw_alignment(P: np.ndarray, Q: np.ndarray) -> List[tuple]:
        """Simple DTW alignment returning index pairs."""
        n, m = len(P), len(Q)
        cost = np.full((n + 1, m + 1), np.inf)
        cost[0, 0] = 0.0

        for i in range(1, n + 1):
            for j in range(1, m + 1):
                d = 1.0 - float(np.dot(P[i - 1], Q[j - 1]))
                cost[i, j] = d + min(cost[i - 1, j], cost[i, j - 1], cost[i - 1, j - 1])

        # Backtrack
        path = []
        i, j = n, m
        while i > 0 and j > 0:
            path.append((i - 1, j - 1))
            candidates = [
                (cost[i - 1, j - 1], i - 1, j - 1),
                (cost[i - 1, j], i - 1, j),
                (cost[i, j - 1], i, j - 1),
            ]
            _, i, j = min(candidates, key=lambda x: x[0])

        path.reverse()
        return path

    # ─── Method 6: topology_summary ───────────────────────────────────

    def topology_summary(
        self,
        n_settling_samples: int = 500,
        max_settling_steps: int = 400,
        settling_eps: float = 1e-3,
        score_channel: Optional[str] = None,
    ) -> dict:
        """Summarize the topology of the learned semantic flow field.

        Computes attractors, basin sizes, mean settling time, transition
        turbulence (curl / speed) and Jacobian-based stability over the
        velocity grid.  All metrics are derived from the existing
        ``velocity_grid`` — no new training is performed.

        These are *analysis-layer* topology metrics extracted from the flow
        field that TraceScope already learned.  They do not guarantee any
        property of the underlying language model itself.

        Args:
            n_settling_samples: Number of random initial points sampled
                inside the bounding box for the mean-settling-time
                estimate. Higher = more accurate but slower.
            max_settling_steps: Cap on integration steps per sample
                before declaring the trajectory non-converging.
            settling_eps: Velocity magnitude below which a trajectory
                is considered "settled" into an attractor.
            score_channel: Optional score channel for per-attractor
                ``mean_score`` (forwarded to ``find_attractors``).

        Returns:
            dict with keys:
              attractors (list of attractor dicts, ``basin_mask`` removed
                  for JSON-friendliness; original is preserved under
                  ``basin_mask`` if present),
              basin_sizes (list[int]),
              basin_fractions (list[float]),
              n_attractors (int),
              mean_settling_time (float | None),
              median_settling_time (float | None),
              fraction_converged (float),
              transition_turbulence (float | None),
              jacobian_stability (dict with spectral_radius_*/divergence_*),
              unstable_regions (list of {position_3d, spectral_radius}),
              has_flow (bool).
        """
        r = self._result
        out = {
            "has_flow": r.velocity_grid is not None,
            "attractors": [],
            "basin_sizes": [],
            "basin_fractions": [],
            "n_attractors": 0,
            "mean_settling_time": None,
            "median_settling_time": None,
            "fraction_converged": 0.0,
            "transition_turbulence": None,
            "jacobian_stability": None,
            "unstable_regions": [],
        }
        if r.velocity_grid is None or r.axis_min is None or r.axis_max is None:
            return out

        # ── Attractors + basin info ─────────────────────────────────
        try:
            attractors = r.find_attractors(score_channel=score_channel)
        except Exception as e:
            logger.warning(f"find_attractors failed in topology_summary: {e}")
            attractors = []

        basin_sizes = []
        basin_fractions = []
        att_serializable = []
        for i, att in enumerate(attractors):
            basin_sizes.append(int(att.get("basin_size", 0)))
            basin_fractions.append(float(att.get("basin_fraction", 0.0)))
            pos = att.get("position")
            att_serializable.append({
                "id": i,
                "position_3d": pos.tolist() if hasattr(pos, "tolist") else list(pos),
                "strength": float(att.get("strength", 0.0)),
                "divergence": float(att.get("divergence", 0.0)),
                "basin_size": int(att.get("basin_size", 0)),
                "basin_fraction": float(att.get("basin_fraction", 0.0)),
                "mean_score": (
                    float(att["mean_score"])
                    if att.get("mean_score") is not None else None
                ),
            })
        out["attractors"] = att_serializable
        out["basin_sizes"] = basin_sizes
        out["basin_fractions"] = basin_fractions
        out["n_attractors"] = len(attractors)

        # ── Mean settling time via short forward integration ────────
        try:
            settling_times = self._estimate_settling_times(
                attractors=attractors,
                n_samples=int(n_settling_samples),
                max_steps=int(max_settling_steps),
                eps=float(settling_eps),
            )
        except Exception as e:
            logger.warning(f"settling-time estimation failed: {e}")
            settling_times = []

        finite = [t for t in settling_times if t is not None]
        if settling_times:
            out["fraction_converged"] = float(len(finite) / len(settling_times))
        if finite:
            out["mean_settling_time"] = float(np.mean(finite))
            out["median_settling_time"] = float(np.median(finite))

        # ── Transition turbulence: mean curl / (speed + eps) ────────
        try:
            vg = r.velocity_grid
            speed = np.linalg.norm(vg, axis=3)
            dvz_dy = np.gradient(vg[:, :, :, 2], axis=1)
            dvy_dz = np.gradient(vg[:, :, :, 1], axis=2)
            dvx_dz = np.gradient(vg[:, :, :, 0], axis=2)
            dvz_dx = np.gradient(vg[:, :, :, 2], axis=0)
            dvy_dx = np.gradient(vg[:, :, :, 1], axis=0)
            dvx_dy = np.gradient(vg[:, :, :, 0], axis=1)
            curl_mag = np.sqrt(
                (dvz_dy - dvy_dz) ** 2
                + (dvx_dz - dvz_dx) ** 2
                + (dvy_dx - dvx_dy) ** 2
            )
            speed_ref = float(np.percentile(speed, 90)) if speed.size else 0.0
            if speed_ref > 1e-12:
                turbulence = float(
                    np.mean(curl_mag / (speed + speed_ref * 1e-3))
                )
                out["transition_turbulence"] = round(turbulence, 6)
        except Exception as e:
            logger.warning(f"turbulence computation failed: {e}")

        # ── Jacobian stability summary over the velocity grid ───────
        try:
            grid = self.stability_grid()
            sr = grid["spectral_radius"]
            dv = grid["divergence"]
            out["jacobian_stability"] = {
                "spectral_radius_mean": float(np.mean(sr)),
                "spectral_radius_median": float(np.median(sr)),
                "spectral_radius_p90": float(np.percentile(sr, 90)),
                "spectral_radius_max": float(np.max(sr)),
                "divergence_mean": float(np.mean(dv)),
                "divergence_negative_fraction": float(
                    np.mean(dv < 0.0)
                ),
            }
            # Top-K unstable regions by spectral radius
            G = sr.shape[0]
            flat = sr.reshape(-1)
            k = min(10, flat.size)
            if k > 0:
                top_idx = np.argpartition(-flat, k - 1)[:k]
                top_idx = top_idx[np.argsort(-flat[top_idx])]
                axis_min = r.axis_min
                axis_max = r.axis_max
                span = axis_max - axis_min
                unstable = []
                for fi in top_idx:
                    ix, iy, iz = np.unravel_index(int(fi), sr.shape)
                    pos = np.array([
                        axis_min[0] + ix / max(G - 1, 1) * span[0],
                        axis_min[1] + iy / max(G - 1, 1) * span[1],
                        axis_min[2] + iz / max(G - 1, 1) * span[2],
                    ], dtype=np.float32)
                    unstable.append({
                        "position_3d": pos.tolist(),
                        "spectral_radius": float(sr[ix, iy, iz]),
                        "divergence": float(dv[ix, iy, iz]),
                    })
                out["unstable_regions"] = unstable
        except Exception as e:
            logger.warning(f"jacobian stability computation failed: {e}")

        return out

    # ─── Method 7: integrate_flow ────────────────────────────────────

    def integrate_flow(
        self,
        text: str,
        method: str = "rk45",
        max_time: float = 25.0,
        convergence_eps: float = 1e-3,
        max_steps: int = 5000,
    ) -> dict:
        """ODE-style probing of the learned semantic flow field.

        Integrates ``dx/dt = velocity_field(x)`` from the embedding of
        ``text`` until either the speed drops below ``convergence_eps``
        (settled into an attractor), the trajectory leaves the bounding
        box, or ``max_time`` is reached.

        This is *analysis-layer* probing of the velocity field that
        TraceScope already learned.  It is **not** training a Neural
        ODE — TraceScope's flow model (MDN or RBF) is not parameterised
        as a continuous-depth neural net.  Neural ODE language only
        applies at the analysis layer (we use an ODE solver to walk
        the existing learned field).

        Args:
            text: Input text — embedded and projected into 3D.
            method: One of {"rk45", "euler"}. ``rk45`` uses
                ``scipy.integrate.solve_ivp`` (RK45 with adaptive step).
                ``euler`` is a lightweight fixed-step fallback.
            max_time: Integration horizon in flow-time units.
            convergence_eps: Speed threshold below which the trajectory
                is considered settled.
            max_steps: Hard cap on Euler steps (ignored for rk45).

        Returns:
            dict with keys:
                trajectory_3d (list[[x,y,z]]),
                times (list[float]),
                final_position (list[float]),
                final_speed (float),
                settling_time (float | None),
                attractor_id (int | None),
                converged (bool),
                method (str),
                escaped (bool),
                source ("integrate_flow").
        """
        r = self._result
        if r.velocity_grid is None or r.axis_min is None or r.axis_max is None:
            raise RuntimeError(
                "No velocity grid available. integrate_flow needs a flow "
                "field — run pipeline.analyze() with train_flow=True."
            )

        embeddings = self._embed_texts([text])
        projected = self._project_to_3d(embeddings)
        x0 = np.asarray(projected[0], dtype=np.float64)

        axis_min = np.asarray(r.axis_min, dtype=np.float64)
        axis_max = np.asarray(r.axis_max, dtype=np.float64)
        span = axis_max - axis_min
        # Pad bounding box very slightly so a probe starting on the
        # boundary doesn't immediately get flagged as "escaped".
        pad = np.where(span > 0, span * 0.05, 1e-6)
        bb_lo = axis_min - pad
        bb_hi = axis_max + pad

        method_lc = (method or "rk45").lower()

        def _vel(x: np.ndarray) -> np.ndarray:
            v = self._sample_velocity(np.asarray(x, dtype=np.float32))
            if v is None:
                return np.zeros(3, dtype=np.float64)
            return np.asarray(v, dtype=np.float64)

        traj = [x0.copy()]
        times_out = [0.0]
        settling_time: Optional[float] = None
        escaped = False
        converged = False

        if method_lc in ("rk45", "rk23", "dop853", "lsoda", "radau", "bdf"):
            try:
                from scipy.integrate import solve_ivp
            except ImportError as e:
                raise RuntimeError(
                    "scipy.integrate.solve_ivp is required for method='rk45'."
                ) from e

            def rhs(t, x):
                return _vel(x)

            # Event: speed below threshold (settled)
            def settled_event(t, x):
                v = _vel(x)
                return float(np.linalg.norm(v)) - convergence_eps
            settled_event.terminal = True
            settled_event.direction = -1

            # Event: left bounding box
            def escaped_event(t, x):
                d = np.concatenate([x - bb_lo, bb_hi - x])
                return float(np.min(d))
            escaped_event.terminal = True
            escaped_event.direction = -1

            method_map = {
                "rk45": "RK45", "rk23": "RK23", "dop853": "DOP853",
                "lsoda": "LSODA", "radau": "Radau", "bdf": "BDF",
            }
            sol = solve_ivp(
                rhs, (0.0, float(max_time)), x0,
                method=method_map[method_lc],
                events=[settled_event, escaped_event],
                max_step=float(max_time) / 50.0 if max_time > 0 else 0.1,
                dense_output=False,
                rtol=1e-4, atol=1e-6,
            )
            ys = sol.y.T  # (n_points, 3)
            ts = sol.t
            traj = [row.copy() for row in ys]
            times_out = ts.tolist()

            # Determine which event fired (if any)
            settled_events = sol.t_events[0] if len(sol.t_events) > 0 else np.array([])
            escape_events = sol.t_events[1] if len(sol.t_events) > 1 else np.array([])
            if settled_events.size > 0:
                converged = True
                settling_time = float(settled_events[0])
            if escape_events.size > 0:
                escaped = True

        else:
            # Fixed-step Euler fallback
            x = x0.copy()
            t = 0.0
            dt = float(max_time) / float(max_steps) if max_steps > 0 else 0.02
            for _step in range(int(max_steps)):
                v = _vel(x)
                speed = float(np.linalg.norm(v))
                if speed < convergence_eps:
                    converged = True
                    settling_time = t
                    break
                x = x + v * dt
                t += dt
                if np.any(x < bb_lo) or np.any(x > bb_hi):
                    escaped = True
                    break
                traj.append(x.copy())
                times_out.append(t)
                if t >= max_time:
                    break

        final_pos = np.asarray(traj[-1], dtype=np.float64)
        final_speed = float(np.linalg.norm(_vel(final_pos)))

        # Attractor matching: find the nearest attractor whose basin
        # contains the final position (if any).  If basin_mask is
        # available we use it; otherwise fall back to nearest-position.
        attractor_id: Optional[int] = None
        if converged:
            try:
                attractors = r.find_attractors()
            except Exception as e:
                logger.warning(f"attractor lookup failed: {e}")
                attractors = []
            if attractors:
                # Try basin-mask membership first
                G = r.velocity_grid.shape[0]
                gi = np.zeros(3, dtype=int)
                for a in range(3):
                    if span[a] > 0:
                        gi[a] = int(np.clip(
                            (final_pos[a] - axis_min[a]) / span[a] * (G - 1),
                            0, G - 1,
                        ))
                for ai, att in enumerate(attractors):
                    bm = att.get("basin_mask")
                    if bm is not None and bm.shape == (G, G, G):
                        if bool(bm[gi[0], gi[1], gi[2]]):
                            attractor_id = ai
                            break
                if attractor_id is None:
                    # Fallback: nearest attractor position
                    dists = [
                        float(np.linalg.norm(final_pos - np.asarray(a["position"])))
                        for a in attractors
                    ]
                    attractor_id = int(np.argmin(dists))

        return {
            "text": text,
            "method": method_lc,
            "trajectory_3d": [list(map(float, p)) for p in traj],
            "times": [float(t) for t in times_out],
            "final_position": final_pos.tolist(),
            "final_speed": round(final_speed, 8),
            "settling_time": settling_time,
            "attractor_id": attractor_id,
            "converged": bool(converged),
            "escaped": bool(escaped),
            "source": "integrate_flow",
        }

    # ─── Method 8: numerical Jacobian diagnostics ────────────────────

    def _numerical_jacobian(
        self,
        point_3d: np.ndarray,
        h: Optional[float] = None,
    ) -> Optional[np.ndarray]:
        """Estimate the local 3×3 Jacobian of the velocity field via
        central differences.  Returns None when no flow field is
        available."""
        r = self._result
        if r.velocity_grid is None or r.axis_min is None or r.axis_max is None:
            return None
        span = np.asarray(r.axis_max - r.axis_min, dtype=np.float64)
        # Default step = a small fraction of the grid cell so we sample
        # *across* a cell, not within numerical noise.
        if h is None:
            G = r.velocity_grid.shape[0]
            cell = span / max(G - 1, 1)
            h_vec = np.maximum(cell * 0.5, 1e-6)
        else:
            h_vec = np.full(3, float(h), dtype=np.float64)
        x = np.asarray(point_3d, dtype=np.float64)
        J = np.zeros((3, 3), dtype=np.float64)
        for j in range(3):
            ej = np.zeros(3, dtype=np.float64)
            ej[j] = h_vec[j]
            v_plus = self._sample_velocity((x + ej).astype(np.float32))
            v_minus = self._sample_velocity((x - ej).astype(np.float32))
            if v_plus is None or v_minus is None:
                return None
            J[:, j] = (np.asarray(v_plus, dtype=np.float64)
                       - np.asarray(v_minus, dtype=np.float64)) / (2.0 * h_vec[j])
        return J

    def local_stability_at(self, text: str) -> dict:
        """Estimate local stability of the flow field at the position
        of ``text`` via a numerical Jacobian.

        Returns spectral radius (max |eigenvalue|), divergence
        (trace of J), and the eigenvalues themselves.  Higher
        spectral radius → faster local divergence/convergence;
        negative trace → on-net contracting (attractor-like).
        """
        r = self._result
        if r.velocity_grid is None or r.axis_min is None or r.axis_max is None:
            raise RuntimeError(
                "No velocity grid available. local_stability_at needs a "
                "flow field — run pipeline.analyze() with train_flow=True."
            )
        embeddings = self._embed_texts([text])
        projected = self._project_to_3d(embeddings)
        pt = projected[0]
        J = self._numerical_jacobian(pt)
        if J is None:
            raise RuntimeError(
                "Failed to compute numerical Jacobian (no flow field)."
            )
        eigvals = np.linalg.eigvals(J)
        spectral_radius = float(np.max(np.abs(eigvals)))
        divergence = float(np.trace(J))
        return {
            "text": text,
            "position_3d": [float(x) for x in pt],
            "jacobian": J.tolist(),
            "eigenvalues_real": [float(v.real) for v in eigvals],
            "eigenvalues_imag": [float(v.imag) for v in eigvals],
            "spectral_radius": round(spectral_radius, 8),
            "divergence": round(divergence, 8),
            "is_locally_attracting": bool(divergence < 0),
            "source": "local_stability_at",
        }

    def stability_grid(self, resolution: Optional[int] = None) -> dict:
        """Return per-cell numerical stability statistics over the
        velocity grid.

        Uses ``np.gradient`` over the velocity grid to build the full
        spatial Jacobian per cell, then returns the spectral radius
        (max |eigenvalue|) and divergence (trace) at each cell.

        Args:
            resolution: Optional output resolution.  If smaller than
                the underlying grid, a stride is applied so callers can
                trade detail for speed.  Defaults to the grid resolution.

        Returns:
            dict with keys:
                spectral_radius (np.ndarray[G,G,G]),
                divergence (np.ndarray[G,G,G]),
                axis_min (list[3]), axis_max (list[3]),
                grid_size (int).
        """
        r = self._result
        if r.velocity_grid is None or r.axis_min is None or r.axis_max is None:
            raise RuntimeError(
                "No velocity grid available. stability_grid needs a "
                "flow field — run pipeline.analyze() with train_flow=True."
            )
        vg = r.velocity_grid.astype(np.float64)
        G = vg.shape[0]
        # Spacing = (axis_max - axis_min) / (G - 1) per axis
        span = np.asarray(r.axis_max - r.axis_min, dtype=np.float64)
        dx = float(span[0] / max(G - 1, 1))
        dy = float(span[1] / max(G - 1, 1))
        dz = float(span[2] / max(G - 1, 1))

        # ∂v_i / ∂x_j  for i,j in {0,1,2}
        dV0 = np.gradient(vg[..., 0], dx, dy, dz)
        dV1 = np.gradient(vg[..., 1], dx, dy, dz)
        dV2 = np.gradient(vg[..., 2], dx, dy, dz)
        # J[i, j] = dV_i / dx_j
        J = np.stack([
            np.stack([dV0[0], dV0[1], dV0[2]], axis=-1),  # i=0
            np.stack([dV1[0], dV1[1], dV1[2]], axis=-1),  # i=1
            np.stack([dV2[0], dV2[1], dV2[2]], axis=-1),  # i=2
        ], axis=-2)  # shape (G,G,G,3,3)

        # Vectorised eigenvalue computation
        eigvals = np.linalg.eigvals(J)        # (G,G,G,3) complex
        spectral_radius = np.max(np.abs(eigvals), axis=-1).astype(np.float32)
        divergence = (J[..., 0, 0] + J[..., 1, 1] + J[..., 2, 2]).astype(np.float32)

        if resolution is not None and 0 < int(resolution) < G:
            stride = max(1, G // int(resolution))
            spectral_radius = spectral_radius[::stride, ::stride, ::stride]
            divergence = divergence[::stride, ::stride, ::stride]

        return {
            "spectral_radius": spectral_radius,
            "divergence": divergence,
            "axis_min": [float(v) for v in r.axis_min],
            "axis_max": [float(v) for v in r.axis_max],
            "grid_size": int(spectral_radius.shape[0]),
        }