1use super::AdjustTimetableEntry;
4use super::entries::{self, TravelMode};
5use crate::graph::{self, Graph, Station};
6use crate::units::distance::Distance;
7use crate::units::time::{Duration, TimetableTime};
8use crate::vehicles::entries::{TimeEstimate, TimetableEntry, TimetableEntryCache};
9use bevy::prelude::*;
10use moonshine_core::kind::Instance;
11
12enum UnwindParams {
13 At {
14 previous: Option<(TimetableTime, Instance<Station>)>,
15 time_offset: Option<Duration>,
16 },
17 For {
18 previous: Option<(TimetableTime, Instance<Station>)>,
19 current: (TimetableTime, Instance<Station>),
20 },
21}
22
23fn clear_estimates(
24 entries: &mut Populated<(&TimetableEntry, &mut TimetableEntryCache)>,
25 stack: &mut Vec<(Entity, Duration)>,
26) {
27 for (timetable_entry_entity, _) in stack.iter() {
28 if let Ok((_, mut cache)) = entries.get_mut(*timetable_entry_entity) {
29 cache.estimate = None;
30 }
31 }
32 stack.clear();
33}
34
35fn build_distances(
36 entries: &mut Populated<(&TimetableEntry, &mut TimetableEntryCache)>,
37 stack: &Vec<(Entity, Duration)>,
38 distances: &mut Vec<Option<Distance>>,
39 graph: &Graph,
40 intervals: &Query<&graph::Interval>,
41 mut previous_station: Instance<Station>,
42 include_last_edge: bool,
43 current_station: Instance<Station>,
44) -> bool {
45 distances.clear();
46 for (timetable_entry_entity, _) in stack.iter() {
47 let station = {
48 let Ok((tte, _)) = entries.get(*timetable_entry_entity) else {
49 return false;
50 };
51 tte.station()
52 };
53 let interval_distance = if previous_station == station {
54 None
55 } else {
56 match graph.edge_weight(previous_station, station) {
57 Some(w) => intervals.get(w.entity()).ok().map(|i| i.length),
58 None => None,
59 }
60 };
61 if previous_station != station && interval_distance.is_none() {
62 return false;
63 }
64 distances.push(interval_distance);
65 previous_station = station;
66 }
67 if include_last_edge {
68 let last = match graph.edge_weight(previous_station, current_station) {
69 Some(w) => intervals.get(w.entity()).ok().map(|i| i.length),
70 None => None,
71 };
72 if last.is_none() {
73 return false;
74 }
75 distances.push(last);
76 } else {
77 distances.push(None);
78 }
79 true
80}
81
82fn process_schedule(
83 entries: &mut Populated<(&TimetableEntry, &mut TimetableEntryCache)>,
84 intervals: &Query<&graph::Interval>,
85 graph: &Graph,
86 schedule: &entries::VehicleScheduleCache,
87 stack: &mut Vec<(Entity, Duration)>,
88 distances: &mut Vec<Option<Distance>>,
89) {
90 stack.clear();
91 let mut stable_time_and_station: Option<(TimetableTime, Instance<Station>)> = None;
92 let mut pending_time_and_station: Option<(TimetableTime, Instance<Station>)> = None;
93 let mut unwind_params: Option<UnwindParams> = None;
94 for timetable_entry_entity in schedule.actual_route.iter().flatten().map(|e| e.inner()) {
95 let (arrival, departure, station) = {
96 let Ok((tte, _)) = entries.get(timetable_entry_entity) else {
97 continue;
98 };
99 (tte.arrival, tte.departure, tte.station())
100 };
101
102 if let Some(v) = pending_time_and_station.take() {
103 stable_time_and_station = Some(v);
104 }
105 match (arrival, departure.unwrap_or(TravelMode::Flexible)) {
106 (TravelMode::At(at), TravelMode::At(dt)) => {
107 if let Ok((_, mut cache)) = entries.get_mut(timetable_entry_entity) {
108 cache.estimate = Some(TimeEstimate {
109 arrival: at,
110 departure: dt,
111 });
112 }
113 unwind_params = Some(UnwindParams::At {
114 previous: stable_time_and_station,
115 time_offset: None,
116 });
117 stable_time_and_station = Some((at, station));
118 pending_time_and_station = Some((dt, station));
119 }
120 (TravelMode::At(at), TravelMode::For(dd)) => {
121 if let Ok((_, mut cache)) = entries.get_mut(timetable_entry_entity) {
122 cache.estimate = Some(TimeEstimate {
123 arrival: at,
124 departure: at + dd,
125 });
126 }
127 unwind_params = Some(UnwindParams::At {
128 previous: stable_time_and_station,
129 time_offset: None,
130 });
131 stable_time_and_station = Some((at, station));
132 pending_time_and_station = Some((at + dd, station));
133 }
134 (TravelMode::At(at), TravelMode::Flexible) => {
135 if let Ok((_, mut cache)) = entries.get_mut(timetable_entry_entity) {
136 cache.estimate = Some(TimeEstimate {
137 arrival: at,
138 departure: at,
139 });
140 }
141 unwind_params = Some(UnwindParams::At {
142 previous: stable_time_and_station,
143 time_offset: None,
144 });
145 stable_time_and_station = Some((at, station));
146 }
147 (TravelMode::For(ad), TravelMode::At(dt)) => {
148 let Some((stable_time, _)) = stable_time_and_station else {
149 clear_estimates(entries, stack);
150 continue;
151 };
152 if let Ok((_, mut cache)) = entries.get_mut(timetable_entry_entity) {
153 cache.estimate = Some(TimeEstimate {
154 arrival: stable_time + ad,
155 departure: dt,
156 });
157 }
158 unwind_params = Some(UnwindParams::For {
159 previous: stable_time_and_station,
160 current: (stable_time + ad, station),
161 });
162 stable_time_and_station = Some((stable_time + ad, station));
163 pending_time_and_station = Some((dt, station));
164 }
165 (TravelMode::For(ad), TravelMode::For(dd)) => {
166 let Some((stable_time, _)) = stable_time_and_station else {
167 clear_estimates(entries, stack);
168 continue;
169 };
170 if let Ok((_, mut cache)) = entries.get_mut(timetable_entry_entity) {
171 cache.estimate = Some(TimeEstimate {
172 arrival: stable_time + ad,
173 departure: stable_time + ad + dd,
174 });
175 }
176 unwind_params = Some(UnwindParams::For {
177 previous: stable_time_and_station,
178 current: (stable_time + ad, station),
179 });
180 stable_time_and_station = Some((stable_time + ad, station));
181 pending_time_and_station = Some((stable_time + ad + dd, station));
182 }
183 (TravelMode::For(ad), TravelMode::Flexible) => {
184 let Some((stable_time, _)) = stable_time_and_station else {
185 clear_estimates(entries, stack);
186 continue;
187 };
188 if let Ok((_, mut cache)) = entries.get_mut(timetable_entry_entity) {
189 cache.estimate = Some(TimeEstimate {
190 arrival: stable_time + ad,
191 departure: stable_time + ad,
192 });
193 }
194 unwind_params = Some(UnwindParams::For {
195 previous: stable_time_and_station,
196 current: (stable_time + ad, station),
197 });
198 stable_time_and_station = Some((stable_time + ad, station));
199 }
200 (TravelMode::Flexible, TravelMode::At(at)) => {
201 if let Ok((_, mut cache)) = entries.get_mut(timetable_entry_entity) {
202 cache.estimate = Some(TimeEstimate {
203 arrival: at,
204 departure: at,
205 });
206 }
207 unwind_params = Some(UnwindParams::At {
208 previous: stable_time_and_station,
209 time_offset: None,
210 });
211 stable_time_and_station = Some((at, station));
212 }
213 (TravelMode::Flexible, TravelMode::For(dd)) => {
214 stack.push((timetable_entry_entity, dd));
215 }
216 (TravelMode::Flexible, TravelMode::Flexible) => {
217 stack.push((timetable_entry_entity, Duration(0)));
218 }
219 };
220 let Some(unwind_params) = unwind_params.take() else {
221 continue;
222 };
223 match unwind_params {
224 UnwindParams::For {
225 previous: previous_time_and_station,
226 current: current_time_and_station,
227 } => {
228 if stack.is_empty() {
229 continue;
230 }
231 let Some((previous_time, previous_station)) = previous_time_and_station else {
232 clear_estimates(entries, stack);
233 continue;
234 };
235 let (mut current_time, current_station) = current_time_and_station;
236 let mut total_time = current_time - previous_time;
237 for (_, dep_dur) in stack.iter() {
238 total_time -= *dep_dur;
239 }
240 if !build_distances(
241 entries,
242 stack,
243 distances,
244 graph,
245 intervals,
246 previous_station,
247 true,
248 current_station,
249 ) {
250 clear_estimates(entries, stack);
251 continue;
252 }
253 let total_distance = distances
254 .iter()
255 .filter_map(|d| *d)
256 .map(|d| d.0)
257 .sum::<i32>();
258 if total_time.0 <= 0 || total_distance <= 0 {
259 clear_estimates(entries, stack);
260 continue;
261 }
262 let velocity = Distance(total_distance) / total_time;
263 debug_assert_eq!(distances.len(), stack.len() + 1);
264 while let (Some((timetable_entry_entity, dep_dur)), Some(distance)) =
265 (stack.pop(), distances.pop())
266 {
267 if let Some(distance) = distance {
268 current_time -= distance / velocity;
269 }
270 let departure_estimate = current_time;
271 current_time = current_time - dep_dur;
272 let arrival_estimate = current_time;
273 if let Ok((_, mut cache)) = entries.get_mut(timetable_entry_entity) {
274 cache.estimate = Some(TimeEstimate {
275 arrival: arrival_estimate,
276 departure: departure_estimate,
277 });
278 }
279 }
280 continue;
281 }
282 UnwindParams::At {
283 previous: previous_time_and_station,
284 time_offset,
285 } => {
286 if stack.is_empty() {
287 continue;
288 }
289 let Some((previous_time, previous_station)) = previous_time_and_station else {
290 clear_estimates(entries, stack);
291 continue;
292 };
293 let Some((mut current_time, current_station)) = stable_time_and_station else {
294 clear_estimates(entries, stack);
295 continue;
296 };
297 let mut total_time = current_time - previous_time;
298 for (_, dep_dur) in stack.iter() {
299 total_time -= *dep_dur;
300 }
301 if !build_distances(
302 entries,
303 stack,
304 distances,
305 graph,
306 intervals,
307 previous_station,
308 time_offset.is_none(),
309 current_station,
310 ) {
311 clear_estimates(entries, stack);
312 continue;
313 }
314 let total_distance = distances
315 .iter()
316 .filter_map(|d| *d)
317 .map(|d| d.0)
318 .sum::<i32>();
319 if let Some(dur) = time_offset {
320 total_time -= dur;
321 }
322 if total_time.0 <= 0 || total_distance <= 0 {
323 clear_estimates(entries, stack);
324 continue;
325 }
326 let velocity = Distance(total_distance) / total_time;
327 if let Some(dur) = time_offset {
328 current_time -= dur;
329 }
330 debug_assert_eq!(distances.len(), stack.len() + 1);
331 while let (Some((timetable_entry_entity, dep_dur)), Some(distance)) =
332 (stack.pop(), distances.pop())
333 {
334 if let Some(distance) = distance {
335 current_time -= distance / velocity;
336 }
337 let departure_estimate = current_time;
338 current_time = current_time - dep_dur;
339 let arrival_estimate = current_time;
340 if let Ok((_, mut cache)) = entries.get_mut(timetable_entry_entity) {
341 cache.estimate = Some(TimeEstimate {
342 arrival: arrival_estimate,
343 departure: departure_estimate,
344 });
345 }
346 }
347 }
348 }
349 }
350 clear_estimates(entries, stack);
351}
352
353pub fn calculate_estimates(
354 mut msg_reader: MessageReader<AdjustTimetableEntry>,
355 mut entries: Populated<(&TimetableEntry, &mut TimetableEntryCache)>,
356 intervals: Query<&graph::Interval>,
357 parents: Populated<&ChildOf>,
358 schedules: Populated<&entries::VehicleScheduleCache>,
359 graph: Res<Graph>,
360 mut stack: Local<Vec<(Entity, Duration)>>,
361 mut distances: Local<Vec<Option<Distance>>>,
362) {
363 let messages = msg_reader.read().collect::<Vec<_>>();
364 if messages.is_empty() {
365 return;
366 }
367
368 for msg in messages {
369 let AdjustTimetableEntry { entity, .. } = msg;
370 let Ok(entry) = parents.get(*entity) else {
371 continue;
372 };
373 let Ok(schedule) = schedules.get(entry.0) else {
374 continue;
375 };
376 process_schedule(
377 &mut entries,
378 &intervals,
379 &graph,
380 schedule,
381 &mut stack,
382 &mut distances,
383 );
384 }
385}