1use super::AdjustTimetableEntry;
2use super::entries::{self, TravelMode};
3use crate::intervals::{self, Graph};
4use crate::units::distance::Distance;
5use crate::units::time::{Duration, TimetableTime};
6use crate::vehicles::entries::{TimeEstimate, TimetableEntry, TimetableEntryCache};
7use bevy::prelude::*;
8
9pub fn calculate_estimates(
10 mut msg_reader: MessageReader<AdjustTimetableEntry>,
11 mut entries: Populated<(&TimetableEntry, &mut TimetableEntryCache)>,
12 intervals: Populated<&intervals::Interval>,
13 parents: Populated<&ChildOf>,
14 schedules: Populated<&entries::VehicleScheduleCache>,
15 graph: Res<Graph>,
16 mut stack: Local<Vec<(Entity, Option<Duration>, Option<Duration>)>>,
17 mut distances: Local<Vec<Option<Distance>>>,
18) {
19 fn clear_estimates(
20 entries: &mut Populated<(&TimetableEntry, &mut TimetableEntryCache)>,
21 stack: &mut Vec<(Entity, Option<Duration>, Option<Duration>)>,
22 ) {
23 for (timetable_entry_entity, _, _) in stack.iter() {
24 if let Ok((_, mut cache)) = entries.get_mut(*timetable_entry_entity) {
25 cache.estimate = None;
26 }
27 }
28 stack.clear();
29 }
30 for msg in msg_reader.read() {
31 let AdjustTimetableEntry { entity, .. } = msg;
32 let Ok(entry) = parents.get(*entity) else {
33 continue;
34 };
35 let Ok(schedule) = schedules.get(entry.0) else {
36 continue;
37 };
38 stack.clear();
39 let mut stable_time_and_station: Option<(TimetableTime, Entity)> = None;
40 let mut pending_time_and_station: Option<(TimetableTime, Entity)> = None;
41 let mut unwind_params: Option<(Option<(TimetableTime, Entity)>, Option<Duration>)> = None;
42 'iter_timetable: for timetable_entry_entity in
43 schedule.actual_route.iter().flatten().map(|e| e.inner())
44 {
45 let (arrival, departure, station) = {
46 let Ok((tte, _)) = entries.get(timetable_entry_entity) else {
47 continue;
48 };
49 (tte.arrival, tte.departure, tte.station)
50 };
51
52 if let Some(v) = pending_time_and_station.take() {
53 stable_time_and_station = Some(v);
54 }
55 match (arrival, departure.unwrap_or(TravelMode::Flexible)) {
56 (TravelMode::At(at), TravelMode::At(dt)) => {
57 if let Ok((_, mut cache)) = entries.get_mut(timetable_entry_entity) {
58 cache.estimate = Some(TimeEstimate {
59 arrival: at,
60 departure: dt,
61 });
62 }
63 unwind_params = Some((stable_time_and_station, None));
64 stable_time_and_station = Some((at, station));
65 pending_time_and_station = Some((dt, station));
66 }
67 (TravelMode::At(at), TravelMode::For(dd)) => {
68 if let Ok((_, mut cache)) = entries.get_mut(timetable_entry_entity) {
69 cache.estimate = Some(TimeEstimate {
70 arrival: at,
71 departure: at + dd,
72 });
73 }
74 unwind_params = Some((stable_time_and_station, None));
75 stable_time_and_station = Some((at, station));
76 pending_time_and_station = Some((at + dd, station));
77 }
78 (TravelMode::At(at), TravelMode::Flexible) => {
79 if let Ok((_, mut cache)) = entries.get_mut(timetable_entry_entity) {
80 cache.estimate = Some(TimeEstimate {
81 arrival: at,
82 departure: at,
83 });
84 }
85 unwind_params = Some((stable_time_and_station, None));
86 stable_time_and_station = Some((at, station));
87 }
88 (TravelMode::For(ad), TravelMode::At(dt)) => {
89 let arrival_estimate = if stack.is_empty()
90 && let Some((stable_time, _)) = stable_time_and_station
91 {
92 stable_time + ad
93 } else {
94 dt
95 };
96 if let Ok((_, mut cache)) = entries.get_mut(timetable_entry_entity) {
97 cache.estimate = Some(TimeEstimate {
98 arrival: arrival_estimate,
99 departure: dt,
100 });
101 }
102 unwind_params = Some((stable_time_and_station, Some(ad)));
103 stable_time_and_station = Some((dt, station));
104 }
105 (TravelMode::For(ad), TravelMode::For(dd)) => {
106 if stack.is_empty()
107 && let Some((stable_time, _)) = stable_time_and_station
108 {
109 if let Ok((_, mut cache)) = entries.get_mut(timetable_entry_entity) {
110 cache.estimate = Some(TimeEstimate {
111 arrival: stable_time + ad,
112 departure: stable_time + ad + dd,
113 });
114 }
115 unwind_params = Some((stable_time_and_station, Some(ad)));
116 stable_time_and_station = Some((stable_time + ad, station));
117 pending_time_and_station = Some((stable_time + ad + dd, station));
118 } else {
119 stack.push((timetable_entry_entity, Some(ad), Some(dd)));
120 }
121 }
122 (TravelMode::For(ad), TravelMode::Flexible) => {
123 if stack.is_empty()
124 && let Some((stable_time, _)) = stable_time_and_station
125 {
126 if let Ok((_, mut cache)) = entries.get_mut(timetable_entry_entity) {
127 cache.estimate = Some(TimeEstimate {
128 arrival: stable_time + ad,
129 departure: stable_time + ad,
130 });
131 }
132 unwind_params = Some((stable_time_and_station, Some(ad)));
133 stable_time_and_station = Some((stable_time + ad, station));
134 } else {
135 stack.push((timetable_entry_entity, Some(ad), None));
136 }
137 }
138 (TravelMode::Flexible, TravelMode::At(at)) => {
139 if let Ok((_, mut cache)) = entries.get_mut(timetable_entry_entity) {
140 cache.estimate = Some(TimeEstimate {
141 arrival: at,
142 departure: at,
143 });
144 }
145 unwind_params = Some((stable_time_and_station, None));
146 stable_time_and_station = Some((at, station));
147 }
148 (TravelMode::Flexible, TravelMode::For(dd)) => {
149 stack.push((timetable_entry_entity, None, Some(dd)));
150 }
151 (TravelMode::Flexible, TravelMode::Flexible) => {
152 stack.push((timetable_entry_entity, None, None));
153 }
154 };
155 let Some((previous_time_and_station, time_offset)) = unwind_params.take() else {
156 continue;
157 };
158 if stack.is_empty() {
159 continue;
160 }
161 let Some((previous_time, mut previous_station)) = previous_time_and_station else {
162 clear_estimates(&mut entries, &mut stack);
163 continue;
164 };
165 let Some((mut current_time, current_station)) = stable_time_and_station else {
166 clear_estimates(&mut entries, &mut stack);
167 continue;
168 };
169 distances.clear();
170 let mut total_time = current_time - previous_time;
171 for (timetable_entry_entity, arr_dur, dep_dur) in stack.iter() {
172 let station = {
173 let Ok((tte, _)) = entries.get(*timetable_entry_entity) else {
174 clear_estimates(&mut entries, &mut stack);
175 continue 'iter_timetable;
176 };
177 tte.station
178 };
179 let interval_distance = if previous_station == station || arr_dur.is_some() {
180 None
181 } else {
182 match graph.edge_weight(previous_station, station) {
183 Some(w) => {
184 if let Ok(interval) = intervals.get(*w) {
185 Some(interval.length)
186 } else {
187 clear_estimates(&mut entries, &mut stack);
188 continue 'iter_timetable;
189 }
190 }
191 None => {
192 clear_estimates(&mut entries, &mut stack);
193 continue 'iter_timetable;
194 }
195 }
196 };
197 if let Some(dur) = arr_dur {
198 total_time -= *dur;
199 }
200 if let Some(dur) = dep_dur {
201 total_time -= *dur;
202 }
203 distances.push(interval_distance);
204 previous_station = station;
205 }
206 distances.push(if time_offset.is_none() {
207 match graph.edge_weight(previous_station, current_station) {
208 Some(w) => {
209 if let Ok(interval) = intervals.get(*w) {
210 Some(interval.length)
211 } else {
212 clear_estimates(&mut entries, &mut stack);
213 continue 'iter_timetable;
214 }
215 }
216 None => {
217 clear_estimates(&mut entries, &mut stack);
218 continue 'iter_timetable;
219 }
220 }
221 } else {
222 None
223 });
224 let total_distance = distances
225 .iter()
226 .filter_map(|d| *d)
227 .map(|d| d.0)
228 .sum::<i32>();
229 if let Some(dur) = time_offset {
230 total_time -= dur;
231 }
232 let velocity = Distance(total_distance) / total_time;
233 if let Some(dur) = time_offset {
234 current_time -= dur;
235 }
236 debug_assert_eq!(distances.len(), stack.len() + 1);
237 while let (Some((timetable_entry_entity, arr_dur, dep_dur)), Some(distance)) =
238 (stack.pop(), distances.pop())
239 {
240 if let Some(distance) = distance {
241 current_time -= distance / velocity;
242 }
243 let departure_estimate = current_time;
244 if let Some(dur) = dep_dur {
245 current_time = current_time - dur;
246 }
247 let arrival_estimate = current_time;
248 if let Some(dur) = arr_dur {
249 current_time = current_time - dur;
250 }
251 if let Ok((_, mut cache)) = entries.get_mut(timetable_entry_entity) {
252 cache.estimate = Some(TimeEstimate {
253 arrival: arrival_estimate,
254 departure: departure_estimate,
255 });
256 }
257 }
258 }
259 clear_estimates(&mut entries, &mut stack);
260 }
261}