strafe-plot 0.1.1

// "Whatever you do, work at it with all your heart, as working for the Lord,
// not for human masters, since you know that you will receive an inheritance
// from the Lord as a reward. It is the Lord Christ you are serving."
// (Col 3:23-24)

use std::{error::Error, iter::once};

use plotters::{
    backend::DrawingBackend,
    chart::SeriesLabelPosition,
    coord::Shift,
    drawing::DrawingArea,
    element::Circle,
    prelude::{Color, LineSeries, Polygon, BLACK},
};

use crate::{
    chart::chart,
    drawing_area::setup_drawing_area,
    drawing_coords::{drawing_coords, DrawingCoords},
    plot_options::PlotOptions,
    plottable::Plottable,
};

pub struct Plot<B: DrawingBackend> {
    pub options: PlotOptions,
    pub plottables: Vec<Box<dyn Plottable<B>>>,
}

impl<B: DrawingBackend> Default for Plot<B> {
    fn default() -> Self {
        Self {
            options: Default::default(),
            plottables: Default::default(),
        }
    }
}

impl<B: DrawingBackend> Plot<B> {
    pub fn new() -> Self {
        Self {
            options: Default::default(),
            plottables: Default::default(),
        }
    }

    pub fn with_options(&mut self, plot_options: PlotOptions) -> &mut Self {
        self.options = plot_options;
        self
    }

    pub fn with_plottable<P: Plottable<B> + 'static>(&mut self, plottable: P) -> &mut Self {
        self.plottables.push(Box::new(plottable));
        self
    }

    pub fn get_drawing_coords(&self, root: &DrawingArea<B, Shift>) -> DrawingCoords {
        let mut xs = Vec::new();
        let mut ys = Vec::new();

        for plottable in &self.plottables {
            if plottable.force_fit() {
                xs.append(&mut plottable.get_x());
                ys.append(&mut plottable.get_y());
            }
        }

        if self.options.x_log {
            for x in xs.iter_mut() {
                *x = x.abs().log10();
            }
        }

        if self.options.y_log {
            for y in ys.iter_mut() {
                *y = y.abs().log10();
            }
        }

        drawing_coords(&xs, &ys, root, &self.options)
    }

    pub fn plot(&self, root: &DrawingArea<B, Shift>) -> Result<(), Box<dyn Error>>
    where
        B::ErrorType: 'static,
    {
        let (width, height) = root.dim_in_pixel();
        let root = &root.split_by_breakpoints(
            [
                width as f64 * self.options.plot_left,
                width as f64 * self.options.plot_right,
            ],
            [
                height as f64 * self.options.plot_top,
                height as f64 * self.options.plot_bottom,
            ],
        )[4];

        let drawing_coords = self.get_drawing_coords(root);

        setup_drawing_area(root, &drawing_coords);

        root.titled(
            &self.options.title,
            (
                self.options.font.as_str(),
                (drawing_coords.chart_width * drawing_coords.chart_height).sqrt()
                    * self.options.title_scalar,
            ),
        )
        .unwrap();

        let mut chart = chart(
            &root,
            &self.options,
            &drawing_coords,
            &self.options.x_axis_label,
            &self.options.y_axis_label,
        );

        let mut legend = false;
        for plottable in &self.plottables {
            let plottable_value = plottable.plot(&self.options, &drawing_coords)?;

            for (coords, size, color, label) in plottable_value.points {
                let c = chart.draw_series(
                    coords
                        .iter()
                        .map(|(x, y)| Circle::new((*x, *y), size, color.filled())),
                )?;
                if let Some(label) = label {
                    c.legend(move |(x, y)| Circle::new((x + 5, y), 2, color.filled()))
                        .label(label);
                }
            }

            for (coords, size, color, label) in plottable_value.lines {
                let c = chart.draw_series(LineSeries::new(coords, color.stroke_width(size)))?;
                if let Some(label) = label {
                    c.legend(move |(x, y)| Circle::new((x + 5, y), 2, color.filled()))
                        .label(label);
                }
            }

            for (coords, color, label) in plottable_value.polygons {
                let c = chart.draw_series(once(Polygon::new(coords, color.mix(0.2))))?;
                if let Some(label) = label {
                    c.legend(move |(x, y)| Circle::new((x + 5, y), 2, color.filled()))
                        .label(label);
                }
            }

            legend = legend || plottable.get_legend();
        }

        if legend {
            chart
                .configure_series_labels()
                .border_style(BLACK)
                .position(SeriesLabelPosition::Coordinate(
                    (self.options.legend_x * (width as f64 * 0.83)) as i32,
                    (self.options.legend_y * (height as f64 * 0.83)) as i32,
                ))
                .label_font((
                    self.options.font.as_str(),
                    (drawing_coords.chart_width * drawing_coords.chart_height).sqrt()
                        * self.options.legend_scalar,
                ))
                .draw()?;
        }

        Ok(())
    }
}