1use ak_core::Structure;
2use std::sync::mpsc::Sender;
3use std::sync::{Arc, Mutex};
4use std::time::Duration;
5
6use crate::ScalarColorMap;
7
8use super::{
9 AppearanceChannel, BondList, FaceList, RenderStyle, SelectedImageAtom, SupercellSettings,
10 ViewerCommand, ViewerReadiness, ViewerSessionClosed, ViewerSnapshot,
11};
12
13#[derive(Clone)]
14pub struct ViewerSessionHandle {
15 sender: Sender<ViewerCommand>,
16 readiness: Arc<ViewerReadiness>,
17 snapshot: Arc<Mutex<ViewerSnapshot>>,
18}
19
20impl ViewerSessionHandle {
21 pub fn new(sender: Sender<ViewerCommand>) -> Self {
22 Self {
23 sender,
24 readiness: Arc::new(ViewerReadiness::new()),
25 snapshot: Arc::new(Mutex::new(ViewerSnapshot::default())),
26 }
27 }
28
29 pub fn with_readiness(sender: Sender<ViewerCommand>, readiness: Arc<ViewerReadiness>) -> Self {
30 Self {
31 sender,
32 readiness,
33 snapshot: Arc::new(Mutex::new(ViewerSnapshot::default())),
34 }
35 }
36
37 pub fn wait_until_ready(&self, timeout: Option<Duration>) -> bool {
38 self.readiness.wait(timeout)
39 }
40
41 pub fn readiness(&self) -> &Arc<ViewerReadiness> {
42 &self.readiness
43 }
44
45 pub fn snapshot(&self) -> &Arc<Mutex<ViewerSnapshot>> {
46 &self.snapshot
47 }
48
49 pub fn load_trajectory(
50 &self,
51 frames: Vec<Structure>,
52 initial_frame: usize,
53 ) -> Result<(), ViewerSessionClosed> {
54 self.sender
55 .send(ViewerCommand::LoadTrajectory {
56 frames,
57 initial_frame,
58 })
59 .map_err(|_| ViewerSessionClosed)
60 }
61
62 pub fn append_frame(&self, frame: Structure) -> Result<(), ViewerSessionClosed> {
63 self.sender
64 .send(ViewerCommand::AppendFrame { frame })
65 .map_err(|_| ViewerSessionClosed)
66 }
67
68 pub fn set_current_frame(&self, index: usize) -> Result<(), ViewerSessionClosed> {
69 self.sender
70 .send(ViewerCommand::SetCurrentFrame { index })
71 .map_err(|_| ViewerSessionClosed)
72 }
73
74 pub fn set_follow_tail(&self, enabled: bool) -> Result<(), ViewerSessionClosed> {
75 self.sender
76 .send(ViewerCommand::SetFollowTail { enabled })
77 .map_err(|_| ViewerSessionClosed)
78 }
79
80 pub fn set_atom_scalars(
81 &self,
82 name: String,
83 values: Vec<f32>,
84 frame_index: Option<usize>,
85 ) -> Result<(), ViewerSessionClosed> {
86 self.sender
87 .send(ViewerCommand::SetAtomScalars {
88 name,
89 values,
90 frame_index,
91 })
92 .map_err(|_| ViewerSessionClosed)
93 }
94
95 pub fn map_appearance_by_scalar(
96 &self,
97 name: String,
98 channel: AppearanceChannel,
99 palette: Option<ScalarColorMap>,
100 min: Option<f32>,
101 max: Option<f32>,
102 append: bool,
103 ) -> Result<(), ViewerSessionClosed> {
104 self.sender
105 .send(ViewerCommand::MapAppearanceByScalar {
106 name,
107 channel,
108 palette,
109 min,
110 max,
111 append,
112 })
113 .map_err(|_| ViewerSessionClosed)
114 }
115
116 pub fn reset_atom_appearance(
117 &self,
118 channel: Option<AppearanceChannel>,
119 ) -> Result<(), ViewerSessionClosed> {
120 self.sender
121 .send(ViewerCommand::ResetAtomAppearance { channel })
122 .map_err(|_| ViewerSessionClosed)
123 }
124
125 pub fn set_bonds(
126 &self,
127 bonds: BondList,
128 frame_index: Option<usize>,
129 ) -> Result<(), ViewerSessionClosed> {
130 self.sender
131 .send(ViewerCommand::SetBonds { bonds, frame_index })
132 .map_err(|_| ViewerSessionClosed)
133 }
134
135 pub fn add_bonds(
136 &self,
137 bonds: BondList,
138 frame_index: Option<usize>,
139 ) -> Result<(), ViewerSessionClosed> {
140 self.sender
141 .send(ViewerCommand::AddBonds { bonds, frame_index })
142 .map_err(|_| ViewerSessionClosed)
143 }
144
145 pub fn remove_bonds(
146 &self,
147 bonds: BondList,
148 frame_index: Option<usize>,
149 ) -> Result<(), ViewerSessionClosed> {
150 self.sender
151 .send(ViewerCommand::RemoveBonds { bonds, frame_index })
152 .map_err(|_| ViewerSessionClosed)
153 }
154
155 pub fn clear_bonds(&self, frame_index: Option<usize>) -> Result<(), ViewerSessionClosed> {
156 self.sender
157 .send(ViewerCommand::ClearBonds { frame_index })
158 .map_err(|_| ViewerSessionClosed)
159 }
160
161 pub fn set_faces(
162 &self,
163 faces: FaceList,
164 frame_index: Option<usize>,
165 ) -> Result<(), ViewerSessionClosed> {
166 self.sender
167 .send(ViewerCommand::SetFaces { faces, frame_index })
168 .map_err(|_| ViewerSessionClosed)
169 }
170
171 pub fn add_faces(
172 &self,
173 faces: FaceList,
174 frame_index: Option<usize>,
175 ) -> Result<(), ViewerSessionClosed> {
176 self.sender
177 .send(ViewerCommand::AddFaces { faces, frame_index })
178 .map_err(|_| ViewerSessionClosed)
179 }
180
181 pub fn remove_faces(
182 &self,
183 faces: FaceList,
184 frame_index: Option<usize>,
185 ) -> Result<(), ViewerSessionClosed> {
186 self.sender
187 .send(ViewerCommand::RemoveFaces { faces, frame_index })
188 .map_err(|_| ViewerSessionClosed)
189 }
190
191 pub fn clear_faces(&self, frame_index: Option<usize>) -> Result<(), ViewerSessionClosed> {
192 self.sender
193 .send(ViewerCommand::ClearFaces { frame_index })
194 .map_err(|_| ViewerSessionClosed)
195 }
196
197 pub fn set_render_style(
198 &self,
199 style: RenderStyle,
200 selection: Vec<bool>,
201 frame_index: Option<usize>,
202 append: bool,
203 ) -> Result<(), ViewerSessionClosed> {
204 self.sender
205 .send(ViewerCommand::SetRenderStyle {
206 style,
207 selection,
208 frame_index,
209 append,
210 })
211 .map_err(|_| ViewerSessionClosed)
212 }
213
214 pub fn reset_render_style(&self) -> Result<(), ViewerSessionClosed> {
215 self.sender
216 .send(ViewerCommand::ResetRenderStyle)
217 .map_err(|_| ViewerSessionClosed)
218 }
219
220 pub fn replace_selection(
221 &self,
222 selection: Vec<bool>,
223 frame_index: Option<usize>,
224 ) -> Result<(), ViewerSessionClosed> {
225 self.sender
226 .send(ViewerCommand::ReplaceSelection {
227 selection,
228 frame_index,
229 })
230 .map_err(|_| ViewerSessionClosed)
231 }
232
233 pub fn add_selection(
234 &self,
235 selection: Vec<bool>,
236 frame_index: Option<usize>,
237 ) -> Result<(), ViewerSessionClosed> {
238 self.sender
239 .send(ViewerCommand::AddSelection {
240 selection,
241 frame_index,
242 })
243 .map_err(|_| ViewerSessionClosed)
244 }
245
246 pub fn remove_selection(
247 &self,
248 selection: Vec<bool>,
249 frame_index: Option<usize>,
250 ) -> Result<(), ViewerSessionClosed> {
251 self.sender
252 .send(ViewerCommand::RemoveSelection {
253 selection,
254 frame_index,
255 })
256 .map_err(|_| ViewerSessionClosed)
257 }
258
259 pub fn clear_selection(&self, frame_index: Option<usize>) -> Result<(), ViewerSessionClosed> {
260 self.sender
261 .send(ViewerCommand::ClearSelection { frame_index })
262 .map_err(|_| ViewerSessionClosed)
263 }
264
265 pub fn replace_image_selection(
266 &self,
267 selection: Vec<SelectedImageAtom>,
268 frame_index: Option<usize>,
269 ) -> Result<(), ViewerSessionClosed> {
270 self.sender
271 .send(ViewerCommand::ReplaceImageSelection {
272 selection,
273 frame_index,
274 })
275 .map_err(|_| ViewerSessionClosed)
276 }
277
278 pub fn add_image_selection(
279 &self,
280 selection: Vec<SelectedImageAtom>,
281 frame_index: Option<usize>,
282 ) -> Result<(), ViewerSessionClosed> {
283 self.sender
284 .send(ViewerCommand::AddImageSelection {
285 selection,
286 frame_index,
287 })
288 .map_err(|_| ViewerSessionClosed)
289 }
290
291 pub fn remove_image_selection(
292 &self,
293 selection: Vec<SelectedImageAtom>,
294 frame_index: Option<usize>,
295 ) -> Result<(), ViewerSessionClosed> {
296 self.sender
297 .send(ViewerCommand::RemoveImageSelection {
298 selection,
299 frame_index,
300 })
301 .map_err(|_| ViewerSessionClosed)
302 }
303
304 pub fn clear_image_selection(
305 &self,
306 frame_index: Option<usize>,
307 ) -> Result<(), ViewerSessionClosed> {
308 self.sender
309 .send(ViewerCommand::ClearImageSelection { frame_index })
310 .map_err(|_| ViewerSessionClosed)
311 }
312
313 pub fn selected_atoms(&self, frame_index: Option<usize>) -> Vec<usize> {
314 let Ok(snapshot) = self.snapshot.lock() else {
315 return Vec::new();
316 };
317 let target_frame = frame_index.unwrap_or(snapshot.current_frame);
318 snapshot.selection.selected_indices(target_frame)
319 }
320
321 pub fn selected_images(&self, frame_index: Option<usize>) -> Vec<SelectedImageAtom> {
322 let Ok(snapshot) = self.snapshot.lock() else {
323 return Vec::new();
324 };
325 let target_frame = frame_index.unwrap_or(snapshot.current_frame);
326 snapshot.image_selection.selected(target_frame)
327 }
328
329 pub fn set_supercell(&self, repeats: [u32; 3]) -> Result<(), ViewerSessionClosed> {
330 self.sender
331 .send(ViewerCommand::SetSupercell { repeats })
332 .map_err(|_| ViewerSessionClosed)
333 }
334
335 pub fn increment_supercell_axis(&self, axis: usize) -> Result<(), ViewerSessionClosed> {
336 self.sender
337 .send(ViewerCommand::IncrementSupercellAxis { axis })
338 .map_err(|_| ViewerSessionClosed)
339 }
340
341 pub fn decrement_supercell_axis(&self, axis: usize) -> Result<(), ViewerSessionClosed> {
342 self.sender
343 .send(ViewerCommand::DecrementSupercellAxis { axis })
344 .map_err(|_| ViewerSessionClosed)
345 }
346
347 pub fn reset_supercell(&self) -> Result<(), ViewerSessionClosed> {
348 self.sender
349 .send(ViewerCommand::ResetSupercell)
350 .map_err(|_| ViewerSessionClosed)
351 }
352
353 pub fn set_ghost_repeated_images(&self, enabled: bool) -> Result<(), ViewerSessionClosed> {
354 self.sender
355 .send(ViewerCommand::SetGhostRepeatedImages { enabled })
356 .map_err(|_| ViewerSessionClosed)
357 }
358
359 pub fn toggle_ghost_repeated_images(&self) -> Result<(), ViewerSessionClosed> {
360 self.sender
361 .send(ViewerCommand::ToggleGhostRepeatedImages)
362 .map_err(|_| ViewerSessionClosed)
363 }
364
365 pub fn supercell(&self) -> SupercellSettings {
366 let Ok(snapshot) = self.snapshot.lock() else {
367 return SupercellSettings::default();
368 };
369 snapshot.supercell
370 }
371
372 pub fn set_camera_view(
373 &self,
374 focus: Option<[f32; 3]>,
375 radius: Option<f32>,
376 yaw: Option<f32>,
377 pitch: Option<f32>,
378 ) -> Result<(), ViewerSessionClosed> {
379 self.sender
380 .send(ViewerCommand::SetCameraView {
381 focus,
382 radius,
383 yaw,
384 pitch,
385 })
386 .map_err(|_| ViewerSessionClosed)
387 }
388
389 pub fn pan_camera(&self, delta: [f32; 3]) -> Result<(), ViewerSessionClosed> {
390 self.sender
391 .send(ViewerCommand::PanCamera { delta })
392 .map_err(|_| ViewerSessionClosed)
393 }
394
395 pub fn zoom_camera(
396 &self,
397 factor: Option<f32>,
398 delta: Option<f32>,
399 ) -> Result<(), ViewerSessionClosed> {
400 self.sender
401 .send(ViewerCommand::ZoomCamera { factor, delta })
402 .map_err(|_| ViewerSessionClosed)
403 }
404
405 pub fn orbit_camera(
406 &self,
407 yaw_delta: f32,
408 pitch_delta: f32,
409 ) -> Result<(), ViewerSessionClosed> {
410 self.sender
411 .send(ViewerCommand::OrbitCamera {
412 yaw_delta,
413 pitch_delta,
414 })
415 .map_err(|_| ViewerSessionClosed)
416 }
417
418 pub fn frame_all(&self) -> Result<(), ViewerSessionClosed> {
419 self.sender
420 .send(ViewerCommand::FrameAll)
421 .map_err(|_| ViewerSessionClosed)
422 }
423
424 pub fn start_orbit(&self, yaw_rate: f32, pitch_rate: f32) -> Result<(), ViewerSessionClosed> {
425 self.sender
426 .send(ViewerCommand::StartOrbit {
427 yaw_rate,
428 pitch_rate,
429 })
430 .map_err(|_| ViewerSessionClosed)
431 }
432
433 pub fn stop_camera_motion(&self) -> Result<(), ViewerSessionClosed> {
434 self.sender
435 .send(ViewerCommand::StopCameraMotion)
436 .map_err(|_| ViewerSessionClosed)
437 }
438
439 pub fn close(&self) -> Result<(), ViewerSessionClosed> {
440 self.sender
441 .send(ViewerCommand::Close)
442 .map_err(|_| ViewerSessionClosed)
443 }
444}