1use bevy::{ecs::entity::EntityHashSet, prelude::*};
6use moonshine_core::prelude::{MapEntities, ReflectMapEntities};
7
8pub struct RoutePlugin;
9impl Plugin for RoutePlugin {
10 fn build(&self, app: &mut App) {
11 app.add_observer(auto_update_length)
12 .add_observer(sort_route_by_direction_trips)
13 .add_systems(Update, (update_route_trips, auto_generate_display_modes));
14 }
15}
16
17use crate::{
18 entry::{EntryMode, EntryQuery},
19 graph::Graph,
20 interval::{Interval, UpdateInterval},
21 station::{ParentStationOrStation, Platform, PlatformEntries, Station, StationQuery},
22 trip::TripQuery,
23};
24
25#[derive(Reflect, Component)]
27#[reflect(Component)]
28pub struct AutoUpdateLength;
29
30#[derive(Reflect, Component, MapEntities)]
31#[reflect(Component, MapEntities)]
32#[require(Name, RouteTrips, RouteByDirectionTrips)]
33pub struct Route {
34 #[entities]
35 pub stops: Vec<Entity>,
36 pub lengths: Vec<f32>,
37}
38
39#[derive(Reflect)]
40pub struct AllTripsDisplayMode {
41 pub departure: bool,
42 pub arrival: bool,
43}
44
45impl AllTripsDisplayMode {
46 pub fn count(&self) -> usize {
47 self.departure as usize + self.arrival as usize
48 }
49}
50
51#[derive(Reflect, Component, Deref, DerefMut)]
52#[reflect(Component)]
53pub struct RouteDisplayModes(Vec<AllTripsDisplayMode>);
54
55fn auto_generate_display_modes(
56 routes: Populated<(Entity, &Route), Without<RouteDisplayModes>>,
57 graph: Res<Graph>,
58 mut commands: Commands,
59) {
60 for (route_entity, route) in routes.iter().filter(|(_, it)| it.stops.len() > 0) {
61 let mut modes: Vec<AllTripsDisplayMode> = Vec::new();
62 modes.resize_with(route.stops.len(), || AllTripsDisplayMode {
63 departure: true,
64 arrival: false,
65 });
66 modes.last_mut().unwrap().arrival = true;
67 modes.last_mut().unwrap().departure = false;
68 for (idx, s) in route.stops.windows(2).enumerate() {
69 let [prev, curr] = s else { unreachable!() };
70 if graph.contains_edge(*prev, *curr) || graph.contains_edge(*curr, *prev) {
71 continue;
72 }
73 modes[idx].departure = false;
74 modes[idx].arrival = true;
75 }
76 commands
77 .entity(route_entity)
78 .insert(RouteDisplayModes(modes));
79 }
80}
81
82#[derive(Default, Reflect, Component, MapEntities, Deref, DerefMut)]
85#[reflect(Component, MapEntities)]
86#[require(Name)]
87pub struct RouteTrips(#[entities] Vec<Entity>);
88
89#[derive(Default, Reflect, Component, MapEntities)]
90#[reflect(Component, MapEntities)]
91pub struct RouteByDirectionTrips {
92 #[entities]
93 pub downward: Vec<Entity>,
94 #[entities]
95 pub upward: Vec<Entity>,
96}
97
98#[derive(EntityEvent)]
99pub struct SortRouteByDirectionTrips {
100 pub entity: Entity,
101}
102
103fn compute_sorted_by_first_entry_estimate(
104 trip_entities: &[Entity],
105 trip_q: &Query<TripQuery>,
106 entry_q: &Query<EntryQuery>,
107) -> Vec<Entity> {
108 let mut out = trip_entities.to_vec();
109 out.sort_unstable_by_key(|trip_entity| {
110 let trip = trip_q.get(*trip_entity).ok();
111 let first_time = trip
112 .and_then(|trip| {
113 entry_q
114 .iter_many(trip.schedule.iter())
115 .find_map(|entry| entry.estimate.map(|it| it.arr.min(it.dep)))
116 })
117 .unwrap_or(crate::units::time::TimetableTime(i32::MAX));
118 (first_time, trip_entity.to_bits())
119 });
120 out
121}
122
123fn compute_directional_members(
124 route: &Route,
125 trip_entities: &[Entity],
126 downwards: bool,
127 trip_q: &Query<TripQuery>,
128 entry_q: &Query<EntryQuery>,
129 parent_station_or_station: &Query<ParentStationOrStation>,
130) -> Vec<Entity> {
131 trip_entities
132 .iter()
133 .copied()
134 .filter_map(|trip_entity| {
135 let trip = trip_q.get(trip_entity).ok()?;
136 let mut stations = if downwards {
137 either::Either::Left(route.stops.iter())
138 } else {
139 either::Either::Right(route.stops.iter().rev())
140 };
141 let mut found_counter = 0;
142 for it in entry_q.iter_many(trip.schedule.iter()) {
143 let station_entity = parent_station_or_station.get(it.stop()).ok()?.parent();
144 if stations.any(|it| *it == station_entity) {
145 found_counter += 1;
146 if found_counter >= 2 {
147 return Some(trip_entity);
148 }
149 }
150 }
151 None
152 })
153 .collect()
154}
155
156fn sync_direction_order(existing: &mut Vec<Entity>, members: &[Entity], fallback_order: &[Entity]) {
157 let member_set: EntityHashSet = members.iter().copied().collect();
158 let mut next = Vec::with_capacity(members.len());
159
160 for entity in existing.iter().copied() {
161 if member_set.contains(&entity) {
162 next.push(entity);
163 }
164 }
165
166 for entity in fallback_order.iter().copied() {
167 if member_set.contains(&entity) && !next.contains(&entity) {
168 next.push(entity);
169 }
170 }
171
172 *existing = next;
173}
174
175impl Route {
176 pub fn iter(&self) -> impl Iterator<Item = (Entity, f32)> {
177 self.stops
178 .iter()
179 .copied()
180 .zip(self.lengths.iter().copied())
181 .scan(0.0_f32, |acc, (stop, len)| {
182 let out = (stop, *acc);
183 *acc += len;
184 Some(out)
185 })
186 }
187}
188
189fn update_route_trips(
190 mut routes: Query<(Entity, &Route, &mut RouteTrips, &mut RouteByDirectionTrips)>,
191 changed_routes: Query<Entity, (With<Route>, Changed<Route>)>,
192 changed_station_entries: Query<Entity, (With<Station>, Changed<PlatformEntries>)>,
193 changed_platform_entries: Query<&ChildOf, (With<Platform>, Changed<PlatformEntries>)>,
194 stations: Query<StationQuery>,
195 platform_entries: Query<&PlatformEntries>,
196 entries: Query<&ChildOf, With<EntryMode>>,
197 trip_q: Query<TripQuery>,
198 entry_q: Query<EntryQuery>,
199 parent_station_or_station: Query<ParentStationOrStation>,
200) {
201 let mut affected_routes = EntityHashSet::default();
202
203 for route_entity in &changed_routes {
204 affected_routes.insert(route_entity);
205 }
206
207 let mut changed_stations = EntityHashSet::default();
208 for station in &changed_station_entries {
209 changed_stations.insert(station);
210 }
211 for parent in &changed_platform_entries {
212 changed_stations.insert(parent.parent());
213 }
214
215 if !changed_stations.is_empty() {
216 for (route_entity, route, _, _) in &routes {
217 if route
218 .stops
219 .iter()
220 .any(|station| changed_stations.contains(station))
221 {
222 affected_routes.insert(route_entity);
223 }
224 }
225 }
226
227 if affected_routes.is_empty() {
228 return;
229 }
230
231 for (route_entity, route, mut route_trips, mut by_direction) in &mut routes {
232 if !affected_routes.contains(&route_entity) {
233 continue;
234 }
235
236 let mut trips = EntityHashSet::default();
237 for station_entity in route.stops.iter().copied() {
238 let Ok(station) = stations.get(station_entity) else {
239 continue;
240 };
241 for entry in station.passing_entries(&platform_entries) {
242 let Ok(parent) = entries.get(entry) else {
243 continue;
244 };
245 trips.insert(parent.parent());
246 }
247 }
248
249 let mut next = route_trips
250 .0
251 .iter()
252 .copied()
253 .filter(|entity| trips.contains(entity))
254 .collect::<Vec<_>>();
255 for entity in trips.iter().copied() {
256 if !next.contains(&entity) {
257 next.push(entity);
258 }
259 }
260 if route_trips.0 != next {
261 route_trips.0 = next.clone();
262 }
263
264 let new_downward = compute_directional_members(
265 route,
266 &next,
267 true,
268 &trip_q,
269 &entry_q,
270 &parent_station_or_station,
271 );
272 let new_upward = compute_directional_members(
273 route,
274 &next,
275 false,
276 &trip_q,
277 &entry_q,
278 &parent_station_or_station,
279 );
280 sync_direction_order(&mut by_direction.downward, &new_downward, &next);
281 sync_direction_order(&mut by_direction.upward, &new_upward, &next);
282 }
283}
284
285fn sort_route_by_direction_trips(
286 trigger: On<SortRouteByDirectionTrips>,
287 mut routes: Query<(&RouteTrips, &mut RouteByDirectionTrips)>,
288 trip_q: Query<TripQuery>,
289 entry_q: Query<EntryQuery>,
290) {
291 let Ok((route_trips, mut by_direction)) = routes.get_mut(trigger.entity) else {
292 return;
293 };
294
295 let sorted_downward =
296 compute_sorted_by_first_entry_estimate(&by_direction.downward, &trip_q, &entry_q);
297 let sorted_upward =
298 compute_sorted_by_first_entry_estimate(&by_direction.upward, &trip_q, &entry_q);
299
300 by_direction.downward = sorted_downward;
301 by_direction.upward = sorted_upward;
302}
303
304fn auto_update_length(
306 updated: On<UpdateInterval>,
307 routes: Populated<&mut Route, With<AutoUpdateLength>>,
308 intervals: Query<&Interval>,
309 graph: Res<Graph>,
310) {
311 for mut route in routes {
312 let Route { stops, lengths } = &mut *route;
313 for (i, w) in stops.windows(2).enumerate() {
314 let [p, c] = w else { unreachable!() };
315 let (p, c) = (*p, *c);
316 if (p == updated.source && c == updated.target)
317 || (p == updated.target && c == updated.source)
318 {
319 } else {
320 continue;
321 }
322 let i1 = graph.edge_weight(p, c).cloned();
323 let i2 = graph.edge_weight(c, p).cloned();
324 match (i1, i2) {
325 (Some(e1), Some(e2)) => {
326 let d1 = intervals.get(e1).unwrap().length;
327 let d2 = intervals.get(e2).unwrap().length;
328 let avg_len = (d1.0 as f32 + d2.0 as f32) / 2.0;
329 lengths[i] = avg_len;
330 }
331 (Some(e), None) | (None, Some(e)) => {
332 let d = intervals.get(e).unwrap().length;
333 lengths[i] = d.0 as f32;
334 }
335 (None, None) => {
336 panic!("Interval disappeared???")
337 }
338 }
339 }
340 }
341}