waybar-collectd/draw/src/stacked.rs

use imageproc::{
    drawing::draw_filled_rect_mut,
    image::{Rgb, RgbImage},
    rect::Rect,
};

use crate::Container;

pub struct StackedConfig {
    pub name: String,
    pub width_px_per_sample: u32,
    pub height_px: u32,
    pub series: Vec<StackedSeriesConfig>,
    pub max_samples: u32,
}

pub struct StackedSeriesConfig {
    pub color: (u8, u8, u8),
}

pub struct StackedSeries {
    pub config: StackedConfig,
    data: Container,
}

impl StackedSeries {
    pub fn new(config: StackedConfig) -> StackedSeries {
        let max_samples = config.max_samples;

        StackedSeries {
            config: config,
            data: Container::new(max_samples),
        }
    }

    pub fn push(&mut self, data: Vec<f32>) {
        self.data.push(data);
    }
}

pub fn stacked(graph: &StackedSeries) -> RgbImage {
    let mut image = RgbImage::new(
        graph.config.width_px_per_sample * graph.config.max_samples,
        graph.config.height_px as u32,
    );
    draw_scaled_graph(&mut image, &graph);

    image
}

/// return scale factor to scale stacked series sum to range 0..1
fn get_scale(series: &Container) -> f32 {
    let max = series
        .iter()
        .map(|s| s.iter().sum())
        .reduce(f32::max)
        .unwrap_or(1f32);

    if max > 1f32 { 1f32 / max } else { 1f32 }
}

fn draw_scaled_graph(image: &mut RgbImage, graph: &StackedSeries) {
    let scale = get_scale(&graph.data);
    let colors = graph
        .config
        .series
        .iter()
        .map(|s| Rgb([s.color.0, s.color.1, s.color.2]))
        .collect::<Vec<Rgb<u8>>>();

    let mut x_offset = 0f32;

    for point_in_time in graph.data.iter() {
        let mut y_offset = graph.config.height_px as f32;

        let mut series_idx = 0usize;
        for sample in point_in_time.iter() {
            let color = colors[series_idx];
            let height_px = (graph.config.height_px as f32) * sample * scale;
            if height_px > 1f32 {
                let rect = Rect::at(x_offset as i32, (y_offset - height_px) as i32)
                    .of_size(graph.config.width_px_per_sample, height_px as u32);
                draw_filled_rect_mut(image, rect, color);
                y_offset -= height_px;
            }
            series_idx += 1;
        }
        x_offset += graph.config.width_px_per_sample as f32;
    }
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn test_get_scale() {
        let mut c = Container::new(10);
        c.push(vec![1., 0., 0.]);
        c.push(vec![2., 1., 0.]);
        c.push(vec![3., 0., 0.]);
        c.push(vec![4., 0., 5.]);
        c.push(vec![5., 0., 5.]);

        assert_eq!(get_scale(&c), 0.1);
    }
}