pav_regression 0.7.0

use crate::coordinate::Coordinate;
use crate::point::{interpolate_two_points, Point};
use eytzinger_interpolation::SliceExt;
use serde::Serialize;
use std::fmt::{Display, Formatter};
use thiserror::Error;

/// Errors that can occur during isotonic regression
#[derive(Error, Debug)]
#[allow(dead_code)]
pub enum IsotonicRegressionError {
    /// Error when a negative point is encountered with intersect_origin set to true
    #[error("With intersect_origin = true, all points must be >= 0 on both x and y axes")]
    NegativePointWithIntersectOrigin,
}

/// A vector of points forming an isotonic regression, along with the
/// centroid point of the original set.
#[derive(Debug, Clone, Serialize)]
pub struct IsotonicRegression<T: Coordinate> {
    direction: Direction,
    points: Vec<Point<T>>,
    centroid_point: Centroid<T>,
    intersect_origin: bool,
}

#[derive(Debug, Clone, PartialEq, Serialize)]
struct Centroid<T: Coordinate> {
    sum_x: T,
    sum_y: T,
    sum_weight: f64,
}

#[derive(Debug, Clone, Serialize, PartialEq)]
#[allow(dead_code)]
/// Specifies the direction of the isotonic regression.
pub enum Direction {
    /// Indicates an ascending (non-decreasing) regression.
    Ascending,
    /// Indicates a descending (non-increasing) regression.
    Descending,
}

impl<T: Coordinate + Display> Display for IsotonicRegression<T> {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        writeln!(f, "IsotonicRegression {{")?;
        writeln!(f, "\tdirection: {:?},", self.direction)?;
        writeln!(f, "\tpoints:")?;
        for point in &self.get_points_sorted() {
            writeln!(
                f,
                "\t\t{}\t{:.2}\t{:.2}",
                point.x(),
                point.y(),
                point.weight()
            )?;
        }
        writeln!(f, "\tcentroid_point:")?;
        writeln!(
            f,
            "\t\t{}\t{:.2}\t{:.2}",
            self.centroid_point.sum_x, self.centroid_point.sum_y, self.centroid_point.sum_weight
        )?;
        write!(f, "}}")
    }
}

#[allow(dead_code)]
impl<T: Coordinate> IsotonicRegression<T> {
    /// Find an ascending isotonic regression from a set of points.
    ///
    /// # Examples
    ///
    /// ```
    /// use pav_regression::{Point, IsotonicRegression};
    ///
    /// let points = vec![
    ///     Point::new(0.0, 1.0),
    ///     Point::new(1.0, 2.0),
    ///     Point::new(2.0, 1.5),
    ///     Point::new(3.0, 3.0),
    /// ];
    /// let regression = IsotonicRegression::new_ascending(&points).unwrap();
    /// assert_eq!(regression.get_points().len(), 3);
    /// ```
    pub fn new_ascending(
        points: &[Point<T>],
    ) -> Result<IsotonicRegression<T>, IsotonicRegressionError> {
        IsotonicRegression::new(points, Direction::Ascending, false)
    }

    /// Find a descending isotonic regression from a set of points.
    ///
    /// # Examples
    ///
    /// ```
    /// use pav_regression::{Point, IsotonicRegression};
    ///
    /// let points = vec![
    ///     Point::new(0.0, 3.0),
    ///     Point::new(1.0, 2.0),
    ///     Point::new(2.0, 2.5),
    ///     Point::new(3.0, 1.0),
    /// ];
    /// let regression = IsotonicRegression::new_descending(&points).unwrap();
    /// assert_eq!(regression.get_points().len(), 3);
    /// ```
    pub fn new_descending(
        points: &[Point<T>],
    ) -> Result<IsotonicRegression<T>, IsotonicRegressionError> {
        IsotonicRegression::new(points, Direction::Descending, false)
    }

    /// Find an isotonic regression in the specified direction.
    ///
    /// If `intersect_origin` is true, the regression will intersect the origin (0,0) and all points must be >= 0 on both axes.
    ///
    /// # Examples
    ///
    /// ```
    /// use pav_regression::{Point, IsotonicRegression};
    /// use pav_regression::isotonic_regression::Direction;
    ///
    /// let points = vec![
    ///     Point::new(0.0, 1.0),
    ///     Point::new(1.0, 2.0),
    ///     Point::new(2.0, 1.5),
    ///     Point::new(3.0, 3.0),
    /// ];
    /// let regression = IsotonicRegression::new(&points, Direction::Ascending, false).unwrap();
    /// assert_eq!(regression.get_points().len(), 3);
    /// ```
    pub fn new(
        points: &[Point<T>],
        direction: Direction,
        intersect_origin: bool,
    ) -> Result<IsotonicRegression<T>, IsotonicRegressionError> {
        let (sum_x, sum_y, sum_weight) =
            points
                .iter()
                .try_fold((T::zero(), T::zero(), 0.0), |(sx, sy, sw), point| {
                    if intersect_origin
                        && (point.x().is_sign_negative() || point.y().is_sign_negative())
                    {
                        Err(IsotonicRegressionError::NegativePointWithIntersectOrigin)
                    } else {
                        Ok((
                            sx + *point.x() * T::from_float(point.weight()),
                            sy + *point.y() * T::from_float(point.weight()),
                            sw + point.weight(),
                        ))
                    }
                })?;

        let mut isotonic_points = isotonic(points, direction.clone());
        isotonic_points.eytzingerize(&mut eytzinger_interpolation::permutation::InplacePermutator);

        Ok(IsotonicRegression {
            direction,
            points: isotonic_points,
            centroid_point: Centroid {
                sum_x,
                sum_y,
                sum_weight,
            },
            intersect_origin,
        })
    }

    /// Find the _y_ point at position `at_x` or None if the regression is empty.
    ///
    /// # Examples
    ///
    /// ```
    /// use pav_regression::{Point, IsotonicRegression};
    ///
    /// let points = vec![
    ///     Point::new(0.0, 1.0),
    ///     Point::new(1.0, 2.0),
    ///     Point::new(2.0, 1.5),
    ///     Point::new(3.0, 3.0),
    /// ];
    /// let regression = IsotonicRegression::new_ascending(&points).unwrap();
    /// let interpolated_y = regression.interpolate(1.5).unwrap();
    /// assert_eq!(interpolated_y, 1.75);
    /// ```
    #[must_use]
    pub fn interpolate(&self, at_x: T) -> Option<T> {
        if self.points.is_empty() {
            return None;
        }

        let interpolation = if self.points.len() == 1 {
            *self.points[0].y()
        } else {
            let (lte, gt) = self
                .points
                .eytzinger_interpolative_search_by(|p| p.x().partial_cmp(&at_x).unwrap());

            match (lte, gt) {
                // Found exact match or need to interpolate between two points
                (Some(lower), Some(upper)) => {
                    interpolate_two_points(&self.points[lower], &self.points[upper], at_x)
                }
                // Requested point meets or exceeds the upper bound
                (Some(upper), None) => {
                    // at_x is beyond the last point - interpolate with centroid
                    interpolate_two_points(
                        &self.get_centroid_point().unwrap(),
                        &self.points[upper],
                        at_x,
                    )
                }
                // Requested point is below the lower bound
                (None, Some(lower)) => {
                    // at_x is before the first point
                    if self.intersect_origin {
                        interpolate_two_points(
                            &Point::new(T::zero(), T::zero()),
                            &self.points[lower],
                            at_x,
                        )
                    } else {
                        interpolate_two_points(
                            &self.points[lower],
                            &self.get_centroid_point().unwrap(),
                            at_x,
                        )
                    }
                }
                // Should never happen - only possible if the slice is empty, which we already checked
                (None, None) => {
                    debug_assert!(
                        false,
                        "Got None, None from eytzinger_interpolative_search_by on non-empty slice"
                    );
                    return None;
                }
            }
        };

        Some(interpolation)
    }

    /// Retrieve the points that make up the isotonic regression. The points are NOT sorted by x value - they are in eytzinger order.
    ///
    /// # Examples
    ///
    /// ```
    /// use pav_regression::{Point, IsotonicRegression};
    ///
    /// let points = vec![
    ///     Point::new(0.0, 1.0),
    ///     Point::new(1.0, 2.0),
    ///     Point::new(2.0, 1.5),
    ///     Point::new(3.0, 3.0),
    /// ];
    /// let regression = IsotonicRegression::new_ascending(&points).unwrap();
    /// assert_eq!(regression.get_points().len(), 3);
    /// ```
    pub fn get_points(&self) -> &[Point<T>] {
        &self.points
    }

    /// Retrieve the points that make up the isotonic regression, sorted by x value.
    pub fn get_points_sorted(&self) -> Vec<Point<T>> {
        let mut points = self.points.clone();
        points.sort_by(|a, b| a.x().partial_cmp(b.x()).unwrap());
        points
    }

    /// Retrieve the mean point of the original point set.
    ///
    /// # Examples
    ///
    /// ```
    /// use pav_regression::{Point, IsotonicRegression};
    ///
    /// let points = vec![
    ///     Point::new(0.0, 1.0),
    ///     Point::new(1.0, 2.0),
    ///     Point::new(2.0, 1.5),
    ///     Point::new(3.0, 3.0),
    /// ];
    /// let regression = IsotonicRegression::new_ascending(&points).unwrap();
    /// let centroid = regression.get_centroid_point().unwrap();
    /// assert_eq!(*centroid.x(), 1.5);
    /// assert_eq!(*centroid.y(), 1.875);
    /// ```
    pub fn get_centroid_point(&self) -> Option<Point<T>> {
        if self.centroid_point.sum_weight == 0.0 {
            None
        } else {
            Some(Point::new_with_weight(
                self.centroid_point.sum_x / T::from_float(self.centroid_point.sum_weight),
                self.centroid_point.sum_y / T::from_float(self.centroid_point.sum_weight),
                1.0,
            ))
        }
    }

    /// Add new points to the regression.
    ///
    /// # Examples
    ///
    /// ```
    /// use pav_regression::{Point, IsotonicRegression};
    ///
    /// let mut regression = IsotonicRegression::new_ascending(&[
    ///     Point::new(0.0, 1.0),
    ///     Point::new(2.0, 2.0),
    /// ]).unwrap();
    /// regression.add_points(&[Point::new(1.0, 1.5)]);
    /// assert_eq!(regression.get_points().len(), 3);
    /// ```
    pub fn add_points(&mut self, points: &[Point<T>]) {
        for point in points {
            assert!(
                !self.intersect_origin
                    || (!point.x().is_sign_negative() && !point.y().is_sign_negative()),
                "With intersect_origin = true, all points must be >= 0 on both x and y axes"
            );
            self.centroid_point.sum_x =
                self.centroid_point.sum_x + *point.x() * T::from_float(point.weight());
            self.centroid_point.sum_y =
                self.centroid_point.sum_y + *point.y() * T::from_float(point.weight());
            self.centroid_point.sum_weight = self.centroid_point.sum_weight + point.weight();
        }

        let mut new_points = self.points.clone();
        new_points.extend_from_slice(points);
        self.points = isotonic(&new_points, self.direction.clone());
        self.points
            .eytzingerize(&mut eytzinger_interpolation::permutation::InplacePermutator);
    }

    /// Remove points from the regression.
    ///
    /// Because PAV merges adjacent points that violate monotonicity, the
    /// original input points may no longer exist verbatim in the internal
    /// point list. This method finds the closest aggregate point (by
    /// x-coordinate) and subtracts the removed point's influence from it.
    /// If the aggregate point's weight drops to zero or below, it is
    /// removed entirely. The remaining points are then re-run through PAV
    /// and eytzingerized.
    ///
    /// # Examples
    ///
    /// ```
    /// use pav_regression::{Point, IsotonicRegression};
    ///
    /// let mut regression = IsotonicRegression::new_ascending(&[
    ///     Point::new(0.0, 1.0),
    ///     Point::new(1.0, 2.0),
    ///     Point::new(2.0, 3.0),
    /// ]).unwrap();
    /// regression.remove_points(&[Point::new(1.0, 2.0)]);
    /// assert_eq!(regression.get_points().len(), 2);
    /// ```
    pub fn remove_points(&mut self, points: &[Point<T>]) {
        // Work with a sorted copy for closest-point lookup
        let mut working_points = self.points.clone();
        working_points.sort_by(|a, b| a.x().partial_cmp(b.x()).unwrap());

        for point in points {
            assert!(
                !self.intersect_origin
                    || (!point.x().is_sign_negative() && !point.y().is_sign_negative()),
                "With intersect_origin = true, all points must be >= 0 on both x and y axes"
            );

            // Update centroid
            self.centroid_point.sum_x =
                self.centroid_point.sum_x - *point.x() * T::from_float(point.weight());
            self.centroid_point.sum_y =
                self.centroid_point.sum_y - *point.y() * T::from_float(point.weight());
            self.centroid_point.sum_weight = self.centroid_point.sum_weight - point.weight();

            if working_points.is_empty() {
                continue;
            }

            // Find the closest aggregate point by x-coordinate distance
            let closest_idx = working_points
                .iter()
                .enumerate()
                .min_by(|(_, a), (_, b)| {
                    a.x()
                        .abs_diff(point.x())
                        .partial_cmp(&b.x().abs_diff(point.x()))
                        .unwrap()
                })
                .map(|(i, _)| i)
                .unwrap();

            let agg = &working_points[closest_idx];
            let new_weight = agg.weight() - point.weight();

            const WEIGHT_EPSILON: f64 = 1e-10;

            if new_weight <= WEIGHT_EPSILON {
                // Aggregate fully consumed — remove it
                working_points.remove(closest_idx);
            } else {
                // Subtract the removed point's influence from the aggregate
                let agg_w = T::from_float(agg.weight());
                let rm_w = T::from_float(point.weight());
                let new_w = T::from_float(new_weight);
                let new_x = (*agg.x() * agg_w - *point.x() * rm_w) / new_w;
                let new_y = (*agg.y() * agg_w - *point.y() * rm_w) / new_w;
                working_points[closest_idx] = Point::new_with_weight(new_x, new_y, new_weight);
            }
        }

        // Re-run PAV and eytzingerize
        self.points = isotonic(&working_points, self.direction.clone());
        self.points
            .eytzingerize(&mut eytzinger_interpolation::permutation::InplacePermutator);
    }

    /// Returns the number of points in the regression.
    ///
    /// # Examples
    ///
    /// ```
    /// use pav_regression::{Point, IsotonicRegression};
    ///
    /// let points = vec![
    ///     Point::new(0.0, 1.0),
    ///     Point::new(1.0, 2.0),
    ///     Point::new(2.0, 1.5),
    ///     Point::new(3.0, 3.0),
    /// ];
    /// let regression = IsotonicRegression::new_ascending(&points).unwrap();
    /// assert_eq!(regression.len(), 4);
    /// ```
    pub fn len(&self) -> usize {
        self.centroid_point.sum_weight.round() as usize
    }

    /// Checks if the regression is empty.
    ///
    /// # Examples
    ///
    /// ```
    /// use pav_regression::IsotonicRegression;
    ///
    /// let regression: IsotonicRegression<f64> = IsotonicRegression::new_ascending(&[]).unwrap();
    /// assert!(regression.is_empty());
    /// ```
    #[must_use]
    pub fn is_empty(&self) -> bool {
        self.centroid_point.sum_weight == 0.0
    }
}

#[allow(dead_code)]
fn isotonic<T: Coordinate>(points: &[Point<T>], direction: Direction) -> Vec<Point<T>> {
    let mut merged_points: Vec<Point<T>> = match direction {
        Direction::Ascending => points.to_vec(),
        Direction::Descending => points
            .iter()
            .map(|p| Point::new_with_weight(*p.x(), *p.y(), p.weight()))
            .collect(),
    };

    // Sort the points by x, and if x is equal, sort by y descending to ensure that points with the same x
    // get merged.
    merged_points.sort_by(|a, b| {
        a.x()
            .partial_cmp(b.x())
            .unwrap_or(std::cmp::Ordering::Equal)
            .then(
                b.y()
                    .partial_cmp(a.y())
                    .unwrap_or(std::cmp::Ordering::Equal),
            )
    });

    let iso_points =
        merged_points
            .into_iter()
            .fold(Vec::new(), |mut acc: Vec<Point<T>>, mut point| {
                while let Some(last) = acc.last() {
                    if (direction == Direction::Ascending && last.y() >= point.y())
                        || (direction == Direction::Descending && last.y() <= point.y())
                    {
                        point.merge_with(&acc.pop().unwrap());
                    } else {
                        break;
                    }
                }
                acc.push(point);
                acc
            });

    match direction {
        Direction::Ascending => iso_points,
        Direction::Descending => iso_points,
    }
}

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

    #[test]
    fn test_ascending_regression() {
        let points = &[
            Point::new(0.0, 1.0),
            Point::new(1.0, 2.0),
            Point::new(2.0, 1.5),
            Point::new(3.0, 3.0),
        ];

        let regression = IsotonicRegression::new_ascending(points).unwrap();
        assert_eq!(regression.get_points_sorted().len(), 3);
        assert_eq!(*regression.get_points_sorted()[0].y(), 1.0);
        assert_eq!(*regression.get_points_sorted()[1].y(), 1.75);
        assert_eq!(*regression.get_points_sorted()[2].y(), 3.0);
    }

    #[test]
    fn test_descending_regression() {
        let points = &[
            Point::new(0.0, 3.0),
            Point::new(1.0, 2.0),
            Point::new(2.0, 2.5),
            Point::new(3.0, 1.0),
        ];

        let regression = IsotonicRegression::new_descending(points).unwrap();
        assert_eq!(regression.get_points_sorted().len(), 3);
        assert_eq!(*regression.get_points_sorted()[0].y(), 3.0);
        assert_eq!(*regression.get_points_sorted()[1].y(), 2.25);
        assert_eq!(*regression.get_points_sorted()[2].y(), 1.0);
    }

    #[test]
    fn test_add_points() {
        let mut regression =
            IsotonicRegression::new_ascending(&[Point::new(0.0, 1.0), Point::new(2.0, 2.0)])
                .unwrap();
        regression.add_points(&[Point::new(1.0, 1.5)]);
        assert_eq!(regression.get_points_sorted().len(), 3);
        assert_eq!(*regression.get_points_sorted()[1].x(), 1.0);
        assert_eq!(*regression.get_points_sorted()[1].y(), 1.5);
    }

    #[test]
    fn test_remove_points() {
        let mut regression = IsotonicRegression::new_ascending(&[
            Point::new(0.0, 1.0),
            Point::new(1.0, 2.0),
            Point::new(2.0, 3.0),
        ])
        .unwrap();
        regression.remove_points(&[Point::new(1.0, 2.0)]);
        assert_eq!(regression.get_points_sorted().len(), 2);
        assert_eq!(*regression.get_points_sorted()[0].x(), 0.0);
        assert_eq!(*regression.get_points_sorted()[1].x(), 2.0);
    }

    #[test]
    fn test_centroid_point() {
        let points = &[
            Point::new(0.0, 1.0),
            Point::new(1.0, 2.0),
            Point::new(2.0, 3.0),
        ];
        let regression = IsotonicRegression::new_ascending(points).unwrap();
        let centroid = regression.get_centroid_point().unwrap();
        assert_eq!(*centroid.x(), 1.0);
        assert_eq!(*centroid.y(), 2.0);
    }

    #[test]
    fn test_empty_regression() {
        let regression: IsotonicRegression<f64> = IsotonicRegression::new_ascending(&[]).unwrap();
        assert!(regression.is_empty());
        assert_eq!(regression.len(), 0);
        assert!(regression.interpolate(1.0).is_none());
    }

    #[test]
    fn test_remove_merged_point() {
        // Points (1.0, 2.0) and (2.0, 1.5) violate ascending order
        // and get merged into a single aggregate point.
        let points = vec![
            Point::new(0.0, 1.0),
            Point::new(1.0, 2.0),
            Point::new(2.0, 1.5),
            Point::new(3.0, 3.0),
        ];
        let mut regression = IsotonicRegression::new_ascending(&points).unwrap();
        assert_eq!(regression.get_points_sorted().len(), 3);

        // Remove (2.0, 1.5) which was merged — should subtract its
        // influence from the aggregate and the regression should adapt.
        regression.remove_points(&[Point::new(2.0, 1.5)]);

        // After removal we should still have a valid regression
        let sorted = regression.get_points_sorted();
        assert!(!sorted.is_empty());

        // The regression should still be monotonically non-decreasing
        for w in sorted.windows(2) {
            assert!(
                w[0].y() <= w[1].y(),
                "Regression not ascending: {:?} > {:?}",
                w[0].y(),
                w[1].y()
            );
        }
    }

    #[test]
    fn test_add_then_remove_returns_to_original() {
        let original_points = vec![
            Point::new(0.0, 1.0),
            Point::new(1.0, 2.0),
            Point::new(2.0, 3.0),
        ];
        let original = IsotonicRegression::new_ascending(&original_points).unwrap();

        let mut regression = IsotonicRegression::new_ascending(&original_points).unwrap();

        // Add some extra points
        let extra = vec![Point::new(0.5, 1.8), Point::new(1.5, 2.2)];
        regression.add_points(&extra);

        // Remove them
        regression.remove_points(&extra);

        // The centroid should be approximately the same
        let orig_centroid = original.get_centroid_point().unwrap();
        let new_centroid = regression.get_centroid_point().unwrap();
        assert!(
            (orig_centroid.x() - new_centroid.x()).abs() < 1e-9,
            "Centroid x mismatch"
        );
        assert!(
            (orig_centroid.y() - new_centroid.y()).abs() < 1e-9,
            "Centroid y mismatch"
        );

        // Interpolated values should be approximately the same
        for x in [0.0, 0.5, 1.0, 1.5, 2.0] {
            let orig_y = original.interpolate(x).unwrap();
            let new_y = regression.interpolate(x).unwrap();
            assert!(
                (orig_y - new_y).abs() < 0.3,
                "At x={x}: original y={orig_y}, after add/remove y={new_y}"
            );
        }
    }

    #[test]
    fn test_remove_reduces_aggregate_to_zero() {
        // Create a regression where two points merge
        let points = vec![
            Point::new(0.0, 2.0),
            Point::new(1.0, 1.0), // These two violate ascending
        ];
        let mut regression = IsotonicRegression::new_ascending(&points).unwrap();
        // Both points get merged into one aggregate (weight 2)
        assert_eq!(regression.get_points_sorted().len(), 1);

        // Remove both original points — aggregate weight goes to zero
        regression.remove_points(&[Point::new(0.0, 2.0), Point::new(1.0, 1.0)]);
        assert!(
            regression.get_points_sorted().is_empty(),
            "Expected empty regression after removing all points"
        );
    }
}