rust-ppm 0.1.1

use crate::{Image, Pixel};

use super::Axes;
use super::raster::{Bounds, draw_axis_arrows, draw_box, draw_line, draw_marker, draw_thick_line};
use super::style::{LineStyle, MarkerStyle};
use super::text::{self, GLYPH_HEIGHT};

const TICK_LENGTH: usize = 4;
const TEXT_GAP: usize = 2;

enum Series {
    Markers(Vec<(f64, f64)>, MarkerStyle),
    Line(Vec<(f64, f64)>, LineStyle),
}

impl Series {
    fn points(&self) -> &[(f64, f64)] {
        match self {
            Self::Markers(points, _) | Self::Line(points, _) => points,
        }
    }
}

/// An image-backed plotting surface.
///
/// `Canvas` keeps the rendered image, the current axis configuration, and any
/// series added to it. It is the lower-level incremental API for custom plotting
/// workflows while `Plot` provides a simpler higher-level builder interface.
pub struct Canvas {
    /// The current axis configuration used to map data coordinates to pixels.
    pub axes: Axes,
    /// The rasterized plot image.
    pub image: Image,
    inset_x: usize,
    inset_y: usize,
    render_scale: usize,
    series: Vec<Series>,
}

impl Canvas {
    /// Creates a new canvas with a given pixel size and axis configuration.
    pub fn new(width: usize, height: usize, axes: Axes) -> Self {
        Self::with_inset(width, height, axes, 0, 0)
    }

    /// Creates a new canvas with a custom inset margin around the drawing bounds.
    pub fn with_inset(
        width: usize,
        height: usize,
        axes: Axes,
        inset_x: usize,
        inset_y: usize,
    ) -> Self {
        Self {
            axes,
            image: Image::new_white(width, height),
            inset_x,
            inset_y,
            render_scale: 1,
            series: Vec::new(),
        }
    }

    /// Returns the canvas width in pixels.
    pub fn width(&self) -> usize {
        self.image.width
    }

    /// Returns the canvas height in pixels.
    pub fn height(&self) -> usize {
        self.image.height
    }

    /// Returns the horizontal and vertical inset values in pixels.
    pub fn insets(&self) -> (usize, usize) {
        (self.inset_x, self.inset_y)
    }

    /// Updates the inset margins used by the plot area.
    pub fn set_inset(&mut self, inset_x: usize, inset_y: usize) {
        self.inset_x = inset_x;
        self.inset_y = inset_y;
    }

    /// Sets the render scale used for higher-resolution output.
    pub fn set_render_scale(&mut self, scale: usize) {
        assert!(scale > 0, "render scale must be greater than zero");
        self.render_scale = scale;
        self.redraw();
    }

    /// Draws the axes, labels, and ticks using the current axes configuration.
    pub fn render(&mut self) {
        let Some(bounds) = self.bounds() else {
            return;
        };
        self.draw_axes(bounds);
    }

    /// Adds a scatter series with a default marker style and explicit marker size.
    pub fn scatter(&mut self, points: &[(f64, f64)], size: usize) {
        self.markers(points, MarkerStyle::new().size(size));
    }

    /// Adds a marker-based series using a custom marker style.
    pub fn markers(&mut self, points: &[(f64, f64)], style: MarkerStyle) {
        self.series.push(Series::Markers(points.to_vec(), style));
        if self.refresh_auto_limits() {
            self.redraw();
        } else {
            self.draw_markers(points, style);
        }
    }

    /// Adds a line series with a default width.
    pub fn plot(&mut self, points: &[(f64, f64)], line_width: usize) {
        self.line(points, LineStyle::new().width(line_width));
    }

    /// Adds a line series using a custom style.
    pub fn line(&mut self, points: &[(f64, f64)], style: LineStyle) {
        self.series.push(Series::Line(points.to_vec(), style));
        if self.refresh_auto_limits() {
            self.redraw();
        } else {
            self.draw_line_series(points, style);
        }
    }

    /// Sets the x-axis limits and redraws the canvas.
    pub fn set_xlim(&mut self, xlim: (f64, f64)) {
        self.axes.set_xlim(xlim);
        self.redraw();
    }

    /// Sets the y-axis limits and redraws the canvas.
    pub fn set_ylim(&mut self, ylim: (f64, f64)) {
        self.axes.set_ylim(ylim);
        self.redraw();
    }

    /// Re-enables automatic x-axis scaling for data-driven limits.
    pub fn use_auto_xlim(&mut self) {
        self.axes.use_auto_xlim();
        self.refresh_auto_limits();
        self.redraw();
    }

    /// Re-enables automatic y-axis scaling for data-driven limits.
    pub fn use_auto_ylim(&mut self) {
        self.axes.use_auto_ylim();
        self.refresh_auto_limits();
        self.redraw();
    }

    fn draw_markers(&mut self, points: &[(f64, f64)], style: MarkerStyle) {
        let Some(bounds) = self.bounds() else {
            return;
        };

        for &(x, y) in points {
            if !self.contains(x, y) {
                continue;
            }

            if let Some(pixel) = self.map_point((x, y), bounds) {
                draw_marker(&mut self.image, pixel, style.size, bounds, style.color);
            }
        }
    }

    fn draw_line_series(&mut self, points: &[(f64, f64)], style: LineStyle) {
        let Some(bounds) = self.bounds() else {
            return;
        };

        let mut previous = None;
        for &(x, y) in points {
            if !self.contains(x, y) {
                previous = None;
                continue;
            }

            if let Some(pixel) = self.map_point((x, y), bounds) {
                if let Some(start) = previous {
                    draw_thick_line(
                        &mut self.image,
                        start,
                        pixel,
                        style.width,
                        bounds,
                        style.color,
                    );
                } else {
                    draw_marker(&mut self.image, pixel, style.width, bounds, style.color);
                }
                previous = Some(pixel);
            }
        }
    }

    fn refresh_auto_limits(&mut self) -> bool {
        let limits = data_limits(self.series.iter().flat_map(|series| series.points()));
        let (xlim, ylim) = limits.unzip();
        self.axes.update_auto_limits(xlim, ylim)
    }

    fn redraw(&mut self) {
        self.image = Image::new_white(self.width(), self.height());
        self.render();

        let series = std::mem::take(&mut self.series);
        for item in &series {
            match item {
                Series::Markers(points, style) => self.draw_markers(points, *style),
                Series::Line(points, style) => self.draw_line_series(points, *style),
            }
        }
        self.series = series;
    }

    /// Consumes the canvas and returns the rendered image.
    pub fn into_image(self) -> Image {
        self.image
    }

    fn bounds(&self) -> Option<Bounds> {
        Bounds::from_insets(self.width(), self.height(), self.inset_x, self.inset_y)
    }

    fn contains(&self, x: f64, y: f64) -> bool {
        x >= self.axes.xlim.0
            && x <= self.axes.xlim.1
            && y >= self.axes.ylim.0
            && y <= self.axes.ylim.1
    }

    fn map_point(&self, point: (f64, f64), bounds: Bounds) -> Option<(usize, usize)> {
        Some((
            map_x(point.0, self.axes.xlim, bounds)?,
            map_y(point.1, self.axes.ylim, bounds)?,
        ))
    }

    fn draw_axes(&mut self, bounds: Bounds) {
        draw_box(&mut self.image, bounds, Pixel::BLACK);

        let y_axis_x = if contains_zero(self.axes.xlim) {
            map_x(0.0, self.axes.xlim, bounds)
        } else {
            Some(bounds.left + (bounds.right - bounds.left) / 2)
        };
        let x_axis_y = if contains_zero(self.axes.ylim) {
            map_y(0.0, self.axes.ylim, bounds)
        } else {
            Some(bounds.top + (bounds.bottom - bounds.top) / 2)
        };

        if let Some(x) = y_axis_x {
            for y in bounds.top..=bounds.bottom {
                self.image.set_pixel(x, y, Pixel::BLACK);
            }
        }
        if let Some(y) = x_axis_y {
            for x in bounds.left..=bounds.right {
                self.image.set_pixel(x, y, Pixel::BLACK);
            }
        }

        self.draw_ticks(bounds);

        if let (Some(y_axis_x), Some(x_axis_y)) = (y_axis_x, x_axis_y) {
            draw_axis_arrows(&mut self.image, y_axis_x, x_axis_y, bounds, Pixel::BLACK);
        }

        self.draw_labels(bounds);
    }

    fn draw_ticks(&mut self, bounds: Bounds) {
        let bottom_margin = self.height() - 1 - bounds.bottom;
        let tick_length = TICK_LENGTH * self.render_scale;
        let text_gap = TEXT_GAP * self.render_scale;
        let glyph_height = GLYPH_HEIGHT * self.render_scale;

        for &tick in &self.axes.xticks {
            let Some(x) = map_x(tick, self.axes.xlim, bounds) else {
                continue;
            };
            let tick_end = (bounds.bottom + tick_length).min(self.height() - 1);
            draw_line(
                &mut self.image,
                (x, bounds.bottom),
                (x, tick_end),
                Pixel::BLACK,
            );

            if bottom_margin >= tick_length + text_gap + glyph_height {
                let label = format_tick(tick);
                let (label_width, _) = text::measure(&label, self.render_scale);
                let label_x = centered_start(x, label_width, self.width());
                text::draw(
                    &mut self.image,
                    (label_x, bounds.bottom + tick_length + text_gap),
                    &label,
                    self.render_scale,
                    Pixel::BLACK,
                );
            }
        }

        for &tick in &self.axes.yticks {
            let Some(y) = map_y(tick, self.axes.ylim, bounds) else {
                continue;
            };
            let tick_start = bounds.left.saturating_sub(tick_length);
            draw_line(
                &mut self.image,
                (tick_start, y),
                (bounds.left, y),
                Pixel::BLACK,
            );

            let label = format_tick(tick);
            let (label_width, label_height) = text::measure(&label, self.render_scale);
            if self.height() >= label_height && bounds.left >= tick_length + text_gap + label_width
            {
                let label_x = tick_start - text_gap - label_width;
                let label_y = y
                    .saturating_sub(label_height / 2)
                    .min(self.height() - label_height);
                text::draw(
                    &mut self.image,
                    (label_x, label_y),
                    &label,
                    self.render_scale,
                    Pixel::BLACK,
                );
            }
        }
    }

    fn draw_labels(&mut self, bounds: Bounds) {
        let bottom_margin = self.height() - 1 - bounds.bottom;
        let tick_length = TICK_LENGTH * self.render_scale;
        let text_gap = TEXT_GAP * self.render_scale;
        let glyph_height = GLYPH_HEIGHT * self.render_scale;

        if !self.axes.title.is_empty() && bounds.top >= glyph_height + text_gap {
            let (width, _) = text::measure(&self.axes.title, self.render_scale);
            let x = centered_start((bounds.left + bounds.right) / 2, width, self.width());
            text::draw(
                &mut self.image,
                (x, bounds.top - glyph_height - text_gap),
                &self.axes.title,
                self.render_scale,
                Pixel::BLACK,
            );
        }

        let xlabel_y = bounds.bottom + tick_length + text_gap + glyph_height + text_gap;
        if !self.axes.xlabel.is_empty() && bottom_margin >= xlabel_y - bounds.bottom + glyph_height
        {
            let (width, _) = text::measure(&self.axes.xlabel, self.render_scale);
            let x = centered_start((bounds.left + bounds.right) / 2, width, self.width());
            text::draw(
                &mut self.image,
                (x, xlabel_y),
                &self.axes.xlabel,
                self.render_scale,
                Pixel::BLACK,
            );
        }

        let max_tick_width = self
            .axes
            .yticks
            .iter()
            .map(|&tick| text::measure(&format_tick(tick), self.render_scale).0)
            .max()
            .unwrap_or(0);
        let required_left = tick_length + text_gap + max_tick_width + text_gap + glyph_height;
        if !self.axes.ylabel.is_empty() && bounds.left >= required_left {
            let (rotated_height, _) = text::measure(&self.axes.ylabel, self.render_scale);
            let x = bounds.left - required_left;
            let y = centered_start(
                (bounds.top + bounds.bottom) / 2,
                rotated_height,
                self.height(),
            );
            text::draw_rotated_counterclockwise(
                &mut self.image,
                (x, y),
                &self.axes.ylabel,
                self.render_scale,
                Pixel::BLACK,
            );
        }
    }
}

fn map_x(value: f64, limits: (f64, f64), bounds: Bounds) -> Option<usize> {
    map_value(value, limits, bounds.width()).map(|pixel| bounds.left + pixel)
}

fn map_y(value: f64, limits: (f64, f64), bounds: Bounds) -> Option<usize> {
    map_value(value, limits, bounds.height()).map(|pixel| bounds.bottom - pixel)
}

fn map_value(value: f64, limits: (f64, f64), size: usize) -> Option<usize> {
    if limits.1 <= limits.0 || size == 0 || value < limits.0 || value > limits.1 {
        return None;
    }
    let scaled = (value - limits.0) / (limits.1 - limits.0);
    Some((scaled * (size - 1) as f64).round() as usize)
}

fn contains_zero(limits: (f64, f64)) -> bool {
    limits.0 <= 0.0 && 0.0 <= limits.1
}

fn centered_start(center: usize, length: usize, extent: usize) -> usize {
    center
        .saturating_sub(length / 2)
        .min(extent.saturating_sub(length))
}

fn format_tick(value: f64) -> String {
    if value == 0.0 {
        return "0".to_owned();
    }

    let magnitude = value.abs();
    if !(0.001..10_000.0).contains(&magnitude) {
        return format!("{value:.1e}");
    }

    format!("{value:.2}")
        .trim_end_matches('0')
        .trim_end_matches('.')
        .to_owned()
}

fn data_limits<'a>(
    points: impl Iterator<Item = &'a (f64, f64)>,
) -> Option<((f64, f64), (f64, f64))> {
    let mut limits: Option<((f64, f64), (f64, f64))> = None;

    for &(x, y) in points.filter(|(x, y)| x.is_finite() && y.is_finite()) {
        limits = Some(match limits {
            None => ((x, x), (y, y)),
            Some(((xmin, xmax), (ymin, ymax))) => {
                ((xmin.min(x), xmax.max(x)), (ymin.min(y), ymax.max(y)))
            }
        });
    }

    limits.map(|(xlim, ylim)| (expand_degenerate(xlim), expand_degenerate(ylim)))
}

fn expand_degenerate(limits: (f64, f64)) -> (f64, f64) {
    if limits.0 != limits.1 {
        return limits;
    }

    let padding = (limits.0.abs() * 0.05).max(1.0);
    (limits.0 - padding, limits.1 + padding)
}

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

    #[test]
    fn coordinates_start_at_bottom_left() {
        let bounds = Bounds::from_insets(7, 7, 0, 0).unwrap();
        assert_eq!(map_x(-2.0, (-2.0, 4.0), bounds), Some(0));
        assert_eq!(map_x(4.0, (-2.0, 4.0), bounds), Some(6));
        assert_eq!(map_y(-2.0, (-2.0, 16.0), bounds), Some(6));
        assert_eq!(map_y(16.0, (-2.0, 16.0), bounds), Some(0));
    }

    #[test]
    fn canvas_respects_horizontal_and_vertical_insets() {
        let mut canvas =
            Canvas::with_inset(11, 11, Axes::from_limits((0.0, 1.0), (0.0, 1.0)), 2, 3);
        canvas.render();

        assert_eq!(canvas.image.get_pixel(0, 0), Some(&Pixel::WHITE));
        assert_eq!(canvas.image.get_pixel(2, 3), Some(&Pixel::BLACK));
        assert_eq!(canvas.image.get_pixel(8, 7), Some(&Pixel::BLACK));
    }

    #[test]
    fn axes_render_origin() {
        let mut canvas = Canvas::new(7, 7, Axes::from_limits((0.0, 2.0), (0.0, 2.0)));
        canvas.render();

        assert_eq!(canvas.image.get_pixel(0, 3), Some(&Pixel::BLACK));
        assert_eq!(canvas.image.get_pixel(3, 6), Some(&Pixel::BLACK));
    }

    #[test]
    fn labels_and_tick_values_render_inside_margins() {
        let axes = Axes::from_limits((0.0, 1.0), (0.0, 1.0))
            .with_labels("time", "value")
            .with_title("demo");
        let mut canvas = Canvas::with_inset(128, 128, axes, 40, 32);
        canvas.render();

        let width = canvas.width();
        let margin_has_text = canvas
            .image
            .pixels()
            .iter()
            .enumerate()
            .any(|(index, pixel)| {
                let x = index % width;
                let y = index / width;
                *pixel == Pixel::BLACK && (x < 36 || !(28..=100).contains(&y))
            });
        assert!(margin_has_text);
    }

    #[test]
    fn tiny_canvas_does_not_panic_while_laying_out_ticks() {
        let mut canvas =
            Canvas::with_inset(64, 1, Axes::from_limits((0.0, 1.0), (0.0, 1.0)), 20, 0);

        canvas.render();
        assert_eq!(canvas.height(), 1);
    }

    #[test]
    fn axes_automatically_fit_all_added_series() {
        let mut canvas = Canvas::new(11, 11, Axes::new());
        canvas.render();
        canvas.scatter(&[(0.0, 0.0)], 1);
        canvas.scatter(&[(3.0, 4.0)], 1);

        assert_eq!(canvas.axes.xlim, (0.0, 3.0));
        assert_eq!(canvas.axes.ylim, (0.0, 4.0));
        assert_eq!(canvas.image.get_pixel(0, 10), Some(&Pixel::rgb(255, 0, 0)));
        assert_eq!(canvas.image.get_pixel(10, 0), Some(&Pixel::rgb(255, 0, 0)));
    }

    #[test]
    fn manual_limits_override_automatic_limits_per_axis() {
        let axes = Axes::new().with_xlim((-10.0, 10.0));
        let mut canvas = Canvas::new(11, 11, axes);
        canvas.scatter(&[(2.0, 3.0), (4.0, 7.0)], 1);

        assert_eq!(canvas.axes.xlim, (-10.0, 10.0));
        assert_eq!(canvas.axes.ylim, (3.0, 7.0));
    }
}