1use std::io::Cursor;
2
3use bevy::platform::collections::HashMap;
4use bevy::prelude::*;
5use moonshine_core::kind::Instance;
6
7use crate::colors::{DisplayedColor, PredefinedColor};
8use crate::entry::{EntryBundle, TravelMode};
9use crate::graph::{Graph, Node, NodeCoor};
10use crate::route::Route;
11use crate::station::{Platform, Station};
12use crate::trip::class::{Class, ClassBundle, DisplayedStroke};
13use crate::trip::{TripBundle, TripClass};
14use crate::units::distance::Distance;
15use crate::units::time::TimetableTime;
16use crate::vehicle::Vehicle;
17
18fn haversine_km(lat1: f64, lon1: f64, lat2: f64, lon2: f64) -> f64 {
19 let r_km = 6371.0_f64;
20 let d_lat = (lat2 - lat1).to_radians();
21 let d_lon = (lon2 - lon1).to_radians();
22 let a = (d_lat / 2.0).sin().powi(2)
23 + lat1.to_radians().cos() * lat2.to_radians().cos() * (d_lon / 2.0).sin().powi(2);
24 let c = 2.0 * a.sqrt().asin();
25 r_km * c
26}
27
28fn route_name(route: Option<>fs_structures::Route>) -> String {
29 route
30 .and_then(|r| {
31 r.long_name
32 .clone()
33 .or_else(|| r.short_name.clone())
34 .or_else(|| Some(r.id.clone()))
35 })
36 .unwrap_or_else(|| "GTFS Route".to_string())
37}
38
39fn class_name(route: Option<>fs_structures::Route>, route_id: &str) -> String {
40 route
41 .and_then(|r| {
42 r.short_name
43 .clone()
44 .or_else(|| r.long_name.clone())
45 .or_else(|| Some(r.id.clone()))
46 })
47 .unwrap_or_else(|| route_id.to_string())
48}
49
50fn class_color(route: Option<>fs_structures::Route>) -> DisplayedColor {
51 if let Some(rgb) = route.and_then(|r| r.color) {
52 return DisplayedColor::Custom(egui::Color32::from_rgb(rgb.r, rgb.g, rgb.b));
53 }
54 DisplayedColor::Predefined(PredefinedColor::Neutral)
55}
56
57fn stop_display_name(stop: >fs_structures::Stop) -> String {
58 stop.name.clone().unwrap_or_else(|| stop.id.clone())
59}
60
61pub fn load_gtfs_static(
62 data: On<super::LoadGTFS>,
63 mut commands: Commands,
64 mut graph: ResMut<Graph>,
65) {
66 info!("Loading GTFS static data...");
67 let reader = Cursor::new(data.content.as_slice());
68 let Ok(gtfs) = gtfs_structures::Gtfs::from_reader(reader) else {
69 warn!("Failed to parse GTFS zip");
70 return;
71 };
72
73 let mut station_entities: HashMap<String, Entity> = HashMap::new();
74 let mut platform_entities: HashMap<String, Entity> = HashMap::new();
75 let mut class_map: HashMap<String, Instance<Class>> = HashMap::new();
76 let mut route_built: std::collections::HashSet<String> = std::collections::HashSet::new();
77 let mut block_to_trips: HashMap<String, Vec<Entity>> = HashMap::new();
78
79 let mut ensure_station =
80 |station_id: &str, station_name: &str, graph: &mut Graph, commands: &mut Commands| {
81 if let Some(&entity) = station_entities.get(station_id) {
82 return entity;
83 }
84 let entity = commands
85 .spawn((Station::default(), Name::new(station_name.to_string())))
86 .id();
87 graph.add_node(entity);
88 station_entities.insert(station_id.to_string(), entity);
89 entity
90 };
91
92 let mut ensure_platform = |platform_id: &str,
93 platform_name: &str,
94 parent_station: Entity,
95 commands: &mut Commands| {
96 if let Some(&entity) = platform_entities.get(platform_id) {
97 return entity;
98 }
99 let entity = commands
100 .spawn((
101 Platform::default(),
102 Name::new(platform_name.to_string()),
103 ChildOf(parent_station),
104 ))
105 .id();
106 platform_entities.insert(platform_id.to_string(), entity);
107 entity
108 };
109
110 for trip in gtfs.trips.values() {
111 if trip.stop_times.is_empty() {
112 continue;
113 }
114
115 let route = gtfs.routes.get(&trip.route_id);
116 let class_name = class_name(route, &trip.route_id);
117 let trip_class =
118 super::make_class(&class_name, &mut class_map, &mut commands, || ClassBundle {
119 class: Class::default(),
120 name: Name::new(class_name.clone()),
121 stroke: DisplayedStroke {
122 color: class_color(route),
123 width: 1.0,
124 },
125 });
126
127 let mut stops_for_trip: Vec<(Entity, Option<f64>, Option<f64>, Option<f32>)> =
128 Vec::with_capacity(trip.stop_times.len());
129 for stop_time in &trip.stop_times {
130 let stop = &stop_time.stop;
131 let stop_name = stop_display_name(stop);
132
133 let (station_entity, platform_entity) =
134 if let Some(parent_station_id) = &stop.parent_station {
135 let parent_stop = gtfs.stops.get(parent_station_id);
136 let parent_name = parent_stop
137 .map(|stop| stop_display_name(stop.as_ref()))
138 .unwrap_or_else(|| parent_station_id.clone());
139 let station_entity =
140 ensure_station(parent_station_id, &parent_name, &mut graph, &mut commands);
141 let platform_entity =
142 ensure_platform(&stop.id, &stop_name, station_entity, &mut commands);
143 (station_entity, platform_entity)
144 } else {
145 let station_entity =
146 ensure_station(&stop.id, &stop_name, &mut graph, &mut commands);
147 (station_entity, station_entity)
148 };
149
150 if let (Some(lat), Some(lon)) = (stop.latitude, stop.longitude) {
151 commands.entity(platform_entity).insert(Node {
152 coor: NodeCoor::new(lon, lat),
153 });
154 if platform_entity != station_entity {
155 commands.entity(station_entity).insert(Node {
156 coor: NodeCoor::new(lon, lat),
157 });
158 }
159 }
160
161 stops_for_trip.push((
162 station_entity,
163 stop.latitude,
164 stop.longitude,
165 stop_time.shape_dist_traveled,
166 ));
167 }
168
169 if !route_built.contains(&trip.route_id) {
170 let mut route_stops: Vec<Entity> = Vec::new();
171 let mut lengths: Vec<f32> = Vec::new();
172 let mut prev_station: Option<Entity> = None;
173 let mut prev_shape_dist: Option<f32> = None;
174 let mut prev_lat_lon: Option<(f64, f64)> = None;
175
176 for (stop, lat, lon, shape_dist) in &stops_for_trip {
177 let curr_station = *stop;
178 if prev_station == Some(curr_station) {
179 continue;
180 }
181
182 route_stops.push(curr_station);
183 if let Some(prev) = prev_station {
184 let mut km = match (shape_dist, prev_shape_dist) {
185 (Some(curr), Some(prev)) => (*curr - prev).abs(),
186 _ => 0.0,
187 };
188 if km <= f32::EPSILON
189 && let (Some((p_lat, p_lon)), Some(c_lat), Some(c_lon)) =
190 (prev_lat_lon, *lat, *lon)
191 {
192 km = haversine_km(p_lat, p_lon, c_lat, c_lon) as f32;
193 }
194 if km <= f32::EPSILON {
195 km = 1.0;
196 }
197 super::add_interval_pair(
198 &mut graph,
199 &mut commands,
200 prev,
201 curr_station,
202 Distance::from_km(km),
203 );
204 lengths.push(km);
205 } else {
206 lengths.push(0.0);
207 }
208
209 prev_station = Some(curr_station);
210 prev_shape_dist = *shape_dist;
211 prev_lat_lon = lat.zip(*lon);
212 }
213
214 if route_stops.len() >= 2 {
215 commands.spawn((
216 Name::new(route_name(route)),
217 Route {
218 stops: route_stops,
219 lengths,
220 },
221 ));
222 }
223 route_built.insert(trip.route_id.clone());
224 }
225
226 let trip_name = trip
227 .trip_short_name
228 .as_ref()
229 .or(trip.trip_headsign.as_ref())
230 .map_or_else(|| trip.id.clone(), std::clone::Clone::clone);
231
232 let mut entry_payloads: Vec<(Entity, Option<TimetableTime>, TimetableTime)> =
233 Vec::with_capacity(trip.stop_times.len());
234 let mut previous_arrival: Option<TimetableTime> = None;
235 for stop_time in &trip.stop_times {
236 let stop = &stop_time.stop;
237 let stop_name = stop_display_name(stop);
238
239 let stop_entity = if let Some(parent_station_id) = &stop.parent_station {
240 let parent_stop = gtfs.stops.get(parent_station_id);
241 let parent_name = parent_stop
242 .map(|stop| stop_display_name(stop.as_ref()))
243 .unwrap_or_else(|| parent_station_id.clone());
244 let parent_station =
245 ensure_station(parent_station_id, &parent_name, &mut graph, &mut commands);
246 ensure_platform(&stop.id, &stop_name, parent_station, &mut commands)
247 } else {
248 ensure_station(&stop.id, &stop_name, &mut graph, &mut commands)
249 };
250
251 if let (Some(lat), Some(lon)) = (stop.latitude, stop.longitude) {
252 commands.entity(stop_entity).insert(Node {
253 coor: NodeCoor::new(lon, lat),
254 });
255 }
256
257 let arr = stop_time
258 .arrival_time
259 .or(stop_time.departure_time)
260 .map(|t| TimetableTime(t as i32));
261 let dep = stop_time
262 .departure_time
263 .or(stop_time.arrival_time)
264 .map(|t| TimetableTime(t as i32));
265
266 let Some(arrival) = arr else {
267 continue;
268 };
269 let departure = dep.unwrap_or(arrival);
270 if let Some(prev) = previous_arrival
271 && arrival < prev
272 {
273 warn!("GTFS trip has non-monotonic time: trip_id={}", trip.id);
274 }
275 previous_arrival = Some(arrival);
276
277 let arr_mode = (departure != arrival).then_some(arrival);
278 entry_payloads.push((stop_entity, arr_mode, departure));
279 }
280
281 let nominal_schedule: Vec<_> = entry_payloads
282 .into_iter()
283 .map(|(stop_entity, arr_mode, departure)| {
284 let arr_mode = arr_mode.map(TravelMode::At);
285 commands
286 .spawn(EntryBundle::new(
287 arr_mode,
288 TravelMode::At(departure),
289 stop_entity,
290 ))
291 .id()
292 })
293 .collect();
294
295 let trip_entity = commands
296 .spawn_empty()
297 .add_children(&nominal_schedule)
298 .insert(TripBundle::new(
299 &trip_name,
300 TripClass(trip_class.entity()),
301 nominal_schedule,
302 ))
303 .id();
304
305 if let Some(block_id) = &trip.block_id {
306 block_to_trips
307 .entry(block_id.clone())
308 .or_default()
309 .push(trip_entity);
310 }
311 }
312
313 for (block_id, trips) in block_to_trips {
314 commands.spawn((
315 Name::new(format!("GTFS block {block_id}")),
316 Vehicle { trips },
317 ));
318 }
319
320 info!(
321 "GTFS import completed: stations={}, classes={}, routes={}, vehicles={}",
322 station_entities.len(),
323 class_map.len(),
324 route_built.len(),
325 gtfs.trips
326 .values()
327 .filter(|t| t.block_id.is_some())
328 .map(|t| t.block_id.as_ref().unwrap())
329 .collect::<std::collections::HashSet<_>>()
330 .len()
331 );
332}