use crate::core::Result;
use crate::plots::traits::{PlotArea, PlotCompute, PlotConfig, PlotData, PlotRender};
use crate::render::skia::SkiaRenderer;
use crate::render::{Color, LineStyle, MarkerStyle, Theme};
#[derive(Debug, Clone)]
pub struct StemConfig {
pub line_color: Option<Color>,
pub line_width: f32,
pub marker_color: Option<Color>,
pub marker_size: f32,
pub marker: StemMarker,
pub baseline: f64,
pub bottom_marker: bool,
pub orientation: StemOrientation,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum StemMarker {
Circle,
Square,
Diamond,
Triangle,
None,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum StemOrientation {
Vertical,
Horizontal,
}
impl Default for StemConfig {
fn default() -> Self {
Self {
line_color: None,
line_width: 1.0,
marker_color: None,
marker_size: 6.0,
marker: StemMarker::Circle,
baseline: 0.0,
bottom_marker: false,
orientation: StemOrientation::Vertical,
}
}
}
impl StemConfig {
pub fn new() -> Self {
Self::default()
}
pub fn line_color(mut self, color: Color) -> Self {
self.line_color = Some(color);
self
}
pub fn line_width(mut self, width: f32) -> Self {
self.line_width = width.max(0.1);
self
}
pub fn marker_color(mut self, color: Color) -> Self {
self.marker_color = Some(color);
self
}
pub fn marker_size(mut self, size: f32) -> Self {
self.marker_size = size.max(0.0);
self
}
pub fn marker(mut self, marker: StemMarker) -> Self {
self.marker = marker;
self
}
pub fn baseline(mut self, baseline: f64) -> Self {
self.baseline = baseline;
self
}
pub fn bottom_marker(mut self, show: bool) -> Self {
self.bottom_marker = show;
self
}
pub fn horizontal(mut self) -> Self {
self.orientation = StemOrientation::Horizontal;
self
}
}
impl PlotConfig for StemConfig {}
pub struct Stem;
#[derive(Debug, Clone, Copy)]
pub struct StemElement {
pub x: f64,
pub y: f64,
pub baseline: f64,
pub index: usize,
}
impl StemElement {
pub fn vertical_line(&self) -> ((f64, f64), (f64, f64)) {
((self.x, self.baseline), (self.x, self.y))
}
pub fn horizontal_line(&self) -> ((f64, f64), (f64, f64)) {
((self.baseline, self.x), (self.y, self.x))
}
pub fn marker_position(&self, orientation: StemOrientation) -> (f64, f64) {
match orientation {
StemOrientation::Vertical => (self.x, self.y),
StemOrientation::Horizontal => (self.y, self.x),
}
}
pub fn baseline_marker_position(&self, orientation: StemOrientation) -> (f64, f64) {
match orientation {
StemOrientation::Vertical => (self.x, self.baseline),
StemOrientation::Horizontal => (self.baseline, self.x),
}
}
}
pub fn compute_stems(x: &[f64], y: &[f64], config: &StemConfig) -> Vec<StemElement> {
let n = x.len().min(y.len());
let mut stems = Vec::with_capacity(n);
for i in 0..n {
stems.push(StemElement {
x: x[i],
y: y[i],
baseline: config.baseline,
index: i,
});
}
stems
}
pub fn stem_range(x: &[f64], y: &[f64], config: &StemConfig) -> ((f64, f64), (f64, f64)) {
if x.is_empty() || y.is_empty() {
return ((0.0, 1.0), (0.0, 1.0));
}
let x_min = x.iter().copied().fold(f64::INFINITY, f64::min);
let x_max = x.iter().copied().fold(f64::NEG_INFINITY, f64::max);
let y_min = y
.iter()
.copied()
.fold(f64::INFINITY, f64::min)
.min(config.baseline);
let y_max = y
.iter()
.copied()
.fold(f64::NEG_INFINITY, f64::max)
.max(config.baseline);
let x_margin = (x_max - x_min) * 0.05;
match config.orientation {
StemOrientation::Vertical => ((x_min - x_margin, x_max + x_margin), (y_min, y_max)),
StemOrientation::Horizontal => ((y_min, y_max), (x_min - x_margin, x_max + x_margin)),
}
}
#[derive(Debug, Clone)]
pub struct StemData {
pub stems: Vec<StemElement>,
pub bounds: ((f64, f64), (f64, f64)),
pub(crate) config: StemConfig,
}
pub struct StemInput<'a> {
pub x: &'a [f64],
pub y: &'a [f64],
}
impl<'a> StemInput<'a> {
pub fn new(x: &'a [f64], y: &'a [f64]) -> Self {
Self { x, y }
}
}
impl PlotCompute for Stem {
type Input<'a> = StemInput<'a>;
type Config = StemConfig;
type Output = StemData;
fn compute(input: Self::Input<'_>, config: &Self::Config) -> Result<Self::Output> {
if input.x.is_empty() || input.y.is_empty() {
return Err(crate::core::PlottingError::EmptyDataSet);
}
let stems = compute_stems(input.x, input.y, config);
let bounds = stem_range(input.x, input.y, config);
Ok(StemData {
stems,
bounds,
config: config.clone(),
})
}
}
impl PlotData for StemData {
fn data_bounds(&self) -> ((f64, f64), (f64, f64)) {
self.bounds
}
fn is_empty(&self) -> bool {
self.stems.is_empty()
}
}
impl PlotRender for StemData {
fn render(
&self,
renderer: &mut SkiaRenderer,
area: &PlotArea,
_theme: &Theme,
color: Color,
) -> Result<()> {
if self.stems.is_empty() {
return Ok(());
}
let config = &self.config;
let line_color = config.line_color.unwrap_or(color);
let marker_color = config.marker_color.unwrap_or(color);
let marker_style = match config.marker {
StemMarker::Circle => Some(MarkerStyle::Circle),
StemMarker::Square => Some(MarkerStyle::Square),
StemMarker::Diamond => Some(MarkerStyle::Diamond),
StemMarker::Triangle => Some(MarkerStyle::Triangle),
StemMarker::None => None,
};
for stem in &self.stems {
let ((x1, y1), (x2, y2)) = match config.orientation {
StemOrientation::Vertical => stem.vertical_line(),
StemOrientation::Horizontal => stem.horizontal_line(),
};
let (sx1, sy1) = area.data_to_screen(x1, y1);
let (sx2, sy2) = area.data_to_screen(x2, y2);
renderer.draw_line(
sx1,
sy1,
sx2,
sy2,
line_color,
config.line_width,
LineStyle::Solid,
)?;
if let Some(style) = marker_style {
let (mx, my) = stem.marker_position(config.orientation);
let (smx, smy) = area.data_to_screen(mx, my);
renderer.draw_marker(smx, smy, config.marker_size, style, marker_color)?;
}
if config.bottom_marker {
if let Some(style) = marker_style {
let (bx, by) = stem.baseline_marker_position(config.orientation);
let (sbx, sby) = area.data_to_screen(bx, by);
renderer.draw_marker(sbx, sby, config.marker_size * 0.7, style, line_color)?;
}
}
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_compute_stems() {
let x = vec![0.0, 1.0, 2.0, 3.0];
let y = vec![2.0, 5.0, 3.0, 4.0];
let config = StemConfig::default();
let stems = compute_stems(&x, &y, &config);
assert_eq!(stems.len(), 4);
assert!((stems[0].x - 0.0).abs() < 1e-10);
assert!((stems[0].y - 2.0).abs() < 1e-10);
assert!((stems[0].baseline - 0.0).abs() < 1e-10);
}
#[test]
fn test_stem_element_vertical_line() {
let stem = StemElement {
x: 1.0,
y: 5.0,
baseline: 0.0,
index: 0,
};
let ((x1, y1), (x2, y2)) = stem.vertical_line();
assert!((x1 - 1.0).abs() < 1e-10);
assert!((y1 - 0.0).abs() < 1e-10);
assert!((x2 - 1.0).abs() < 1e-10);
assert!((y2 - 5.0).abs() < 1e-10);
}
#[test]
fn test_stem_marker_position() {
let stem = StemElement {
x: 1.0,
y: 5.0,
baseline: 0.0,
index: 0,
};
let (mx, my) = stem.marker_position(StemOrientation::Vertical);
assert!((mx - 1.0).abs() < 1e-10);
assert!((my - 5.0).abs() < 1e-10);
let (mx, my) = stem.marker_position(StemOrientation::Horizontal);
assert!((mx - 5.0).abs() < 1e-10);
assert!((my - 1.0).abs() < 1e-10);
}
#[test]
fn test_stem_range() {
let x = vec![0.0, 1.0, 2.0];
let y = vec![1.0, 3.0, 2.0];
let config = StemConfig::default();
let ((x_min, x_max), (y_min, y_max)) = stem_range(&x, &y, &config);
assert!(x_min < 0.0); assert!(x_max > 2.0); assert!((y_min - 0.0).abs() < 1e-10); assert!((y_max - 3.0).abs() < 1e-10);
}
#[test]
fn test_stem_with_baseline() {
let x = vec![0.0, 1.0];
let y = vec![2.0, 4.0];
let config = StemConfig::default().baseline(1.0);
let stems = compute_stems(&x, &y, &config);
assert!((stems[0].baseline - 1.0).abs() < 1e-10);
}
#[test]
fn test_stem_config_implements_plot_config() {
fn assert_plot_config<T: PlotConfig>() {}
assert_plot_config::<StemConfig>();
}
#[test]
fn test_stem_plot_compute_trait() {
use crate::plots::traits::PlotCompute;
let x = vec![0.0, 1.0, 2.0, 3.0];
let y = vec![2.0, 5.0, 3.0, 4.0];
let config = StemConfig::default();
let input = StemInput::new(&x, &y);
let result = Stem::compute(input, &config);
assert!(result.is_ok());
let stem_data = result.unwrap();
assert_eq!(stem_data.stems.len(), 4);
}
#[test]
fn test_stem_plot_compute_empty() {
use crate::plots::traits::PlotCompute;
let x: Vec<f64> = vec![];
let y: Vec<f64> = vec![];
let config = StemConfig::default();
let input = StemInput::new(&x, &y);
let result = Stem::compute(input, &config);
assert!(result.is_err());
}
#[test]
fn test_stem_plot_data_trait() {
use crate::plots::traits::{PlotCompute, PlotData};
let x = vec![0.0, 1.0, 2.0];
let y = vec![1.0, 3.0, 2.0];
let config = StemConfig::default();
let input = StemInput::new(&x, &y);
let stem_data = Stem::compute(input, &config).unwrap();
let ((x_min, x_max), (y_min, y_max)) = stem_data.data_bounds();
assert!(x_min < 0.0); assert!(x_max > 2.0); assert!(y_min <= y_max);
assert!(!stem_data.is_empty());
}
#[test]
fn test_stem_plot_compute_horizontal() {
use crate::plots::traits::PlotCompute;
let x = vec![0.0, 1.0, 2.0];
let y = vec![2.0, 4.0, 3.0];
let config = StemConfig::default().horizontal();
let input = StemInput::new(&x, &y);
let result = Stem::compute(input, &config);
assert!(result.is_ok());
let stem_data = result.unwrap();
assert_eq!(stem_data.config.orientation, StemOrientation::Horizontal);
}
}