use ml_learn::{KNNModel, Model, NaiveBayesClassifier, Perceptron, TestDataSet, TrainDataset};
use ndarray::Array2;

#[test]
fn perceptron_fit_and_predict() {
    // 线性可分数据：x1 + x2 > 2 为正类
    let input_data = Array2::from_shape_vec((6, 2), vec![
        2.0, 2.0, 3.0, 1.0, 2.0, 3.0,   // "pos"
        0.0, 0.0, 1.0, 0.0, 0.0, 1.0,   // "neg"
    ]).unwrap();
    let output_data = vec![
        "pos".to_string(), "pos".to_string(), "pos".to_string(),
        "neg".to_string(), "neg".to_string(), "neg".to_string(),
    ];
    let train = TrainDataset { input_data, output_data };

    let mut model = Perceptron::new(1.0, 1000);
    model.fit(&train);
    assert_eq!(model.score(&train), 1.0);

    let test = TestDataSet {
        input_data: Array2::from_shape_vec((2, 2), vec![4.0, 4.0, -1.0, 0.0]).unwrap(),
    };
    let result = model.predict(&test);
    assert_eq!(result.output_data, vec!["pos".to_string(), "neg".to_string()]);
}

#[test]
fn knn_fit_and_predict() {
    let input_data = Array2::from_shape_vec((6, 2), vec![
        1.0, 1.0, 1.2, 0.8, 0.8, 1.2,   // "A"
        5.0, 5.0, 5.2, 4.8, 4.8, 5.2,   // "B"
    ]).unwrap();
    let output_data = vec![
        "A".to_string(), "A".to_string(), "A".to_string(),
        "B".to_string(), "B".to_string(), "B".to_string(),
    ];
    let train = TrainDataset { input_data, output_data };

    let mut model = KNNModel::new("base", 3);
    model.fit(&train);
    assert_eq!(model.score(&train), 1.0);

    let test = TestDataSet {
        input_data: Array2::from_shape_vec((2, 2), vec![1.1, 0.9, 5.1, 5.1]).unwrap(),
    };
    let result = model.predict(&test);
    assert_eq!(result.output_data, vec!["A".to_string(), "B".to_string()]);
}

#[test]
fn gaussian_naive_bayes_fit_and_predict() {
    // 两个类别，单一连续特征，分布明显分开
    let input_data = Array2::from_shape_vec((6, 1), vec![
        1.0, 1.2, 0.8,   // "A"
        5.0, 5.2, 4.8,   // "B"
    ]).unwrap();
    let output_data = vec![
        "A".to_string(), "A".to_string(), "A".to_string(),
        "B".to_string(), "B".to_string(), "B".to_string(),
    ];
    let train = TrainDataset { input_data, output_data };

    let mut clf = NaiveBayesClassifier::new();
    clf.create_classifier(vec!["GNB"]);
    clf.fit(&train);
    assert_eq!(clf.score(&train), 1.0);

    let test = TestDataSet {
        input_data: Array2::from_shape_vec((2, 1), vec![1.1, 5.1]).unwrap(),
    };
    let result = clf.predict(&test);
    assert_eq!(result.output_data, vec!["A".to_string(), "B".to_string()]);
}

#[test]
fn poisson_naive_bayes_fit_and_predict() {
    // 计数特征：A 类 λ≈1，B 类 λ≈5
    let input_data = Array2::from_shape_vec((6, 1), vec![
        0.0, 1.0, 2.0,   // "A"
        4.0, 5.0, 6.0,   // "B"
    ]).unwrap();
    let output_data = vec![
        "A".to_string(), "A".to_string(), "A".to_string(),
        "B".to_string(), "B".to_string(), "B".to_string(),
    ];
    let train = TrainDataset { input_data, output_data };

    let mut clf = NaiveBayesClassifier::new();
    clf.create_classifier(vec!["PoissonNB"]);
    clf.fit(&train);
    assert_eq!(clf.score(&train), 1.0);

    let test = TestDataSet {
        input_data: Array2::from_shape_vec((2, 1), vec![1.0, 5.0]).unwrap(),
    };
    let result = clf.predict(&test);
    assert_eq!(result.output_data, vec!["A".to_string(), "B".to_string()]);
}

#[test]
fn bernoulli_naive_bayes_fit_and_predict() {
    // 0/1 特征：A 类大概率为 1，B 类大概率为 0
    let input_data = Array2::from_shape_vec((6, 1), vec![
        1.0, 1.0, 1.0,   // "A"
        0.0, 0.0, 0.0,   // "B"
    ]).unwrap();
    let output_data = vec![
        "A".to_string(), "A".to_string(), "A".to_string(),
        "B".to_string(), "B".to_string(), "B".to_string(),
    ];
    let train = TrainDataset { input_data, output_data };

    let mut clf = NaiveBayesClassifier::new();
    clf.create_classifier(vec!["BernoulliNB"]);
    clf.fit(&train);

    let test = TestDataSet {
        input_data: Array2::from_shape_vec((2, 1), vec![1.0, 0.0]).unwrap(),
    };
    let result = clf.predict(&test);
    assert_eq!(result.output_data, vec!["A".to_string(), "B".to_string()]);
}

#[test]
fn naive_bayes_mixed_layers() {
    // 两个特征列分别用不同层：连续特征 GNB + 计数特征 PoissonNB
    let input_data = Array2::from_shape_vec((6, 2), vec![
        1.0, 0.0,
        1.2, 1.0,
        0.8, 2.0,   // "A"
        5.0, 4.0,
        5.2, 5.0,
        4.8, 6.0,   // "B"
    ]).unwrap();
    let output_data = vec![
        "A".to_string(), "A".to_string(), "A".to_string(),
        "B".to_string(), "B".to_string(), "B".to_string(),
    ];
    let train = TrainDataset { input_data, output_data };

    let mut clf = NaiveBayesClassifier::new();
    clf.create_classifier(vec!["GNB", "PoissonNB"]);
    clf.fit(&train);
    assert_eq!(clf.score(&train), 1.0);
}