extern crate savgol_rs;
use nalgebra::{DMatrix, DVector};
use std::convert::TryFrom;

use savgol_rs::savgol_filter;
pub fn savgol(data: Vec<f64>, window: usize, order: usize) -> Vec<f64> {
    let svinput= savgol_rs::SavGolInput{data:&data,window_length:window,poly_order:order,derivative:0};

    savgol_filter(&svinput).unwrap()
}

/// 多项式拟合函数 (f64版本)
/// 
/// 参数:
/// - x: x坐标序列
/// - y: y坐标序列
/// - degree: 多项式次数 (7或8)
/// 
/// 返回: 平滑后的y值序列
pub fn polynomial_fit_f64(x: &[f64], y: &[f64], degree: usize) -> Vec<f64> {
    assert_eq!(x.len(), y.len(), "x和y的长度必须相同");
    if x.len() < degree + 1 {
        panic!("数据点数量必须大于多项式次数");
    }
    let n = x.len();
    let y_vec = DVector::from_vec(y.to_vec());
    // 构建范德蒙矩阵
    let mut vandermonde = DMatrix::zeros(n, degree + 1);
    for i in 0..n {
        for j in 0..=degree {
            vandermonde[(i, j)] = x[i].powi(j as i32);
        }
    }
    // 解最小二乘问题 - 新版本nalgebra的调用方式
    let svd = vandermonde.svd(true, true);
    let coefficients = svd.solve(&y_vec, f64::EPSILON).unwrap();
    // 计算拟合值
    let fitted_y: Vec<f64> = x
        .iter()
        .map(|&xi| {
            (0..=degree).fold(0.0, |acc, j| acc + coefficients[j] * xi.powi(j as i32))
        })
        .collect();
    fitted_y
}
/// 多项式拟合函数 (u32版本)
/// 
/// 参数:
/// - x: x坐标序列
/// - y: y坐标序列
/// - degree: 多项式次数 (7或8)
/// 
/// 返回: 平滑后的y值序列
pub fn polynomial_fit_u32(x: &[u32], y: &[u32], degree: usize) -> Vec<u32> {
    // 转换为f64处理
    let x_f64: Vec<f64> = x.iter().map(|&xi| xi as f64).collect();
    let y_f64: Vec<f64> = y.iter().map(|&yi| yi as f64).collect();
    
    let fitted_f64 = polynomial_fit_f64(&x_f64, &y_f64, degree);
    
    // 转换回u32，处理可能的负值(截断为0)和溢出
    fitted_f64
        .into_iter()
        .map(|y| {
            if y < 0.0 {
                0
            } else {
                u32::try_from(y.round() as i64).unwrap_or(u32::MAX)
            }
        })
        .collect()
}
/// 简化版：当x是等间距时的平滑函数 (u32版本)
pub fn polynomial_smooth_u32(y: &[u32], degree: usize) -> Vec<u32> {
    let x: Vec<u32> = (0..y.len() as u32).collect();
    polynomial_fit_u32(&x, y, degree)
}
/// 多项式拟合函数 (u16版本)
pub fn polynomial_fit_u16(x: &[u16], y: &[u16], degree: usize) -> Vec<u16> {
    let x_f64: Vec<f64> = x.iter().map(|&xi| xi as f64).collect();
    let y_f64: Vec<f64> = y.iter().map(|&yi| yi as f64).collect();
    
    let fitted_f64 = polynomial_fit_f64(&x_f64, &y_f64, degree);
    
    fitted_f64
        .into_iter()
        .map(|y| {
            if y < 0.0 {
                0
            } else if y > u16::MAX as f64 {
                u16::MAX
            } else {
                y.round() as u16
            }
        })
        .collect()
}
/// 简化版：当x是等间距时的平滑函数 (u16版本)
pub fn polynomial_smooth_u16(y: &[u16], degree: usize) -> Vec<u16> {
    let x: Vec<u16> = (0..y.len() as u16).collect();
    polynomial_fit_u16(&x, y, degree)
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn test_polynomial_fit() {
        // 测试数据: 一个简单的二次函数加一些噪声
        let x = vec![0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0];
        let y = vec![1.0, 3.0, 6.0, 10.0, 15.0, 24.0, 35.0, 50.0, 65.0, 82.0];
        
        // 2次多项式拟合应该能很好拟合
        let fitted = polynomial_fit_f64(&x, &y, 2);
        assert_eq!(fitted.len(), y.len());
        
        // 检查拟合结果是否接近原始数据
        for (original, fitted) in y.iter().zip(fitted.iter()) {
            assert!((original - fitted).abs() < 5.0);
        }
    }
    #[test]
    fn test_polynomial_smooth() {
        // 测试数据: 一个简单的上升序列加一些噪声
        let y = vec![
           10,12,11,13,12,14,13,15,14,16,
           15,17,16,18,17,19,18,20,19,21,
        ];
        
        // 7次多项式平滑
        let smoothed_7 = polynomial_smooth_u32(&y, 7);
        print!("smoothed_7: {:?}", smoothed_7);
        assert_eq!(smoothed_7.len(), y.len());
        
        // 8次多项式平滑
        let smoothed_8 = polynomial_smooth_u32(&y, 8);
        assert_eq!(smoothed_8.len(), y.len());
    }
}
#[test]
fn test_savgol() {
    // 示例数据
    let data = vec![1.0, 1.9, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0];
    let window = 5;
    let order = 2;

    // 调用 savgol 函数
    let smoothed_data = savgol(data.clone(), window, order);
    println!("Smoothed data: {:?}", smoothed_data);

}
