game/layer/
pandemic.rs

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
15// TODO Disable drawing unzoomed agents... or alternatively, implement this by asking Sim to
16// return this kind of data instead!
17pub 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        // Faster to grab all agent positions than individually map trips to agent positions.
74        // TODO If we ever revive this simulation, need to also grab transit riders here.
75        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        // Many people are probably in the same building. If we're building a heatmap, we
84        // absolutely care about these repeats! If we're just drawing the simple dot map, avoid
85        // drawing repeat circles.
86        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                // Already covered above
91                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            // It's quite silly to produce triangles for the same circle over and over again. ;)
121            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// TODO This should live in sim
151#[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    // If None, just a dot map
163    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}