mod common;

use anofox_ml::prelude::*;
use common::{assert_close, json_to_array1, json_to_array2, load_golden_data};

#[test]
fn test_golden_hist_gb_classifier_predictions() {
    let cases = load_golden_data("hist_gradient_boosting.json");
    let case = cases
        .iter()
        .find(|c| c["name"] == "hist_gb_classifier")
        .unwrap();

    let x = json_to_array2(&case["X"]);
    let y = json_to_array1(&case["y"]);
    let _expected_preds = json_to_array1(&case["predictions"]);

    let model = HistGradientBoostingClassifier::new()
        .with_n_estimators(20)
        .with_max_depth(3)
        .with_learning_rate(0.1)
        .with_min_samples_leaf(1);

    let fitted = Fit::fit(&model, &x, &y).unwrap();
    let preds = fitted.predict(&x).unwrap();

    // Different binning strategies mean predictions may differ from sklearn.
    // Verify behavioral correctness: high accuracy on well-separated data.
    let correct: usize = preds
        .iter()
        .zip(y.iter())
        .filter(|(&p, &t)| (p - t).abs() < 1e-10)
        .count();
    assert!(
        correct >= 8,
        "hist_gb_clf should classify most correctly, got {}/{}",
        correct,
        y.len()
    );
}

#[test]
fn test_golden_hist_gb_classifier_proba_sums_to_one() {
    let cases = load_golden_data("hist_gradient_boosting.json");
    let case = cases
        .iter()
        .find(|c| c["name"] == "hist_gb_classifier")
        .unwrap();

    let x = json_to_array2(&case["X"]);
    let y = json_to_array1(&case["y"]);

    let model = HistGradientBoostingClassifier::new()
        .with_n_estimators(20)
        .with_max_depth(3)
        .with_min_samples_leaf(1);

    let fitted = Fit::fit(&model, &x, &y).unwrap();
    let proba = fitted.predict_proba(&x).unwrap();

    for i in 0..x.nrows() {
        let row_sum: f64 = (0..proba.ncols()).map(|c| proba[[i, c]]).sum();
        assert_close(
            row_sum,
            1.0,
            1e-10,
            &format!("hist_gb_clf proba row {} sum", i),
        );
    }
}

#[test]
fn test_golden_hist_gb_regressor_predictions() {
    let cases = load_golden_data("hist_gradient_boosting.json");
    let case = cases
        .iter()
        .find(|c| c["name"] == "hist_gb_regressor")
        .unwrap();

    let x = json_to_array2(&case["X"]);
    let y = json_to_array1(&case["y"]);
    let _expected_preds = json_to_array1(&case["predictions"]);

    let model = HistGradientBoostingRegressor::new()
        .with_n_estimators(50)
        .with_max_depth(3)
        .with_learning_rate(0.1)
        .with_min_samples_leaf(1);

    let fitted = Fit::fit(&model, &x, &y).unwrap();
    let preds = fitted.predict(&x).unwrap();

    // Different binning strategies mean exact values differ from sklearn.
    // Verify behavioral correctness: R² should be high on training data.
    let y_mean: f64 = y.iter().sum::<f64>() / y.len() as f64;
    let ss_res: f64 = preds
        .iter()
        .zip(y.iter())
        .map(|(&p, &t)| (p - t).powi(2))
        .sum();
    let ss_tot: f64 = y.iter().map(|&t| (t - y_mean).powi(2)).sum();
    let r2 = 1.0 - ss_res / ss_tot;
    assert!(r2 > 0.8, "hist_gb_reg R² should be > 0.8, got {:.4}", r2);
}
