use approx::assert_relative_eq;
use fleximl_models::ml::linear_model::LinearModel;
use fleximl_models::utils::tasks::Task;
use ndarray::{arr1, arr2};

#[test]
fn test_linear_model_creation() {
    let model = LinearModel::new(3, 1, 0.01, Task::BinaryClassification, 42);
    assert_eq!(model.weights.shape(), &[3, 1]);
    assert_eq!(model.bias.len(), 1);
    assert_eq!(model.learning_rate, 0.01);
    assert_eq!(model.task, Task::BinaryClassification);
}

#[test]
fn test_binary_classification_prediction() {
    let mut model = LinearModel::new(2, 1, 0.01, Task::BinaryClassification, 42);
    model.weights = arr2(&[[1.0], [-1.0]]);
    model.bias = arr1(&[0.5]);

    let x = arr1(&[2.0, 1.0]);
    let prediction = model.predict(&x);
    assert_relative_eq!(prediction[0], 0.8175744761936437, epsilon = 1e-8);
}

#[test]
fn test_regression_prediction() {
    let mut model = LinearModel::new(2, 1, 0.01, Task::Regression, 42);
    model.weights = arr2(&[[1.5], [-0.5]]);
    model.bias = arr1(&[1.0]);

    let x = arr1(&[2.0, 3.0]);
    let prediction = model.predict(&x);
    assert_eq!(prediction[0], 2.5);
}

#[test]
fn test_binary_classification_fit() {
    let mut model = LinearModel::new(2, 1, 0.1, Task::BinaryClassification, 42);
    let x = arr2(&[[1.0, 2.0], [2.0, 3.0], [3.0, 4.0], [4.0, 5.0]]);
    let y = arr2(&[[0.0], [0.0], [1.0], [1.0]]);

    model.fit(&x, &y, 1000);

    let prediction1 = model.predict(&arr1(&[1.0, 2.0]))[0];
    let prediction2 = model.predict(&arr1(&[4.0, 5.0]))[0];

    assert!(prediction1 < 0.1);
    assert!(prediction2 > 0.9);
}

#[test]
fn test_regression_fit() {
    let mut model = LinearModel::new(1, 1, 0.01, Task::Regression, 42);
    let x = arr2(&[[1.0], [2.0], [3.0], [4.0]]);
    let y = arr2(&[[2.0], [4.0], [6.0], [8.0]]);

    model.fit(&x, &y, 1000);

    let prediction = model.predict(&arr1(&[5.0]))[0];
    assert_relative_eq!(prediction, 10.0, epsilon = 0.1);
}

#[test]
fn test_multi_class_classification() {
    let mut model = LinearModel::new(2, 3, 0.1, Task::MultiClassClassification, 42);
    let x = arr2(&[[1.0, 2.0], [2.0, 3.0], [3.0, 4.0], [4.0, 5.0]]);
    let y = arr2(&[
        [1.0, 0.0, 0.0],
        [0.0, 1.0, 0.0],
        [0.0, 0.0, 1.0],
        [1.0, 0.0, 0.0],
    ]);

    model.fit(&x, &y, 1000);

    let prediction = model.predict(&arr1(&[1.0, 2.0]));
    assert_eq!(prediction.len(), 3);
    assert!(prediction.iter().all(|&x| x >= 0.0 && x <= 1.0));
    assert_relative_eq!(prediction.sum(), 1.0, epsilon = 1e-6);
}