1use super::{Plot, PlotError, plot_spec_from, validate_data_lengths, with_plot_str_or_empty};
4use crate::{ItemFlags, StemsFlags, sys};
5
6pub struct StemPlot<'a> {
8 label: &'a str,
9 x_data: &'a [f64],
10 y_data: &'a [f64],
11 y_ref: f64,
12 flags: StemsFlags,
13 item_flags: ItemFlags,
14 offset: i32,
15 stride: i32,
16}
17
18impl<'a> StemPlot<'a> {
19 pub fn new(label: &'a str, x_data: &'a [f64], y_data: &'a [f64]) -> Self {
21 Self {
22 label,
23 x_data,
24 y_data,
25 y_ref: 0.0, flags: StemsFlags::NONE,
27 item_flags: ItemFlags::NONE,
28 offset: 0,
29 stride: std::mem::size_of::<f64>() as i32,
30 }
31 }
32
33 pub fn with_y_ref(mut self, y_ref: f64) -> Self {
36 self.y_ref = y_ref;
37 self
38 }
39
40 pub fn with_flags(mut self, flags: StemsFlags) -> Self {
42 self.flags = flags;
43 self
44 }
45
46 pub fn with_item_flags(mut self, flags: ItemFlags) -> Self {
48 self.item_flags = flags;
49 self
50 }
51
52 pub fn with_offset(mut self, offset: i32) -> Self {
54 self.offset = offset;
55 self
56 }
57
58 pub fn with_stride(mut self, stride: i32) -> Self {
60 self.stride = stride;
61 self
62 }
63
64 pub fn validate(&self) -> Result<(), PlotError> {
66 validate_data_lengths(self.x_data, self.y_data)
67 }
68}
69
70impl<'a> Plot for StemPlot<'a> {
71 fn plot(&self) {
72 if self.validate().is_err() {
73 return;
74 }
75 let Ok(count) = i32::try_from(self.x_data.len()) else {
76 return;
77 };
78
79 with_plot_str_or_empty(self.label, |label_ptr| unsafe {
80 let spec = plot_spec_from(
81 self.flags.bits() | self.item_flags.bits(),
82 self.offset,
83 self.stride,
84 );
85 sys::ImPlot_PlotStems_doublePtrdoublePtr(
86 label_ptr,
87 self.x_data.as_ptr(),
88 self.y_data.as_ptr(),
89 count,
90 self.y_ref,
91 spec,
92 );
93 })
94 }
95
96 fn label(&self) -> &str {
97 self.label
98 }
99}
100
101pub struct SimpleStemPlot<'a> {
103 label: &'a str,
104 values: &'a [f64],
105 y_ref: f64,
106 x_scale: f64,
107 x_start: f64,
108}
109
110impl<'a> SimpleStemPlot<'a> {
111 pub fn new(label: &'a str, values: &'a [f64]) -> Self {
113 Self {
114 label,
115 values,
116 y_ref: 0.0,
117 x_scale: 1.0,
118 x_start: 0.0,
119 }
120 }
121
122 pub fn with_y_ref(mut self, y_ref: f64) -> Self {
124 self.y_ref = y_ref;
125 self
126 }
127
128 pub fn with_x_scale(mut self, scale: f64) -> Self {
130 self.x_scale = scale;
131 self
132 }
133
134 pub fn with_x_start(mut self, start: f64) -> Self {
136 self.x_start = start;
137 self
138 }
139}
140
141impl<'a> Plot for SimpleStemPlot<'a> {
142 fn plot(&self) {
143 if self.values.is_empty() {
144 return;
145 }
146 let Ok(count) = i32::try_from(self.values.len()) else {
147 return;
148 };
149
150 let x_data: Vec<f64> = (0..self.values.len())
152 .map(|i| self.x_start + i as f64 * self.x_scale)
153 .collect();
154
155 with_plot_str_or_empty(self.label, |label_ptr| unsafe {
156 let spec = plot_spec_from(0, 0, std::mem::size_of::<f64>() as i32);
157 sys::ImPlot_PlotStems_doublePtrdoublePtr(
158 label_ptr,
159 x_data.as_ptr(),
160 self.values.as_ptr(),
161 count,
162 self.y_ref,
163 spec,
164 );
165 })
166 }
167
168 fn label(&self) -> &str {
169 self.label
170 }
171}
172
173impl<'ui> crate::PlotUi<'ui> {
175 pub fn stem_plot(&self, label: &str, x_data: &[f64], y_data: &[f64]) -> Result<(), PlotError> {
177 let plot = StemPlot::new(label, x_data, y_data);
178 plot.validate()?;
179 plot.plot();
180 Ok(())
181 }
182
183 pub fn stem_plot_with_ref(
185 &self,
186 label: &str,
187 x_data: &[f64],
188 y_data: &[f64],
189 y_ref: f64,
190 ) -> Result<(), PlotError> {
191 let plot = StemPlot::new(label, x_data, y_data).with_y_ref(y_ref);
192 plot.validate()?;
193 plot.plot();
194 Ok(())
195 }
196
197 pub fn simple_stem_plot(&self, label: &str, values: &[f64]) -> Result<(), PlotError> {
199 if values.is_empty() {
200 return Err(PlotError::EmptyData);
201 }
202 let plot = SimpleStemPlot::new(label, values);
203 plot.plot();
204 Ok(())
205 }
206
207 pub fn simple_stem_plot_with_ref(
209 &self,
210 label: &str,
211 values: &[f64],
212 y_ref: f64,
213 ) -> Result<(), PlotError> {
214 if values.is_empty() {
215 return Err(PlotError::EmptyData);
216 }
217 let plot = SimpleStemPlot::new(label, values).with_y_ref(y_ref);
218 plot.plot();
219 Ok(())
220 }
221}
222
223#[cfg(test)]
224mod tests {
225 use super::*;
226
227 #[test]
228 fn test_stem_plot_creation() {
229 let x_data = [1.0, 2.0, 3.0, 4.0];
230 let y_data = [1.0, 4.0, 2.0, 3.0];
231
232 let plot = StemPlot::new("test", &x_data, &y_data);
233 assert_eq!(plot.label(), "test");
234 assert!(plot.validate().is_ok());
235 }
236
237 #[test]
238 fn test_stem_plot_validation() {
239 let x_data = [1.0, 2.0, 3.0];
240 let y_data = [1.0, 4.0]; let plot = StemPlot::new("test", &x_data, &y_data);
243 assert!(plot.validate().is_err());
244 }
245
246 #[test]
247 fn test_simple_stem_plot() {
248 let values = [1.0, 2.0, 3.0, 4.0];
249 let plot = SimpleStemPlot::new("test", &values);
250 assert_eq!(plot.label(), "test");
251 }
252
253 #[test]
254 fn test_stem_plot_with_ref() {
255 let x_data = [1.0, 2.0, 3.0, 4.0];
256 let y_data = [1.0, 4.0, 2.0, 3.0];
257
258 let plot = StemPlot::new("test", &x_data, &y_data).with_y_ref(1.0);
259 assert_eq!(plot.y_ref, 1.0);
260 assert!(plot.validate().is_ok());
261 }
262}