geographdb-core 0.5.4

//! Geometric graph inference: walker and decoder utilities.
//!
//! The graph topology is the model. A walker traverses sense-nodes in 3D space
//! using spatial proximity, velocity alignment, sequential bias, edge weights,
//! repetition penalty, and optional A* planning. This module provides a reusable
//! Rust API for both greedy and beam-search decoding.

use crate::algorithms::parallel_transport::{graph_attention_scores, transport_score};
use crate::spatial::octree::{BoundingBox, Octree};
use crate::storage::data_structures::NodePoint;
use crate::{astar_find_path_4d, GraphNode4D, GraphPath4D, TemporalWindow, TraversalContext4D};
use glam::Vec3;
use std::collections::{HashMap, HashSet};

/// How the walker selects the next node from the candidate set.
#[derive(Debug, Clone, Copy, Default)]
pub enum TransitionMode {
    /// Original distance/velocity/alignment scoring.
    #[default]
    DistanceKnn,
    /// Geometric parallel transport: Rodrigues rotation from the incoming edge
    /// tangent to the outgoing edge tangent, optionally modulated by Ricci
    /// curvature. `lambda` controls decay with rotation angle.
    RodriguesTransport { lambda: f32 },
    /// Graph attention over transported neighbor features, with a UCB exploration
    /// bonus. For the first test W_Q and W_K are identity, so the geometry alone
    /// drives the scores.
    GraphAttentionTransport { ucb_c: f32 },
}

/// Configuration for a geometric walker.
#[derive(Debug, Clone, Copy)]
pub struct WalkerConfig {
    /// Number of spatial nearest neighbors to consider each step.
    pub knn: usize,
    /// Softmax temperature for candidate selection.
    pub temperature: f32,
    /// Momentum (retention) for the velocity update.
    pub momentum: f32,
    /// Step size for the position update.
    pub step_size: f32,
    /// Penalty applied to recently visited token IDs.
    pub repetition_penalty: f32,
    /// Number of recent token IDs to remember for repetition penalty.
    pub recent_window: usize,
    /// Number of steps between A* plan replenishments (0 disables planning).
    pub plan_interval: usize,
    /// Weight for the sliding-window context centroid coherence term.
    pub context_weight: f32,
    /// Optional goal position to steer the walker toward (e.g. a tool region).
    pub goal_position: Option<Vec3>,
    /// Weight for goal alignment.
    pub goal_weight: f32,
    /// Transition mechanism used to rank successor candidates.
    pub transition_mode: TransitionMode,
}

impl Default for WalkerConfig {
    fn default() -> Self {
        Self {
            knn: 20,
            temperature: 0.05,
            momentum: 0.7,
            step_size: 0.3,
            repetition_penalty: 0.3,
            recent_window: 8,
            plan_interval: 8,
            context_weight: 0.0,
            goal_position: None,
            goal_weight: 1.5,
            transition_mode: TransitionMode::default(),
        }
    }
}

/// Decode a sense-node ID to its base token ID.
#[inline]
pub fn decode_node_id(node_id: u64) -> u64 {
    node_id / 1000
}

/// Build a lookup from node ID to graph index.
pub fn build_node_index(graph: &[GraphNode4D]) -> HashMap<u64, usize> {
    graph.iter().enumerate().map(|(i, n)| (n.id, i)).collect()
}

/// Build a lookup from directed edge to its weight.
pub fn build_edge_weights(graph: &[GraphNode4D]) -> HashMap<(u64, u64), f32> {
    let mut weights = HashMap::new();
    for node in graph {
        for edge in &node.successors {
            weights.insert((node.id, edge.dst), edge.weight);
        }
    }
    weights
}

/// Compute the geometric centroid of sense-nodes for a tokenized prompt.
///
/// This centroid acts as the walker's starting cognitive state: a point in graph
/// space that summarizes the prompt meaning. The walker then routes toward nearby
/// semantic regions (text, code, math, tool) naturally.
pub fn prompt_centroid(
    graph: &[GraphNode4D],
    tokenizer: &tokenizers::Tokenizer,
    prompt: &str,
) -> Option<Vec3> {
    let encoding = tokenizer.encode(prompt.to_string(), false).ok()?;
    let token_ids: Vec<u32> = encoding.get_ids().to_vec();
    if token_ids.is_empty() {
        return None;
    }

    let node_index = build_node_index(graph);
    let mut points = Vec::new();
    for tid in token_ids {
        let base = (tid as u64) * 1000;
        for sense_offset in 0..1000 {
            if let Some(&idx) = node_index.get(&(base + sense_offset)) {
                points.push(graph[idx].position());
            }
        }
    }

    if points.is_empty() {
        return None;
    }
    let sum: Vec3 = points.iter().copied().sum();
    Some(sum / points.len() as f32)
}

/// Build an octree over the graph nodes for fast spatial nearest-neighbor queries.
pub fn build_octree(graph: &[GraphNode4D]) -> Octree {
    let mut min = graph[0].position();
    let mut max = graph[0].position();
    for node in graph {
        min = min.min(node.position());
        max = max.max(node.position());
    }
    let span = (max - min).length().max(1.0);
    let pad = Vec3::splat(span * 0.25);
    let bounds = BoundingBox::new(min - pad, max + pad);
    let mut octree = Octree::new(bounds);
    for node in graph {
        octree.insert(NodePoint {
            id: node.id,
            x: node.x,
            y: node.y,
            z: node.z,
        });
    }
    octree
}

/// A single geometric walker state.
pub struct GeometricWalker {
    config: WalkerConfig,
    current_node: u64,
    previous_node: Option<u64>,
    position: Vec3,
    velocity: Vec3,
    time_step: u64,
    recent_tokens: Vec<u64>,
    recent_positions: Vec<Vec3>,
    planned_path: Vec<u64>,
    path_index: usize,
    cum_score: f32,
    trajectory: Vec<u64>,
    visit_counts: HashMap<u64, u32>,
}

impl GeometricWalker {
    /// Create a walker starting at the given graph node.
    pub fn new(start: &GraphNode4D, config: WalkerConfig) -> Self {
        let mut visit_counts = HashMap::new();
        visit_counts.insert(start.id, 1);
        Self {
            config,
            current_node: start.id,
            previous_node: None,
            position: start.position(),
            velocity: Vec3::ZERO,
            time_step: 0,
            recent_tokens: Vec::new(),
            recent_positions: Vec::new(),
            planned_path: Vec::new(),
            path_index: 0,
            cum_score: 0.0,
            trajectory: vec![start.id],
            visit_counts,
        }
    }

    /// Current node ID.
    pub fn current_node(&self) -> u64 {
        self.current_node
    }

    /// Current floating-point position in graph space.
    pub fn position(&self) -> Vec3 {
        self.position
    }

    /// Set the walker position (useful to start from a prompt centroid rather than a node).
    pub fn set_position(&mut self, position: Vec3) {
        self.position = position;
    }

    /// Cumulative trajectory score.
    pub fn cum_score(&self) -> f32 {
        self.cum_score
    }

    /// Recent token IDs (oldest first).
    pub fn recent_tokens(&self) -> &[u64] {
        &self.recent_tokens
    }

    /// Sliding-window centroid of recent selected node positions.
    pub fn context_centroid(&self) -> Vec3 {
        if self.recent_positions.is_empty() {
            return self.position;
        }
        let n = self.recent_positions.len() as f32;
        let sum: Vec3 = self.recent_positions.iter().copied().sum();
        sum / n
    }

    /// UCB exploration bonus for a candidate node.
    ///
    /// Bonus = c * sqrt(ln(N) / n), where N is total visits across all nodes and
    /// n is the number of times this candidate has been visited.
    fn ucb_bonus(&self, node_id: u64, ucb_c: f32) -> f32 {
        let total: u32 = self.visit_counts.values().sum();
        if total == 0 || ucb_c <= 0.0 {
            return 0.0;
        }
        let visits = self.visit_counts.get(&node_id).copied().unwrap_or(0).max(1);
        ucb_c * ((total as f32).ln() / visits as f32).sqrt()
    }

    /// Score a candidate node for the next step.
    ///
    /// The base term depends on `transition_mode`:
    /// - `DistanceKnn`: spatial proximity plus velocity alignment.
    /// - `RodriguesTransport`: parallel-transport score between the incoming edge,
    ///   the candidate edge, and the incoming Ricci curvature.
    ///
    /// All modes share sequential, edge-weight, plan, context, goal, and
    /// repetition terms.
    #[rustfmt::skip]
    #[allow(clippy::too_many_arguments, reason = "candidate scoring combines many walker context pieces")]
    fn candidate_score(
        &self,
        graph: &[GraphNode4D],
        node_index: &HashMap<u64, usize>,
        edge_weights: &HashMap<(u64, u64), f32>,
        cid: u64,
        node_pos: Vec3,
        current_successors: &HashSet<u64>,
        next_planned: Option<u64>,
        context_centroid: Vec3,
        recent_set: &HashSet<u64>,
        curvature: Option<&HashMap<(u64, u64), f32>>,
    ) -> f32 {
        let base_score = match self.config.transition_mode {
            TransitionMode::DistanceKnn => {
                let dir = node_pos - self.position;
                let dist_sq = dir.length_squared();
                let dist = dist_sq.sqrt().max(1e-6);

                // Spatial proximity.
                let spatial_score = (-dist_sq / self.config.temperature).exp();

                // Velocity alignment.
                let vel_norm = self.velocity.length().max(1e-6);
                let alignment = if vel_norm > 1e-6 {
                    self.velocity.dot(dir) / (vel_norm * dist)
                } else {
                    0.0
                };
                let alignment_score = (alignment * 0.5 + 0.5).max(0.0);

                spatial_score * (1.0 + alignment_score)
            }
            TransitionMode::RodriguesTransport { lambda } => {
                if let Some(prev_id) = self.previous_node {
                    let prev_idx = node_index[&prev_id];
                    let curr_idx = node_index[&self.current_node];
                    let next_idx = node_index[&cid];
                    transport_score(graph, prev_idx, curr_idx, next_idx, lambda, curvature)
                } else {
                    // No previous node yet: fall back to spatial proximity.
                    let dist_sq = node_pos.distance_squared(self.position);
                    (-dist_sq / self.config.temperature).exp()
                }
            }
            TransitionMode::GraphAttentionTransport { ucb_c } => {
                if let Some(prev_id) = self.previous_node {
                    let prev_idx = node_index[&prev_id];
                    let curr_idx = node_index[&self.current_node];
                    let next_idx = node_index[&cid];
                    let scores = graph_attention_scores(
                        graph,
                        curr_idx,
                        Some(prev_idx),
                        &[next_idx],
                    );
                    let (raw_score, _alpha) = scores[&next_idx];

                    // Modulate the UCB bonus by Ollivier-Ricci curvature.
                    // Negative κ (bridge edges) amplifies exploration;
                    // positive κ (cluster interiors) dampens it.
                    let kappa = curvature
                        .and_then(|c| c.get(&(self.current_node, cid)).copied())
                        .unwrap_or(0.0f32);
                    let kappa_weight = (1.0f32 - kappa).max(0.1f32);

                    raw_score + self.ucb_bonus(cid, ucb_c) * kappa_weight
                } else {
                    // No previous node yet: fall back to spatial proximity.
                    let dist_sq = node_pos.distance_squared(self.position);
                    (-dist_sq / self.config.temperature).exp()
                }
            }
        };

        // Sequential bias: direct graph successors.
        let sequential_bonus = if current_successors.contains(&cid) {
            0.5
        } else {
            0.0
        };

        // Edge weight boost.
        let edge_bonus = edge_weights
            .get(&(self.current_node, cid))
            .copied()
            .unwrap_or(0.0)
            * 0.05;

        // Repetition penalty.
        let token_id = decode_node_id(cid);
        let repetition_penalty = if recent_set.contains(&token_id) {
            self.config.repetition_penalty
                * recent_set.iter().filter(|&&t| t == token_id).count() as f32
        } else {
            0.0
        };

        // Plan bonus.
        let plan_bonus = if next_planned == Some(cid) { 2.0 } else { 0.0 };

        // Context centroid coherence.
        let context_dist = node_pos.distance(context_centroid);
        let context_score = if self.config.context_weight > 0.0 {
            self.config.context_weight * (-context_dist / self.config.temperature).exp()
        } else {
            0.0
        };

        // Goal alignment.
        let goal_score = if let Some(goal_pos) = self.config.goal_position {
            let goal_dist = node_pos.distance(goal_pos);
            self.config.goal_weight * (-goal_dist / self.config.temperature).exp()
        } else {
            0.0
        };

        (base_score + sequential_bonus + edge_bonus + plan_bonus + context_score + goal_score
            - repetition_penalty)
            .max(1e-6)
    }

    /// Take one greedy step and return the chosen node ID.
    pub fn step(
        &mut self,
        graph: &[GraphNode4D],
        node_index: &HashMap<u64, usize>,
        octree: &Octree,
        edge_weights: &HashMap<(u64, u64), f32>,
        curvature: Option<&HashMap<(u64, u64), f32>>,
    ) -> u64 {
        self.replenish_plan(graph);

        let current_idx = node_index[&self.current_node];
        let current_successors: HashSet<u64> = graph[current_idx]
            .successors
            .iter()
            .map(|e| e.dst)
            .collect();

        // Spatial KNN + direct successors form the candidate set.
        let mut candidate_ids: HashSet<u64> = current_successors.clone();
        let knn = octree.query_knn(self.position, self.config.knn);
        for (np, _) in &knn {
            if np.id != self.current_node {
                candidate_ids.insert(np.id);
            }
        }

        let mut candidate_positions: HashMap<u64, Vec3> = HashMap::new();
        for (np, _) in &knn {
            candidate_positions.insert(np.id, Vec3::new(np.x, np.y, np.z));
        }
        for &sid in &current_successors {
            candidate_positions.entry(sid).or_insert_with(|| {
                let snode = &graph[node_index[&sid]];
                snode.position()
            });
        }

        let recent_set: HashSet<u64> = self
            .recent_tokens
            .iter()
            .rev()
            .take(self.config.recent_window)
            .copied()
            .collect();

        let next_planned = self.planned_path.get(self.path_index).copied();
        let context_centroid = self.context_centroid();

        let mut candidates: Vec<(u64, f32)> = Vec::new();
        for &cid in &candidate_ids {
            let node_pos = candidate_positions[&cid];
            let score = self.candidate_score(
                graph,
                node_index,
                edge_weights,
                cid,
                node_pos,
                &current_successors,
                next_planned,
                context_centroid,
                &recent_set,
                curvature,
            );
            candidates.push((cid, score));
        }

        let next_id = if candidates.is_empty() {
            self.current_node
        } else {
            softmax_sample(&candidates)
        };

        self.update_state(
            graph,
            node_index,
            next_id,
            candidates
                .iter()
                .find(|(id, _)| *id == next_id)
                .map(|(_, s)| *s)
                .unwrap_or(0.0),
        );
        next_id
    }

    /// Greedy walk for `steps` iterations, returning the node IDs visited (including start).
    pub fn walk(
        &mut self,
        graph: &[GraphNode4D],
        node_index: &HashMap<u64, usize>,
        octree: &Octree,
        edge_weights: &HashMap<(u64, u64), f32>,
        curvature: Option<&HashMap<(u64, u64), f32>>,
        steps: usize,
    ) -> Vec<u64> {
        for _ in 0..steps {
            self.step(graph, node_index, octree, edge_weights, curvature);
        }
        self.trajectory.clone()
    }

    /// Beam-search walk. Returns the highest-scoring trajectory (including start).
    #[allow(
        clippy::too_many_arguments,
        reason = "beam search exposes the same walker context pieces as the rest of the API"
    )]
    pub fn walk_beam(
        graph: &[GraphNode4D],
        node_index: &HashMap<u64, usize>,
        octree: &Octree,
        edge_weights: &HashMap<(u64, u64), f32>,
        curvature: Option<&HashMap<(u64, u64), f32>>,
        start: &GraphNode4D,
        steps: usize,
        beam_width: usize,
        config: &WalkerConfig,
    ) -> Vec<u64> {
        if beam_width == 0 {
            return vec![start.id];
        }

        let mut beams: Vec<GeometricWalker> = vec![GeometricWalker::new(start, *config)];

        for _ in 0..steps {
            let mut next_beams: Vec<GeometricWalker> = Vec::new();
            for beam in &mut beams {
                let current_idx = node_index[&beam.current_node];
                let current_successors: HashSet<u64> = graph[current_idx]
                    .successors
                    .iter()
                    .map(|e| e.dst)
                    .collect();

                let mut candidate_ids: HashSet<u64> = current_successors.clone();
                let knn = octree.query_knn(beam.position, config.knn);
                for (np, _) in &knn {
                    if np.id != beam.current_node {
                        candidate_ids.insert(np.id);
                    }
                }

                let mut candidate_positions: HashMap<u64, Vec3> = HashMap::new();
                for (np, _) in &knn {
                    candidate_positions.insert(np.id, Vec3::new(np.x, np.y, np.z));
                }
                for &sid in &current_successors {
                    candidate_positions.entry(sid).or_insert_with(|| {
                        let snode = &graph[node_index[&sid]];
                        snode.position()
                    });
                }

                let recent_set: HashSet<u64> = beam
                    .recent_tokens
                    .iter()
                    .rev()
                    .take(config.recent_window)
                    .copied()
                    .collect();

                let next_planned = beam.planned_path.get(beam.path_index).copied();
                let context_centroid = beam.context_centroid();

                for &cid in &candidate_ids {
                    let Some(&node_pos) = candidate_positions.get(&cid) else {
                        continue;
                    };
                    let mut child = clone_for_beam(beam, graph, node_index);

                    let step_score = beam.candidate_score(
                        graph,
                        node_index,
                        edge_weights,
                        cid,
                        node_pos,
                        &current_successors,
                        next_planned,
                        context_centroid,
                        &recent_set,
                        curvature,
                    );

                    child.update_state(graph, node_index, cid, step_score);
                    next_beams.push(child);
                }
            }

            // Keep top beams by cumulative score.
            next_beams.sort_by(|a, b| b.cum_score.partial_cmp(&a.cum_score).unwrap());
            next_beams.truncate(beam_width);

            // If no expansion happened, stop early.
            if next_beams.is_empty() {
                break;
            }
            beams = next_beams;
        }

        beams
            .into_iter()
            .max_by(|a, b| a.cum_score.partial_cmp(&b.cum_score).unwrap())
            .map(|b| b.trajectory)
            .unwrap_or_else(|| vec![start.id])
    }

    fn replenish_plan(&mut self, graph: &[GraphNode4D]) {
        if self.config.plan_interval == 0 {
            return;
        }
        if !self.planned_path.is_empty() && self.path_index < self.planned_path.len() {
            return;
        }
        if graph.len() < 2 {
            return;
        }

        const PLAN_SPATIAL_RADIUS: f32 = 2.0;
        let current_pos = self.position;
        let mut candidates: Vec<u64> = graph
            .iter()
            .filter(|n| {
                n.id != self.current_node
                    && n.position().distance(current_pos) < PLAN_SPATIAL_RADIUS * 2.0
            })
            .map(|n| n.id)
            .collect();
        if candidates.is_empty() {
            candidates = graph
                .iter()
                .filter(|n| n.id != self.current_node)
                .map(|n| n.id)
                .collect();
        }
        if candidates.is_empty() {
            return;
        }

        let goal_idx =
            ((self.current_node.wrapping_add(self.time_step)) as usize) % candidates.len();
        let goal_id = candidates[goal_idx];

        let ctx = TraversalContext4D {
            time_window: Some(TemporalWindow {
                start: self.time_step,
                end: self.time_step + self.config.plan_interval as u64 * 2,
            }),
            spatial_region: None,
            spatial_candidates: None,
            graph_weight: 1.0,
            spatial_weight: 0.0,
            temporal_weight: 0.5,
        };

        if let Some(GraphPath4D { node_ids, .. }) =
            astar_find_path_4d(graph, self.current_node, goal_id, &ctx)
        {
            self.planned_path = node_ids
                .into_iter()
                .skip_while(|&id| id == self.current_node)
                .collect();
            self.path_index = 0;
        }
    }

    fn update_state(
        &mut self,
        graph: &[GraphNode4D],
        node_index: &HashMap<u64, usize>,
        next_id: u64,
        step_score: f32,
    ) {
        let next_node = &graph[node_index[&next_id]];
        let target = next_node.position();
        let dir = target - self.position;
        let dir_norm = dir.length().max(1e-6);
        let dir_unit = dir / dir_norm;

        self.previous_node = Some(self.current_node);
        self.velocity =
            self.velocity * self.config.momentum + dir_unit * (1.0 - self.config.momentum);
        self.position += self.velocity * self.config.step_size;
        self.current_node = next_id;
        self.time_step += 1;
        self.cum_score += step_score;
        self.trajectory.push(next_id);
        *self.visit_counts.entry(next_id).or_insert(0) += 1;

        let token_id = decode_node_id(next_id);
        self.recent_tokens.push(token_id);
        self.recent_positions.push(next_node.position());
        if self.recent_tokens.len() > self.config.recent_window {
            self.recent_tokens.remove(0);
            self.recent_positions.remove(0);
        }

        if self.config.plan_interval > 0 {
            self.path_index += 1;
        }
    }
}

/// Softmax sample from scored candidates and return the highest-probability ID.
fn softmax_sample(candidates: &[(u64, f32)]) -> u64 {
    let max_score = candidates.iter().map(|(_, s)| *s).fold(0.0f32, f32::max);
    let exp_scores: Vec<f32> = candidates
        .iter()
        .map(|(_, s)| (s - max_score).exp())
        .collect();
    let sum_exp: f32 = exp_scores.iter().sum();
    let probs: Vec<f32> = exp_scores.iter().map(|e| e / sum_exp).collect();

    probs
        .iter()
        .enumerate()
        .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap())
        .map(|(i, _)| candidates[i].0)
        .unwrap_or(candidates[0].0)
}

/// Clone a walker for beam expansion, resetting the planned-path pointer to share the parent plan.
fn clone_for_beam(
    walker: &GeometricWalker,
    _graph: &[GraphNode4D],
    _node_index: &HashMap<u64, usize>,
) -> GeometricWalker {
    GeometricWalker {
        config: walker.config,
        current_node: walker.current_node,
        previous_node: walker.previous_node,
        position: walker.position,
        velocity: walker.velocity,
        time_step: walker.time_step,
        recent_tokens: walker.recent_tokens.clone(),
        recent_positions: walker.recent_positions.clone(),
        planned_path: walker.planned_path.clone(),
        path_index: walker.path_index,
        cum_score: walker.cum_score,
        trajectory: walker.trajectory.clone(),
        visit_counts: walker.visit_counts.clone(),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::algorithms::four_d::{GraphProperties, TemporalEdge};
    use crate::GraphNode4D;

    fn make_node(id: u64, x: f32, y: f32, z: f32) -> GraphNode4D {
        GraphNode4D {
            id,
            x,
            y,
            z,
            begin_ts: 0,
            end_ts: u64::MAX,
            properties: GraphProperties::default(),
            successors: Vec::new(),
        }
    }

    fn make_graph_line() -> Vec<GraphNode4D> {
        let mut a = make_node(1000, 0.0, 0.0, 0.0);
        let mut b = make_node(2000, 1.0, 0.0, 0.0);
        let c = make_node(3000, 2.0, 0.0, 0.0);
        a.successors.push(TemporalEdge {
            dst: 2000,
            weight: 1.0,
            begin_ts: 0,
            end_ts: u64::MAX,
        });
        b.successors.push(TemporalEdge {
            dst: 3000,
            weight: 1.0,
            begin_ts: 0,
            end_ts: u64::MAX,
        });
        vec![a, b, c]
    }

    #[test]
    fn decode_node_id_maps_sense_to_token() {
        assert_eq!(decode_node_id(1234001), 1234);
        assert_eq!(decode_node_id(0), 0);
    }

    #[test]
    fn build_node_index_and_edge_weights() {
        let graph = make_graph_line();
        let idx = build_node_index(&graph);
        assert_eq!(idx[&1000], 0);
        assert_eq!(idx[&2000], 1);
        assert_eq!(idx[&3000], 2);

        let weights = build_edge_weights(&graph);
        assert_eq!(weights.get(&(1000, 2000)), Some(&1.0));
        assert_eq!(weights.get(&(2000, 3000)), Some(&1.0));
    }

    #[test]
    fn walker_follows_successors() {
        let graph = make_graph_line();
        let idx = build_node_index(&graph);
        let octree = build_octree(&graph);
        let weights = build_edge_weights(&graph);

        let config = WalkerConfig {
            knn: 0, // only direct successors so the line is deterministic
            temperature: 0.01,
            plan_interval: 0,
            ..Default::default()
        };
        let mut walker = GeometricWalker::new(&graph[0], config);

        let traj = walker.walk(&graph, &idx, &octree, &weights, None, 2);
        assert_eq!(traj, vec![1000, 2000, 3000]);
    }

    #[test]
    fn context_centroid_averages_recent_positions() {
        let graph = make_graph_line();
        let idx = build_node_index(&graph);
        let octree = build_octree(&graph);
        let weights = build_edge_weights(&graph);

        let config = WalkerConfig {
            knn: 0,
            recent_window: 2,
            plan_interval: 0,
            ..Default::default()
        };
        let mut walker = GeometricWalker::new(&graph[0], config);
        walker.walk(&graph, &idx, &octree, &weights, None, 2);

        let centroid = walker.context_centroid();
        assert!((centroid.x - 1.5).abs() < 1e-6);
    }

    #[test]
    fn beam_search_returns_a_trajectory() {
        let graph = make_graph_line();
        let idx = build_node_index(&graph);
        let octree = build_octree(&graph);
        let weights = build_edge_weights(&graph);

        let config = WalkerConfig {
            knn: 0,
            plan_interval: 0,
            ..Default::default()
        };
        let traj = GeometricWalker::walk_beam(
            &graph, &idx, &octree, &weights, None, &graph[0], 2, 2, &config,
        );
        assert_eq!(traj, vec![1000, 2000, 3000]);
    }

    #[test]
    fn rodrigues_transport_prefers_straight_line() {
        // a -- b -- c is straight; a -- b -- d is a 90° turn.
        let mut a = make_node(1000, 0.0, 0.0, 0.0);
        let mut b = make_node(2000, 1.0, 0.0, 0.0);
        let c = make_node(3000, 2.0, 0.0, 0.0);
        let d = make_node(4000, 1.0, 1.0, 0.0);
        a.successors.push(TemporalEdge {
            dst: 2000,
            weight: 1.0,
            begin_ts: 0,
            end_ts: u64::MAX,
        });
        b.successors.push(TemporalEdge {
            dst: 3000,
            weight: 1.0,
            begin_ts: 0,
            end_ts: u64::MAX,
        });
        b.successors.push(TemporalEdge {
            dst: 4000,
            weight: 1.0,
            begin_ts: 0,
            end_ts: u64::MAX,
        });
        let graph = vec![a, b, c, d];
        let idx = build_node_index(&graph);
        let octree = build_octree(&graph);
        let weights = build_edge_weights(&graph);

        let config = WalkerConfig {
            knn: 0,
            plan_interval: 0,
            temperature: 0.01,
            transition_mode: TransitionMode::RodriguesTransport { lambda: 1.0 },
            ..Default::default()
        };

        // First step a -> b uses fallback distance scoring.
        let mut walker = GeometricWalker::new(&graph[0], config);
        walker.step(&graph, &idx, &octree, &weights, None);
        assert_eq!(walker.current_node(), 2000);

        // Second step: from b, with history a -> b, the straight successor c
        // must outrank the 90° turn d.
        let next = walker.step(&graph, &idx, &octree, &weights, None);
        assert_eq!(next, 3000);
    }
}