Skip to main content

ak_vis/viewer/
app.rs

1use ak_core::{Structure, Trajectory};
2use bevy::{input_focus::InputDispatchPlugin, prelude::*, ui_widgets::UiWidgetsPlugins};
3use bevy_panorbit_camera::{PanOrbitCameraPlugin, PanOrbitCameraSystemSet};
4use std::sync::{Arc, Mutex, mpsc};
5use std::thread;
6
7use crate::ui::{
8    handle_playback_buttons, set_camera_input_enabled, setup_viewer_ui, sync_inspector_camera,
9    sync_inspector_state, sync_inspector_text, sync_playback_camera, sync_playback_slider_value,
10    sync_playback_state, sync_playback_text, sync_playback_visibility, toggle_hints_visibility,
11};
12use crate::viewer::ViewerConfig;
13use crate::viewer::controls::{
14    keyboard_controls, navigate_frames, screenshot_on_spacebar, screenshot_saving,
15    toggle_ui_visibility, toggle_view,
16};
17use crate::viewer::orientation_widget::{
18    setup_orientation_widget, sync_orientation_widget, sync_orientation_widget_letter_strokes,
19    update_orientation_widget_viewport,
20};
21use crate::viewer::runtime::{configure_shared_app, register_viewer_fonts};
22use crate::viewer::session::{ViewerCommand, ViewerReadiness, ViewerSessionHandle};
23use crate::viewer::systems::{
24    MarqueeSelectionState, handle_atom_clicks, handle_marquee_selection, setup_marquee_overlay,
25    sync_marquee_overlay,
26};
27
28#[derive(Clone, Debug, Default)]
29pub struct ViewerAppOptions {
30    pub canvas_selector: Option<String>,
31    pub fit_canvas_to_parent: bool,
32}
33
34#[derive(Resource, Clone)]
35struct ViewerLifecycle {
36    readiness: Arc<ViewerReadiness>,
37    ready_signaled: bool,
38}
39
40impl Drop for ViewerLifecycle {
41    fn drop(&mut self) {
42        self.readiness.mark_closed();
43    }
44}
45
46pub fn build_app(
47    trajectory: Trajectory,
48    config: ViewerConfig,
49    receiver: Option<mpsc::Receiver<ViewerCommand>>,
50    readiness: Arc<ViewerReadiness>,
51    snapshot: Arc<Mutex<crate::viewer::session::ViewerSnapshot>>,
52) -> App {
53    build_app_with_options(
54        trajectory,
55        config,
56        receiver,
57        readiness,
58        snapshot,
59        ViewerAppOptions::default(),
60    )
61}
62
63pub fn build_app_with_options(
64    trajectory: Trajectory,
65    config: ViewerConfig,
66    receiver: Option<mpsc::Receiver<ViewerCommand>>,
67    readiness: Arc<ViewerReadiness>,
68    snapshot: Arc<Mutex<crate::viewer::session::ViewerSnapshot>>,
69    options: ViewerAppOptions,
70) -> App {
71    let mut app = App::new();
72    let mut plugins = DefaultPlugins.build().disable::<bevy::audio::AudioPlugin>();
73
74    if let Some(primary_window) = primary_window_override(&config, &options) {
75        plugins = plugins.set(WindowPlugin {
76            primary_window: Some(primary_window),
77            ..default()
78        });
79    }
80
81    app.add_plugins((
82        plugins,
83        MeshPickingPlugin,
84        PanOrbitCameraPlugin,
85        UiWidgetsPlugins,
86        InputDispatchPlugin,
87    ));
88    register_viewer_fonts(&mut app);
89    configure_shared_app(&mut app, trajectory, config, receiver, snapshot);
90    app.insert_resource(MarqueeSelectionState::default());
91    app.insert_resource(ViewerLifecycle {
92        readiness,
93        ready_signaled: false,
94    })
95    .add_systems(Startup, (setup_orientation_widget, setup_marquee_overlay))
96    .add_systems(
97        Update,
98        (
99            (
100                handle_marquee_selection,
101                handle_atom_clicks,
102                sync_marquee_overlay,
103            )
104                .chain(),
105            toggle_view,
106            keyboard_controls,
107            screenshot_on_spacebar,
108            screenshot_saving,
109            navigate_frames,
110            sync_orientation_widget,
111            sync_orientation_widget_letter_strokes,
112            update_orientation_widget_viewport,
113            signal_viewer_ready,
114        ),
115    );
116
117    if app.world().resource::<ViewerConfig>().render.show_ui {
118        app.add_systems(Startup, setup_viewer_ui);
119        app.add_systems(
120            Update,
121            (
122                toggle_ui_visibility,
123                toggle_hints_visibility,
124                handle_playback_buttons,
125                sync_playback_slider_value,
126            ),
127        );
128        app.add_systems(
129            PostUpdate,
130            set_camera_input_enabled.before(PanOrbitCameraSystemSet),
131        );
132        app.add_systems(
133            PostUpdate,
134            (
135                sync_inspector_camera,
136                sync_inspector_state,
137                sync_inspector_text,
138                sync_playback_camera,
139                sync_playback_state,
140                sync_playback_visibility,
141                sync_playback_text,
142            )
143                .chain(),
144        );
145    }
146
147    app
148}
149
150fn primary_window_override(config: &ViewerConfig, options: &ViewerAppOptions) -> Option<Window> {
151    let needs_size_override = config.window_width.is_some() || config.window_height.is_some();
152    let needs_canvas_override = options.canvas_selector.is_some() || options.fit_canvas_to_parent;
153
154    if !needs_size_override && !needs_canvas_override {
155        return None;
156    }
157
158    let mut window = Window::default();
159
160    if needs_size_override {
161        let width = config.window_width.unwrap_or(750);
162        let height = config.window_height.unwrap_or(750);
163        window.resolution = (width, height).into();
164    }
165
166    window.canvas.clone_from(&options.canvas_selector);
167    window.fit_canvas_to_parent = options.fit_canvas_to_parent;
168
169    Some(window)
170}
171
172fn run_app(
173    trajectory: Trajectory,
174    config: ViewerConfig,
175    receiver: Option<mpsc::Receiver<ViewerCommand>>,
176    readiness: Arc<ViewerReadiness>,
177    snapshot: Arc<Mutex<crate::viewer::session::ViewerSnapshot>>,
178) {
179    let mut app = build_app(trajectory, config, receiver, readiness, snapshot);
180    app.run();
181}
182
183fn signal_viewer_ready(
184    mut lifecycle: ResMut<'_, ViewerLifecycle>,
185    windows: Query<'_, '_, Entity, With<Window>>,
186) {
187    if lifecycle.ready_signaled || windows.is_empty() {
188        return;
189    }
190
191    lifecycle.ready_signaled = true;
192    lifecycle.readiness.mark_ready();
193}
194
195pub fn run_prepared(
196    trajectory: Trajectory,
197    config: ViewerConfig,
198    receiver: mpsc::Receiver<ViewerCommand>,
199    readiness: Arc<ViewerReadiness>,
200    snapshot: Arc<Mutex<crate::viewer::session::ViewerSnapshot>>,
201) {
202    run_app(trajectory, config, Some(receiver), readiness, snapshot);
203}
204
205pub fn launch(trajectory: Trajectory, config: ViewerConfig) -> ViewerSessionHandle {
206    let (sender, receiver) = mpsc::channel();
207    let handle = ViewerSessionHandle::new(sender);
208    let launch_readiness = handle.readiness().clone();
209    let launch_snapshot = handle.snapshot().clone();
210
211    thread::Builder::new()
212        .name("ak-viewer-session".to_string())
213        .spawn(move || {
214            run_app(
215                trajectory,
216                config,
217                Some(receiver),
218                launch_readiness,
219                launch_snapshot,
220            );
221        })
222        .expect("failed to launch viewer session thread");
223
224    handle
225}
226
227pub fn run_with_session<F>(trajectory: Trajectory, config: ViewerConfig, driver: F)
228where
229    F: FnOnce(ViewerSessionHandle) + Send + 'static,
230{
231    let (sender, receiver) = mpsc::channel();
232    let handle = ViewerSessionHandle::new(sender);
233    let readiness = handle.readiness().clone();
234    let snapshot = handle.snapshot().clone();
235
236    thread::Builder::new()
237        .name("ak-viewer-session-driver".to_string())
238        .spawn(move || {
239            driver(handle);
240        })
241        .expect("failed to launch viewer session driver thread");
242
243    run_app(trajectory, config, Some(receiver), readiness, snapshot);
244}
245
246pub fn run(trajectory: Trajectory, config: ViewerConfig) {
247    run_app(
248        trajectory,
249        config,
250        None,
251        Arc::new(ViewerReadiness::new()),
252        Arc::new(Mutex::new(crate::viewer::session::ViewerSnapshot::default())),
253    );
254}
255
256pub fn run_default(trajectory: Trajectory) {
257    let config = ViewerConfig::default();
258    run(trajectory, config)
259}
260
261pub fn run_structure(structure: Structure, config: ViewerConfig) {
262    let trajectory = Trajectory::new(vec![structure]);
263    run(trajectory, config)
264}
265
266pub fn run_structure_default(structure: Structure) {
267    let trajectory = Trajectory::new(vec![structure]);
268    run_default(trajectory);
269}
270
271#[cfg(test)]
272mod tests {
273    use super::{ViewerAppOptions, primary_window_override};
274    use crate::viewer::ViewerConfig;
275
276    #[test]
277    fn default_app_options_do_not_override_primary_window() {
278        let config = ViewerConfig::default();
279        let options = ViewerAppOptions::default();
280
281        assert!(primary_window_override(&config, &options).is_none());
282    }
283
284    #[test]
285    fn canvas_selector_overrides_primary_window_canvas() {
286        let config = ViewerConfig::default();
287        let options = ViewerAppOptions {
288            canvas_selector: Some("#viewer-canvas".to_string()),
289            fit_canvas_to_parent: true,
290        };
291
292        let window = primary_window_override(&config, &options).unwrap();
293
294        assert_eq!(window.canvas.as_deref(), Some("#viewer-canvas"));
295        assert!(window.fit_canvas_to_parent);
296    }
297
298    #[test]
299    fn width_and_height_config_still_override_primary_window_size() {
300        let config = ViewerConfig {
301            window_width: Some(640),
302            window_height: Some(480),
303            ..Default::default()
304        };
305        let options = ViewerAppOptions::default();
306
307        let window = primary_window_override(&config, &options).unwrap();
308
309        assert_eq!(window.resolution.width(), 640.0);
310        assert_eq!(window.resolution.height(), 480.0);
311        assert_eq!(window.canvas, None);
312        assert!(!window.fit_canvas_to_parent);
313    }
314}