rustyml 0.17.0

//! Integration tests for the KNN (K-Nearest Neighbors) classifier: constructor
//! validation, fit/predict error paths, distance metrics, weighting, parallel
//! paths, and save/load round-trips

use ndarray::{Array1, Array2, array};
use rustyml::error::Error;
use rustyml::machine_learning::DistanceCalculationMetric as Metric;
use rustyml::machine_learning::{KNN, WeightingStrategy};

/// k=0 must return Error::InvalidParameter
#[test]
fn constructor_k_zero_returns_invalid_parameter() {
    let err = KNN::<i32>::new(0).unwrap_err();
    assert!(
        matches!(err, Error::InvalidParameter { .. }),
        "expected InvalidParameter, got {err:?}"
    );
}

/// k=1 (minimum legal) must succeed
#[test]
fn constructor_k_one_succeeds() {
    KNN::<i32>::new(1)
        .unwrap()
        .with_weighting_strategy(WeightingStrategy::Uniform)
        .with_metric(Metric::Euclidean)
        .unwrap();
}

/// Minkowski p < 1 is not a valid metric and must return Error::InvalidParameter
#[test]
fn constructor_rejects_minkowski_p_below_one() {
    let err = KNN::<i32>::new(1)
        .unwrap()
        .with_metric(Metric::Minkowski(0.5))
        .unwrap_err();
    assert!(
        matches!(err, Error::InvalidParameter { .. }),
        "expected InvalidParameter for Minkowski(0.5), got {err:?}"
    );
}

/// Default constructor must expose k=5, Uniform, Euclidean
#[test]
fn constructor_default_values() {
    let knn = KNN::<i32>::default();
    assert_eq!(knn.get_k(), 5);
    assert_eq!(knn.get_weighting_strategy(), WeightingStrategy::Uniform);
    assert_eq!(knn.get_metric(), Metric::Euclidean);
    assert!(knn.get_x_train().is_none());
}

/// new() must store and expose the exact parameters supplied
#[test]
fn constructor_stores_parameters() {
    let knn = KNN::<i32>::new(3)
        .unwrap()
        .with_weighting_strategy(WeightingStrategy::Distance)
        .with_metric(Metric::Manhattan)
        .unwrap();
    assert_eq!(knn.get_k(), 3);
    assert_eq!(knn.get_weighting_strategy(), WeightingStrategy::Distance);
    assert_eq!(knn.get_metric(), Metric::Manhattan);
}

/// Fitting with fewer samples than k returns Error::InvalidInput (k=5, 3 rows)
#[test]
fn fit_fewer_samples_than_k_returns_invalid_input() {
    let mut knn = KNN::<i32>::new(5)
        .unwrap()
        .with_weighting_strategy(WeightingStrategy::Uniform)
        .with_metric(Metric::Euclidean)
        .unwrap();
    let x = array![[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]; // 3 rows < k=5
    let y = array![0, 1, 0];
    let err = knn.fit(&x, &y).unwrap_err();
    assert!(
        matches!(err, Error::InvalidInput(_)),
        "expected InvalidInput, got {err:?}"
    );
}

/// Fitting with y length != x.nrows() returns DimensionMismatch (x has 3 rows, y has 2 labels)
#[test]
fn fit_mismatched_y_length_returns_dimension_mismatch() {
    let mut knn = KNN::<i32>::new(1)
        .unwrap()
        .with_weighting_strategy(WeightingStrategy::Uniform)
        .with_metric(Metric::Euclidean)
        .unwrap();
    let x = array![[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]; // 3 rows
    let y = array![0, 1]; // only 2 labels
    let err = knn.fit(&x, &y).unwrap_err();
    assert!(
        matches!(err, Error::DimensionMismatch { .. }),
        "expected DimensionMismatch, got {err:?}"
    );
}

/// predict with wrong feature count must return Error::DimensionMismatch (train 2, predict 3)
#[test]
fn predict_wrong_feature_count_returns_dimension_mismatch() {
    let mut knn = KNN::<i32>::new(1)
        .unwrap()
        .with_weighting_strategy(WeightingStrategy::Uniform)
        .with_metric(Metric::Euclidean)
        .unwrap();
    let x_train = array![[0.0, 0.0], [10.0, 0.0]];
    let y_train = array![0, 1];
    knn.fit(&x_train, &y_train).unwrap();

    // supply 3 features instead of 2
    let x_test = array![[1.0, 2.0, 3.0]];
    let err = knn.predict(&x_test).unwrap_err();
    assert!(
        matches!(err, Error::DimensionMismatch { .. }),
        "expected DimensionMismatch, got {err:?}"
    );
}

/// k=1 Euclidean Uniform picks the single nearest neighbor per query
#[test]
fn predict_k1_euclidean_correctness() {
    let mut knn = KNN::<i32>::new(1)
        .unwrap()
        .with_weighting_strategy(WeightingStrategy::Uniform)
        .with_metric(Metric::Euclidean)
        .unwrap();
    let x_train = array![[0.0, 0.0], [10.0, 0.0], [0.0, 10.0]];
    let y_train = array![0, 1, 1];
    knn.fit(&x_train, &y_train).unwrap();

    let x_test = array![[0.5, 0.0], [9.5, 0.0], [0.0, 9.5]];
    let predictions = knn.predict(&x_test).unwrap();

    assert_eq!(predictions[0], 0, "Q0=(0.5,0) should map to class 0");
    assert_eq!(predictions[1], 1, "Q1=(9.5,0) should map to class 1");
    assert_eq!(predictions[2], 1, "Q2=(0,9.5) should map to class 1");
}

/// k=1 Manhattan Uniform picks the single nearest neighbor per query
#[test]
fn predict_k1_manhattan_correctness() {
    let mut knn = KNN::<i32>::new(1)
        .unwrap()
        .with_weighting_strategy(WeightingStrategy::Uniform)
        .with_metric(Metric::Manhattan)
        .unwrap();
    let x_train = array![[0.0, 0.0], [10.0, 0.0], [0.0, 10.0]];
    let y_train = array![0, 1, 1];
    knn.fit(&x_train, &y_train).unwrap();

    let x_test = array![[0.5, 0.0], [9.5, 0.0], [0.0, 9.5]];
    let predictions = knn.predict(&x_test).unwrap();

    assert_eq!(predictions[0], 0, "Q0=(0.5,0) should map to class 0");
    assert_eq!(predictions[1], 1, "Q1=(9.5,0) should map to class 1");
    assert_eq!(predictions[2], 1, "Q2=(0,9.5) should map to class 1");
}

/// Minkowski(p=1) predictions must equal Manhattan predictions on the same data
#[test]
fn predict_minkowski_p1_equals_manhattan() {
    // Manhattan reference model
    let mut knn_man = KNN::<i32>::new(1)
        .unwrap()
        .with_weighting_strategy(WeightingStrategy::Uniform)
        .with_metric(Metric::Manhattan)
        .unwrap();
    let x_train = array![[0.0, 0.0], [10.0, 0.0], [0.0, 10.0]];
    let y_train = array![0, 1, 1];
    knn_man.fit(&x_train, &y_train).unwrap();

    // Minkowski(p=1) model
    let mut knn_mink = KNN::<i32>::new(1)
        .unwrap()
        .with_weighting_strategy(WeightingStrategy::Uniform)
        .with_metric(Metric::Minkowski(1.0))
        .unwrap();
    knn_mink.fit(&x_train, &y_train).unwrap();

    let x_test = array![[0.5, 0.0], [9.5, 0.0], [0.0, 9.5]];
    let preds_man = knn_man.predict(&x_test).unwrap();
    let preds_mink = knn_mink.predict(&x_test).unwrap();

    assert_eq!(
        preds_man, preds_mink,
        "Minkowski(p=1) predictions must equal Manhattan"
    );
}

/// Minkowski(p=2) predictions must equal Euclidean predictions on the same data
#[test]
fn predict_minkowski_p2_equals_euclidean() {
    // Euclidean reference model
    let mut knn_euc = KNN::<i32>::new(1)
        .unwrap()
        .with_weighting_strategy(WeightingStrategy::Uniform)
        .with_metric(Metric::Euclidean)
        .unwrap();
    let x_train = array![[0.0, 0.0], [10.0, 0.0], [0.0, 10.0]];
    let y_train = array![0, 1, 1];
    knn_euc.fit(&x_train, &y_train).unwrap();

    // Minkowski(p=2) model
    let mut knn_mink = KNN::<i32>::new(1)
        .unwrap()
        .with_weighting_strategy(WeightingStrategy::Uniform)
        .with_metric(Metric::Minkowski(2.0))
        .unwrap();
    knn_mink.fit(&x_train, &y_train).unwrap();

    let x_test = array![[0.5, 0.0], [9.5, 0.0], [0.0, 9.5]];
    let preds_euc = knn_euc.predict(&x_test).unwrap();
    let preds_mink = knn_mink.predict(&x_test).unwrap();

    assert_eq!(
        preds_euc, preds_mink,
        "Minkowski(p=2) predictions must equal Euclidean"
    );
}

/// k=3 Euclidean Uniform resolves a clean 3-vs-0 majority per query
#[test]
fn predict_k3_majority_clean() {
    let mut knn = KNN::<i32>::new(3)
        .unwrap()
        .with_weighting_strategy(WeightingStrategy::Uniform)
        .with_metric(Metric::Euclidean)
        .unwrap();
    let x_train = array![
        [0.0, 0.0],
        [1.0, 0.0],
        [2.0, 0.0],
        [10.0, 0.0],
        [11.0, 0.0],
        [12.0, 0.0]
    ];
    let y_train = array![0, 0, 0, 1, 1, 1];
    knn.fit(&x_train, &y_train).unwrap();

    let x_test = array![[0.5, 0.0], [10.5, 0.0]];
    let predictions = knn.predict(&x_test).unwrap();

    assert_eq!(predictions[0], 0, "Q0=(0.5,0) → 3 votes for class 0");
    assert_eq!(predictions[1], 1, "Q1=(10.5,0) → 3 votes for class 1");
}

/// Above 8 features, KNN uses a brute-force search.
/// Both prediction methods must return the expected labels.
#[test]
fn predict_high_dimensional_falls_back_to_brute_force() {
    let n_features = 18; // above the kd-tree dimensionality cutoff (8)
    let mut x_train = Array2::<f64>::zeros((6, n_features));
    for i in 0..3 {
        for j in 0..n_features {
            x_train[[i, j]] = 0.1 * i as f64; // class 0 near the origin
        }
    }
    for i in 3..6 {
        for j in 0..n_features {
            x_train[[i, j]] = 10.0 + 0.1 * i as f64; // class 1 near 10
        }
    }
    let y_train = array![0, 0, 0, 1, 1, 1];

    let mut knn = KNN::<i32>::new(3)
        .unwrap()
        .with_weighting_strategy(WeightingStrategy::Uniform)
        .with_metric(Metric::Euclidean)
        .unwrap();
    knn.fit(&x_train, &y_train).unwrap();

    let mut x_test = Array2::<f64>::zeros((2, n_features));
    for j in 0..n_features {
        x_test[[0, j]] = 0.15; // near class 0
        x_test[[1, j]] = 10.15; // near class 1
    }

    let seq = knn.predict(&x_test).unwrap();
    let par = knn.predict_parallel(&x_test).unwrap();
    assert_eq!(seq[0], 0, "high-dim query near class 0 must predict 0");
    assert_eq!(seq[1], 1, "high-dim query near class 1 must predict 1");
    assert_eq!(
        seq, par,
        "sequential and parallel must agree on the brute-force path"
    );
}

/// k=3 Euclidean Uniform resolves a 2:1 split (2 of 3 neighbors decide the class)
#[test]
fn predict_k3_majority_two_to_one_split() {
    let mut knn = KNN::<i32>::new(3)
        .unwrap()
        .with_weighting_strategy(WeightingStrategy::Uniform)
        .with_metric(Metric::Euclidean)
        .unwrap();
    // 4 training points (exactly k+1) keep the problem minimal
    let x_train = array![[0.0, 0.0], [1.0, 0.0], [5.0, 0.0], [6.0, 0.0]];
    let y_train = array![0, 0, 1, 1];
    knn.fit(&x_train, &y_train).unwrap();

    // At Q=(0.5,0), neighbors P0,P1,P2 give class 0 a 2-to-1 majority.
    // At Q=(5.5,0), neighbors P2,P3,P1 give class 1 a 2-to-1 majority.
    let x_test = array![[0.5, 0.0], [5.5, 0.0]];
    let predictions = knn.predict(&x_test).unwrap();

    assert_eq!(
        predictions[0], 0,
        "Q0=(0.5,0): 2 votes class 0, 1 vote class 1 → class 0"
    );
    assert_eq!(
        predictions[1], 1,
        "Q1=(5.5,0): 2 votes class 1, 1 vote class 0 → class 1"
    );
}

/// k=2 Distance weighting: the closer neighbor decides the class
#[test]
fn predict_distance_weighting_closer_wins() {
    let mut knn = KNN::<i32>::new(2)
        .unwrap()
        .with_weighting_strategy(WeightingStrategy::Distance)
        .with_metric(Metric::Euclidean)
        .unwrap();
    let x_train = array![[0.0, 0.0], [10.0, 0.0]];
    let y_train = array![0, 1];
    knn.fit(&x_train, &y_train).unwrap();

    let x_test = array![[1.0, 0.0], [9.0, 0.0]];
    let predictions = knn.predict(&x_test).unwrap();

    assert_eq!(
        predictions[0], 0,
        "Q=(1,0): weight(P0)=1 > weight(P1)=1/9 → class 0"
    );
    assert_eq!(
        predictions[1], 1,
        "Q=(9,0): weight(P1)=1 > weight(P0)=1/9 → class 1"
    );
}

/// Distance weighting: an exact-match query (distance 0) returns the matched point's class
#[test]
fn predict_distance_zero_exact_match_returns_immediately() {
    let mut knn = KNN::<i32>::new(2)
        .unwrap()
        .with_weighting_strategy(WeightingStrategy::Distance)
        .with_metric(Metric::Euclidean)
        .unwrap();
    // at least k=2 training samples needed to fit
    let x_train = array![[3.0, 4.0], [100.0, 100.0]];
    let y_train = array![7, 99];
    knn.fit(&x_train, &y_train).unwrap();

    let x_test = array![[3.0, 4.0]]; // exact match for P0
    let predictions = knn.predict(&x_test).unwrap();

    // exact match wins regardless of the other neighbor
    assert_eq!(predictions[0], 7, "exact-match query must return class 7");
}

/// KNN with T=String returns the nearest neighbor's string label
#[test]
fn predict_string_labels() {
    let mut knn = KNN::<String>::new(1)
        .unwrap()
        .with_weighting_strategy(WeightingStrategy::Uniform)
        .with_metric(Metric::Euclidean)
        .unwrap();
    let x_train = array![[0.0, 0.0], [10.0, 0.0], [0.0, 10.0]];
    let y_train: Array1<String> = Array1::from(vec![
        "cat".to_string(),
        "dog".to_string(),
        "dog".to_string(),
    ]);
    knn.fit(&x_train, &y_train).unwrap();

    let x_test = array![[0.5, 0.0], [9.5, 0.0]];
    let predictions = knn.predict(&x_test).unwrap();

    assert_eq!(predictions[0], "cat");
    assert_eq!(predictions[1], "dog");
}

/// predict and predict_parallel must produce identical results (Uniform, small k)
#[test]
fn predict_and_predict_parallel_agree_small_k() {
    let mut knn = KNN::<i32>::new(3)
        .unwrap()
        .with_weighting_strategy(WeightingStrategy::Uniform)
        .with_metric(Metric::Euclidean)
        .unwrap();
    let x_train = array![
        [0.0, 0.0],
        [1.0, 0.0],
        [2.0, 0.0],
        [10.0, 0.0],
        [11.0, 0.0],
        [12.0, 0.0]
    ];
    let y_train = array![0, 0, 0, 1, 1, 1];
    knn.fit(&x_train, &y_train).unwrap();

    let x_test = array![[0.5, 0.0], [10.5, 0.0], [1.0, 0.0], [11.0, 0.0]];
    let preds_seq = knn.predict(&x_test).unwrap();
    let preds_par = knn.predict_parallel(&x_test).unwrap();

    assert_eq!(
        preds_seq, preds_par,
        "predict and predict_parallel must agree"
    );
}

/// predict and predict_parallel must also agree with Distance weighting
#[test]
fn predict_and_predict_parallel_agree_distance_weights() {
    let mut knn = KNN::<i32>::new(2)
        .unwrap()
        .with_weighting_strategy(WeightingStrategy::Distance)
        .with_metric(Metric::Euclidean)
        .unwrap();
    let x_train = array![[0.0, 0.0], [10.0, 0.0], [0.0, 10.0]];
    let y_train = array![0, 1, 1];
    knn.fit(&x_train, &y_train).unwrap();

    let x_test = array![[1.0, 0.0], [9.0, 0.0], [0.0, 1.0]];
    let preds_seq = knn.predict(&x_test).unwrap();
    let preds_par = knn.predict_parallel(&x_test).unwrap();

    assert_eq!(
        preds_seq, preds_par,
        "predict and predict_parallel must agree with Distance weighting"
    );
}

/// Checks predict and predict_parallel agree at k=100 with Uniform weighting.
#[test]
fn predict_parallel_large_k_exercises_parallel_voting_branch() {
    const N: usize = 200; // 100 per class

    // 100 class-0 at x=0..99, 100 class-1 at x=200..299 (all on y=0)
    let mut x_rows: Vec<[f64; 2]> = Vec::with_capacity(N);
    let mut y_vals: Vec<i32> = Vec::with_capacity(N);
    for i in 0..100_usize {
        x_rows.push([i as f64, 0.0]);
        y_vals.push(0);
    }
    for i in 200..300_usize {
        x_rows.push([i as f64, 0.0]);
        y_vals.push(1);
    }
    let x_train =
        Array2::from_shape_vec((N, 2), x_rows.iter().flatten().copied().collect()).unwrap();
    let y_train = Array1::from_vec(y_vals);

    let mut knn = KNN::<i32>::new(100)
        .unwrap()
        .with_weighting_strategy(WeightingStrategy::Uniform)
        .with_metric(Metric::Euclidean)
        .unwrap();
    knn.fit(&x_train, &y_train).unwrap();

    let x_test = array![[50.0, 0.0], [250.0, 0.0]];

    // sequential path
    let preds_seq = knn.predict(&x_test).unwrap();
    assert_eq!(
        preds_seq[0], 0,
        "Q=(50,0): 100 class-0 neighbours → class 0"
    );
    assert_eq!(
        preds_seq[1], 1,
        "Q=(250,0): 100 class-1 neighbours → class 1"
    );

    let preds_par = knn.predict_parallel(&x_test).unwrap();
    assert_eq!(
        preds_seq, preds_par,
        "parallel path must agree with sequential at k=100"
    );
}

/// Checks predict and predict_parallel agree at k=100 with Distance weighting.
#[test]
fn predict_parallel_large_k_distance_weighting_exercises_parallel_weight_branch() {
    const N: usize = 200;
    let mut x_rows: Vec<[f64; 2]> = Vec::with_capacity(N);
    let mut y_vals: Vec<i32> = Vec::with_capacity(N);
    for i in 0..100_usize {
        x_rows.push([i as f64 + 1.0, 0.0]); // avoid x=0 to prevent the dist=0 edge case
        y_vals.push(0);
    }
    for i in 200..300_usize {
        x_rows.push([i as f64, 0.0]);
        y_vals.push(1);
    }
    let x_train =
        Array2::from_shape_vec((N, 2), x_rows.iter().flatten().copied().collect()).unwrap();
    let y_train = Array1::from_vec(y_vals);

    let mut knn = KNN::<i32>::new(100)
        .unwrap()
        .with_weighting_strategy(WeightingStrategy::Distance)
        .with_metric(Metric::Euclidean)
        .unwrap();
    knn.fit(&x_train, &y_train).unwrap();

    let x_test = array![[50.5, 0.0]];
    let preds_seq = knn.predict(&x_test).unwrap();
    let preds_par = knn.predict_parallel(&x_test).unwrap();

    // class-0 points are all within distance <= 100. class-1 points start at distance >= 149.5,
    // so class-0's aggregate weight far exceeds class-1's even with 100 neighbors.
    assert_eq!(
        preds_seq[0], 0,
        "distance-weighted: class-0 should win near x=50.5"
    );
    assert_eq!(preds_seq, preds_par, "parallel must agree with sequential");
}

/// fit_predict with k=1 on the training data returns every training label unchanged
#[test]
fn fit_predict_returns_training_labels() {
    let mut knn = KNN::<i32>::new(1)
        .unwrap()
        .with_weighting_strategy(WeightingStrategy::Uniform)
        .with_metric(Metric::Euclidean)
        .unwrap();
    let x_train = array![[0.0, 0.0], [10.0, 0.0]];
    let y_train = array![0, 1];
    let result = knn.fit_predict(&x_train, &y_train).unwrap();

    assert_eq!(result[0], 0);
    assert_eq!(result[1], 1);
}

/// k=1 Minkowski(p=3) finds the correct nearest neighbor under the Lp-3 metric
#[test]
fn predict_minkowski_p3_correct_nearest_neighbour() {
    let mut knn = KNN::<i32>::new(1)
        .unwrap()
        .with_weighting_strategy(WeightingStrategy::Uniform)
        .with_metric(Metric::Minkowski(3.0))
        .unwrap();
    let x_train = array![[3.0, 0.0], [0.0, 4.0]];
    let y_train = array![0, 1];
    knn.fit(&x_train, &y_train).unwrap();

    // At Q=(0,0), distances to P0 and P1 are 3 and 4, so P0 gives class 0.
    // At Q=(0,3), distances are 54^(1/3) and 1, so P1 gives class 1.
    let x_test = array![[0.0, 0.0], [0.0, 3.0]];
    let predictions = knn.predict(&x_test).unwrap();

    assert_eq!(
        predictions[0], 0,
        "Q=(0,0): Minkowski-3 nearest is P0=(3,0)"
    );
    assert_eq!(
        predictions[1], 1,
        "Q=(0,3): Minkowski-3 nearest is P1=(0,4)"
    );
}

/// A reloaded KNN model must produce predictions identical to the original
#[test]
fn save_and_load_produces_identical_predictions() {
    let mut knn = KNN::<i32>::new(3)
        .unwrap()
        .with_weighting_strategy(WeightingStrategy::Uniform)
        .with_metric(Metric::Euclidean)
        .unwrap();
    let x_train = array![
        [0.0, 0.0],
        [1.0, 0.0],
        [2.0, 0.0],
        [10.0, 0.0],
        [11.0, 0.0],
        [12.0, 0.0]
    ];
    let y_train = array![0, 0, 0, 1, 1, 1];
    knn.fit(&x_train, &y_train).unwrap();

    let x_test = array![[0.5, 0.0], [5.0, 0.0], [10.5, 0.0]];
    let preds_before = knn.predict(&x_test).unwrap();

    let path = "/tmp/rustyml_knn_save_load_test.bin";
    knn.save_to_path(path).unwrap();

    let knn_loaded = KNN::<i32>::load_from_path(path).unwrap();
    let preds_after = knn_loaded.predict(&x_test).unwrap();

    assert_eq!(
        preds_before, preds_after,
        "predictions must be identical after save/load round-trip"
    );

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

/// Save/load preserves the Distance weighting strategy, metric, and k
#[test]
fn save_and_load_preserves_model_metadata() {
    let mut knn = KNN::<i32>::new(7)
        .unwrap()
        .with_weighting_strategy(WeightingStrategy::Distance)
        .with_metric(Metric::Manhattan)
        .unwrap();
    // 7 training points to satisfy nrows >= k
    let x_train = Array2::from_shape_vec(
        (7, 2),
        vec![
            0.0, 0.0, 1.0, 0.0, 2.0, 0.0, 3.0, 0.0, 4.0, 0.0, 5.0, 0.0, 6.0, 0.0,
        ],
    )
    .unwrap();
    let y_train = array![0, 0, 0, 1, 1, 1, 1];
    knn.fit(&x_train, &y_train).unwrap();

    let path = "/tmp/rustyml_knn_metadata_test.bin";
    knn.save_to_path(path).unwrap();

    let knn_loaded = KNN::<i32>::load_from_path(path).unwrap();
    assert_eq!(knn_loaded.get_k(), 7);
    assert_eq!(
        knn_loaded.get_weighting_strategy(),
        WeightingStrategy::Distance
    );
    assert_eq!(knn_loaded.get_metric(), Metric::Manhattan);

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