1use crate::direction::{NeighborDirection, NodeNeighbor};
2use kiddo::{KdTree, SquaredEuclidean};
3use nalgebra::{Point2, Vector2};
4
5pub struct NeighborCandidate {
7 pub index: usize,
9 pub offset: Vector2<f32>,
11 pub distance: f32,
13}
14
15pub trait NeighborValidator {
20 type PointData;
23
24 fn validate(
28 &self,
29 source_index: usize,
30 source_data: &Self::PointData,
31 candidate: &NeighborCandidate,
32 candidate_data: &Self::PointData,
33 ) -> Option<(NeighborDirection, f32)>;
34}
35
36#[derive(Clone, Debug)]
38pub struct GridGraphParams {
39 pub k_neighbors: usize,
41 pub max_distance: f32,
43}
44
45impl Default for GridGraphParams {
46 fn default() -> Self {
47 Self {
48 k_neighbors: 8,
49 max_distance: f32::MAX,
50 }
51 }
52}
53
54pub struct GridGraph {
59 pub neighbors: Vec<Vec<NodeNeighbor>>,
62}
63
64impl GridGraph {
65 pub fn build<V: NeighborValidator>(
72 positions: &[Point2<f32>],
73 point_data: &[V::PointData],
74 validator: &V,
75 params: &GridGraphParams,
76 ) -> Self {
77 assert_eq!(
78 positions.len(),
79 point_data.len(),
80 "positions and point_data must have the same length"
81 );
82
83 let coords: Vec<[f32; 2]> = positions.iter().map(|p| [p.x, p.y]).collect();
84 let tree: KdTree<f32, 2> = (&coords).into();
85 let max_dist_sq = params.max_distance * params.max_distance;
86
87 let mut neighbors = Vec::with_capacity(positions.len());
88
89 for (i, pos) in positions.iter().enumerate() {
90 let query = [pos.x, pos.y];
91 let results = tree.nearest_n::<SquaredEuclidean>(&query, params.k_neighbors);
92
93 let mut candidates = Vec::new();
94
95 for nn in results {
96 let j = nn.item as usize;
97 if j == i {
98 continue;
99 }
100
101 let dist_sq = nn.distance;
102 if dist_sq > max_dist_sq {
103 continue;
104 }
105
106 let neighbor_pos = positions[j];
107 let offset = Vector2::new(neighbor_pos.x - pos.x, neighbor_pos.y - pos.y);
108 let distance = dist_sq.sqrt();
109
110 let candidate = NeighborCandidate {
111 index: j,
112 offset,
113 distance,
114 };
115
116 if let Some((direction, score)) =
117 validator.validate(i, &point_data[i], &candidate, &point_data[j])
118 {
119 candidates.push(NodeNeighbor {
120 direction,
121 index: j,
122 distance,
123 score,
124 });
125 }
126 }
127
128 neighbors.push(select_neighbors(candidates));
129 }
130
131 Self { neighbors }
132 }
133}
134
135fn select_neighbors(candidates: Vec<NodeNeighbor>) -> Vec<NodeNeighbor> {
137 let mut best: [Option<NodeNeighbor>; 4] = [None, None, None, None];
138
139 for candidate in candidates {
140 let slot = match candidate.direction {
141 NeighborDirection::Right => &mut best[0],
142 NeighborDirection::Left => &mut best[1],
143 NeighborDirection::Up => &mut best[2],
144 NeighborDirection::Down => &mut best[3],
145 };
146
147 let replace = match slot {
148 None => true,
149 Some(current) => {
150 candidate.score < current.score
151 || (candidate.score == current.score && candidate.distance < current.distance)
152 }
153 };
154
155 if replace {
156 *slot = Some(candidate);
157 }
158 }
159
160 best.into_iter().flatten().collect()
161}