ggplot_rs/scale/

continuous.rs

use crate::aes::Aesthetic;
use crate::data::Value;

use super::format::LabelFormatter;
use super::sec_axis::SecAxis;
use super::transform::ScaleTransform;
use super::util::{extended_breaks, format_number};
use super::Scale;

/// Continuous linear scale.
#[derive(Clone)]
pub struct ScaleContinuous {
    aesthetic: Aesthetic,
    name: String,
    min: f64,
    max: f64,
    trained: bool,
    filter_oob: bool,
    // Per-side expansion: (mult_lower, add_lower, mult_upper, add_upper).
    expand: (f64, f64, f64, f64),
    // Axis position: false = default (bottom for x, left for y), true = opposite
    // side (top for x, right for y).
    position_opposite: bool,
    pub(crate) scale_transform: ScaleTransform,
    custom_breaks: Option<Vec<f64>>,
    custom_labels: Option<Vec<String>>,
    pub(crate) sec_axis: Option<SecAxis>,
    label_formatter: Option<LabelFormatter>,
}

impl ScaleContinuous {
    pub fn new() -> Self {
        ScaleContinuous {
            aesthetic: Aesthetic::X,
            name: String::new(),
            min: f64::INFINITY,
            max: f64::NEG_INFINITY,
            trained: false,
            filter_oob: false,
            expand: (0.05, 0.0, 0.05, 0.0),
            position_opposite: false,
            scale_transform: ScaleTransform::Identity,
            custom_breaks: None,
            custom_labels: None,
            sec_axis: None,
            label_formatter: None,
        }
    }

    pub fn for_aesthetic(mut self, aes: Aesthetic) -> Self {
        self.aesthetic = aes;
        self
    }

    pub fn with_name(mut self, name: &str) -> Self {
        self.name = name.to_string();
        self
    }

    pub fn with_limits(mut self, min: f64, max: f64) -> Self {
        self.min = min;
        self.max = max;
        self.trained = true;
        self.filter_oob = true;
        self
    }

    pub fn with_transform(mut self, transform: ScaleTransform) -> Self {
        self.scale_transform = transform;
        self
    }

    /// Set custom break positions (data values where ticks appear).
    pub fn with_breaks(mut self, breaks: Vec<f64>) -> Self {
        self.custom_breaks = Some(breaks);
        self
    }

    /// Set custom labels for breaks. Must match the number of breaks.
    pub fn with_labels(mut self, labels: Vec<String>) -> Self {
        self.custom_labels = Some(labels);
        self
    }

    /// Set the expansion multiplier and additive constant.
    /// Like R's `expand = c(mult, add)`, applied symmetrically. Default `(0.05, 0.0)`.
    pub fn with_expand(mut self, mult: f64, add: f64) -> Self {
        self.expand = (mult, add, mult, add);
        self
    }

    /// Per-side expansion (R's `expansion(mult = c(l, u), add = c(l, u))`):
    /// separate multiplicative/additive expansion for the lower and upper ends.
    pub fn with_expand_sides(
        mut self,
        mult_lower: f64,
        add_lower: f64,
        mult_upper: f64,
        add_upper: f64,
    ) -> Self {
        self.expand = (mult_lower, add_lower, mult_upper, add_upper);
        self
    }

    /// Place this axis on the opposite side (x → top, y → right).
    pub fn with_position_opposite(mut self) -> Self {
        self.position_opposite = true;
        self
    }

    /// Set a label formatter. Accepts a plain `fn` (e.g. `label_comma`) or a
    /// configurable formatter such as `label_si()` / `label_number(...)`.
    pub fn with_label_formatter<F>(mut self, f: F) -> Self
    where
        F: Fn(f64) -> String + Send + Sync + 'static,
    {
        self.label_formatter = Some(std::sync::Arc::new(f));
        self
    }

    /// Add a secondary axis with a transformation function.
    pub fn with_sec_axis(mut self, sec: SecAxis) -> Self {
        self.sec_axis = Some(sec);
        self
    }

    /// Get the secondary axis, if any.
    pub fn sec_axis(&self) -> Option<&SecAxis> {
        self.sec_axis.as_ref()
    }

    fn format_label(&self, v: f64) -> String {
        if let Some(f) = &self.label_formatter {
            f(v)
        } else {
            format_number(v)
        }
    }

    fn expanded_range(&self) -> (f64, f64) {
        let range = self.max - self.min;
        let (ml, al, mu, au) = self.expand;
        (self.min - range * ml - al, self.max + range * mu + au)
    }

    /// Whether this axis should be drawn on the opposite side (x → top, y → right).
    pub fn is_opposite(&self) -> bool {
        self.position_opposite
    }
}

impl Default for ScaleContinuous {
    fn default() -> Self {
        Self::new()
    }
}

impl Scale for ScaleContinuous {
    fn aesthetic(&self) -> Aesthetic {
        self.aesthetic.clone()
    }

    fn train(&mut self, values: &[Value]) {
        for v in values {
            if let Some(f) = v.as_f64() {
                if f.is_finite() {
                    if f < self.min {
                        self.min = f;
                    }
                    if f > self.max {
                        self.max = f;
                    }
                }
            }
        }
        self.trained = true;
    }

    fn map(&self, value: &Value) -> f64 {
        let f = match value.as_f64() {
            Some(f) => f,
            None => return 0.0,
        };
        // ±Inf means "panel edge" (ggplot2): -Inf → lower edge, +Inf → upper.
        if f == f64::INFINITY {
            return 1.0;
        }
        if f == f64::NEG_INFINITY {
            return 0.0;
        }
        // An untrained scale (empty data) has no domain — centre everything
        // rather than producing NaN.
        if !self.trained || !self.min.is_finite() || !self.max.is_finite() {
            return 0.5;
        }
        let (emin, emax) = self.expanded_range();
        let range = emax - emin;
        if range.abs() < f64::EPSILON {
            0.5
        } else {
            (f - emin) / range
        }
    }

    fn expanded_domain(&self) -> Option<(f64, f64)> {
        if !self.trained || !self.min.is_finite() || !self.max.is_finite() {
            return None;
        }
        let (a, b) = self.expanded_range();
        if !(a.is_finite() && b.is_finite()) {
            return None;
        }
        if (b - a).abs() < f64::EPSILON {
            // A degenerate domain maps everything to the panel centre.
            return Some((a - 0.5, b + 0.5));
        }
        Some((a, b))
    }

    fn breaks(&self) -> Vec<(f64, String)> {
        if !self.trained || self.min > self.max {
            return vec![];
        }

        // Use custom breaks if provided
        if let Some(ref custom) = self.custom_breaks {
            return custom
                .iter()
                .enumerate()
                .map(|(i, &v)| {
                    let pos = self.map(&Value::Float(v));
                    let label = if let Some(ref labels) = self.custom_labels {
                        labels
                            .get(i)
                            .cloned()
                            .unwrap_or_else(|| self.format_label(v))
                    } else {
                        self.format_label(self.scale_transform.inverse(v))
                    };
                    (pos, label)
                })
                .collect();
        }

        if super::util::is_degenerate_range(self.min, self.max) {
            let label = self.format_label(self.scale_transform.inverse(self.min));
            return vec![(0.5, label)];
        }

        // Extended-Wilkinson breaks over the data range (matching ggplot2's
        // scales::extended_breaks), keeping those within the expanded (visible)
        // range. Labels show the original (inverse-transformed) value.
        let (emin, emax) = self.expanded_range();
        let tol = (emax - emin).abs() * 1e-9;
        let breaks: Vec<f64> = extended_breaks(self.min, self.max, 5)
            .into_iter()
            .filter(|&v| v >= emin - tol && v <= emax + tol)
            .collect();
        let shown: Vec<f64> = breaks
            .iter()
            .map(|&v| self.scale_transform.inverse(v))
            .collect();
        // Without a custom formatter, label the breaks jointly so neighbours
        // stay distinguishable on narrow ranges.
        let labels: Vec<String> = match &self.label_formatter {
            Some(f) => shown.iter().map(|&v| f(v)).collect(),
            None => super::util::format_numbers(&shown),
        };
        breaks
            .iter()
            .zip(labels)
            .map(|(&v, label)| (self.map(&Value::Float(v)), label))
            .collect()
    }

    fn name(&self) -> &str {
        &self.name
    }

    fn set_name(&mut self, name: &str) {
        self.name = name.to_string();
    }

    fn transform(&self, value: &Value) -> Value {
        self.scale_transform.transform_value(value)
    }

    fn sec_axis(&self) -> Option<&SecAxis> {
        self.sec_axis.as_ref()
    }

    fn set_limits(&mut self, min: f64, max: f64) {
        self.min = min;
        self.max = max;
        self.trained = true;
    }

    fn filter_limits(&self) -> Option<(f64, f64)> {
        if self.filter_oob && self.trained {
            Some((self.min, self.max))
        } else {
            None
        }
    }

    fn domain(&self) -> Option<(f64, f64)> {
        if self.trained {
            Some((self.min, self.max))
        } else {
            None
        }
    }

    fn axis_position_opposite(&self) -> bool {
        self.position_opposite
    }

    fn clone_box(&self) -> Box<dyn Scale> {
        Box::new(self.clone())
    }

    fn reset_training(&mut self) {
        self.min = f64::INFINITY;
        self.max = f64::NEG_INFINITY;
        self.trained = false;
    }
}