//! Various helper functions that work on vectors in Rust

use num_traits::Float;

use crate::prelude::BixverseFloat;

////////////
// Consts //
////////////

/// MAD scaling constant. R's `mad()` applies this factor by default.
pub const MAD_SCALE: f64 = 1.482_602_218_505_602;

///////////////
// Functions //
///////////////

/// Generate the rank of a vector with tie correction.
///
/// ### Params
///
/// * `vec` - The slice of numericals to rank.
///
/// ### Returns
///
/// The ranked vector (also f64)
pub fn rank_vector<T>(vec: &[T]) -> Vec<T>
where
    T: Float,
{
    let n = vec.len();
    if n == 0 {
        return Vec::new();
    }
    let mut indexed_values: Vec<(T, usize)> = vec
        .iter()
        .copied()
        .enumerate()
        .map(|(i, v)| (v, i))
        .collect();
    indexed_values
        .sort_unstable_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(std::cmp::Ordering::Equal));
    let mut ranks = vec![T::zero(); n];
    let mut i = 0;
    while i < n {
        let current_value = indexed_values[i].0;
        let start = i;
        while i < n && indexed_values[i].0 == current_value {
            i += 1;
        }
        let avg_rank = (start + i + 1) as f64 / 2.0;
        let rank_value = T::from(avg_rank).unwrap();
        for j in start..i {
            ranks[indexed_values[j].1] = rank_value;
        }
    }
    ranks
}

/// Get the median
///
/// ### Params
///
/// * `x` - The slice for which to calculate the median for.
///
/// ### Results
///
/// The median (if the vector is not empty)
pub fn median<T>(x: &[T]) -> Option<T>
where
    T: BixverseFloat,
{
    if x.is_empty() {
        return None;
    }
    let mut data = x.to_vec();
    let len = data.len();
    if len.is_multiple_of(2) {
        let (_, median1, right) =
            data.select_nth_unstable_by(len / 2 - 1, |a, b| a.partial_cmp(b).unwrap());
        let median2 = right
            .iter()
            .min_by(|a, b| a.partial_cmp(b).unwrap())
            .unwrap();
        Some((*median1 + *median2) / T::from_f64(2.0).unwrap())
    } else {
        let (_, median, _) = data.select_nth_unstable_by(len / 2, |a, b| a.partial_cmp(b).unwrap());
        Some(*median)
    }
}

/// Calculate the MAD
///
/// ### Params
///
/// * `x` - Slice for which to calculate the MAD for
/// * `scale` - Optional scaling factor. Pass `Some(1.4826)` for consistency with
///   the standard deviation under normality (R's default). `None` returns the raw MAD.
///
/// ### Results
///
/// The (optionally scaled) MAD of the slice.
pub fn mad<T>(x: &[T], scale: Option<T>) -> Option<T>
where
    T: BixverseFloat,
{
    if x.is_empty() {
        return None;
    }
    let median_val = median(x)?;
    let deviations: Vec<T> = x.iter().map(|&val| (val - median_val).abs()).collect();
    let raw = median(&deviations)?;
    Some(match scale {
        Some(k) => raw * k,
        None => raw,
    })
}

/// Standard deviation
///
/// ### Params
///
/// * `x` Slice of `f64`
///
/// ### Returns
///
/// The standard deviation
pub fn standard_deviation<T>(x: &[T]) -> T
where
    T: BixverseFloat + std::iter::Sum,
{
    let n = T::from_usize(x.len()).unwrap();
    let mean: T = x.iter().copied().sum::<T>() / n;
    let variance = x.iter().map(|&val| (val - mean).powi(2)).sum::<T>() / (n - T::one());
    variance.sqrt()
}

/// Calculate the mean while removing NaNs
///
/// ### Params
///
/// * `x` - The slice of floats for which to calculate the mean (while
///   ignoring `NaN`'s)
///
/// ### Returns
///
/// The mean of the slice without `NaN`s
pub fn mean_nan<T>(x: &[T]) -> T
where
    T: BixverseFloat + std::iter::Sum,
{
    let finite: Vec<T> = x.iter().copied().filter(|x| x.is_finite()).collect();
    if finite.is_empty() {
        T::nan()
    } else {
        finite.iter().copied().sum::<T>() / T::from_usize(finite.len()).unwrap()
    }
}

/// Linearly interpolated quantile over a pre-sorted slice.
///
/// ### Params
///
/// * `sorted` - Sorted slice of values.
/// * `q` - Quantile in `[0.0, 1.0]`.
///
/// ### Returns
///
/// The interpolated quantile value, or `0.0` if `sorted` is empty.
pub fn quantile_sorted<T>(sorted: &[T], q: T) -> T
where
    T: BixverseFloat,
{
    if sorted.is_empty() {
        return T::zero();
    }
    let pos = q.clamp(T::zero(), T::one()) * T::from_usize(sorted.len() - 1).unwrap();
    let lo = pos.floor().to_usize().unwrap();
    let hi = (lo + 1).min(sorted.len() - 1);
    let frac = pos - T::from_usize(lo).unwrap();
    sorted[lo] * (T::one() - frac) + sorted[hi] * frac
}

/// Linearly interpolated quantile over an unsorted slice, matching numpy's
/// default.
///
/// ### Params
///
/// * `values` - Unsorted slice of values.
/// * `q` - Quantile in `[0.0, 1.0]`.
///
/// ### Returns
///
/// The interpolated quantile value, or `0.0` if `values` is empty.
pub fn quantile<T>(values: &[T], q: T) -> T
where
    T: BixverseFloat,
{
    let mut sorted: Vec<T> = values.to_vec();
    sorted.sort_unstable_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
    quantile_sorted(&sorted, q)
}

///////////
// Tests //
///////////

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

    #[test]
    fn test_rank_vector() {
        let vec = vec![3.0, 1.0, 2.0, 3.0];
        // Ranks should account for ties:
        // sorted: 1.0 (idx 1), 2.0 (idx 2), 3.0 (idx 0), 3.0 (idx 3)
        // ranks: 1.0 -> 1.0, 2.0 -> 2.0, 3.0 -> 3.5, 3.0 -> 3.5
        let ranks = rank_vector(&vec);
        assert_eq!(ranks, vec![3.5, 1.0, 2.0, 3.5]);

        let empty: Vec<f64> = vec![];
        assert_eq!(rank_vector(&empty), Vec::<f64>::new());
    }

    #[test]
    fn test_median() {
        let vec_odd = vec![1.0, 3.0, 2.0];
        assert_eq!(median(&vec_odd), Some(2.0));

        let vec_even = vec![1.0, 4.0, 2.0, 3.0];
        assert_eq!(median(&vec_even), Some(2.5));

        let empty: Vec<f64> = vec![];
        assert_eq!(median(&empty), None);
    }

    #[test]
    fn test_mad() {
        let vec = vec![1.0, 1.0, 2.0, 2.0, 4.0, 6.0, 9.0];
        // Median is 2.0
        // Absolute deviations: [1.0, 1.0, 0.0, 0.0, 2.0, 4.0, 7.0]
        // Sorted deviations: [0.0, 0.0, 1.0, 1.0, 2.0, 4.0, 7.0]
        // MAD (median of deviations) is 1.0
        assert_eq!(mad(&vec, None), Some(1.0));
    }

    #[test]
    fn test_standard_deviation() {
        let vec = vec![2.0, 4.0, 4.0, 4.0, 5.0, 5.0, 7.0, 9.0];
        let std = standard_deviation(&vec);
        assert!((std - 2.1380899352).abs() < 1e-6);
    }
}
