1use std::collections::{HashMap, HashSet, VecDeque};
2use std::sync::Arc;
3use std::sync::atomic::{AtomicUsize, Ordering};
4
5use bevy::ecs::entity::EntityHashMap;
6use bevy::prelude::*;
7use bevy::tasks::futures_lite::future::poll_once;
8use bevy::tasks::{AsyncComputeTaskPool, Task, block_on};
9use petgraph::graph::NodeIndex;
10use serde::Deserialize;
11use visgraph::layout::force_directed::force_directed_layout;
12
13use super::{Graph, Node, NodeCoor};
14
15#[derive(Clone, Copy, Debug, PartialEq, Eq)]
16pub enum GraphLayoutKind {
17 ForceDirected,
18 OSM,
19}
20
21#[derive(Resource)]
22pub struct GraphLayoutTask {
23 pub task: Task<Vec<(Entity, NodeCoor)>>,
24 finished: Arc<AtomicUsize>,
25 queued_for_retry: Arc<AtomicUsize>,
26 pub total: usize,
27 pub kind: GraphLayoutKind,
28}
29
30impl GraphLayoutTask {
31 fn new(
32 task: Task<Vec<(Entity, NodeCoor)>>,
33 finished: Arc<AtomicUsize>,
34 queued_for_retry: Arc<AtomicUsize>,
35 total: usize,
36 kind: GraphLayoutKind,
37 ) -> Self {
38 Self {
39 task,
40 finished,
41 queued_for_retry,
42 total,
43 kind,
44 }
45 }
46
47 pub fn progress(&self) -> (usize, usize, usize) {
48 (
49 self.finished.load(Ordering::Relaxed),
50 self.total,
51 self.queued_for_retry.load(Ordering::Relaxed),
52 )
53 }
54}
55
56pub fn apply_graph_layout_task(
57 mut commands: Commands,
58 task: Option<ResMut<GraphLayoutTask>>,
59 mut nodes: Query<&mut Node>,
60) {
61 let Some(mut task) = task else {
62 return;
63 };
64 let Some(found) = block_on(poll_once(&mut task.task)) else {
65 return;
66 };
67 for (entity, coor) in found {
68 let Ok(mut node) = nodes.get_mut(entity) else {
69 continue;
70 };
71 node.coor = coor;
72 }
73 let (finished, total, queued_for_retry) = task.progress();
74 info!(
75 "Graph arrange completed: mode={:?}, mapped={finished}/{total}, retry_queued={queued_for_retry}",
76 task.kind
77 );
78 commands.remove_resource::<GraphLayoutTask>();
79}
80
81pub fn apply_force_directed_layout(
82 In(iterations): In<u32>,
83 graph_map: Res<Graph>,
84 mut nodes: Query<&mut Node>,
85) {
86 let graph: petgraph::Graph<_, _, _, usize> = graph_map.map.clone().into_graph();
87 let binding = &graph;
88 let entity_map: EntityHashMap<NodeIndex<usize>> = graph
89 .node_indices()
90 .map(|idx| (*graph.node_weight(idx).unwrap(), idx))
91 .collect();
92 let layout = force_directed_layout(&binding, iterations, 0.1);
93
94 for node_entity in graph_map.nodes() {
95 let Some(&idx) = entity_map.get(&node_entity) else {
96 continue;
97 };
98 let Ok(mut node) = nodes.get_mut(node_entity) else {
99 continue;
100 };
101 let (nx, ny) = layout(idx);
102 node.coor = NodeCoor::from_xy(nx as f64, ny as f64);
103 }
104}
105
106pub fn auto_arrange_graph(
107 (In(ctx), In(iterations)): (In<egui::Context>, In<u32>),
108 mut commands: Commands,
109 graph_map: Res<Graph>,
110) {
111 let graph: petgraph::Graph<_, _, _, usize> = graph_map.map.clone().into_graph();
112 let total = graph.node_count();
113 let finished = Arc::new(AtomicUsize::new(0));
114 let queued_for_retry = Arc::new(AtomicUsize::new(0));
115 let finished_in_task = Arc::clone(&finished);
116
117 info!(
118 "Starting force-directed arrange: nodes={}, iterations={}",
119 total, iterations
120 );
121
122 let task = AsyncComputeTaskPool::get().spawn(async move {
123 let binding = &graph;
124 let layout = force_directed_layout(&binding, iterations, 0.1);
125 let out: Vec<(Entity, NodeCoor)> = graph
126 .node_indices()
127 .map(|idx| {
128 let (x, y) = layout(idx);
129 (
130 *graph.node_weight(idx).unwrap(),
131 NodeCoor::from_xy(x as f64 * 10000.0, y as f64 * 10000.0),
132 )
133 })
134 .collect();
135 finished_in_task.store(total, Ordering::Relaxed);
136 ctx.request_repaint();
137 out
138 });
139
140 commands.insert_resource(GraphLayoutTask::new(
141 task,
142 finished,
143 queued_for_retry,
144 total,
145 GraphLayoutKind::ForceDirected,
146 ));
147}
148
149#[derive(Deserialize)]
150struct OSMResponse {
151 elements: Vec<OSMElement>,
152}
153
154#[derive(Deserialize)]
155struct OSMElement {
156 lat: Option<f64>,
157 lon: Option<f64>,
158 center: Option<OSMCenter>,
159 #[serde(default)]
160 tags: std::collections::HashMap<String, String>,
161}
162
163#[derive(Deserialize)]
164struct OSMCenter {
165 lat: f64,
166 lon: f64,
167}
168
169impl OSMElement {
170 fn coor(&self) -> Option<NodeCoor> {
171 match (self.lon, self.lat, self.center.as_ref()) {
172 (Some(lon), Some(lat), _) => Some(NodeCoor::new(lon, lat)),
173 (_, _, Some(center)) => Some(NodeCoor::new(center.lon, center.lat)),
174 _ => None,
175 }
176 }
177}
178
179fn escape_overpass_regex(input: &str) -> String {
180 let mut out = String::with_capacity(input.len());
181 for c in input.chars() {
182 match c {
183 '\\' | '.' | '+' | '*' | '?' | '^' | '$' | '(' | ')' | '[' | ']' | '{' | '}' | '|' => {
184 out.push('\\');
185 out.push(c);
186 }
187 _ => out.push(c),
188 }
189 }
190 out
191}
192
193fn name_tag_weight(key: &str) -> f64 {
194 match key {
195 "name" => 0.06,
196 _ if key.starts_with("name:") => 0.05,
197 "official_name" => 0.04,
198 _ if key.starts_with("official_name:") => 0.04,
199 "short_name" => 0.03,
200 _ if key.starts_with("short_name:") => 0.03,
201 "loc_name" => 0.02,
202 _ if key.starts_with("loc_name:") => 0.02,
203 "alt_name" => 0.01,
204 _ if key.starts_with("alt_name:") => 0.01,
205 "old_name" => 0.0,
206 _ if key.starts_with("old_name:") => 0.0,
207 _ => -1.0,
208 }
209}
210
211fn station_kind_weight(tags: &HashMap<String, String>) -> f64 {
212 let railway_weight: f64 = match tags.get("railway").map(String::as_str) {
213 Some("station") => 0.60,
214 Some("halt") => 0.55,
215 Some("tram_stop") => 0.45,
216 Some("stop") => 0.40,
217 Some("light_rail") | Some("subway") | Some("monorail_station") => 0.40,
218 Some("stop_position") => 0.20,
219 Some("platform") => 0.15,
220 Some("disused_station") | Some("preserved") => 0.10,
221 Some(_) | None => 0.0,
222 };
223 let public_transport_weight: f64 = match tags.get("public_transport").map(String::as_str) {
224 Some("station") => 0.50,
225 Some("stop_area") => 0.35,
226 Some("platform") => 0.20,
227 Some("stop_position") => 0.15,
228 Some(_) | None => 0.0,
229 };
230 let station_weight: f64 = match tags.get("station").map(String::as_str) {
231 Some("subway") | Some("light_rail") => 0.20,
232 Some(_) | None => 0.0,
233 };
234 railway_weight
235 .max(public_transport_weight)
236 .max(station_weight)
237}
238
239fn best_name_match<'a>(elements: &'a [OSMElement], station_name: &str) -> Option<&'a OSMElement> {
240 let mut best: Option<(&OSMElement, f64)> = None;
241 for element in elements {
242 if element.coor().is_none() {
243 continue;
244 }
245 let base_weight = station_kind_weight(&element.tags);
246 for (key, value) in &element.tags {
247 let name_weight = name_tag_weight(key);
248 if name_weight < 0.0 {
249 continue;
250 }
251
252 let score = if value == station_name {
253 2.0 + base_weight + name_weight
254 } else {
255 let similarity = strsim::jaro_winkler(station_name, value);
256 if similarity <= 0.9 {
257 continue;
258 }
259 similarity + base_weight + name_weight
260 };
261
262 if best
263 .as_ref()
264 .is_none_or(|(_, best_score)| score > *best_score)
265 {
266 best = Some((element, score));
267 }
268 }
269 }
270 best.map(|(element, _)| element)
271}
272
273fn fill_unmatched_via_neighbors(
274 graph: &petgraph::Graph<Entity, Entity, petgraph::Directed, usize>,
275 known_positions: &mut HashMap<Entity, NodeCoor>,
276 all_stations: &[Entity],
277) -> usize {
278 let entity_to_index: HashMap<Entity, NodeIndex<usize>> = graph
279 .node_indices()
280 .map(|idx| (*graph.node_weight(idx).unwrap(), idx))
281 .collect();
282
283 let mut fallback_count = 0usize;
284 for &station in all_stations {
285 if known_positions.contains_key(&station) {
286 continue;
287 }
288 let Some(&start_idx) = entity_to_index.get(&station) else {
289 continue;
290 };
291
292 let mut queue = VecDeque::new();
293 let mut visited = HashSet::new();
294 let mut found_neighbor_positions = Vec::new();
295
296 queue.push_back(start_idx);
297 visited.insert(start_idx);
298
299 while let Some(current) = queue.pop_front() {
300 for neighbor in graph.neighbors_undirected(current) {
301 if !visited.insert(neighbor) {
302 continue;
303 }
304 let neighbor_entity = *graph.node_weight(neighbor).unwrap();
305 if let Some(coor) = known_positions.get(&neighbor_entity) {
306 found_neighbor_positions.push(*coor);
307 } else {
308 queue.push_back(neighbor);
309 }
310 }
311 }
312
313 if found_neighbor_positions.is_empty() {
314 continue;
315 }
316
317 let count = found_neighbor_positions.len() as f64;
318 let avg_lon = found_neighbor_positions.iter().map(|p| p.lon).sum::<f64>() / count;
319 let avg_lat = found_neighbor_positions.iter().map(|p| p.lat).sum::<f64>() / count;
320 known_positions.insert(station, NodeCoor::new(avg_lon, avg_lat));
321 fallback_count += 1;
322 }
323
324 fallback_count
325}
326
327pub fn arrange_via_osm(
328 (In(ctx), In(area_name)): (In<egui::Context>, In<Option<String>>),
329 mut commands: Commands,
330 graph_map: Res<Graph>,
331 station_names: Query<(Entity, &Name), With<crate::station::Station>>,
332) {
333 const MAX_RETRY_COUNT: usize = 3;
334 const OVERPASS_ENDPOINTS: [&str; 2] = [
335 "https://maps.mail.ru/osm/tools/overpass/api/interpreter",
336 "https://overpass-api.de/api/interpreter",
337 ];
338 let stations: Vec<(Entity, String)> = station_names
339 .iter()
340 .map(|(entity, name)| (entity, name.to_string()))
341 .collect();
342 let total = stations.len();
343 let station_entities: Vec<Entity> = stations.iter().map(|(entity, _)| *entity).collect();
344 let graph: petgraph::Graph<_, _, _, usize> = graph_map.map.clone().into_graph();
345
346 info!(
347 "Starting OSM arrange: stations={}, area={}",
348 total,
349 area_name.as_deref().unwrap_or("<global>")
350 );
351
352 let finished = Arc::new(AtomicUsize::new(0));
353 let queued_for_retry = Arc::new(AtomicUsize::new(0));
354 let finished_in_task = Arc::clone(&finished);
355 let queued_in_task = Arc::clone(&queued_for_retry);
356
357 let mut task_queue: VecDeque<(Vec<(Entity, String)>, usize)> = stations
358 .chunks(100)
359 .map(|chunk| (chunk.to_vec(), 0))
360 .collect();
361
362 let (area_def, area_filter) = match area_name.as_ref() {
363 Some(area) => {
364 if area.len() == 2 && area.chars().all(|c| c.is_ascii_alphabetic()) {
366 let country_code = area.to_uppercase();
367 info!(?country_code);
368 (
369 format!(r#"area["ISO3166-1"="{country_code}"]->.searchArea;"#),
370 "(area.searchArea)",
371 )
372 } else {
373 info!(?area);
374 (
375 format!(r#"area[name="{}"]->.searchArea;"#, area),
376 "(area.searchArea)",
377 )
378 }
379 }
380 None => (String::new(), ""),
381 };
382
383 let task = AsyncComputeTaskPool::get().spawn(async move {
384 let mut known_positions: HashMap<Entity, NodeCoor> = HashMap::new();
385
386 while let Some((chunk, retry_count)) = task_queue.pop_front() {
387 if retry_count >= MAX_RETRY_COUNT {
388 finished_in_task.fetch_add(chunk.len(), Ordering::Relaxed);
389 continue;
390 }
391
392 let names_regex = chunk
393 .iter()
394 .map(|(_, name)| escape_overpass_regex(name))
395 .collect::<Vec<_>>()
396 .join("|");
397
398 let query = format!(
399 r#"[out:json];{area_def}(node[~"^(railway|public_transport|station|subway|light_rail)$"~"^(station|halt|stop|tram_stop|subway_entrance|monorail_station|light_rail_station|narrow_gauge_station|funicular_station|preserved|disused_station|stop_position|platform|stop_area|subway|railway|tram|yes)$"][~"name(:.*)?"~"^({names_regex})$"]{area_filter};);out;"#,
400 );
401
402 let mut osm_data: Option<OSMResponse> = None;
403 for endpoint in OVERPASS_ENDPOINTS {
404 let request = ehttp::Request::post(
405 endpoint,
406 format!("data={}", urlencoding::encode(&query)).into_bytes(),
407 );
408
409 let response = match ehttp::fetch_async(request).await {
410 Ok(resp) => resp,
411 Err(e) => {
412 warn!(
413 "OSM request failed: endpoint={}, chunk(size={}), retry={}/{} ({:?})",
414 endpoint,
415 chunk.len(),
416 retry_count + 1,
417 MAX_RETRY_COUNT,
418 e
419 );
420 continue;
421 }
422 };
423
424 if !response.ok {
425 let body_preview = response
426 .text()
427 .map(|t| t.chars().take(200).collect::<String>())
428 .unwrap_or_else(|| "<non-utf8>".to_string());
429 warn!(
430 "OSM bad response: endpoint={}, status={} {}, content_type={:?}, body_preview={:?}",
431 endpoint,
432 response.status,
433 response.status_text,
434 response.content_type(),
435 body_preview
436 );
437 continue;
438 }
439
440 match response.json() {
441 Ok(data) => {
442 info!(
443 "OSM chunk fetched: endpoint={}, chunk(size={}), retry={}/{}",
444 endpoint,
445 chunk.len(),
446 retry_count,
447 MAX_RETRY_COUNT
448 );
449 osm_data = Some(data);
450 break;
451 }
452 Err(e) => {
453 let body_preview = response
454 .text()
455 .map(|t| t.chars().take(200).collect::<String>())
456 .unwrap_or_else(|| "<non-utf8>".to_string());
457 warn!(
458 "OSM response parse failed: endpoint={}, chunk(size={}), retry={}/{} ({:?}), content_type={:?}, body_preview={:?}",
459 endpoint,
460 chunk.len(),
461 retry_count + 1,
462 MAX_RETRY_COUNT,
463 e,
464 response.content_type(),
465 body_preview
466 );
467 }
468 }
469 }
470
471 let Some(osm_data) = osm_data else {
472 queued_in_task.fetch_add(chunk.len(), Ordering::Relaxed);
473 task_queue.push_back((chunk, retry_count + 1));
474 continue;
475 };
476
477 let chunk_size = chunk.len();
478 let mut matched_count = 0usize;
479 for (entity, name) in chunk {
480 if let Some(element) = best_name_match(&osm_data.elements, &name) {
481 if let Some(coor) = element.coor() {
482 known_positions.insert(entity, coor);
483 matched_count += 1;
484 }
485 }
486 finished_in_task.fetch_add(1, Ordering::Relaxed);
487 }
488 info!(
489 "OSM chunk processed: matched={}/{}, progress={}/{}",
490 matched_count,
491 chunk_size,
492 finished_in_task.load(Ordering::Relaxed),
493 total
494 );
495 ctx.request_repaint();
496 }
497
498 let fallback_count = fill_unmatched_via_neighbors(&graph, &mut known_positions, &station_entities);
499 info!(
500 "OSM neighbour fallback applied: fallback_mapped={}, total_mapped={}/{}",
501 fallback_count,
502 known_positions.len(),
503 total
504 );
505
506 known_positions.into_iter().collect()
507 });
508
509 commands.insert_resource(GraphLayoutTask::new(
510 task,
511 finished,
512 queued_for_retry,
513 total,
514 GraphLayoutKind::OSM,
515 ));
516}