1use ak_core::Trajectory;
2use bevy::prelude::{Resource, Vec3};
3use std::collections::HashMap;
4
5use super::snapshot::{clamp_frame, scaled_cell_translation, supercell_offsets};
6use super::{
7 AtomAppearanceRule, BondFrames, DisplayAtom, FaceFrames, ImageSelectionFrames, RenderStyleRule,
8 SelectedImageAtom, SelectionFrames, SupercellSettings,
9};
10
11#[derive(Clone, Copy, Debug, PartialEq)]
12pub struct OrbitMotion {
13 pub yaw_rate: f32,
14 pub pitch_rate: f32,
15}
16
17#[derive(Clone, Copy, Debug, PartialEq)]
18pub struct CameraView {
19 pub focus: Vec3,
20 pub radius: f32,
21 pub yaw: f32,
22 pub pitch: f32,
23}
24
25#[derive(Resource, Clone, Debug, PartialEq)]
26pub struct CameraState {
27 pub focus: Vec3,
28 pub radius: f32,
29 pub yaw: f32,
30 pub pitch: f32,
31 pub needs_apply: bool,
32 pub motion: Option<OrbitMotion>,
33}
34
35#[derive(Resource)]
36pub struct ViewerState {
37 pub traj: Trajectory,
38 pub current: usize,
39 pub follow_tail: bool,
40 pub atom_scalars: HashMap<String, Vec<Option<Vec<f32>>>>,
41 pub atom_appearance_rules: Vec<AtomAppearanceRule>,
42 pub bonds: BondFrames,
43 pub faces: FaceFrames,
44 pub render_style_rules: Vec<RenderStyleRule>,
45 pub selection: SelectionFrames,
46 pub image_selection: ImageSelectionFrames,
47 pub supercell: SupercellSettings,
48 pub needs_render: bool,
49 pub needs_camera_reset: bool,
50}
51
52impl ViewerState {
53 pub fn new(traj: Trajectory, initial_frame: usize) -> Self {
54 let frame_count = traj.len();
55 let current = clamp_frame(initial_frame, traj.len());
56 let selection = SelectionFrames::new(&traj);
57 let image_selection = ImageSelectionFrames::new(frame_count);
58 Self {
59 traj,
60 current,
61 follow_tail: false,
62 atom_scalars: HashMap::new(),
63 atom_appearance_rules: Vec::new(),
64 bonds: BondFrames::new(frame_count),
65 faces: FaceFrames::new(frame_count),
66 render_style_rules: Vec::new(),
67 selection,
68 image_selection,
69 supercell: SupercellSettings::default(),
70 needs_render: true,
71 needs_camera_reset: true,
72 }
73 }
74
75 pub fn trajectory_len(&self) -> usize {
76 self.traj.len()
77 }
78
79 pub fn has_frames(&self) -> bool {
80 !self.traj.is_empty()
81 }
82
83 pub fn step_frame(&mut self, delta: isize) -> bool {
84 if self.traj.is_empty() {
85 return false;
86 }
87
88 let next = self.current as isize + delta;
89 if next < 0 || next >= self.traj.len() as isize {
90 return false;
91 }
92
93 let next = next as usize;
94 if next == self.current {
95 return false;
96 }
97
98 self.set_current_frame_index(next)
99 }
100
101 pub fn set_current_frame_index(&mut self, index: usize) -> bool {
102 if index >= self.traj.len() || index == self.current {
103 return false;
104 }
105
106 let previous = self.current;
107 self.copy_selection_between_compatible_frames(previous, index);
108 self.current = index;
109 self.needs_render = true;
110 self.needs_camera_reset = self.cell_changed(previous, self.current);
111 true
112 }
113
114 pub fn is_at_last_frame(&self) -> bool {
115 self.has_frames() && self.current + 1 >= self.traj.len()
116 }
117
118 pub fn selected_atoms(&self, frame_index: usize) -> Vec<usize> {
119 self.selection.selected_indices(frame_index)
120 }
121
122 pub fn selected_images(&self, frame_index: usize) -> Vec<SelectedImageAtom> {
123 let selected = self.image_selection.selected(frame_index);
124 if selected.is_empty() {
125 self.selection.selected_main_images(frame_index)
126 } else {
127 selected
128 }
129 }
130
131 pub fn current_selection(&self) -> &[bool] {
132 self.selection.get(self.current).unwrap_or(&[])
133 }
134
135 pub fn current_image_selection(&self) -> &[SelectedImageAtom] {
136 self.image_selection.get(self.current).unwrap_or(&[])
137 }
138
139 pub fn current_supercell(&self) -> SupercellSettings {
140 self.supercell
141 }
142
143 pub fn replace_atom_selection(&mut self, atom: SelectedImageAtom) -> bool {
144 let changed = self.image_selection.replace_single(self.current, atom);
145 if changed {
146 self.sync_main_selection_from_images(self.current);
147 }
148 changed
149 }
150
151 pub fn toggle_atom_selection(&mut self, atom: SelectedImageAtom) -> bool {
152 let changed = self.image_selection.toggle(self.current, atom);
153 if changed {
154 self.sync_main_selection_from_images(self.current);
155 }
156 changed
157 }
158
159 pub fn display_atoms(&self, frame_index: usize) -> Vec<DisplayAtom> {
160 if frame_index >= self.traj.len() {
161 return Vec::new();
162 }
163
164 let view = self.traj.view(frame_index);
165 let offsets = supercell_offsets(self.supercell.repeats);
166 let mut atoms = Vec::with_capacity(view.positions.len() * offsets.len());
167 for image_offset in offsets {
168 let shift = scaled_cell_translation(view.cell, image_offset);
169 for (atom_index, position) in view.positions.iter().copied().enumerate() {
170 atoms.push(DisplayAtom {
171 identity: SelectedImageAtom {
172 atom_index,
173 image_offset,
174 },
175 position: [
176 position[0] + shift[0],
177 position[1] + shift[1],
178 position[2] + shift[2],
179 ],
180 is_main_cell: image_offset == [0, 0, 0],
181 });
182 }
183 }
184 atoms
185 }
186
187 pub(super) fn cell_changed(&self, old_index: usize, new_index: usize) -> bool {
188 self.traj.view(old_index).cell != self.traj.view(new_index).cell
189 }
190
191 fn frames_have_matching_atom_identity(&self, source_index: usize, target_index: usize) -> bool {
192 if source_index >= self.traj.len() || target_index >= self.traj.len() {
193 return false;
194 }
195
196 let source = self.traj.view(source_index);
197 let target = self.traj.view(target_index);
198 source.positions.len() == target.positions.len() && source.numbers == target.numbers
199 }
200
201 pub(super) fn copy_selection_between_compatible_frames(
202 &mut self,
203 source_index: usize,
204 target_index: usize,
205 ) {
206 if !self.frames_have_matching_atom_identity(source_index, target_index) {
207 return;
208 }
209
210 let selection = self.selected_images(source_index);
211 if !self.validate_image_selection(target_index, &selection) {
212 return;
213 }
214
215 self.image_selection.replace(target_index, selection);
216 self.sync_main_selection_from_images(target_index);
217 }
218
219 pub(super) fn validate_scalar_values(&self, frame_index: usize, values: &[f32]) -> bool {
220 frame_index < self.traj.len() && self.traj.view(frame_index).positions.len() == values.len()
221 }
222
223 pub(super) fn validate_selection(&self, frame_index: usize, selection: &[bool]) -> bool {
224 frame_index < self.traj.len()
225 && self.traj.view(frame_index).positions.len() == selection.len()
226 }
227
228 pub(super) fn validate_image_selection(
229 &self,
230 frame_index: usize,
231 selection: &[SelectedImageAtom],
232 ) -> bool {
233 if frame_index >= self.traj.len() {
234 return false;
235 }
236 let atom_count = self.traj.view(frame_index).positions.len();
237 let allowed_offsets = supercell_offsets(self.supercell.repeats);
238 selection.iter().all(|atom| {
239 atom.atom_index < atom_count && allowed_offsets.contains(&atom.image_offset)
240 })
241 }
242
243 pub(super) fn store_scalars(&mut self, name: String, frame_index: usize, values: Vec<f32>) {
244 let frame_count = self.traj.len();
245 let entries = self
246 .atom_scalars
247 .entry(name)
248 .or_insert_with(|| vec![None; frame_count]);
249 if entries.len() < frame_count {
250 entries.resize(frame_count, None);
251 }
252 entries[frame_index] = Some(values);
253 }
254
255 pub(super) fn resize_scalar_storage(&mut self, frame_count: usize) {
256 for values in self.atom_scalars.values_mut() {
257 values.resize(frame_count, None);
258 }
259 }
260
261 pub(super) fn sync_main_selection_from_images(&mut self, frame_index: usize) {
262 let Some(selected) = self.image_selection.get(frame_index) else {
263 return;
264 };
265 let atom_count = self.traj.view(frame_index).positions.len();
266 self.selection.replace(
267 frame_index,
268 SelectionFrames::mask_from_images(atom_count, selected),
269 );
270 self.selection.set_order(
271 frame_index,
272 SelectionFrames::ordered_atoms_from_images(selected),
273 );
274 }
275}
276
277#[derive(Default)]
278pub struct CommandOutcome {
279 pub should_close: bool,
280}