graph.rs

use plotters::{coord::AsRangedCoord, prelude::*};
use std::error::Error;
use std::fmt::Debug;
use std::ops::Range;
use std::path::Path;

const IMAGE_SIZE: (u32, u32) = (900, 600);

pub struct Graph<X, Y> {
    bounds: (Range<X>, Range<Y>),
    data_sets: Vec<DataSeries<X, Y>>,
}

struct DataSeries<X, Y> {
    name: String,
    style: ShapeStyle,
    data: Vec<(X, Y)>,
}

impl<X, Y> Graph<X, Y> {
    pub fn new(x_bounds: Range<X>, y_bounds: Range<Y>) -> Self {
        Graph {
            bounds: (x_bounds, y_bounds),
            data_sets: Vec::new(),
        }
    }

    pub fn add_series<S, I>(&mut self, name: &str, style: S, series: I)
    where
        S: Into<ShapeStyle>,
        I: IntoIterator<Item = (X, Y)>,
    {
        let series = DataSeries {
            name: name.into(),
            style: style.into(),
            data: series.into_iter().collect(),
        };

        self.data_sets.push(series);
    }

    pub fn render<P: AsRef<Path>>(&self, filename: P) -> Result<(), Box<dyn Error>>
    where
        X: Debug + Copy + 'static,
        Y: Debug + Copy + 'static,
        Range<X>: AsRangedCoord<Value = X>,
        Range<Y>: AsRangedCoord<Value = Y>,
    {
        let ctx = BitMapBackend::new(&filename, IMAGE_SIZE).into_drawing_area();
        ctx.fill(&WHITE)?;

        let mut chart = ChartBuilder::on(&ctx)
            .margin(30)
            .x_label_area_size(20)
            .y_label_area_size(20)
            .build_ranged(self.bounds.0.clone(), self.bounds.1.clone())?;

        chart.configure_mesh().draw()?;

        for series in &self.data_sets {
            let data = series.data.iter().cloned();
            let line = LineSeries::new(data, series.style.clone());
            let anno = chart.draw_series(line)?;

            anno.label(&series.name);

            anno.legend(move |(x, y)| {
                let coords = vec![(x, y), (x + 20, y)];
                PathElement::new(coords, series.style.clone())
            });
        }

        chart
            .configure_series_labels()
            .background_style(&WHITE.mix(0.8))
            .border_style(&BLACK)
            .draw()?;

        Ok(())
    }
}