box-plotters 0.1.0

use std::marker::PhantomData;

use plotters::{element::{Drawable, PointCollection}, style::{BLACK, ShapeStyle, TRANSPARENT}};
use plotters_backend::{DrawingBackend, DrawingErrorKind};

use crate::{box_plot::orientation::{Horizontal, Orientation, Vertical}, quartiles::Quartiles};

pub mod orientation;

const DEFAULT_WIDTH: f64 = 10.0;
const DEFAULT_WHISKER_WIDTH: f64 = 10.0;

pub struct BoxPlot<K, O> {
    line_style: ShapeStyle,
    box_style: ShapeStyle,
    width: f64,
    whisker_width: f64,
    key: K,
    quartiles: Quartiles,
    _phantom: PhantomData<O>,
}

impl<K> BoxPlot<K, Vertical<K, f64>> {
    /// Constructs a new vertical box plot based on a key and quartiles.
    pub fn vertical_from_key_quartiles(key: K, quartiles: Quartiles) -> Self {
        Self {
            line_style: BLACK.into(),
            box_style: TRANSPARENT.into(),
            width: DEFAULT_WIDTH,
            whisker_width: DEFAULT_WHISKER_WIDTH,
            key,
            quartiles,
            _phantom: PhantomData,
        }
    }
}

impl<K> BoxPlot<K, Horizontal<K, f64>> {
    /// Constructs a new horizontal box plot based on a key and quartiles.
    pub fn horizontal_from_key_quartiles(key: K, quartiles: Quartiles) -> Self {
        Self {
            line_style: BLACK.into(),
            box_style: TRANSPARENT.into(),
            width: DEFAULT_WIDTH,
            whisker_width: DEFAULT_WHISKER_WIDTH,
            key,
            quartiles,
            _phantom: PhantomData,
        }
    }
}

impl<K, O> BoxPlot<K, O> {
    /// Sets the line style used by the box plot
    pub fn with_line_style<S: Into<ShapeStyle>>(&mut self, style: S) -> &mut Self {
        self.line_style = style.into();
        self
    }

    /// Sets the box style used to fill the box plot
    pub fn with_box_style<S: Into<ShapeStyle>>(&mut self, style: S) -> &mut Self {
        self.box_style = style.into();
        self
    }

    /// Sets the width of the box.
    pub fn with_width(&mut self, width: f64) -> &mut Self {
        self.width = width;
        self
    }

    /// Sets the width of the whiskers.
    pub fn with_whisker_width(&mut self, width: f64) -> &mut Self {
        self.whisker_width = width;
        self
    }
}

impl<'a, K: Clone, O: Orientation<K, f64>> PointCollection<'a, (O::XType, O::YType)>
    for &'a BoxPlot<K, O>
{
    type Point = (O::XType, O::YType);
    type IntoIter = Vec<Self::Point>;
    fn point_iter(self) -> Self::IntoIter {
        self.quartiles
            .values()
            .iter()
            .map(|v| O::make_coord(self.key.clone(), *v))
            .collect()
    }
}

impl<K, DB: DrawingBackend, O: Orientation<K, f64>> Drawable<DB> for BoxPlot<K, O> {
    fn draw<I: Iterator<Item = <plotters::element::BackendCoordOnly as plotters::element::CoordMapper>::Output>>(
        &self,
        pos: I,
        backend: &mut DB,
        _parent_dim: (u32, u32),
    ) -> Result<(), DrawingErrorKind<DB::ErrorType>> {
        let points = pos.collect::<Vec<_>>();
        let whisker_start_offset = |p| O::with_offset(p, (-self.whisker_width / 2.0) as i32);
        let whisker_end_offset   = |p| O::with_offset(p, ( self.whisker_width / 2.0) as i32);
        let bar_start_offset     = |p| O::with_offset(p, (-self.width / 2.0) as i32);
        let bar_end_offset       = |p| O::with_offset(p, ( self.width / 2.0) as i32);

        let lower_whisker_start = whisker_start_offset(points[0]);
        let lower_whisker_end = whisker_end_offset(points[0]);
        let lower_quartile_start = bar_start_offset(points[1]);
        let lower_quartile_end = bar_end_offset(points[1]);
        let median_start = bar_start_offset(points[2]);
        let median_end = bar_end_offset(points[2]);
        let upper_quartile_start = bar_start_offset(points[3]);
        let upper_quartile_end = bar_end_offset(points[3]);
        let upper_whisker_start = whisker_start_offset(points[4]);
        let upper_whisker_end = whisker_end_offset(points[4]);

        backend.draw_line(lower_whisker_start, lower_whisker_end, &self.line_style)?;
        backend.draw_line(points[0], points[1], &self.line_style)?;
        
        let corners = [
            lower_quartile_start,
            lower_quartile_end,
            upper_quartile_start,
            upper_quartile_end
        ];

        let x_min = corners.iter().map(|p| p.0).min().unwrap();
        let x_max = corners.iter().map(|p| p.0).max().unwrap();
        let y_min = corners.iter().map(|p| p.1).min().unwrap();
        let y_max = corners.iter().map(|p| p.1).max().unwrap();

        backend.draw_rect((x_min, y_min), (x_max, y_max), &self.box_style, true)?;
        backend.draw_rect((x_min, y_min), (x_max, y_max), &self.line_style, false)?;
        backend.draw_line(median_start, median_end, &self.line_style)?;

        backend.draw_line(points[3], points[4], &self.line_style)?;
        backend.draw_line(upper_whisker_start, upper_whisker_end, &self.line_style)?;

        Ok(())
    }
}