rialo-aggregators-utils 0.2.0

// Copyright (c) Subzero Labs, Inc.
// SPDX-License-Identifier: Apache-2.0

use base64::DecodeError;
use rialo_events_core::RialoEvent;
use thiserror::Error;

pub const PRICE_KEY: &str = "price";

/// Defines different methods for calculating the median when the number of elements is even
#[derive(Clone, Copy, Debug, Default, PartialEq)]
pub enum MedianMethod {
    /// Average the two middle values (mathematical median): (a + b) / 2
    Average,
    /// Take the lower of the two middle values
    Lower,
    /// Take the upper of the two middle values (also equivalent to NearestRank)
    #[default]
    Upper,
}

/// Represents errors that can occur when extracting values from oracle updates
#[derive(Debug, Error)]
pub enum ValueExtractionError {
    /// Error decoding Base64 data
    #[error("Base64 decode error: {0}")]
    DecodeError(#[from] DecodeError),
    /// Error deserializing JSON data
    #[error("JSON parse error: {0}")]
    JsonError(#[from] serde_json::Error),
    /// Missing key in JSON data
    #[error("Missing key in JSON data: {0}")]
    MissingKey(String),
    /// Value is not a number
    #[error("Value is not a number: {0}")]
    NotANumber(String),
}

/// Result of median calculation including potential errors
#[derive(Debug)]
pub struct MedianResult {
    /// The calculated median
    pub median: Median,
    /// Errors encountered during value extraction, mapped to their indices
    pub errors: Vec<(usize, ValueExtractionError)>,
}

impl MedianResult {
    /// Get errors from the result
    ///
    /// # Returns
    ///
    /// A slice of tuples containing the index and the error
    pub fn errors(&self) -> &[(usize, ValueExtractionError)] {
        &self.errors
    }
}

// TODO event regarding aggregation method should be defined for user program and voted through by governance
#[derive(Clone, Debug, Default, PartialEq, RialoEvent)]
pub struct Median {
    pub median: f64,
    pub sample_size: usize,
}

impl Median {
    /// Creates a new Median with the specified method for even-sized datasets,
    /// collecting errors encountered during processing
    pub fn from_with_method(mut values: Vec<f64>, method: MedianMethod) -> Self {
        let sample_size = values.len();

        // Calculate median or default to 0 if empty
        if sample_size == 0 {
            return Self {
                median: 0.0,
                sample_size: 0,
            };
        }

        // NaN values are filtered out in extract_value and returned as errors.
        // This comparison function is defensive programming in case any NaNs slip through.
        //
        // Standard floating-point comparison using partial_cmp can return None when comparing with NaN,
        // but sorting algorithms require a total ordering (must always return Some(Ordering)).
        // This function provides a total ordering by treating NaN values specially.
        let safe_cmp = |a: &f64, b: &f64| {
            match (a.is_nan(), b.is_nan()) {
                // Normal case: both values are valid numbers, use standard comparison
                (false, false) => a.partial_cmp(b).unwrap(),

                // If 'a' is NaN but 'b' is not, put NaN at the end (consider 'a' > 'b')
                // This ensures NaNs don't interfere with finding valid median values
                (true, false) => std::cmp::Ordering::Greater,

                // If 'b' is NaN but 'a' is not, put NaN at the end (consider 'a' < 'b')
                (false, true) => std::cmp::Ordering::Less,

                // If both are NaN, consider them equal for sorting stability
                // The order between NaNs doesn't matter for median calculation
                (true, true) => std::cmp::Ordering::Equal,
            }
        };

        // Calculate median based on the method
        #[allow(clippy::manual_is_multiple_of)]
        let median_value = if sample_size % 2 == 0 {
            // Even number of elements
            match method {
                MedianMethod::Average => {
                    // Sort once and take the two middle elements
                    values.sort_by(safe_cmp);
                    let mid_lower = sample_size / 2 - 1;
                    let mid_upper = sample_size / 2;
                    (values[mid_lower] + values[mid_upper]) / 2.0
                }
                MedianMethod::Lower => {
                    let mid_lower = sample_size / 2 - 1;
                    values.select_nth_unstable_by(mid_lower, safe_cmp);
                    values[mid_lower]
                }
                MedianMethod::Upper => {
                    let mid_upper = sample_size / 2;
                    values.select_nth_unstable_by(mid_upper, safe_cmp);
                    values[mid_upper]
                }
            }
        } else {
            // Odd number of elements
            let mid = sample_size / 2;
            values.select_nth_unstable_by(mid, safe_cmp);
            values[mid]
        };

        Self {
            median: median_value,
            sample_size,
        }
    }
}

impl From<MedianResult> for Median {
    fn from(result: MedianResult) -> Self {
        result.median
    }
}

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

    #[test]
    fn test_different_median_methods_even_elements() {
        // After sorting: [10.0, 20.0, 30.0, 40.0]
        let values = vec![20.0, 10.0, 40.0, 30.0];

        // Test Average method (default)
        let median_average = Median::from_with_method(values.clone(), MedianMethod::Average);
        assert_eq!(median_average.median, 25.0); // (20.0 + 30.0) / 2

        // Test Lower method
        let median_lower = Median::from_with_method(values.clone(), MedianMethod::Lower);
        assert_eq!(median_lower.median, 20.0); // Lower of the two middle values

        // Test Upper method
        let median_upper = Median::from_with_method(values, MedianMethod::Upper);
        assert_eq!(median_upper.median, 30.0); // Upper of the two middle values
    }

    #[test]
    fn test_different_median_methods_odd_elements() {
        let values = vec![10.0, 20.0, 30.0, 40.0, 50.0];
        // After sorting: [10.0, 20.0, 30.0, 40.0, 50.0]

        // All methods should return the same value for odd number of elements
        for method in [
            MedianMethod::Average,
            MedianMethod::Lower,
            MedianMethod::Upper,
        ] {
            let median = Median::from_with_method(values.clone(), method);
            assert_eq!(median.median, 30.0); // Middle element
        }
    }

    #[test]
    fn test_empty_vector() {
        let values = vec![];

        // Using the new API that returns MedianResult
        let median_price = Median::from_with_method(values.clone(), MedianMethod::default());

        assert_eq!(median_price.sample_size, 0);
        assert_eq!(median_price.median, 0.0);
    }

    #[test]
    fn test_single_element() {
        let values = vec![10.0];

        // Using the new API that returns MedianResult
        let median_price = Median::from_with_method(values, MedianMethod::default());

        assert_eq!(median_price.sample_size, 1);
        assert_eq!(median_price.median, 10.0);
    }

    #[test]
    fn test_odd_number_of_elements() {
        let values = vec![10.0, 30.0, 20.0, 50.0, 40.0];

        // Using the new API that returns MedianResult
        let median_price = Median::from_with_method(values, MedianMethod::default());

        assert_eq!(median_price.sample_size, 5);
        assert_eq!(median_price.median, 30.0);
    }

    #[test]
    fn test_even_number_of_elements() {
        let values = vec![10.0, 20.0, 30.0, 40.0];

        // Using the new API that returns MedianResult
        let median_price = Median::from_with_method(values, MedianMethod::default());

        assert_eq!(median_price.sample_size, 4);
        assert_eq!(median_price.median, 30.0); // Default is Upper method, so 30.0
    }

    #[test]
    fn test_many_elements() {
        // Create 101 updates with prices 0.0 through 100.0
        let values: Vec<f64> = (0..=100).map(|x| x as f64).collect();

        let median_price = Median::from_with_method(values, MedianMethod::default());

        assert_eq!(median_price.sample_size, 101);
        assert_eq!(median_price.median, 50.0);
    }

    #[test]
    fn test_different_original_update_keys() {
        let values = vec![10.0];

        let median_price1 = Median::from_with_method(values.clone(), MedianMethod::default());
        let median_price2 = Median::from_with_method(values, MedianMethod::default());

        assert_eq!(median_price1.median, median_price2.median);
        assert_eq!(median_price1.sample_size, median_price2.sample_size);
    }

    #[test]
    fn test_average_method_comprehensive() {
        // Test with various even-sized datasets
        let test_cases = vec![
            // Simple case
            (vec![10.0, 20.0], 15.0),
            // Four elements
            (vec![10.0, 20.0, 30.0, 40.0], 25.0), // (20 + 30) / 2
            // Six elements
            (vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], 3.5), // (3 + 4) / 2
            // With decimals
            (vec![1.5, 2.5, 3.5, 4.5], 3.0), // (2.5 + 3.5) / 2
            // Negative numbers
            (vec![-10.0, -5.0, 5.0, 10.0], 0.0), // (-5 + 5) / 2
            // Large numbers
            (vec![1000.0, 2000.0, 3000.0, 4000.0], 2500.0),
        ];

        for (values, expected) in test_cases {
            let result = Median::from_with_method(values.clone(), MedianMethod::Average);
            assert_eq!(result.median, expected);
            assert_eq!(result.sample_size, values.len());
        }
    }

    #[test]
    fn test_lower_method_comprehensive() {
        // Test with various even-sized datasets
        let test_cases = vec![
            // Simple case
            (vec![10.0, 20.0], 10.0), // Lower of [10, 20]
            // Four elements
            (vec![10.0, 20.0, 30.0, 40.0], 20.0), // Lower of [20, 30]
            // Six elements
            (vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], 3.0), // Lower of [3, 4]
            // Unordered input
            (vec![40.0, 10.0, 30.0, 20.0], 20.0), // Should sort first
            // With decimals
            (vec![1.1, 2.2, 3.3, 4.4], 2.2), // Lower of [2.2, 3.3]
            // Negative numbers
            (vec![-10.0, -5.0, 5.0, 10.0], -5.0), // Lower of [-5, 5]
        ];

        for (values, expected) in test_cases {
            let result = Median::from_with_method(values.clone(), MedianMethod::Lower);
            assert_eq!(result.median, expected);
            assert_eq!(result.sample_size, values.len());
        }
    }

    #[test]
    fn test_all_methods_with_same_odd_dataset() {
        // All methods should give the same result for odd number of elements
        let values = vec![5.0, 1.0, 9.0, 3.0, 7.0];
        // Sorted: [1.0, 3.0, 5.0, 7.0, 9.0] -> median should be 5.0

        for method in [
            MedianMethod::Average,
            MedianMethod::Lower,
            MedianMethod::Upper,
        ] {
            let result = Median::from_with_method(values.clone(), method);
            assert_eq!(result.median, 5.0, "Method {method:?} failed");
            assert_eq!(result.sample_size, 5);
        }
    }

    #[test]
    fn test_all_methods_with_duplicates() {
        // Test with duplicate values in even-sized dataset
        let values = vec![10.0, 20.0, 20.0, 30.0];
        // Sorted: [10.0, 20.0, 20.0, 30.0]
        // Middle elements are 20.0 and 20.0

        let result_avg = Median::from_with_method(values.clone(), MedianMethod::Average);
        assert_eq!(result_avg.median, 20.0); // (20 + 20) / 2

        let result_lower = Median::from_with_method(values.clone(), MedianMethod::Lower);
        assert_eq!(result_lower.median, 20.0); // Lower of [20, 20]

        let result_upper = Median::from_with_method(values, MedianMethod::Upper);
        assert_eq!(result_upper.median, 20.0); // Upper of [20, 20]

        // All should be the same when middle elements are identical
        assert_eq!(result_avg.median, result_lower.median);
        assert_eq!(result_avg.median, result_upper.median);
    }

    #[test]
    fn test_single_element_all_methods() {
        let values = vec![42.0];

        // All methods should return the same value for single element
        for method in [
            MedianMethod::Average,
            MedianMethod::Lower,
            MedianMethod::Upper,
        ] {
            let result = Median::from_with_method(values.clone(), method);
            assert_eq!(result.median, 42.0, "Method {method:?} failed");
            assert_eq!(result.sample_size, 1);
        }
    }

    #[test]
    fn test_two_elements_all_methods() {
        // Sorted: [10.0, 30.0]
        let values = vec![10.0, 30.0];

        let result_avg = Median::from_with_method(values.clone(), MedianMethod::Average);
        assert_eq!(result_avg.median, 20.0); // (10 + 30) / 2

        let result_lower = Median::from_with_method(values.clone(), MedianMethod::Lower);
        assert_eq!(result_lower.median, 10.0); // Lower element

        let result_upper = Median::from_with_method(values, MedianMethod::Upper);
        assert_eq!(result_upper.median, 30.0); // Upper element
    }

    #[test]
    fn test_large_dataset_all_methods() {
        // Test with larger dataset to ensure performance is reasonable
        // Values: 0.0, 1.0, 2.0, ..., 999.0
        // Middle elements: 499.0 and 500.0 (at indices 499 and 500)
        let values: Vec<f64> = (0..1000).map(|x| x as f64).collect();

        let result_avg = Median::from_with_method(values.clone(), MedianMethod::Average);
        assert_eq!(result_avg.median, 499.5); // (499 + 500) / 2
        assert_eq!(result_avg.sample_size, 1000);

        let result_lower = Median::from_with_method(values.clone(), MedianMethod::Lower);
        assert_eq!(result_lower.median, 499.0);

        let result_upper = Median::from_with_method(values, MedianMethod::Upper);
        assert_eq!(result_upper.median, 500.0);
    }

    #[test]
    fn test_extreme_values_all_methods() {
        let values = vec![f64::MIN, -1000.0, 1000.0, f64::MAX];
        // Sorted: [f64::MIN, -1000.0, 1000.0, f64::MAX]
        // Middle elements: -1000.0 and 1000.0

        let result_avg = Median::from_with_method(values.clone(), MedianMethod::Average);
        assert_eq!(result_avg.median, 0.0); // (-1000 + 1000) / 2

        let result_lower = Median::from_with_method(values.clone(), MedianMethod::Lower);
        assert_eq!(result_lower.median, -1000.0);

        let result_upper = Median::from_with_method(values, MedianMethod::Upper);
        assert_eq!(result_upper.median, 1000.0);
    }
}