1use rand::prelude::SliceRandom;
2use rand::{Rng, SeedableRng};
3use rand_xorshift::XorShiftRng;
4
5use abstutil::{prettyprint_usize, Timer};
6use geom::{Distance, Duration, FindClosest};
7use map_model::{AmenityType, BuildingID, Map};
8use synthpop::{IndividTrip, Scenario, ScenarioModifier, TripEndpoint, TripMode, TripPurpose};
9
10pub fn run(
11 input_scenario: String,
12 should_add_return_trips: bool,
13 should_add_lunch_trips: bool,
14 modifiers: Vec<ScenarioModifier>,
15 should_delete_cancelled_trips: bool,
16 rng_seed: u64,
17) {
18 let mut rng = XorShiftRng::seed_from_u64(rng_seed);
19 let mut timer = Timer::new("augment scenario");
20
21 let mut scenario: Scenario = abstio::must_read_object(input_scenario, &mut timer);
22 let map = Map::load_synchronously(scenario.map_name.path(), &mut timer);
23
24 if should_add_return_trips {
25 add_return_trips(&mut scenario, &mut rng);
26 }
27 if should_add_lunch_trips {
28 add_lunch_trips(&mut scenario, &map, &mut rng, &mut timer);
29 }
30
31 for m in modifiers {
32 scenario = m.apply(&map, scenario, &mut rng);
33 }
34
35 if should_delete_cancelled_trips {
36 scenario = delete_cancelled_trips(scenario);
37 }
38
39 scenario.save();
40}
41
42fn add_return_trips(scenario: &mut Scenario, rng: &mut XorShiftRng) {
43 let mut cnt = 0;
44 for person in &mut scenario.people {
45 if person.trips.len() != 1 {
46 continue;
47 }
48
49 let depart =
51 person.trips[0].depart + rand_duration(rng, Duration::hours(4), Duration::hours(12));
52 person.trips.push(IndividTrip::new(
53 depart,
54 TripPurpose::Home,
55 person.trips[0].destination,
56 person.trips[0].origin,
57 person.trips[0].mode,
58 ));
59 cnt += 1;
60 }
61 info!("Added return trips to {} people", prettyprint_usize(cnt));
62}
63
64fn rand_duration(rng: &mut XorShiftRng, low: Duration, high: Duration) -> Duration {
65 Duration::seconds(rng.gen_range(low.inner_seconds()..high.inner_seconds()))
66}
67
68fn add_lunch_trips(scenario: &mut Scenario, map: &Map, rng: &mut XorShiftRng, timer: &mut Timer) {
69 timer.start("index lunch spots");
71 let mut closest_spots: FindClosest<BuildingID> = FindClosest::new();
72 for b in map.all_buildings() {
73 if b.amenities
74 .iter()
75 .any(|a| AmenityType::categorize(&a.amenity_type) == Some(AmenityType::Food))
76 {
77 closest_spots.add_polygon(b.id, &b.polygon);
78 }
79 }
80 timer.stop("index lunch spots");
81
82 timer.start_iter("add lunch trips", scenario.people.len());
83 let mut cnt = 0;
84 for person in &mut scenario.people {
85 timer.next();
86 let num_trips = person.trips.len();
87 if num_trips <= 1 || person.trips[num_trips - 1].destination != person.trips[0].origin {
89 continue;
90 }
91
92 let work = match person.trips[num_trips - 2].destination {
93 TripEndpoint::Building(b) => b,
94 _ => continue,
95 };
96 let has_bike = person.trips[num_trips - 2].mode == TripMode::Bike;
97 let (restaurant, mode) =
98 if let Some(pair) = pick_lunch_spot(work, has_bike, &closest_spots, map, rng) {
99 pair
100 } else {
101 continue;
102 };
103
104 let t1 = person.trips[num_trips - 2].depart;
106 let t2 = person.trips[num_trips - 1].depart;
107 let depart = t1 + (t2 - t1) / 2.0;
108 let return_home = person.trips.pop().unwrap();
109 person.trips.push(IndividTrip::new(
110 depart,
111 TripPurpose::Meal,
112 TripEndpoint::Building(work),
113 TripEndpoint::Building(restaurant),
114 mode,
115 ));
116 person.trips.push(IndividTrip::new(
117 depart + Duration::minutes(30),
118 TripPurpose::Work,
119 TripEndpoint::Building(restaurant),
120 TripEndpoint::Building(work),
121 mode,
122 ));
123 person.trips.push(return_home);
124 cnt += 1;
125 }
126 info!("Added lunch trips to {} people", prettyprint_usize(cnt));
127}
128
129fn pick_lunch_spot(
130 work: BuildingID,
131 has_bike: bool,
132 closest_spots: &FindClosest<BuildingID>,
133 map: &Map,
134 rng: &mut XorShiftRng,
135) -> Option<(BuildingID, TripMode)> {
136 let choices =
139 closest_spots.all_close_pts(map.get_b(work).polygon.center(), Distance::miles(10.0));
140 let (b, _, dist) = choices
141 .choose_weighted(rng, |(_, _, dist)| dist.inner_meters())
142 .ok()?;
143 let mode = if *dist <= Distance::miles(1.0) {
145 TripMode::Walk
146 } else if has_bike {
147 TripMode::Bike
148 } else {
149 TripMode::Drive
150 };
151 Some((*b, mode))
152}
153
154fn delete_cancelled_trips(mut scenario: Scenario) -> Scenario {
155 for person in &mut scenario.people {
156 person.trips.retain(|trip| !trip.cancelled);
157 }
158 scenario.remove_weird_schedules(false)
159}