paiagram/graph/
arrange.rs

1use crate::graph::{Graph, Station};
2use bevy::ecs::entity::EntityHashSet;
3use bevy::platform::collections::{HashMap, HashSet};
4use bevy::prelude::*;
5use bevy::tasks::{AsyncComputeTaskPool, Task};
6use egui::{Context, Pos2};
7use moonshine_core::prelude::*;
8use serde::Deserialize;
9use std::collections::VecDeque;
10use std::sync::atomic::{AtomicUsize, Ordering};
11use visgraph::layout::force_directed::force_directed_layout;
12
13#[derive(Resource)]
14pub struct GraphLayoutTask {
15    pub task: Task<(Vec<(Instance<Station>, Pos2)>, Vec<Instance<Station>>)>,
16    pub finished: AtomicUsize,
17    pub queued_for_retry: AtomicUsize,
18    pub total: usize,
19}
20
21impl GraphLayoutTask {
22    pub fn in_progress(&self) -> usize {
23        let finished = self.finished.load(Ordering::Relaxed);
24        self.total.saturating_sub(finished)
25    }
26}
27
28impl GraphLayoutTask {
29    pub fn new(
30        task: Task<(Vec<(Instance<Station>, Pos2)>, Vec<Instance<Station>>)>,
31        total: usize,
32    ) -> Self {
33        Self {
34            task,
35            finished: AtomicUsize::new(0),
36            queued_for_retry: AtomicUsize::new(0),
37            total,
38        }
39    }
40
41    pub fn finish(&self, amount: usize) {
42        self.finished.fetch_add(amount, Ordering::Relaxed);
43    }
44
45    pub fn finished_count(&self) -> usize {
46        self.finished.load(Ordering::Relaxed)
47    }
48
49    pub fn queue_retry(&self, amount: usize) {
50        self.queued_for_retry.fetch_add(amount, Ordering::Relaxed);
51    }
52
53    pub fn queued_for_retry_count(&self) -> usize {
54        self.queued_for_retry.load(Ordering::Relaxed)
55    }
56}
57
58/// Apply the graph layout once the task is complete
59/// This system should only be run if the [`GraphLayoutTask`] resource exists
60pub(super) fn apply_graph_layout(
61    mut task: ResMut<GraphLayoutTask>,
62    mut stations: Query<(NameOrEntity, &mut Station)>,
63    mut commands: Commands,
64    graph: Res<Graph>,
65    mut settings: ResMut<crate::settings::ApplicationSettings>,
66) {
67    let Some((found, not_found)) =
68        bevy::tasks::block_on(bevy::tasks::futures_lite::future::poll_once(&mut task.task))
69    else {
70        return;
71    };
72    let mut mapped = false;
73    for (station_instance, pos) in found {
74        if let Ok((_, mut station)) = stations.get_mut(station_instance.entity()) {
75            station.0 = pos;
76            mapped = true;
77        }
78    }
79    let not_found_entities: EntityHashSet = not_found.iter().map(|s| s.entity()).collect();
80    // find the connecting edges for stations that were not found
81    // then find the average position of their connected stations
82    // then assign that position to the not found station
83    for station_instance in not_found.iter().copied() {
84        info!(
85            "Station {:?} is not found in database, arranging via neighbors",
86            stations
87                .get(station_instance.entity())
88                .map_or("<Unknown>".to_string(), |(name, _)| name.to_string())
89        );
90        let Some(node_index) = graph.node_index(station_instance) else {
91            error!("Station {:?} not found in graph", station_instance);
92            continue;
93        };
94
95        let mut valid_neighbor_positions = Vec::new();
96        let mut visited = HashSet::new();
97        visited.insert(node_index);
98
99        let mut queue = VecDeque::new();
100        queue.push_back(node_index);
101
102        while let Some(current_node) = queue.pop_front() {
103            for neighbor_index in graph.inner().neighbors_undirected(current_node) {
104                if !visited.insert(neighbor_index) {
105                    continue;
106                }
107                let Some(stn_instance) = graph.entity(neighbor_index) else {
108                    continue;
109                };
110                if not_found_entities.contains(&stn_instance.entity()) {
111                    queue.push_back(neighbor_index);
112                } else if let Ok((_, stn)) = stations.get(stn_instance.entity()) {
113                    valid_neighbor_positions.push(stn.0);
114                }
115            }
116        }
117
118        let average_pos = if valid_neighbor_positions.is_empty() {
119            Pos2::new(0.0, 0.0)
120        } else {
121            let sum_x: f32 = valid_neighbor_positions.iter().map(|p| p.x).sum();
122            let sum_y: f32 = valid_neighbor_positions.iter().map(|p| p.y).sum();
123            Pos2::new(
124                sum_x / valid_neighbor_positions.len() as f32,
125                sum_y / valid_neighbor_positions.len() as f32,
126            )
127        };
128        if let Ok((_, mut station)) = stations.get_mut(station_instance.entity()) {
129            station.0 = average_pos;
130        }
131    }
132    // cleanup the task resource
133    commands.remove_resource::<GraphLayoutTask>();
134    info!("Finished applying graph layout");
135}
136
137// TODO: move layout algorithms to a separate module
138pub fn auto_arrange_graph(
139    (In(ctx), In(iterations)): (In<Context>, In<u32>),
140    mut commands: Commands,
141    graph: Res<Graph>,
142) {
143    info!("Auto arranging graph with {} iterations", iterations);
144    let inner = graph.inner().clone();
145    let thread_pool = AsyncComputeTaskPool::get();
146    let task = thread_pool.spawn(async move {
147        let graph_ref = &inner;
148        let layout = force_directed_layout(&graph_ref, iterations, 0.1);
149        let results = inner
150            .node_indices()
151            .map(|node| {
152                let pos = layout(node);
153                (inner[node], Pos2::new(pos.0 * 500.0, pos.1 * 500.0))
154            })
155            .collect::<Vec<_>>();
156        ctx.request_repaint();
157        // No stations are "not found" in this method
158        (results, Vec::new())
159    });
160    commands.insert_resource(GraphLayoutTask::new(task, graph.inner().node_count()));
161}
162
163#[derive(Deserialize)]
164struct OSMResponse {
165    elements: Vec<OSMElement>,
166}
167
168impl OSMResponse {
169    fn get_element_by_name(&self, name: &str) -> Option<&OSMElement> {
170        // find the element with the closest matching name
171        let mut best_match: Option<(&OSMElement, f64, &str)> = None;
172
173        for element in &self.elements {
174            for (k, v) in &element.tags {
175                if !k.starts_with("name") {
176                    continue;
177                }
178
179                if v == name {
180                    return Some(element);
181                }
182
183                let score = strsim::jaro_winkler(name, v);
184                if score > 0.9 {
185                    if best_match.as_ref().map_or(true, |&(_, s, _)| score > s) {
186                        best_match = Some((element, score, v));
187                    }
188                }
189            }
190        }
191
192        if let Some((element, score, matched_name)) = best_match {
193            info!(
194                "Fuzzy matched '{}' to '{:?}' (score: {:.2})",
195                name, matched_name, score
196            );
197            Some(element)
198        } else {
199            None
200        }
201    }
202}
203
204#[derive(Deserialize)]
205struct OSMElement {
206    lat: f64,
207    lon: f64,
208    tags: HashMap<String, String>,
209}
210
211impl OSMElement {
212    fn to_pos2(&self) -> Pos2 {
213        // Web Mercator projection (EPSG:3857)
214        // This preserves angles and local shapes, making the map look "natural".
215        let lat_rad = self.lat.to_radians();
216        let lon_rad = self.lon.to_radians();
217
218        const EARTH_RADIUS: f64 = 6378137.0;
219
220        let x = EARTH_RADIUS * lon_rad;
221        let y = EARTH_RADIUS * ((lat_rad / 2.0) + (std::f64::consts::PI / 4.0)).tan().ln();
222
223        // In Egui, Y increases downwards. Mapping North to smaller Y (Up)
224        // and East to larger X (Right).
225        Pos2::new(x as f32, -y as f32)
226    }
227}
228
229// TODO: move all OSM reading related stuff into a separate module
230pub fn arrange_via_osm(
231    (In(ctx), In(area_name)): (In<Context>, In<Option<String>>),
232    mut commands: Commands,
233    station_names: Query<(Instance<Station>, &Name)>,
234) {
235    const MAX_RETRY_COUNT: usize = 3;
236    info!("Arranging graph via OSM with parameters...");
237    info!(?area_name);
238    let mut task_queue: VecDeque<(_, usize)> = station_names
239        .iter()
240        .map(|(instance, name)| (instance, name.to_string()))
241        .collect::<Vec<_>>()
242        .chunks(50)
243        .map(|chunk| (chunk.to_vec(), 0))
244        .collect();
245    let thread_pool = AsyncComputeTaskPool::get();
246
247    let async_task = async move {
248        let mut found: Vec<(Instance<Station>, Pos2)> = Vec::new();
249        let mut not_found: Vec<Instance<Station>> = Vec::new();
250        while let Some((task, retry_count)) = task_queue.pop_front() {
251            if retry_count >= MAX_RETRY_COUNT {
252                error!("Max retry count reached for chunk: {:?}", task);
253                for (instance, _) in task {
254                    not_found.push(instance);
255                }
256                continue;
257            }
258            // Build Overpass Query for the chunk
259            let names_regex = task
260                .iter()
261                .map(|(_, name)| name.as_str())
262                .collect::<Vec<_>>()
263                .join("|");
264
265            let (area_def, area_filter) = match area_name.as_ref() {
266                Some(area) => (
267                    format!(r#"area[name="{}"]->.searchArea;"#, area),
268                    "(area.searchArea)",
269                ),
270                None => ("".to_string(), ""),
271            };
272
273            let query = format!(
274                r#"[out:json];{}(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(:.*)?"~"^({})$"]{};);out;"#,
275                area_def, names_regex, area_filter
276            );
277
278            // 2. Fetch data from Overpass API using a POST request to handle large queries
279            let url = "https://maps.mail.ru/osm/tools/overpass/api/interpreter";
280            let request = ehttp::Request::post(
281                url,
282                format!("data={}", urlencoding::encode(&query)).into_bytes(),
283            );
284
285            let response = match ehttp::fetch_async(request).await {
286                Ok(resp) => resp,
287                Err(e) => {
288                    error!("Failed to fetch OSM data for chunk: {}", e);
289                    task_queue.push_back((task, retry_count + 1));
290                    continue;
291                }
292            };
293
294            let osm_data: OSMResponse = match response.json() {
295                Ok(data) => data,
296                Err(e) => {
297                    error!(
298                        "Failed to parse OSM data: {}, response: {:?}",
299                        e,
300                        response.text()
301                    );
302                    task_queue.push_back((task, retry_count + 1));
303                    continue;
304                }
305            };
306
307            // 3. Match stations and get positions for this chunk
308            for (instance, name) in task {
309                if let Some(osm_element) = osm_data.get_element_by_name(&name) {
310                    let pos = osm_element.to_pos2();
311                    found.push((instance, pos));
312                    info!(
313                        "Matched station '{}' to OSM element at position {:?}",
314                        name, pos
315                    );
316                } else {
317                    warn!("No matching OSM element found for station: {}", name);
318                    not_found.push(instance);
319                }
320            }
321        }
322        ctx.request_repaint();
323        (found, not_found)
324    };
325
326    let task = thread_pool.spawn(async_task);
327    commands.insert_resource(GraphLayoutTask::new(task, station_names.iter().count()));
328}
paiagram/graph/arrange.rs

paiagram/graph/
arrange.rs
