paiagram/vehicles/
calculate_estimates.rs

1//! Calculate the time estimates for each entry in a timetable
2
3use 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}