1use ak_core::PERIODIC_TABLE;
2use bevy::prelude::*;
3
4use crate::components::{
5 InspectorHintsContainer, InspectorHintsToggle, InspectorMeasurementBody,
6 InspectorMeasurementSection, InspectorPanelRoot, InspectorSelectionBody,
7 InspectorSelectionSection, MainSceneCamera,
8};
9use crate::viewer::ViewerState;
10
11#[derive(Clone, Debug, PartialEq)]
12pub enum MeasurementStatus {
13 Empty,
14 NeedOneMoreAtom,
15 Distance { atoms: [usize; 2], angstrom: f64 },
16 Angle { atoms: [usize; 3], degrees: f64 },
17 UnsupportedCount { count: usize },
18}
19
20#[derive(Resource, Clone, Debug, PartialEq, Default)]
21pub struct InspectorState {
22 pub selected_atoms: Vec<SelectedAtomSummary>,
23 pub measurement: MeasurementStatus,
24 pub hints_expanded: bool,
25}
26
27#[derive(Clone, Debug, PartialEq, Eq)]
28pub struct SelectedAtomSummary {
29 pub index: usize,
30 pub symbol: String,
31 pub image_offset: [i32; 3],
32}
33
34impl Default for MeasurementStatus {
35 fn default() -> Self {
36 Self::Empty
37 }
38}
39
40pub fn sync_inspector_state(viewer: Res<ViewerState>, mut inspector: ResMut<InspectorState>) {
41 let hints_expanded = inspector.hints_expanded;
42 *inspector = derive_inspector_state(viewer.as_ref());
43 inspector.hints_expanded = hints_expanded;
44}
45
46pub fn sync_inspector_camera(
47 mut commands: Commands,
48 main_camera: Query<Entity, With<MainSceneCamera>>,
49 roots: Query<(Entity, Option<&UiTargetCamera>), With<InspectorPanelRoot>>,
50) {
51 let Ok(main_camera) = main_camera.single() else {
52 return;
53 };
54
55 for (entity, target_camera) in roots.iter() {
56 if target_camera.is_none() {
57 commands.entity(entity).insert(UiTargetCamera(main_camera));
58 }
59 }
60}
61
62pub fn sync_inspector_text(
63 inspector: Res<InspectorState>,
64 mut sections: ParamSet<(
65 Query<&mut Node, With<InspectorSelectionSection>>,
66 Query<&mut Node, With<InspectorMeasurementSection>>,
67 )>,
68 mut bodies: ParamSet<(
69 Query<&mut Text, With<InspectorSelectionBody>>,
70 Query<&mut Text, With<InspectorMeasurementBody>>,
71 Query<&mut Text, With<InspectorHintsToggle>>,
72 )>,
73) {
74 if !inspector.is_changed() {
75 return;
76 }
77
78 if let Ok(mut node) = sections.p0().single_mut() {
79 node.display = if inspector.selected_atoms.is_empty() {
80 Display::None
81 } else {
82 Display::Flex
83 };
84 }
85 if let Ok(mut node) = sections.p1().single_mut() {
86 node.display = if shows_measurement(inspector.as_ref()) {
87 Display::Flex
88 } else {
89 Display::None
90 };
91 }
92 if let Ok(mut text) = bodies.p0().single_mut() {
93 *text = Text::new(selection_text(inspector.as_ref()));
94 }
95 if let Ok(mut text) = bodies.p1().single_mut() {
96 *text = Text::new(measurement_text(inspector.as_ref()));
97 }
98 if let Ok(mut text) = bodies.p2().single_mut() {
99 *text = Text::new(hints_toggle_text(inspector.as_ref()));
100 }
101}
102
103pub fn derive_inspector_state(viewer: &ViewerState) -> InspectorState {
104 let selected_indices = viewer.selected_images(viewer.current);
105 let frame = viewer.traj.view(viewer.current);
106 let selected_atoms = selected_indices
107 .iter()
108 .map(|selected| SelectedAtomSummary {
109 index: selected.atom_index,
110 symbol: frame
111 .numbers
112 .get(selected.atom_index)
113 .map(|&number| PERIODIC_TABLE.get(number).symbol.clone())
114 .unwrap_or_else(|| "?".to_string()),
115 image_offset: selected.image_offset,
116 })
117 .collect();
118
119 InspectorState {
120 measurement: measurement_for_selection(
121 &viewer.display_atoms(viewer.current),
122 &selected_indices,
123 ),
124 selected_atoms,
125 hints_expanded: false,
126 }
127}
128
129fn measurement_for_selection(
130 positions: &[crate::viewer::DisplayAtom],
131 selected_indices: &[crate::viewer::SelectedImageAtom],
132) -> MeasurementStatus {
133 match selected_indices {
134 [] => MeasurementStatus::Empty,
135 [_] => MeasurementStatus::NeedOneMoreAtom,
136 [a, b] => MeasurementStatus::Distance {
137 atoms: [a.atom_index, b.atom_index],
138 angstrom: distance_between(positions, *a, *b).unwrap_or(0.0),
139 },
140 [a, b, c] => MeasurementStatus::Angle {
141 atoms: [a.atom_index, b.atom_index, c.atom_index],
142 degrees: angle_between(positions, *a, *b, *c).unwrap_or(0.0),
143 },
144 many => MeasurementStatus::UnsupportedCount { count: many.len() },
145 }
146}
147
148fn distance_between(
149 positions: &[crate::viewer::DisplayAtom],
150 a: crate::viewer::SelectedImageAtom,
151 b: crate::viewer::SelectedImageAtom,
152) -> Option<f64> {
153 let pa = positions.iter().find(|atom| atom.identity == a)?.position;
154 let pb = positions.iter().find(|atom| atom.identity == b)?.position;
155 let dx = pa[0] - pb[0];
156 let dy = pa[1] - pb[1];
157 let dz = pa[2] - pb[2];
158 Some((dx * dx + dy * dy + dz * dz).sqrt())
159}
160
161fn angle_between(
162 positions: &[crate::viewer::DisplayAtom],
163 a: crate::viewer::SelectedImageAtom,
164 b: crate::viewer::SelectedImageAtom,
165 c: crate::viewer::SelectedImageAtom,
166) -> Option<f64> {
167 let pa = positions.iter().find(|atom| atom.identity == a)?.position;
168 let pb = positions.iter().find(|atom| atom.identity == b)?.position;
169 let pc = positions.iter().find(|atom| atom.identity == c)?.position;
170
171 let ba = [pa[0] - pb[0], pa[1] - pb[1], pa[2] - pb[2]];
172 let bc = [pc[0] - pb[0], pc[1] - pb[1], pc[2] - pb[2]];
173 let ba_norm = vector_norm(ba);
174 let bc_norm = vector_norm(bc);
175 if ba_norm <= f64::EPSILON || bc_norm <= f64::EPSILON {
176 return None;
177 }
178
179 let cosine = (dot(ba, bc) / (ba_norm * bc_norm)).clamp(-1.0, 1.0);
180 Some(cosine.acos().to_degrees())
181}
182
183fn dot(a: [f64; 3], b: [f64; 3]) -> f64 {
184 a[0] * b[0] + a[1] * b[1] + a[2] * b[2]
185}
186
187fn vector_norm(v: [f64; 3]) -> f64 {
188 dot(v, v).sqrt()
189}
190
191fn selection_text(inspector: &InspectorState) -> String {
192 let count = inspector.selected_atoms.len();
193 if count == 0 {
194 return "Selected atoms: 0\nClick an atom to inspect it.".to_string();
195 }
196
197 let summary = if count <= 3 {
198 inspector
199 .selected_atoms
200 .iter()
201 .map(|atom| {
202 if atom.image_offset == [0, 0, 0] {
203 format!("#{} {}", atom.index, atom.symbol)
204 } else {
205 format!(
206 "#{} {} [{}, {}, {}]",
207 atom.index,
208 atom.symbol,
209 atom.image_offset[0],
210 atom.image_offset[1],
211 atom.image_offset[2]
212 )
213 }
214 })
215 .collect::<Vec<_>>()
216 .join(", ")
217 } else {
218 inspector
219 .selected_atoms
220 .iter()
221 .map(|atom| {
222 if atom.image_offset == [0, 0, 0] {
223 format!("#{}", atom.index)
224 } else {
225 format!(
226 "#{} [{}, {}, {}]",
227 atom.index,
228 atom.image_offset[0],
229 atom.image_offset[1],
230 atom.image_offset[2]
231 )
232 }
233 })
234 .collect::<Vec<_>>()
235 .join(", ")
236 };
237
238 format!("Selected atoms: {count}\n{summary}")
239}
240
241fn measurement_text(inspector: &InspectorState) -> String {
242 match &inspector.measurement {
243 MeasurementStatus::Empty => {
244 "No measurement yet.\nSelect 2 atoms for a distance or 3 atoms for an angle."
245 .to_string()
246 }
247 MeasurementStatus::NeedOneMoreAtom => {
248 "Select one more atom to compute a distance.".to_string()
249 }
250 MeasurementStatus::Distance { atoms, angstrom } => {
251 format!("Distance (#{} -> #{})\n{angstrom:.3} A", atoms[0], atoms[1])
252 }
253 MeasurementStatus::Angle { atoms, degrees } => {
254 format!(
255 "Angle (#{}-#{}-#{})\n{degrees:.2} deg",
256 atoms[0], atoms[1], atoms[2]
257 )
258 }
259 MeasurementStatus::UnsupportedCount { count } => {
260 format!("Selected atoms: {count}\nMeasurements are shown for 2 or 3 atoms.")
261 }
262 }
263}
264
265fn hints_toggle_text(inspector: &InspectorState) -> String {
266 let _ = inspector;
267 ")".to_string()
268}
269
270fn shows_measurement(inspector: &InspectorState) -> bool {
271 matches!(
272 inspector.measurement,
273 MeasurementStatus::Distance { .. } | MeasurementStatus::Angle { .. }
274 )
275}
276
277pub fn toggle_hints_visibility(
278 keys: Res<ButtonInput<KeyCode>>,
279 mut inspector: ResMut<InspectorState>,
280 mut hint_bodies: Query<&mut Node, With<InspectorHintsContainer>>,
281) {
282 if !keys.just_pressed(KeyCode::KeyH) {
283 return;
284 }
285
286 inspector.hints_expanded = !inspector.hints_expanded;
287 let display = if inspector.hints_expanded {
288 Display::Flex
289 } else {
290 Display::None
291 };
292 for mut node in &mut hint_bodies {
293 node.display = display;
294 }
295}