1use std::collections::HashSet;
2
3use abstutil::prettyprint_usize;
4use geom::{Circle, Distance, Pt2D, Time};
5use map_gui::tools::{make_heatmap, HeatmapOptions};
6use sim::PersonState;
7use widgetry::mapspace::ToggleZoomed;
8use widgetry::{
9 Choice, Color, EventCtx, GfxCtx, Line, Outcome, Panel, Text, TextExt, Toggle, Widget,
10};
11
12use crate::app::App;
13use crate::layer::{header, Layer, LayerOutcome, PANEL_PLACEMENT};
14
15pub struct Pandemic {
18 time: Time,
19 opts: Options,
20 draw: ToggleZoomed,
21 panel: Panel,
22}
23
24impl Layer for Pandemic {
25 fn name(&self) -> Option<&'static str> {
26 Some("pandemic model")
27 }
28 fn event(&mut self, ctx: &mut EventCtx, app: &mut App) -> Option<LayerOutcome> {
29 if app.primary.sim.time() != self.time {
30 let mut new = Pandemic::new(ctx, app, self.opts.clone());
31 new.panel.restore(ctx, &self.panel);
32 *self = new;
33 }
34
35 match self.panel.event(ctx) {
36 Outcome::Clicked(x) => match x.as_ref() {
37 "close" => {
38 return Some(LayerOutcome::Close);
39 }
40 _ => unreachable!(),
41 },
42 _ => {
43 let new_opts = self.options();
44 if self.opts != new_opts {
45 *self = Pandemic::new(ctx, app, new_opts);
46 }
47 }
48 }
49 None
50 }
51 fn draw(&self, g: &mut GfxCtx, _: &App) {
52 self.panel.draw(g);
53 self.draw.draw(g);
54 }
55 fn draw_minimap(&self, g: &mut GfxCtx) {
56 g.redraw(&self.draw.unzoomed);
57 }
58}
59
60impl Pandemic {
61 pub fn new(ctx: &mut EventCtx, app: &App, opts: Options) -> Pandemic {
62 let model = app.primary.sim.get_pandemic_model().unwrap();
63
64 let filter = |p| match opts.state {
65 Seir::Sane => model.is_sane(p),
66 Seir::Exposed => model.is_exposed(p),
67 Seir::Infected => model.is_exposed(p),
68 Seir::Recovered => model.is_recovered(p),
69 Seir::Dead => model.is_dead(p),
70 };
71
72 let mut pts = Vec::new();
73 for a in app.primary.sim.get_unzoomed_agents(&app.primary.map) {
76 if let Some(p) = a.person {
77 if filter(p) {
78 pts.push(a.pos);
79 }
80 }
81 }
82
83 let mut seen_bldgs = HashSet::new();
87 let mut repeat_pts = Vec::new();
88 for person in app.primary.sim.get_all_people() {
89 match person.state {
90 PersonState::Trip(_) => {}
92 PersonState::Inside(b) => {
93 if !filter(person.id) {
94 continue;
95 }
96
97 let pt = app.primary.map.get_b(b).polygon.center();
98 if seen_bldgs.contains(&b) {
99 repeat_pts.push(pt);
100 } else {
101 seen_bldgs.insert(b);
102 pts.push(pt);
103 }
104 }
105 PersonState::OffMap => {}
106 }
107 }
108
109 let mut draw = ToggleZoomed::builder();
110 let legend = if let Some(ref o) = opts.heatmap {
111 pts.extend(repeat_pts);
112 Some(make_heatmap(
113 ctx,
114 &mut draw.unzoomed,
115 app.primary.map.get_bounds(),
116 pts,
117 o,
118 ))
119 } else {
120 let circle = Circle::new(Pt2D::new(0.0, 0.0), Distance::meters(10.0)).to_polygon();
122 for pt in pts {
123 draw.unzoomed
124 .push(Color::RED.alpha(0.8), circle.translate(pt.x(), pt.y()));
125 }
126 None
127 };
128 let controls = make_controls(ctx, app, &opts, legend);
129 Pandemic {
130 time: app.primary.sim.time(),
131 opts,
132 draw: draw.build(ctx),
133 panel: controls,
134 }
135 }
136
137 fn options(&self) -> Options {
138 let heatmap = if self.panel.is_checked("Show heatmap") {
139 Some(HeatmapOptions::from_controls(&self.panel))
140 } else {
141 None
142 };
143 Options {
144 heatmap,
145 state: self.panel.dropdown_value("seir"),
146 }
147 }
148}
149
150#[derive(Clone, Copy, PartialEq, Debug)]
152pub enum Seir {
153 Sane,
154 Exposed,
155 Infected,
156 Recovered,
157 Dead,
158}
159
160#[derive(Clone, PartialEq)]
161pub struct Options {
162 pub heatmap: Option<HeatmapOptions>,
164 pub state: Seir,
165}
166
167fn make_controls(ctx: &mut EventCtx, app: &App, opts: &Options, legend: Option<Widget>) -> Panel {
168 let model = app.primary.sim.get_pandemic_model().unwrap();
169 let pct = 100.0 / (model.count_total() as f64);
170
171 let mut col = vec![
172 header(ctx, "Pandemic model"),
173 Text::from_multiline(vec![
174 Line(format!(
175 "{} Sane ({:.1}%)",
176 prettyprint_usize(model.count_sane()),
177 (model.count_sane() as f64) * pct
178 )),
179 Line(format!(
180 "{} Exposed ({:.1}%)",
181 prettyprint_usize(model.count_exposed()),
182 (model.count_exposed() as f64) * pct
183 )),
184 Line(format!(
185 "{} Infected ({:.1}%)",
186 prettyprint_usize(model.count_infected()),
187 (model.count_infected() as f64) * pct
188 )),
189 Line(format!(
190 "{} Recovered ({:.1}%)",
191 prettyprint_usize(model.count_recovered()),
192 (model.count_recovered() as f64) * pct
193 )),
194 Line(format!(
195 "{} Dead ({:.1}%)",
196 prettyprint_usize(model.count_dead()),
197 (model.count_dead() as f64) * pct
198 )),
199 ])
200 .into_widget(ctx),
201 Widget::row(vec![
202 "Filter:".text_widget(ctx),
203 Widget::dropdown(
204 ctx,
205 "seir",
206 opts.state,
207 vec![
208 Choice::new("sane", Seir::Sane),
209 Choice::new("exposed", Seir::Exposed),
210 Choice::new("infected", Seir::Infected),
211 Choice::new("recovered", Seir::Recovered),
212 Choice::new("dead", Seir::Dead),
213 ],
214 ),
215 ]),
216 ];
217
218 col.push(Toggle::switch(
219 ctx,
220 "Show heatmap",
221 None,
222 opts.heatmap.is_some(),
223 ));
224 if let Some(ref o) = opts.heatmap {
225 col.extend(o.to_controls(ctx, legend.unwrap()));
226 }
227
228 Panel::new_builder(Widget::col(col))
229 .aligned_pair(PANEL_PLACEMENT)
230 .build(ctx)
231}