1use std::collections::{BTreeSet, HashSet};
9
10use crate::{
11 ControlTrafficSignal, DrivingSide, Intersection, IntersectionCluster, IntersectionID, Map,
12 MapConfig, MovementID, RoadID, Stage, StageType, TurnPriority, TurnType,
13};
14use geom::Duration;
15
16mod lagging_green;
17
18pub fn get_possible_policies(map: &Map, id: IntersectionID) -> Vec<(String, ControlTrafficSignal)> {
22 let mut results = Vec::new();
23
24 let i = map.get_i(id);
25
26 if let Some(ts) = four_way_two_stage(map, i) {
29 results.push(("two-stage".to_string(), ts));
30 }
31 if let Some(ts) = three_way(map, i) {
32 results.push(("three-stage".to_string(), ts));
33 }
34 if let Some(ts) = four_way_four_stage(map, i) {
35 results.push(("four-stage".to_string(), ts));
36 }
37 if let Some(ts) = half_signal(i) {
38 results.push(("half signal (2 roads with crosswalk)".to_string(), ts));
39 }
40 if let Some(ts) = degenerate(map, i) {
41 results.push(("degenerate (2 roads)".to_string(), ts));
42 }
43 if let Some(ts) = lagging_green::make_traffic_signal(map, i) {
44 results.push(("lagging green".to_string(), ts));
45 }
46 results.push(("stage per road".to_string(), stage_per_road(i)));
47 results.push(("arbitrary assignment".to_string(), greedy_assignment(i)));
48 results.push((
49 "all walk, then free-for-all yield".to_string(),
50 all_walk_all_yield(i),
51 ));
52
53 for (_, signal) in &mut results {
55 for stage in &mut signal.stages {
56 let crosswalks: Vec<MovementID> = stage
57 .protected_movements
58 .iter()
59 .filter(|id| id.crosswalk)
60 .cloned()
61 .collect();
62 for id in crosswalks {
63 stage.enforce_minimum_crosswalk_time(&i.movements[&id]);
64 }
65 }
66 }
67
68 results.retain(|pair| pair.1.validate(i).is_ok());
69 results
70}
71
72fn new(id: IntersectionID) -> ControlTrafficSignal {
73 ControlTrafficSignal {
74 id,
75 stages: Vec::new(),
76 offset: Duration::ZERO,
77 }
78}
79
80fn greedy_assignment(i: &Intersection) -> ControlTrafficSignal {
81 let mut ts = new(i.id);
82
83 let mut remaining_movements: Vec<MovementID> = i.movements.keys().cloned().collect();
85 let mut current_stage = Stage::new();
86 loop {
87 let add = remaining_movements
88 .iter()
89 .position(|&g| current_stage.could_be_protected(g, i));
90 match add {
91 Some(idx) => {
92 current_stage
93 .protected_movements
94 .insert(remaining_movements.remove(idx));
95 }
96 None => {
97 assert!(!current_stage.protected_movements.is_empty());
98 ts.stages.push(current_stage);
99 current_stage = Stage::new();
100 if remaining_movements.is_empty() {
101 break;
102 }
103 }
104 }
105 }
106
107 expand_all_stages(&mut ts, i);
108
109 ts
110}
111
112fn degenerate(map: &Map, i: &Intersection) -> Option<ControlTrafficSignal> {
113 let roads = i.get_sorted_incoming_roads(map);
114 if roads.len() != 2 {
115 return None;
116 }
117 let (r1, r2) = (roads[0], roads[1]);
118
119 let mut ts = new(i.id);
120 make_stages(
121 &mut ts,
122 &map.config,
123 i,
124 vec![vec![(vec![r1, r2], TurnType::Straight, PROTECTED)]],
125 );
126 Some(ts)
127}
128
129fn half_signal(i: &Intersection) -> Option<ControlTrafficSignal> {
130 if i.roads.len() != 2 {
131 return None;
132 }
133
134 let mut ts = new(i.id);
135 let mut vehicle_stage = Stage::new();
136 let mut ped_stage = Stage::new();
137 for (id, movement) in &i.movements {
138 if id.crosswalk {
139 ped_stage.edit_movement(movement, TurnPriority::Protected);
140 } else {
141 vehicle_stage.edit_movement(movement, TurnPriority::Protected);
142 }
143 }
144 vehicle_stage.stage_type = StageType::Fixed(Duration::minutes(1));
145 ped_stage.stage_type = StageType::Fixed(Duration::seconds(10.0));
146
147 ts.stages = vec![vehicle_stage, ped_stage];
148 Some(ts)
149}
150
151fn three_way(map: &Map, i: &Intersection) -> Option<ControlTrafficSignal> {
152 let roads = i.get_sorted_incoming_roads(map);
153 if roads.len() != 3 {
154 return None;
155 }
156 let mut ts = new(i.id);
157
158 let straight = i
160 .movements
161 .values()
162 .find(|g| g.turn_type == TurnType::Straight)?;
163 let (north, south) = (straight.id.from.road, straight.id.to.road);
164 let east = roads
165 .into_iter()
166 .find(|r| *r != north && *r != south)
167 .unwrap();
168
169 make_stages(
171 &mut ts,
172 &map.config,
173 i,
174 vec![
175 vec![
176 (vec![north, south], TurnType::Straight, PROTECTED),
177 (vec![north, south], TurnType::Right, YIELD),
178 (vec![north, south], TurnType::Left, YIELD),
179 (vec![east], TurnType::Right, YIELD),
180 (vec![east], TurnType::Crosswalk, PROTECTED),
181 (vec![east], TurnType::UnmarkedCrossing, PROTECTED),
183 ],
184 vec![
185 (vec![east], TurnType::Straight, PROTECTED),
186 (vec![east], TurnType::Right, YIELD),
187 (vec![east], TurnType::Left, YIELD),
188 (vec![north, south], TurnType::Right, YIELD),
189 (vec![north, south], TurnType::Crosswalk, PROTECTED),
190 (vec![north, south], TurnType::UnmarkedCrossing, PROTECTED),
191 ],
192 ],
193 );
194
195 Some(ts)
196}
197
198fn four_way_four_stage(map: &Map, i: &Intersection) -> Option<ControlTrafficSignal> {
199 let roads = i.get_sorted_incoming_roads(map);
200 if roads.len() != 4 {
201 return None;
202 }
203
204 let (north, west, south, east) = (roads[0], roads[1], roads[2], roads[3]);
206
207 let mut ts = new(i.id);
210 make_stages(
211 &mut ts,
212 &map.config,
213 i,
214 vec![
215 vec![
216 (vec![north, south], TurnType::Straight, PROTECTED),
217 (vec![north, south], TurnType::Right, YIELD),
218 (vec![east, west], TurnType::Right, YIELD),
219 ],
220 vec![(vec![north, south], TurnType::Left, PROTECTED)],
221 vec![
222 (vec![east, west], TurnType::Straight, PROTECTED),
223 (vec![east, west], TurnType::Right, YIELD),
224 (vec![north, south], TurnType::Right, YIELD),
225 ],
226 vec![(vec![east, west], TurnType::Left, PROTECTED)],
227 ],
228 );
229 Some(ts)
230}
231
232fn four_way_two_stage(map: &Map, i: &Intersection) -> Option<ControlTrafficSignal> {
233 let roads = i.get_sorted_incoming_roads(map);
234 if roads.len() != 4 {
235 return None;
236 }
237
238 let (north, west, south, east) = (roads[0], roads[1], roads[2], roads[3]);
240
241 let mut ts = new(i.id);
243 make_stages(
244 &mut ts,
245 &map.config,
246 i,
247 vec![
248 vec![
249 (vec![north, south], TurnType::Straight, PROTECTED),
250 (vec![north, south], TurnType::Right, YIELD),
251 (vec![north, south], TurnType::Left, YIELD),
252 (vec![east, west], TurnType::Right, YIELD),
253 ],
254 vec![
255 (vec![east, west], TurnType::Straight, PROTECTED),
256 (vec![east, west], TurnType::Right, YIELD),
257 (vec![east, west], TurnType::Left, YIELD),
258 (vec![north, south], TurnType::Right, YIELD),
259 ],
260 ],
261 );
262 Some(ts)
263}
264
265fn all_walk_all_yield(i: &Intersection) -> ControlTrafficSignal {
266 let mut ts = new(i.id);
267
268 let mut all_walk = Stage::new();
269 let mut all_yield = Stage::new();
270
271 for movement in i.movements.values() {
272 if movement.turn_type.pedestrian_crossing() {
273 all_walk.protected_movements.insert(movement.id);
274 } else {
275 all_yield.yield_movements.insert(movement.id);
276 }
277 }
278
279 ts.stages = vec![all_walk, all_yield];
280 ts
281}
282
283fn stage_per_road(i: &Intersection) -> ControlTrafficSignal {
284 let mut ts = new(i.id);
285
286 let sorted_roads = i.roads.clone();
287 for idx in 0..sorted_roads.len() {
288 let r = sorted_roads[idx];
289 let adj1 = *abstutil::wraparound_get(&sorted_roads, (idx as isize) - 1);
290 let adj2 = *abstutil::wraparound_get(&sorted_roads, (idx as isize) + 1);
291
292 let mut stage = Stage::new();
293 for movement in i.movements.values() {
294 if movement.turn_type.pedestrian_crossing() {
295 if movement.id.from.road == adj1 || movement.id.from.road == adj2 {
296 stage.protected_movements.insert(movement.id);
297 }
298 } else if movement.id.from.road == r {
299 stage.yield_movements.insert(movement.id);
300 }
301 }
302 if !stage.yield_movements.is_empty() {
304 ts.stages.push(stage);
305 }
306 }
307 ts
308}
309
310fn expand_all_stages(ts: &mut ControlTrafficSignal, i: &Intersection) {
312 for stage in ts.stages.iter_mut() {
313 for g in i.movements.keys() {
314 if stage.could_be_protected(*g, i) {
315 stage.protected_movements.insert(*g);
316 }
317 }
318 }
319}
320
321const PROTECTED: bool = true;
322const YIELD: bool = false;
323
324fn make_stages(
325 ts: &mut ControlTrafficSignal,
326 map_config: &MapConfig,
327 i: &Intersection,
328 stage_specs: Vec<Vec<(Vec<RoadID>, TurnType, bool)>>,
329) {
330 for specs in stage_specs {
331 let mut stage = Stage::new();
332 let mut explicit_crosswalks = false;
333 for (roads, mut turn_type, protected) in specs.iter() {
334 if map_config.driving_side == DrivingSide::Left {
337 if turn_type == TurnType::Right {
338 turn_type = TurnType::Left;
339 } else if turn_type == TurnType::Left {
340 turn_type = TurnType::Right;
341 }
342 }
343 if turn_type.pedestrian_crossing() {
344 explicit_crosswalks = true;
345 }
346
347 for movement in i.movements.values() {
348 if !roads.contains(&movement.id.from.road) || turn_type != movement.turn_type {
349 continue;
350 }
351
352 if !map_config.turn_on_red
355 && !specs.iter().any(|(other_roads, _, other_protected)| {
356 *other_protected && other_roads.contains(&movement.id.from.road)
357 })
358 {
359 continue;
360 }
361
362 stage.edit_movement(
363 movement,
364 if *protected {
365 TurnPriority::Protected
366 } else {
367 TurnPriority::Yield
368 },
369 );
370 }
371 }
372
373 if !explicit_crosswalks {
377 for movement in i.movements.values() {
381 if movement.turn_type.pedestrian_crossing()
382 && stage.could_be_protected(movement.id, i)
383 {
384 stage.edit_movement(movement, TurnPriority::Protected);
385 }
386 }
387 }
388
389 if stage.protected_movements.is_empty() && stage.yield_movements.is_empty() {
391 continue;
392 }
393
394 ts.stages.push(stage);
395 }
396
397 if ts.stages.len() > 1 {
398 let smallest = ts
401 .stages
402 .iter()
403 .min_by_key(|p| p.protected_movements.len() + p.yield_movements.len())
404 .cloned()
405 .unwrap();
406 if ts.stages.iter().any(|p| {
407 p != &smallest
408 && smallest
409 .protected_movements
410 .is_subset(&p.protected_movements)
411 && smallest.yield_movements.is_subset(&p.yield_movements)
412 }) {
413 ts.stages.retain(|p| p != &smallest);
414 }
415 }
416}
417
418pub fn synchronize(map: &mut Map) {
424 let mut seen = HashSet::new();
425 let mut pairs = Vec::new();
426 for i in map.all_intersections() {
427 if !i.is_traffic_signal() || seen.contains(&i.id) {
428 continue;
429 }
430 if let Some(list) = IntersectionCluster::autodetect(i.id, map) {
431 let list = list.into_iter().collect::<Vec<_>>();
432 if list.len() == 2
433 && map.get_traffic_signal(list[0]).stages.len() == 2
434 && map.get_traffic_signal(list[1]).stages.len() == 2
435 {
436 pairs.push((list[0], list[1]));
437 seen.insert(list[0]);
438 seen.insert(list[1]);
439 }
440 }
441 }
442
443 for (i1, i2) in pairs {
444 let ts1 = map.get_traffic_signal(i1);
445 let ts2 = map.get_traffic_signal(i2);
446 let flip1 = ts1.stages[0].protected_movements.iter().any(|m1| {
447 !m1.crosswalk
448 && ts2.stages[1]
449 .protected_movements
450 .iter()
451 .any(|m2| !m2.crosswalk && (m1.to == m2.from || m1.from == m2.to))
452 });
453 let flip2 = ts1.stages[1].protected_movements.iter().any(|m1| {
454 !m1.crosswalk
455 && ts2.stages[0]
456 .protected_movements
457 .iter()
458 .any(|m2| !m2.crosswalk && (m1.to == m2.from || m1.from == m2.to))
459 });
460 if flip1 || flip2 {
461 info!(
462 "Flipping stage order of {} and {} to synchronize them",
463 i1, i2
464 );
465 map.traffic_signals.get_mut(&i1).unwrap().stages.swap(0, 1);
466 }
467 }
468}