rustyml 0.14.0

//! Integration tests for [`rustyml::machine_learning::LinearRegression`].
//!
//! Every expected value comes from the problem design or a closed-form analytic result,
//! never from recording the model's output.

use approx::assert_abs_diff_eq;
use ndarray::{Array1, Array2, array};
use rustyml::error::Error;
use rustyml::machine_learning::linear_model::LeastSquaresSolver;
use rustyml::machine_learning::{LinearRegression, RegularizationType};

use crate::common::assert_allclose;

// Constructor validation

/// A non-positive or non-finite learning_rate (0.0 / negative / NaN / +inf) -> InvalidParameter
#[test]
fn with_solver_rejects_invalid_learning_rate() {
    for lr in [0.0, -0.01, f64::NAN, f64::INFINITY] {
        let result = LinearRegression::new(true).with_solver(LeastSquaresSolver::GradientDescent {
            learning_rate: lr,
            max_iter: 100,
            tol: 1e-6,
        });
        assert!(
            matches!(result, Err(Error::InvalidParameter { .. })),
            "expected InvalidParameter for learning_rate={lr:?}, got {result:?}"
        );
    }
}

/// max_iterations = 0 -> InvalidParameter
#[test]
fn with_solver_rejects_zero_max_iter() {
    let result = LinearRegression::new(true).with_solver(LeastSquaresSolver::GradientDescent {
        learning_rate: 0.01,
        max_iter: 0,
        tol: 1e-6,
    });
    assert!(
        matches!(result, Err(Error::InvalidParameter { .. })),
        "expected InvalidParameter, got {:?}",
        result
    );
}

/// A non-positive or non-finite tolerance (0.0 / negative / NaN / +inf) -> InvalidParameter
#[test]
fn with_solver_rejects_invalid_tolerance() {
    for tol in [0.0, -1e-6, f64::NAN, f64::INFINITY] {
        let result = LinearRegression::new(true).with_solver(LeastSquaresSolver::GradientDescent {
            learning_rate: 0.01,
            max_iter: 100,
            tol,
        });
        assert!(
            matches!(result, Err(Error::InvalidParameter { .. })),
            "expected InvalidParameter for tolerance={tol:?}, got {result:?}"
        );
    }
}

/// A negative or non-finite regularization alpha (L2(-0.1) / L1(-0.5) / L2(NaN))
/// -> InvalidParameter
#[test]
fn constructor_invalid_regularization_alpha_is_invalid() {
    for reg in [
        RegularizationType::L2(-0.1),
        RegularizationType::L1(-0.5),
        RegularizationType::L2(f64::NAN),
    ] {
        let result = LinearRegression::new(true)
            .with_solver(LeastSquaresSolver::GradientDescent {
                learning_rate: 0.01,
                max_iter: 100,
                tol: 1e-6,
            })
            .unwrap()
            .with_regularization(reg);
        assert!(
            matches!(result, Err(Error::InvalidParameter { .. })),
            "expected InvalidParameter for regularization={reg:?}, got {result:?}"
        );
    }
}

/// Valid constructor with all legal parameters -> Ok
#[test]
fn constructor_valid_parameters_succeeds() {
    let result = LinearRegression::new(true).with_solver(LeastSquaresSolver::GradientDescent {
        learning_rate: 0.01,
        max_iter: 1000,
        tol: 1e-6,
    });
    assert!(result.is_ok(), "expected Ok, got {:?}", result);
}

/// Getters on a freshly constructed model return the supplied values
#[test]
fn constructor_getters_round_trip() {
    let model = LinearRegression::new(false)
        .with_solver(LeastSquaresSolver::GradientDescent {
            learning_rate: 0.05,
            max_iter: 500,
            tol: 1e-4,
        })
        .unwrap();
    assert!(!model.get_fit_intercept());
    // The solver variant carries the gradient-descent settings, not the model
    match model.get_solver() {
        LeastSquaresSolver::GradientDescent {
            learning_rate,
            max_iter,
            tol,
        } => {
            assert_abs_diff_eq!(learning_rate, 0.05, epsilon = 1e-15);
            assert_eq!(max_iter, 500);
            assert_abs_diff_eq!(tol, 1e-4, epsilon = 1e-20);
        }
        LeastSquaresSolver::Normal => panic!("expected the gradient-descent solver"),
    }
    assert!(model.get_coefficients().is_none());
    assert!(model.get_intercept().is_none());
    assert!(model.get_actual_iterations().is_none());
}

// NotFitted errors before fit

/// predict() on an unfitted model -> NotFitted
#[test]
fn predict_before_fit_returns_not_fitted() {
    let model = LinearRegression::new(true)
        .with_solver(LeastSquaresSolver::GradientDescent {
            learning_rate: 0.01,
            max_iter: 100,
            tol: 1e-6,
        })
        .unwrap();
    let x = array![[1.0, 2.0]];
    let result = model.predict(&x);
    assert!(
        matches!(result, Err(Error::NotFitted(_))),
        "expected NotFitted, got {:?}",
        result
    );
}

// fit() input-validation errors

/// fit() with empty X -> EmptyInput
#[test]
fn fit_empty_x_returns_empty_input() {
    let mut model = LinearRegression::new(true)
        .with_solver(LeastSquaresSolver::GradientDescent {
            learning_rate: 0.01,
            max_iter: 100,
            tol: 1e-6,
        })
        .unwrap();
    let x: Array2<f64> = Array2::zeros((0, 2));
    let y: Array1<f64> = Array1::zeros(0);
    let result = model.fit(&x, &y);
    assert!(
        matches!(result, Err(Error::EmptyInput(_))),
        "expected EmptyInput, got {:?}",
        result
    );
}

/// fit() with a non-finite sentinel (NaN / +Inf) in X -> NonFinite
#[test]
fn fit_non_finite_in_x_returns_non_finite() {
    for sentinel in [f64::NAN, f64::INFINITY] {
        let mut model = LinearRegression::new(true)
            .with_solver(LeastSquaresSolver::GradientDescent {
                learning_rate: 0.01,
                max_iter: 100,
                tol: 1e-6,
            })
            .unwrap();
        let x = array![[1.0, sentinel], [2.0, 3.0]];
        let y = array![1.0, 2.0];
        let result = model.fit(&x, &y);
        assert!(
            matches!(result, Err(Error::NonFinite(_))),
            "expected NonFinite for sentinel={sentinel:?}, got {result:?}"
        );
    }
}

/// fit() with mismatched y length -> DimensionMismatch
#[test]
fn fit_y_length_mismatch_returns_dimension_mismatch() {
    let mut model = LinearRegression::new(true)
        .with_solver(LeastSquaresSolver::GradientDescent {
            learning_rate: 0.01,
            max_iter: 100,
            tol: 1e-6,
        })
        .unwrap();
    // 3 rows in x but 2 elements in y
    let x = array![[1.0], [2.0], [3.0]];
    let y = array![1.0, 2.0];
    let result = model.fit(&x, &y);
    assert!(
        matches!(result, Err(Error::DimensionMismatch { .. })),
        "expected DimensionMismatch, got {:?}",
        result
    );
}

// predict() input-validation errors (after fit)

/// predict() with empty matrix -> EmptyInput
#[test]
fn predict_empty_matrix_returns_empty_input() {
    let mut model = LinearRegression::new(true)
        .with_solver(LeastSquaresSolver::GradientDescent {
            learning_rate: 0.01,
            max_iter: 5000,
            tol: 1e-8,
        })
        .unwrap();
    // train on y = 2x + 1
    let x_train = array![[1.0], [2.0], [3.0], [4.0], [5.0]];
    let y_train = array![3.0, 5.0, 7.0, 9.0, 11.0];
    model.fit(&x_train, &y_train).unwrap();

    let x_empty: Array2<f64> = Array2::zeros((0, 1));
    let result = model.predict(&x_empty);
    assert!(
        matches!(result, Err(Error::EmptyInput(_))),
        "expected EmptyInput, got {:?}",
        result
    );
}

/// predict() with wrong number of columns -> DimensionMismatch
#[test]
fn predict_wrong_feature_count_returns_dimension_mismatch() {
    let mut model = LinearRegression::new(true)
        .with_solver(LeastSquaresSolver::GradientDescent {
            learning_rate: 0.01,
            max_iter: 5000,
            tol: 1e-8,
        })
        .unwrap();
    // trained on 1 feature
    let x_train = array![[1.0], [2.0], [3.0], [4.0], [5.0]];
    let y_train = array![3.0, 5.0, 7.0, 9.0, 11.0];
    model.fit(&x_train, &y_train).unwrap();

    // predict with 2 features, mismatched
    let x_wrong = array![[1.0, 2.0]];
    let result = model.predict(&x_wrong);
    assert!(
        matches!(result, Err(Error::DimensionMismatch { .. })),
        "expected DimensionMismatch, got {:?}",
        result
    );
}

/// predict() with a non-finite sentinel (NaN / +Inf) in X -> NonFinite
#[test]
fn predict_non_finite_in_x_returns_non_finite() {
    for sentinel in [f64::NAN, f64::INFINITY] {
        let mut model = LinearRegression::new(true)
            .with_solver(LeastSquaresSolver::GradientDescent {
                learning_rate: 0.01,
                max_iter: 5000,
                tol: 1e-8,
            })
            .unwrap();
        let x_train = array![[1.0], [2.0], [3.0], [4.0], [5.0]];
        let y_train = array![3.0, 5.0, 7.0, 9.0, 11.0];
        model.fit(&x_train, &y_train).unwrap();

        let x_bad = array![[sentinel]];
        let result = model.predict(&x_bad);
        assert!(
            matches!(result, Err(Error::NonFinite(_))),
            "expected NonFinite for sentinel={sentinel:?}, got {result:?}"
        );
    }
}

// Correctness: univariate y = 2x + 1
// OLS on x=[1..5], y=[3,5,7,9,11] gives slope 2.0, intercept 1.0

/// After fit, coefficient ~= 2.0 and intercept ~= 1.0 (tight tolerance)
#[test]
fn univariate_y_equals_2x_plus_1_coefficient_and_intercept() {
    // small learning rate and many iterations so gradient descent converges
    let mut model = LinearRegression::new(true)
        .with_solver(LeastSquaresSolver::GradientDescent {
            learning_rate: 0.01,
            max_iter: 10_000,
            tol: 1e-10,
        })
        .unwrap();
    let x = array![[1.0], [2.0], [3.0], [4.0], [5.0]];
    let y = array![3.0, 5.0, 7.0, 9.0, 11.0];
    model.fit(&x, &y).unwrap();

    // iterative solver reaches the OLS solution (coef 2, intercept 1) to ~1e-3, so assert
    // within iterative-solver tolerance rather than exact-OLS precision
    let coeff = model.get_coefficients().unwrap();
    assert_abs_diff_eq!(coeff[0], 2.0, epsilon = 3e-3);

    let intercept = model.get_intercept().unwrap();
    assert_abs_diff_eq!(intercept, 1.0, epsilon = 3e-3);
}

/// predict on x=6 -> 13.0, and predict on x=0 -> 1.0
#[test]
fn univariate_y_equals_2x_plus_1_predictions() {
    let mut model = LinearRegression::new(true)
        .with_solver(LeastSquaresSolver::GradientDescent {
            learning_rate: 0.01,
            max_iter: 10_000,
            tol: 1e-10,
        })
        .unwrap();
    let x = array![[1.0], [2.0], [3.0], [4.0], [5.0]];
    let y = array![3.0, 5.0, 7.0, 9.0, 11.0];
    model.fit(&x, &y).unwrap();

    let preds = model.predict(&array![[6.0], [0.0]]).unwrap();
    let expected = array![13.0, 1.0];
    assert_allclose(&preds, &expected, 1e-3);
}

/// After fit, n_iter is set (model ran at least 1 iteration)
#[test]
fn fit_sets_n_iter() {
    let mut model = LinearRegression::new(true)
        .with_solver(LeastSquaresSolver::GradientDescent {
            learning_rate: 0.01,
            max_iter: 10_000,
            tol: 1e-10,
        })
        .unwrap();
    let x = array![[1.0], [2.0], [3.0], [4.0], [5.0]];
    let y = array![3.0, 5.0, 7.0, 9.0, 11.0];
    model.fit(&x, &y).unwrap();

    let n_iter = model.get_actual_iterations();
    assert!(n_iter.is_some(), "n_iter should be set after fit");
    assert!(n_iter.unwrap() >= 1, "n_iter must be at least 1");
}

// Correctness: multivariate y = 2*x1 + 3*x2 + 1
// Coefficients converge to [2.0, 3.0], intercept to 1.0

/// Multivariate coefficients and intercept converge to known values
#[test]
fn multivariate_y_equals_2x1_plus_3x2_plus_1_coefficients() {
    let mut model = LinearRegression::new(true)
        .with_solver(LeastSquaresSolver::GradientDescent {
            learning_rate: 0.01,
            max_iter: 20_000,
            tol: 1e-10,
        })
        .unwrap();
    // 6 training points spanning the feature space
    let x = Array2::from_shape_vec(
        (6, 2),
        vec![
            1.0, 1.0, // y = 2+3+1 = 6
            2.0, 1.0, // y = 4+3+1 = 8
            1.0, 2.0, // y = 2+6+1 = 9
            3.0, 2.0, // y = 6+6+1 = 13
            2.0, 3.0, // y = 4+9+1 = 14
            4.0, 1.0, // y = 8+3+1 = 12
        ],
    )
    .unwrap();
    let y = array![6.0, 8.0, 9.0, 13.0, 14.0, 12.0];
    model.fit(&x, &y).unwrap();

    let coeff = model.get_coefficients().unwrap();
    assert_abs_diff_eq!(coeff[0], 2.0, epsilon = 3e-3);
    assert_abs_diff_eq!(coeff[1], 3.0, epsilon = 3e-3);

    let intercept = model.get_intercept().unwrap();
    assert_abs_diff_eq!(intercept, 1.0, epsilon = 3e-3);
}

/// Multivariate predictions match closed-form y = 2*x1 + 3*x2 + 1
#[test]
fn multivariate_predictions_match_closed_form() {
    let mut model = LinearRegression::new(true)
        .with_solver(LeastSquaresSolver::GradientDescent {
            learning_rate: 0.01,
            max_iter: 20_000,
            tol: 1e-10,
        })
        .unwrap();
    let x = Array2::from_shape_vec(
        (6, 2),
        vec![1.0, 1.0, 2.0, 1.0, 1.0, 2.0, 3.0, 2.0, 2.0, 3.0, 4.0, 1.0],
    )
    .unwrap();
    let y = array![6.0, 8.0, 9.0, 13.0, 14.0, 12.0];
    model.fit(&x, &y).unwrap();

    let x_new = array![[1.0, 1.0], [2.0, 3.0]];
    let preds = model.predict(&x_new).unwrap();
    let expected = array![6.0, 14.0];
    assert_allclose(&preds, &expected, 5e-3);
}

// fit_intercept = false
// Data y = 2x through the origin: OLS slope 2.0, stored intercept 0.0 by contract

/// With fit_intercept=false the stored intercept is exactly 0.0
#[test]
fn no_intercept_stored_intercept_is_zero() {
    let mut model = LinearRegression::new(false)
        .with_solver(LeastSquaresSolver::GradientDescent {
            learning_rate: 0.01,
            max_iter: 10_000,
            tol: 1e-10,
        })
        .unwrap();
    let x = array![[1.0], [2.0], [3.0], [4.0], [5.0]];
    // y = 2x, passes through the origin
    let y = array![2.0, 4.0, 6.0, 8.0, 10.0];
    model.fit(&x, &y).unwrap();

    let intercept = model.get_intercept().unwrap();
    assert_abs_diff_eq!(intercept, 0.0, epsilon = 1e-15);
}

/// With fit_intercept=false the coefficient converges to slope ~= 2.0
#[test]
fn no_intercept_coefficient_converges_to_slope() {
    let mut model = LinearRegression::new(false)
        .with_solver(LeastSquaresSolver::GradientDescent {
            learning_rate: 0.01,
            max_iter: 10_000,
            tol: 1e-10,
        })
        .unwrap();
    let x = array![[1.0], [2.0], [3.0], [4.0], [5.0]];
    let y = array![2.0, 4.0, 6.0, 8.0, 10.0];
    model.fit(&x, &y).unwrap();

    let coeff = model.get_coefficients().unwrap();
    assert_abs_diff_eq!(coeff[0], 2.0, epsilon = 1e-4);
}

/// With fit_intercept=false get_fit_intercept() returns false
#[test]
fn no_intercept_getter_returns_false() {
    let model = LinearRegression::new(false)
        .with_solver(LeastSquaresSolver::GradientDescent {
            learning_rate: 0.01,
            max_iter: 1000,
            tol: 1e-6,
        })
        .unwrap();
    assert!(!model.get_fit_intercept());
}

// Analytic OLS sanity: y = 3x + 2
// OLS on x=[1..5], y=[5,8,11,14,17] gives slope 3.0, intercept 2.0

/// OLS converges to slope=3.0, intercept=2.0 on y=3x+2
#[test]
fn ols_sanity_y_equals_3x_plus_2_parameters() {
    let mut model = LinearRegression::new(true)
        .with_solver(LeastSquaresSolver::GradientDescent {
            learning_rate: 0.01,
            max_iter: 10_000,
            tol: 1e-10,
        })
        .unwrap();
    let x = array![[1.0], [2.0], [3.0], [4.0], [5.0]];
    let y = array![5.0, 8.0, 11.0, 14.0, 17.0];
    model.fit(&x, &y).unwrap();

    let coeff = model.get_coefficients().unwrap();
    // analytic slope = 3.0 (gradient descent converges to ~1e-3)
    assert_abs_diff_eq!(coeff[0], 3.0, epsilon = 3e-3);

    // analytic intercept = 2.0
    let intercept = model.get_intercept().unwrap();
    assert_abs_diff_eq!(intercept, 2.0, epsilon = 3e-3);
}

/// OLS prediction at x=6 -> 20.0, x=10 -> 32.0 on y=3x+2
#[test]
fn ols_sanity_y_equals_3x_plus_2_predictions() {
    let mut model = LinearRegression::new(true)
        .with_solver(LeastSquaresSolver::GradientDescent {
            learning_rate: 0.01,
            max_iter: 10_000,
            tol: 1e-10,
        })
        .unwrap();
    let x = array![[1.0], [2.0], [3.0], [4.0], [5.0]];
    let y = array![5.0, 8.0, 11.0, 14.0, 17.0];
    model.fit(&x, &y).unwrap();

    let preds = model.predict(&array![[6.0], [10.0]]).unwrap();
    let expected = array![20.0, 32.0];
    assert_allclose(&preds, &expected, 3e-3);
}

// fit_predict matches separate fit + predict

/// fit_predict() returns the same predictions as fit() + predict() on training data
#[test]
fn fit_predict_matches_fit_then_predict() {
    let x = array![[1.0], [2.0], [3.0], [4.0], [5.0]];
    let y = array![3.0, 5.0, 7.0, 9.0, 11.0];

    let mut model_a = LinearRegression::new(true)
        .with_solver(LeastSquaresSolver::GradientDescent {
            learning_rate: 0.01,
            max_iter: 10_000,
            tol: 1e-10,
        })
        .unwrap();
    let preds_a = model_a.fit_predict(&x, &y).unwrap();

    let mut model_b = LinearRegression::new(true)
        .with_solver(LeastSquaresSolver::GradientDescent {
            learning_rate: 0.01,
            max_iter: 10_000,
            tol: 1e-10,
        })
        .unwrap();
    model_b.fit(&x, &y).unwrap();
    let preds_b = model_b.predict(&x).unwrap();

    assert_allclose(&preds_a, &preds_b, 1e-12);
}

/// fit_predict() predictions match known true values from y=2x+1
#[test]
fn fit_predict_values_match_known_true_values() {
    let x = array![[1.0], [2.0], [3.0], [4.0], [5.0]];
    let y = array![3.0, 5.0, 7.0, 9.0, 11.0];

    let mut model = LinearRegression::new(true)
        .with_solver(LeastSquaresSolver::GradientDescent {
            learning_rate: 0.01,
            max_iter: 10_000,
            tol: 1e-10,
        })
        .unwrap();
    let preds = model.fit_predict(&x, &y).unwrap();

    // Trained on exactly this data, so predictions on it closely match y = 2x + 1
    let expected = array![3.0, 5.0, 7.0, 9.0, 11.0];
    assert_allclose(&preds, &expected, 5e-3);
}

// Regularization: L2 shrinks the coefficient
// Ridge alpha > 0 makes |w_ridge| strictly smaller than |w_ols|.
// The test checks only that inequality.

/// L2 regularization shrinks the L2-norm of coefficients below the unregularized value
#[test]
fn l2_regularization_shrinks_coefficient_norm() {
    let x = array![[1.0], [2.0], [3.0], [4.0], [5.0]];
    let y = array![3.0, 5.0, 7.0, 9.0, 11.0];

    let mut unregularized = LinearRegression::new(true)
        .with_solver(LeastSquaresSolver::GradientDescent {
            learning_rate: 0.01,
            max_iter: 10_000,
            tol: 1e-10,
        })
        .unwrap();
    unregularized.fit(&x, &y).unwrap();
    let coeff_unreg = unregularized.get_coefficients().unwrap()[0].abs();

    let mut ridge = LinearRegression::new(true)
        .with_solver(LeastSquaresSolver::GradientDescent {
            learning_rate: 0.01,
            max_iter: 10_000,
            tol: 1e-10,
        })
        .unwrap()
        .with_regularization(RegularizationType::L2(5.0))
        .unwrap();
    ridge.fit(&x, &y).unwrap();
    let coeff_ridge = ridge.get_coefficients().unwrap()[0].abs();

    assert!(
        coeff_ridge < coeff_unreg,
        "Ridge coefficient {coeff_ridge} should be smaller than unregularized {coeff_unreg}"
    );
}

// Regularization: L1 shrinks the coefficient
// Lasso alpha drives the coefficient toward zero, so |w_lasso| < |w_ols|

/// L1 regularization shrinks the coefficient below the unregularized value
#[test]
fn l1_regularization_shrinks_coefficient() {
    let x = array![[1.0], [2.0], [3.0], [4.0], [5.0]];
    let y = array![3.0, 5.0, 7.0, 9.0, 11.0];

    let mut unregularized = LinearRegression::new(true)
        .with_solver(LeastSquaresSolver::GradientDescent {
            learning_rate: 0.01,
            max_iter: 10_000,
            tol: 1e-10,
        })
        .unwrap();
    unregularized.fit(&x, &y).unwrap();
    let coeff_unreg = unregularized.get_coefficients().unwrap()[0].abs();

    let mut lasso = LinearRegression::new(true)
        .with_solver(LeastSquaresSolver::GradientDescent {
            learning_rate: 0.01,
            max_iter: 10_000,
            tol: 1e-10,
        })
        .unwrap()
        .with_regularization(RegularizationType::L1(5.0))
        .unwrap();
    lasso.fit(&x, &y).unwrap();
    let coeff_lasso = lasso.get_coefficients().unwrap()[0].abs();

    assert!(
        coeff_lasso < coeff_unreg,
        "Lasso coefficient {coeff_lasso} should be smaller than unregularized {coeff_unreg}"
    );
}

/// With fit_intercept=true and moderate L2 alpha the intercept stays close to 1.0
#[test]
fn l2_regularization_intercept_within_reasonable_range() {
    let x = array![[1.0], [2.0], [3.0], [4.0], [5.0]];
    let y = array![3.0, 5.0, 7.0, 9.0, 11.0];

    let mut ridge = LinearRegression::new(true)
        .with_solver(LeastSquaresSolver::GradientDescent {
            learning_rate: 0.01,
            max_iter: 10_000,
            tol: 1e-10,
        })
        .unwrap()
        .with_regularization(RegularizationType::L2(0.1))
        .unwrap();
    ridge.fit(&x, &y).unwrap();

    // alpha=0.1 is small, so it should only mildly shrink the intercept away from 1.0
    let intercept = ridge.get_intercept().unwrap();
    assert!(
        (intercept - 1.0).abs() < 0.5,
        "Intercept {intercept} deviates too far from 1.0 under weak L2 regularization"
    );
}

// Determinism: same data -> identical results
// Gradient descent has no internal randomness, so identical runs are bit-identical

/// 2 identical LinearRegression models trained on the same data produce identical predictions
#[test]
fn determinism_same_data_identical_predictions() {
    let x = array![[1.0], [2.0], [3.0], [4.0], [5.0]];
    let y = array![3.0, 5.0, 7.0, 9.0, 11.0];
    let x_test = array![[6.0], [7.0], [8.0]];

    let mut model_a = LinearRegression::new(true)
        .with_solver(LeastSquaresSolver::GradientDescent {
            learning_rate: 0.01,
            max_iter: 10_000,
            tol: 1e-10,
        })
        .unwrap();
    model_a.fit(&x, &y).unwrap();
    let preds_a = model_a.predict(&x_test).unwrap();

    let mut model_b = LinearRegression::new(true)
        .with_solver(LeastSquaresSolver::GradientDescent {
            learning_rate: 0.01,
            max_iter: 10_000,
            tol: 1e-10,
        })
        .unwrap();
    model_b.fit(&x, &y).unwrap();
    let preds_b = model_b.predict(&x_test).unwrap();

    // Gradient descent is deterministic, so results must be bit-identical
    assert_allclose(&preds_a, &preds_b, 0.0);
}

// Save / load round-trip

/// save_to_path + load_from_path round-trip yields identical predictions
#[test]
fn save_load_round_trip_identical_predictions() {
    let x_train = array![[1.0], [2.0], [3.0], [4.0], [5.0]];
    let y_train = array![3.0, 5.0, 7.0, 9.0, 11.0];
    let x_test = array![[6.0], [7.0], [0.5]];

    let mut model = LinearRegression::new(true)
        .with_solver(LeastSquaresSolver::GradientDescent {
            learning_rate: 0.01,
            max_iter: 10_000,
            tol: 1e-10,
        })
        .unwrap();
    model.fit(&x_train, &y_train).unwrap();
    let preds_before = model.predict(&x_test).unwrap();

    // Tmp path unique to this test to avoid collisions
    let path = "/tmp/rustyml_linear_regression_test_round_trip.bin";
    model.save_to_path(path).unwrap();

    let loaded = LinearRegression::load_from_path(path).unwrap();
    let preds_after = loaded.predict(&x_test).unwrap();

    assert_allclose(&preds_before, &preds_after, 0.0);

    let _ = std::fs::remove_file(path);
}

/// After load, getter values match those of the original model
#[test]
fn save_load_preserves_model_state() {
    let x_train = array![[1.0], [2.0], [3.0], [4.0], [5.0]];
    let y_train = array![3.0, 5.0, 7.0, 9.0, 11.0];

    let mut model = LinearRegression::new(false)
        .with_solver(LeastSquaresSolver::GradientDescent {
            learning_rate: 0.005,
            max_iter: 8_000,
            tol: 1e-9,
        })
        .unwrap();
    model.fit(&x_train, &y_train).unwrap();

    let path = "/tmp/rustyml_linear_regression_test_state.bin";
    model.save_to_path(path).unwrap();
    let loaded = LinearRegression::load_from_path(path).unwrap();

    assert_eq!(loaded.get_fit_intercept(), model.get_fit_intercept());
    // The solver, with its settings, round-trips as 1 value
    assert_eq!(loaded.get_solver(), model.get_solver());

    let orig_coeff = model.get_coefficients().unwrap();
    let load_coeff = loaded.get_coefficients().unwrap();
    assert_allclose(orig_coeff, load_coeff, 0.0);

    let _ = std::fs::remove_file(path);
}

// Default constructor

/// `default()` is `new(true)`: fit_intercept = true and the closed-form solver
#[test]
fn default_constructor_has_expected_hyperparameters() {
    let model = LinearRegression::default();
    assert!(model.get_fit_intercept());
    assert_eq!(model.get_solver(), LeastSquaresSolver::Normal);
    // The 2 constructors must agree, so that neither can silently build a different algorithm
    assert_eq!(model.get_solver(), LinearRegression::new(true).get_solver());
    assert!(model.get_coefficients().is_none());
    assert!(model.get_intercept().is_none());
    assert!(model.get_actual_iterations().is_none());
}

/// Default model can be fit and predict without error
#[test]
fn default_constructor_can_fit_and_predict() {
    let mut model = LinearRegression::default();
    let x = array![[1.0], [2.0], [3.0]];
    let y = array![3.0, 5.0, 7.0];
    model.fit(&x, &y).unwrap();
    let preds = model.predict(&array![[4.0]]).unwrap();
    // y = 2x + 1, predict(4) ~= 9.0
    assert_abs_diff_eq!(preds[0], 9.0, epsilon = 5e-2);
}

// Clone

/// Clone of a fitted model makes identical predictions
#[test]
fn clone_of_fitted_model_makes_identical_predictions() {
    let x_train = array![[1.0], [2.0], [3.0], [4.0], [5.0]];
    let y_train = array![3.0, 5.0, 7.0, 9.0, 11.0];
    let x_test = array![[6.0], [0.0]];

    let mut model = LinearRegression::new(true)
        .with_solver(LeastSquaresSolver::GradientDescent {
            learning_rate: 0.01,
            max_iter: 10_000,
            tol: 1e-10,
        })
        .unwrap();
    model.fit(&x_train, &y_train).unwrap();
    let preds_orig = model.predict(&x_test).unwrap();

    let cloned = model.clone();
    let preds_clone = cloned.predict(&x_test).unwrap();

    assert_allclose(&preds_orig, &preds_clone, 0.0);
}

// In-loop NonFinite divergence guards
// A huge learning_rate overshoots to +/- inf inside the loop, tripping the in-loop finiteness guard

/// fit() on finite, valid data with a huge learning_rate diverges to +/- inf inside the loop.
/// This trips the in-loop guard and returns Error::NonFinite, not the up-front check on x.
#[test]
fn fit_huge_learning_rate_diverges_returns_non_finite() {
    // learning_rate = 1e8 is positive and finite, so the constructor accepts it
    let mut model = LinearRegression::new(true)
        .with_solver(LeastSquaresSolver::GradientDescent {
            learning_rate: 1e8,
            max_iter: 1000,
            tol: 1e-10,
        })
        .unwrap();

    // clean, finite data: y = 2x + 1 on x = [1..5]
    let x = array![[1.0], [2.0], [3.0], [4.0], [5.0]];
    let y = array![3.0, 5.0, 7.0, 9.0, 11.0];

    let result = model.fit(&x, &y);
    assert!(
        matches!(result, Err(Error::NonFinite(_))),
        "expected NonFinite from in-loop divergence guard, got {:?}",
        result
    );
}

// L1 regularization with many features (n_features = 200)
// Column 0 carries the signal (y=3*x0), the remaining 199 columns are noise

/// L1 regularization with 200 features: the 1 informative feature (column 0, y = 3*x0)
/// ends up with the dominant coefficient over the noise columns
#[test]
fn l1_regularization_many_features_recovers_informative_feature() {
    let n_samples = 12usize;
    let n_features = 200usize;

    // column 0 = centered signal, columns 1.. = tiny noise
    let x = Array2::from_shape_fn((n_samples, n_features), |(i, j)| {
        if j == 0 {
            // centered, varying signal in [-5.5, 5.5]
            (i as f64) - 5.5
        } else {
            // bounded zig-zag noise in {-0.03, ..., 0.03}, uncorrelated with y
            0.01 * (((i * 31 + j * 17) % 7) as f64 - 3.0)
        }
    });

    // y depends only on column 0: y = 3 * x0 (no intercept needed)
    let y = Array1::from_shape_fn(n_samples, |i| 3.0 * ((i as f64) - 5.5));

    // weak L1 shrinks the dominant coefficient only slightly. fit_intercept = false
    let mut model = LinearRegression::new(false)
        .with_solver(LeastSquaresSolver::GradientDescent {
            learning_rate: 0.01,
            max_iter: 20_000,
            tol: 1e-12,
        })
        .unwrap()
        .with_regularization(RegularizationType::L1(1e-3))
        .unwrap();
    model
        .fit(&x, &y)
        .expect("fit with 200 features and L1 should succeed");

    let coeffs = model.get_coefficients().unwrap();
    assert_eq!(
        coeffs.len(),
        n_features,
        "coefficient vector length must equal feature count"
    );

    let c0 = coeffs[0];
    // (a) the informative coefficient is substantial and positive (true slope is 3.0)
    assert!(
        c0 > 1.0,
        "informative coefficient[0] = {c0} should be a large positive value (true slope 3.0)"
    );

    // (b) every noise coefficient stays small. The 0.5 bound sits above what the bounded
    // (|x_j| <= 0.03) noise columns can earn, yet far below coefficient[0]
    let max_other = coeffs
        .iter()
        .skip(1)
        .fold(0.0_f64, |acc, &w| acc.max(w.abs()));
    assert!(
        max_other < 0.5,
        "uninformative coefficients should stay small; largest |other| = {max_other}"
    );

    // (c) coefficient[0] strictly dominates every other |coefficient|
    assert!(
        c0.abs() > max_other,
        "|coefficient[0]| = {} should dominate the largest other |coefficient| = {}",
        c0.abs(),
        max_other
    );
}

// score (coefficient of determination R^2)

/// On exactly-linear data (y = 3x0 - 2x1 + 5) a converged model achieves R^2 about 1.
#[test]
fn score_is_one_on_perfectly_linear_data() {
    let x = array![
        [1.0, 1.0],
        [2.0, 0.0],
        [0.0, 3.0],
        [4.0, 2.0],
        [3.0, 1.0],
        [1.0, 4.0]
    ];
    let y = array![6.0, 11.0, -1.0, 13.0, 12.0, 0.0];
    let mut model = LinearRegression::new(true)
        .with_solver(LeastSquaresSolver::GradientDescent {
            learning_rate: 0.02,
            max_iter: 300_000,
            tol: 1e-13,
        })
        .unwrap();
    model.fit(&x, &y).unwrap();
    let r2 = model.score(&x, &y).unwrap();
    assert!(r2 <= 1.0 + 1e-9, "R² must not exceed 1, got {r2}");
    assert!(
        r2 > 0.999,
        "expected R² ≈ 1 on exactly-linear data, got {r2}"
    );
}

/// score equals the textbook R^2 definition computed independently from predict()
#[test]
fn score_matches_r2_definition() {
    let x = array![[1.0], [2.0], [3.0], [4.0], [5.0]];
    let y = array![2.1, 3.9, 6.2, 7.8, 10.1]; // noisy linear, R^2 strictly < 1
    let mut model = LinearRegression::new(true)
        .with_solver(LeastSquaresSolver::GradientDescent {
            learning_rate: 0.01,
            max_iter: 100_000,
            tol: 1e-12,
        })
        .unwrap();
    model.fit(&x, &y).unwrap();

    let preds = model.predict(&x).unwrap();
    let y_mean = y.iter().sum::<f64>() / y.len() as f64;
    let ss_res: f64 = y
        .iter()
        .zip(preds.iter())
        .map(|(yi, pi)| (yi - pi).powi(2))
        .sum();
    let ss_tot: f64 = y.iter().map(|yi| (yi - y_mean).powi(2)).sum();
    let expected = 1.0 - ss_res / ss_tot;

    let r2 = model.score(&x, &y).unwrap();
    assert_abs_diff_eq!(r2, expected, epsilon = 1e-12);
    assert!(r2 < 1.0, "noisy data must score strictly below 1, got {r2}");
}

/// A model that predicts the mean of y scores R^2 = 0. Here zero coefficients plus an
/// intercept fitted on mean-centered features converges to predicting y_bar.
#[test]
fn score_mean_predictor_is_about_zero() {
    // Feature is uninformative about y (y alternates independently of x), so the best
    // linear fit is y_hat about y_bar and R^2 about 0
    let x = array![[1.0], [2.0], [3.0], [4.0], [5.0], [6.0]];
    let y = array![1.0, 0.0, 1.0, 0.0, 1.0, 0.0];
    let mut model = LinearRegression::new(true)
        .with_solver(LeastSquaresSolver::GradientDescent {
            learning_rate: 0.001,
            max_iter: 200_000,
            tol: 1e-13,
        })
        .unwrap();
    model.fit(&x, &y).unwrap();
    let r2 = model.score(&x, &y).unwrap();
    assert!(
        r2.abs() < 0.1,
        "an uninformative feature should give R² near 0, got {r2}"
    );
    assert!(r2 <= 1.0 + 1e-9);
}

/// score on an unfitted model returns NotFitted
#[test]
fn score_not_fitted_errors() {
    let model = LinearRegression::new(true)
        .with_solver(LeastSquaresSolver::GradientDescent {
            learning_rate: 0.01,
            max_iter: 100,
            tol: 1e-6,
        })
        .unwrap();
    let x = array![[1.0], [2.0]];
    let y = array![1.0, 2.0];
    assert!(matches!(
        model.score(&x, &y),
        Err(Error::NotFitted("LinearRegression"))
    ));
}

/// score with a y of the wrong length returns DimensionMismatch
#[test]
fn score_y_length_mismatch_errors() {
    let x = array![[1.0], [2.0], [3.0]];
    let y = array![1.0, 2.0, 3.0];
    let mut model = LinearRegression::new(true)
        .with_solver(LeastSquaresSolver::GradientDescent {
            learning_rate: 0.01,
            max_iter: 1000,
            tol: 1e-6,
        })
        .unwrap();
    model.fit(&x, &y).unwrap();
    let y_wrong = array![1.0, 2.0];
    assert!(matches!(
        model.score(&x, &y_wrong),
        Err(Error::DimensionMismatch { .. })
    ));
}

// Closed-form (normal-equation) solver

/// On exactly-linear data the closed-form solver recovers the true coefficients and
/// intercept exactly (no iteration / learning-rate tuning), unlike gradient descent which
/// only approaches them
#[test]
fn normal_solver_recovers_exact_coefficients() {
    // y = 3*x0 - 2*x1 + 5, exactly
    let x = array![
        [1.0, 1.0],
        [2.0, 0.0],
        [0.0, 3.0],
        [4.0, 2.0],
        [3.0, 1.0],
        [1.0, 4.0]
    ];
    let y = array![6.0, 11.0, -1.0, 13.0, 12.0, 0.0];

    // Selecting the closed form leaves no iteration settings to carry.
    // `LeastSquaresSolver::Normal` has no fields, so it cannot even accept a learning rate.
    let mut model = LinearRegression::new(true)
        .with_solver(LeastSquaresSolver::Normal)
        .unwrap();
    model.fit(&x, &y).unwrap();

    let coefs = model.get_coefficients().unwrap();
    assert_abs_diff_eq!(coefs[0], 3.0, epsilon = 1e-9);
    assert_abs_diff_eq!(coefs[1], -2.0, epsilon = 1e-9);
    assert_abs_diff_eq!(model.get_intercept().unwrap(), 5.0, epsilon = 1e-9);
    // Closed form performs no gradient-descent iterations
    assert_eq!(model.get_actual_iterations(), Some(0));
}

/// The closed-form L2 (ridge) solution matches what gradient descent converges to on the
/// same objective. GD uses penalty (alpha/2)||w||^2, so the closed form uses lambda = n*alpha.
#[test]
fn normal_solver_l2_matches_gradient_descent() {
    let x = array![
        [1.0, 0.5],
        [2.0, -1.0],
        [3.0, 0.0],
        [-1.0, 2.0],
        [0.5, 1.5],
        [2.5, -0.5],
        [1.0, 1.0],
        [-2.0, 0.5]
    ];
    let y = array![2.0, 1.0, 3.5, -0.5, 1.0, 2.2, 1.8, -1.5];

    let alpha = 0.3;
    let mut gd = LinearRegression::new(true)
        .with_solver(LeastSquaresSolver::GradientDescent {
            learning_rate: 0.03,
            max_iter: 400_000,
            tol: 1e-13,
        })
        .unwrap()
        .with_regularization(RegularizationType::L2(alpha))
        .unwrap();
    gd.fit(&x, &y).unwrap();

    let mut normal = LinearRegression::new(true)
        .with_solver(LeastSquaresSolver::GradientDescent {
            learning_rate: 0.01,
            max_iter: 1,
            tol: 1e-6,
        })
        .unwrap()
        .with_regularization(RegularizationType::L2(alpha))
        .unwrap()
        .with_solver(LeastSquaresSolver::Normal)
        .unwrap();
    normal.fit(&x, &y).unwrap();

    let gd_c = gd.get_coefficients().unwrap();
    let nm_c = normal.get_coefficients().unwrap();
    for i in 0..gd_c.len() {
        assert_abs_diff_eq!(gd_c[i], nm_c[i], epsilon = 1e-3);
    }
    assert_abs_diff_eq!(
        gd.get_intercept().unwrap(),
        normal.get_intercept().unwrap(),
        epsilon = 1e-3
    );
}

/// Closed-form OLS matches the hand-evaluated normal equation w = (X^T X)^-1 X^T y on a
/// fit_intercept=false problem
#[test]
fn normal_solver_no_intercept_matches_normal_equation() {
    // Simple 2-feature, no-intercept system
    let x = array![[1.0, 2.0], [3.0, 1.0], [2.0, 4.0], [0.0, 1.0]];
    let y = array![5.0, 5.0, 10.0, 2.0]; // y = 1*x0 + 2*x1 (exact)
    let mut model = LinearRegression::new(false)
        .with_solver(LeastSquaresSolver::GradientDescent {
            learning_rate: 0.01,
            max_iter: 1,
            tol: 1e-6,
        })
        .unwrap()
        .with_solver(LeastSquaresSolver::Normal)
        .unwrap();
    model.fit(&x, &y).unwrap();
    let c = model.get_coefficients().unwrap();
    assert_abs_diff_eq!(c[0], 1.0, epsilon = 1e-9);
    assert_abs_diff_eq!(c[1], 2.0, epsilon = 1e-9);
    assert_eq!(model.get_intercept().unwrap(), 0.0);
}

/// The Normal solver rejects L1 regularization, which has no closed form
#[test]
fn normal_solver_rejects_l1_regularization() {
    let x = array![[1.0], [2.0], [3.0]];
    let y = array![1.0, 2.0, 3.0];
    let mut model = LinearRegression::new(true)
        .with_solver(LeastSquaresSolver::GradientDescent {
            learning_rate: 0.01,
            max_iter: 1,
            tol: 1e-6,
        })
        .unwrap()
        .with_regularization(RegularizationType::L1(0.5))
        .unwrap()
        .with_solver(LeastSquaresSolver::Normal)
        .unwrap();
    let result = model.fit(&x, &y);
    assert!(
        matches!(result, Err(Error::InvalidInput(_))),
        "Normal solver + L1 must error, got {result:?}"
    );
}

/// Ridge shrinks coefficients relative to unregularized OLS (closed form)
#[test]
fn normal_solver_ridge_shrinks_coefficients() {
    let x = array![[1.0, 0.9], [2.0, 2.1], [3.0, 2.9], [4.0, 4.2], [5.0, 5.1]];
    let y = array![1.0, 2.0, 3.0, 4.0, 5.0];

    let mut ols = LinearRegression::new(true)
        .with_solver(LeastSquaresSolver::GradientDescent {
            learning_rate: 0.01,
            max_iter: 1,
            tol: 1e-6,
        })
        .unwrap()
        .with_solver(LeastSquaresSolver::Normal)
        .unwrap();
    ols.fit(&x, &y).unwrap();

    let mut ridge = LinearRegression::new(true)
        .with_solver(LeastSquaresSolver::GradientDescent {
            learning_rate: 0.01,
            max_iter: 1,
            tol: 1e-6,
        })
        .unwrap()
        .with_regularization(RegularizationType::L2(1.0))
        .unwrap()
        .with_solver(LeastSquaresSolver::Normal)
        .unwrap();
    ridge.fit(&x, &y).unwrap();

    let ols_norm: f64 = ols.get_coefficients().unwrap().iter().map(|c| c * c).sum();
    let ridge_norm: f64 = ridge
        .get_coefficients()
        .unwrap()
        .iter()
        .map(|c| c * c)
        .sum();
    assert!(
        ridge_norm < ols_norm,
        "ridge ||w||^2 ({ridge_norm}) must be smaller than OLS ({ols_norm})"
    );
}

// scikit-learn parity

/// `LinearRegression::default()` is exact OLS, like Python's `LinearRegression()`. scikit-learn
/// 1.9.0 gives `coef_ = [1.11666667, 0.93333333]` and `intercept_ = 0.05` on this data.
#[test]
fn default_solver_is_exact_ols_matching_scikit_learn() {
    let x: Array2<f64> = array![[1.0, 1.0], [2.0, 1.0], [3.0, 2.0], [4.0, 3.0], [5.0, 5.0]];
    let y: Array1<f64> = array![2.0, 3.5, 5.1, 7.2, 10.4];

    let mut model = LinearRegression::default();
    model.fit(&x, &y).expect("fit should succeed");

    let coefficients = model.get_coefficients().expect("fit sets coefficients");
    assert_abs_diff_eq!(coefficients[0], 1.116_666_666_666_667, epsilon = 1e-12);
    assert_abs_diff_eq!(coefficients[1], 0.933_333_333_333_333_3, epsilon = 1e-12);
    assert_abs_diff_eq!(
        model.get_intercept().expect("fit sets the intercept"),
        0.05,
        epsilon = 1e-12
    );
}

/// L1 drives an unsupported coefficient to exactly `0.0`, not merely close to it.
///
/// The proximal (soft-thresholding) step makes L1 a feature selector by reaching exact zero.
#[test]
fn l1_regularization_produces_exact_zeros() {
    // Feature 0 explains y exactly. Feature 1 is noise.
    let x: Array2<f64> = array![
        [1.0, 0.3],
        [2.0, -0.7],
        [3.0, 0.5],
        [4.0, -0.2],
        [5.0, 0.9],
        [6.0, -0.4],
        [7.0, 0.1],
        [8.0, -0.8],
        [9.0, 0.6],
        [10.0, -0.3]
    ];
    let y: Array1<f64> = array![2.0, 4.0, 6.0, 8.0, 10.0, 12.0, 14.0, 16.0, 18.0, 20.0];

    let mut model = LinearRegression::new(true)
        .with_solver(LeastSquaresSolver::GradientDescent {
            learning_rate: 0.01,
            max_iter: 20_000,
            tol: 1e-12,
        })
        .expect("valid params")
        .with_regularization(RegularizationType::L1(0.5))
        .expect("valid alpha");
    model.fit(&x, &y).expect("fit should succeed");

    let coefficients = model.get_coefficients().expect("fit sets coefficients");
    assert_eq!(
        coefficients[1], 0.0,
        "the noise coefficient must be exactly zero, got {}",
        coefficients[1]
    );
    assert!(
        coefficients[0] > 1.0,
        "the informative coefficient must survive, got {}",
        coefficients[0]
    );
}