plotters-statistical 0.1.0

//! Violin plot series — a mirrored KDE outline, optionally with an embedded box
//! plot, built on [`crate::stats::kde`] and reusing [`BoxPlot`](crate::BoxPlot)'s geometry.

use plotters::element::{Drawable, PointCollection};
use plotters::style::{RGBColor, ShapeStyle};
use plotters_backend::{BackendCoord, DrawingBackend, DrawingErrorKind};

use crate::series::box_plot::BoxStyle;
use crate::stats::{kde_curve, quartiles, StatsError};
use crate::style::{stroke_style, translucent_fill};

const DEFAULT_VIOLIN_FILL: RGBColor = RGBColor(0, 158, 115); // Okabe–Ito green
const DEFAULT_VIOLIN_WIDTH: u32 = 60;
const DEFAULT_KDE_POINTS: usize = 200;
const EMBEDDED_BOX_WIDTH: u32 = 10;
const N_BOX_POINTS: usize = 5;

/// Styling for a violin plot.
#[derive(Debug, Clone)]
pub struct ViolinStyle {
    /// Fill of the violin body.
    pub fill: ShapeStyle,
    /// Outline of the violin body.
    pub outline: ShapeStyle,
    /// Whether to draw the embedded box plot inside the violin.
    pub show_box: bool,
}

impl Default for ViolinStyle {
    fn default() -> Self {
        Self {
            fill: translucent_fill(&DEFAULT_VIOLIN_FILL, 0.4),
            outline: stroke_style(&DEFAULT_VIOLIN_FILL, 2),
            show_box: false,
        }
    }
}

/// A single violin at one coordinate on the category axis.
///
/// Like [`BoxPlot`](crate::BoxPlot), it is generic over the full chart coordinate `C`: a
/// vertical violin lives on a `(X, f64)` chart, a horizontal one on `(f64, Y)`.
/// The half-width scales with the kernel density, so the widest part of the
/// violin corresponds to the mode of the sample.
#[derive(Debug, Clone)]
pub struct ViolinPlot<C> {
    // Layout: [n_kde density points][5 box points]. The KDE points carry the
    // outline; the 5 box points feed the optional embedded box.
    points: Vec<C>,
    density: Vec<f64>,
    max_density: f64,
    n_kde: usize,
    width: u32,
    horizontal: bool,
    style: ViolinStyle,
    box_style: BoxStyle,
}

impl<X: Clone> ViolinPlot<(X, f64)> {
    /// A vertical violin at category position `x` from a raw `data` sample,
    /// using a Silverman-rule bandwidth. Errors if the sample is empty or
    /// degenerate (see [`crate::stats::kde`]).
    pub fn vertical(x: X, data: &[f64]) -> Result<Self, StatsError> {
        Self::vertical_with_bandwidth(x, data, None)
    }

    /// A vertical violin with an explicit KDE bandwidth override.
    pub fn vertical_with_bandwidth(
        x: X,
        data: &[f64],
        bandwidth: Option<f64>,
    ) -> Result<Self, StatsError> {
        let curve = kde_curve(data, bandwidth, DEFAULT_KDE_POINTS, 3.0)?;
        let q = quartiles(data)?;
        let mut points: Vec<(X, f64)> = curve.xs.iter().map(|&y| (x.clone(), y)).collect();
        points.push((x.clone(), q.upper_whisker));
        points.push((x.clone(), q.q3));
        points.push((x.clone(), q.median));
        points.push((x.clone(), q.q1));
        points.push((x.clone(), q.lower_whisker));
        Ok(Self {
            n_kde: curve.xs.len(),
            max_density: curve.max_density(),
            density: curve.density,
            points,
            width: DEFAULT_VIOLIN_WIDTH,
            horizontal: false,
            style: ViolinStyle::default(),
            box_style: BoxStyle::default(),
        })
    }
}

impl<Y: Clone> ViolinPlot<(f64, Y)> {
    /// A horizontal violin at category position `y` from a raw `data` sample.
    pub fn horizontal(y: Y, data: &[f64]) -> Result<Self, StatsError> {
        Self::horizontal_with_bandwidth(y, data, None)
    }

    /// A horizontal violin with an explicit KDE bandwidth override.
    pub fn horizontal_with_bandwidth(
        y: Y,
        data: &[f64],
        bandwidth: Option<f64>,
    ) -> Result<Self, StatsError> {
        let curve = kde_curve(data, bandwidth, DEFAULT_KDE_POINTS, 3.0)?;
        let q = quartiles(data)?;
        let mut points: Vec<(f64, Y)> = curve.xs.iter().map(|&x| (x, y.clone())).collect();
        points.push((q.upper_whisker, y.clone()));
        points.push((q.q3, y.clone()));
        points.push((q.median, y.clone()));
        points.push((q.q1, y.clone()));
        points.push((q.lower_whisker, y.clone()));
        Ok(Self {
            n_kde: curve.xs.len(),
            max_density: curve.max_density(),
            density: curve.density,
            points,
            width: DEFAULT_VIOLIN_WIDTH,
            horizontal: true,
            style: ViolinStyle::default(),
            box_style: BoxStyle::default(),
        })
    }
}

impl<C> ViolinPlot<C> {
    /// Set the full violin width in **pixels** (default 60).
    pub fn width(mut self, width: u32) -> Self {
        self.width = width;
        self
    }

    /// Enable/disable the embedded box plot overlay.
    pub fn show_box(mut self, show: bool) -> Self {
        self.style.show_box = show;
        self
    }

    /// Replace the violin style block.
    pub fn style(mut self, style: ViolinStyle) -> Self {
        self.style = style;
        self
    }

    /// Replace the embedded-box style block.
    pub fn box_style(mut self, style: BoxStyle) -> Self {
        self.box_style = style;
        self
    }
}

impl<'a, C: 'a> PointCollection<'a, C> for &'a ViolinPlot<C> {
    type Point = &'a C;
    type IntoIter = &'a [C];
    fn point_iter(self) -> &'a [C] {
        &self.points
    }
}

impl<C, DB: DrawingBackend> Drawable<DB> for ViolinPlot<C> {
    fn draw<I: Iterator<Item = BackendCoord>>(
        &self,
        points: I,
        backend: &mut DB,
        _parent_dim: (u32, u32),
    ) -> Result<(), DrawingErrorKind<DB::ErrorType>> {
        let pix: Vec<BackendCoord> = points.collect();
        if pix.len() < self.n_kde + N_BOX_POINTS || self.max_density <= 0.0 {
            return Ok(());
        }
        let half = self.width as f64 / 2.0;
        let kde = &pix[..self.n_kde];

        // Build the mirrored outline. `offset(i)` is the half-width in pixels
        // for the i-th density sample.
        let offset =
            |i: usize| -> i32 { ((self.density[i] / self.max_density) * half).round() as i32 };

        let mut outline: Vec<BackendCoord> = Vec::with_capacity(self.n_kde * 2 + 1);
        if self.horizontal {
            let py = |i: usize| kde[i].1;
            let px = |i: usize| kde[i].0;
            for i in 0..self.n_kde {
                outline.push((px(i), py(i) - offset(i)));
            }
            for i in (0..self.n_kde).rev() {
                outline.push((px(i), py(i) + offset(i)));
            }
        } else {
            let px = |i: usize| kde[i].0;
            let py = |i: usize| kde[i].1;
            for i in 0..self.n_kde {
                outline.push((px(i) - offset(i), py(i)));
            }
            for i in (0..self.n_kde).rev() {
                outline.push((px(i) + offset(i), py(i)));
            }
        }

        backend.fill_polygon(outline.iter().copied(), &self.style.fill)?;
        // Close the outline for the stroke.
        if let Some(&first) = outline.first() {
            outline.push(first);
        }
        backend.draw_path(outline, &self.style.outline)?;

        // Optional embedded box plot, using the 5 trailing box points.
        if self.style.show_box {
            let b = &pix[self.n_kde..self.n_kde + N_BOX_POINTS];
            draw_embedded_box(backend, b, self.horizontal, &self.box_style)?;
        }
        Ok(())
    }
}

/// Draw a thin box (IQR box + median + whiskers) from the five box points, for
/// the embedded-box overlay. Mirrors [`BoxPlot`]'s geometry at a fixed narrow
/// width so it sits inside the violin.
fn draw_embedded_box<DB: DrawingBackend>(
    backend: &mut DB,
    b: &[BackendCoord],
    horizontal: bool,
    style: &BoxStyle,
) -> Result<(), DrawingErrorKind<DB::ErrorType>> {
    let (upper_w, q3, median, q1, lower_w) = (b[0], b[1], b[2], b[3], b[4]);
    let half = (EMBEDDED_BOX_WIDTH / 2) as i32;
    if horizontal {
        let cy = q1.1;
        let (bx1, bx2) = (q1.0.min(q3.0), q1.0.max(q3.0));
        backend.draw_line((lower_w.0, cy), (upper_w.0, cy), &style.whisker)?;
        backend.draw_rect((bx1, cy - half), (bx2, cy + half), &style.box_fill, true)?;
        backend.draw_rect((bx1, cy - half), (bx2, cy + half), &style.box_border, false)?;
        backend.draw_line((median.0, cy - half), (median.0, cy + half), &style.median)?;
    } else {
        let cx = q1.0;
        let (by1, by2) = (q3.1.min(q1.1), q3.1.max(q1.1));
        backend.draw_line((cx, lower_w.1), (cx, upper_w.1), &style.whisker)?;
        backend.draw_rect((cx - half, by1), (cx + half, by2), &style.box_fill, true)?;
        backend.draw_rect((cx - half, by1), (cx + half, by2), &style.box_border, false)?;
        backend.draw_line((cx - half, median.1), (cx + half, median.1), &style.median)?;
    }
    Ok(())
}

/// A group of violins laid out side by side — mirrors [`BoxPlotSeries`](crate::BoxPlotSeries).
#[derive(Debug, Clone)]
pub struct ViolinPlotSeries<C> {
    violins: Vec<ViolinPlot<C>>,
}

impl<X: Clone> ViolinPlotSeries<(X, f64)> {
    /// Build a vertical multi-violin series from `(position, sample)` pairs.
    pub fn from_samples<I, S>(groups: I) -> Result<Self, StatsError>
    where
        I: IntoIterator<Item = (X, S)>,
        S: AsRef<[f64]>,
    {
        let violins = groups
            .into_iter()
            .map(|(x, s)| ViolinPlot::vertical(x, s.as_ref()))
            .collect::<Result<Vec<_>, _>>()?;
        Ok(Self { violins })
    }
}

impl<C> ViolinPlotSeries<C> {
    /// Apply a common width (pixels) to every violin.
    pub fn width(mut self, width: u32) -> Self {
        self.violins = self.violins.into_iter().map(|v| v.width(width)).collect();
        self
    }

    /// Enable/disable the embedded box overlay on every violin.
    pub fn show_box(mut self, show: bool) -> Self {
        self.violins = self.violins.into_iter().map(|v| v.show_box(show)).collect();
        self
    }

    /// Apply a common violin style to every violin.
    pub fn style(mut self, style: ViolinStyle) -> Self {
        self.violins = self
            .violins
            .into_iter()
            .map(|v| v.style(style.clone()))
            .collect();
        self
    }
}

impl<C> IntoIterator for ViolinPlotSeries<C> {
    type Item = ViolinPlot<C>;
    type IntoIter = std::vec::IntoIter<ViolinPlot<C>>;
    fn into_iter(self) -> Self::IntoIter {
        self.violins.into_iter()
    }
}