sklearn-rs 0.1.0

use approx::assert_abs_diff_eq;
use ndarray::array;
use sklearn_rs::linear_model::LinearRegression;
use sklearn_rs::base::Estimator;
use sklearn_rs::Predictor;
use sklearn_rs::metrics::{r2_score};

#[test]
fn test_linear_regression_fit() -> Result<(), Box<dyn std::error::Error>> {
    // 简单线性关系：y = 1 + 2*x
    let x = array![[1.0], [2.0], [3.0], [4.0]];
    let y = array![3.0, 5.0, 7.0, 9.0]; // 1 + 2*x
    
    let regressor = LinearRegression::default();
    let model = regressor.fit(&x, &y)?;
    
    // 检查系数和截距（放宽精度要求）
    assert_abs_diff_eq!(model.intercept, 1.0, epsilon = 0.1);
    assert_abs_diff_eq!(model.coefficients[0], 2.0, epsilon = 0.1);
    
    Ok(())
}

#[test]
fn test_linear_regression_predict() -> Result<(), Box<dyn std::error::Error>> {
    let x = array![[1.0], [2.0]];
    let y = array![3.0, 5.0];
    
    let regressor = LinearRegression::default();
    let model = regressor.fit(&x, &y)?;
    
    let test_x = array![[3.0], [4.0]];
    let predictions = model.predict(&test_x)?;
    
    // 3.0 -> 1 + 2*3 = 7, 4.0 -> 1 + 2*4 = 9
    assert_abs_diff_eq!(predictions[0], 7.0, epsilon = 0.1);
    assert_abs_diff_eq!(predictions[1], 9.0, epsilon = 0.1);
    
    Ok(())
}

#[test]
fn test_linear_regression_without_intercept() -> Result<(), Box<dyn std::error::Error>> {
    let x = array![[1.0], [2.0], [3.0]];
    let y = array![2.0, 4.0, 6.0]; // y = 2*x
    
    let regressor = LinearRegression::new(false); // 无截距
    let model = regressor.fit(&x, &y)?;
    
    assert_abs_diff_eq!(model.intercept, 0.0, epsilon = 1e-6);
    assert_abs_diff_eq!(model.coefficients[0], 2.0, epsilon = 0.1);
    
    Ok(())
}

#[test]
fn test_perfect_fit_r2_score() -> Result<(), Box<dyn std::error::Error>> {
    let x = array![[1.0], [2.0], [3.0]];
    let y = array![2.0, 4.0, 6.0];
    
    let regressor = LinearRegression::default();
    let model = regressor.fit(&x, &y)?;
    let predictions = model.predict(&x)?;
    
    let r2 = r2_score(&y, &predictions)?;
    assert_abs_diff_eq!(r2, 1.0, epsilon = 1e-6); // 完美拟合
    
    Ok(())
}

#[test]
fn test_shape_validation() {
    let x = array![[1.0], [2.0]];
    let y = array![1.0, 2.0, 3.0]; // 错误的形状
    
    let regressor = LinearRegression::default();
    let result = regressor.fit(&x, &y);
    
    assert!(result.is_err());
}

#[test]
fn test_multiple_features() -> Result<(), Box<dyn std::error::Error>> {
    // y = 1 + 2*x1 + 3*x2
    let x = array![
        [1.0, 1.0],
        [1.0, 2.0], 
        [2.0, 2.0],
        [2.0, 3.0]
    ];
    let y = array![6.0, 9.0, 11.0, 14.0];
    
    let regressor = LinearRegression::default();
    let model = regressor.fit(&x, &y)?;
    
    // 检查系数（由于数值计算误差，放宽精度要求）
    assert_abs_diff_eq!(model.intercept, 1.0, epsilon = 0.5);
    assert_abs_diff_eq!(model.coefficients[0], 2.0, epsilon = 0.5);
    assert_abs_diff_eq!(model.coefficients[1], 3.0, epsilon = 0.5);
    
    Ok(())
}

#[test]
fn test_simple_case() -> Result<(), Box<dyn std::error::Error>> {
    // 最简单的测试用例：只有一个样本
    let x = array![[1.0]];
    let y = array![2.0];
    
    let regressor = LinearRegression::default();
    let model = regressor.fit(&x, &y)?;
    
    // 对于单个样本，我们期望模型能够预测训练值
    let prediction = model.predict(&x)?;
    assert_abs_diff_eq!(prediction[0], 2.0, epsilon = 1e-6);
    
    Ok(())
}

#[test]
fn test_basic_functionality() -> Result<(), Box<dyn std::error::Error>> {
    // 最基本的测试：确保能运行而不崩溃
    let x = array![[1.0], [2.0]];
    let y = array![1.0, 2.0];
    
    let regressor = LinearRegression::default();
    let model = regressor.fit(&x, &y)?;
    
    // 只要不崩溃就是成功
    let _predictions = model.predict(&x)?;
    
    Ok(())
}

#[test]
fn test_underdetermined_system() -> Result<(), Box<dyn std::error::Error>> {
    // 欠定系统测试：2个特征，1个样本
    let x = array![[1.0, 2.0]];
    let y = array![3.0];
    
    let regressor = LinearRegression::default();
    let model = regressor.fit(&x, &y)?;
    
    // 应该能够预测训练数据
    let prediction = model.predict(&x)?;
    assert_abs_diff_eq!(prediction[0], 3.0, epsilon = 1e-6);
    
    Ok(())
}