//! Golden tests for ensemble models against sklearn 1.8.0.
//!
//! ExtraTrees (C/R), AdaBoost (C/R), Bagging (C/R), Voting (C/R),
//! StackingRegressor, RandomForest OOB.

mod common;

use anofox_ml::prelude::*;
use anofox_ml_trees::{DecisionTreeClassifier, DecisionTreeRegressor};
use common::{json_to_array1, json_to_array2, load_golden_data};
use ndarray::Array1;

fn accuracy(pred: &Array1<f64>, target: &Array1<f64>) -> f64 {
    let correct: usize = pred
        .iter()
        .zip(target.iter())
        .filter(|(&p, &t)| (p - t).abs() < 1e-10)
        .count();
    correct as f64 / pred.len() as f64
}

fn r2(pred: &Array1<f64>, target: &Array1<f64>) -> f64 {
    let mean = target.iter().sum::<f64>() / target.len() as f64;
    let ss_res: f64 = pred
        .iter()
        .zip(target.iter())
        .map(|(&p, &t)| (p - t).powi(2))
        .sum();
    let ss_tot: f64 = target.iter().map(|&t| (t - mean).powi(2)).sum();
    if ss_tot > 0.0 {
        1.0 - ss_res / ss_tot
    } else {
        0.0
    }
}

#[test]
fn test_golden_extra_trees_classifier_accuracy() {
    let cases = load_golden_data("ensemble_batch.json");
    let case = cases
        .iter()
        .find(|c| c["name"] == "extra_trees_classifier")
        .unwrap();

    let x = json_to_array2(&case["X"]);
    let y = json_to_array1(&case["y"]);
    let sklearn_acc = case["accuracy"].as_f64().unwrap();

    let model = ExtraTreesClassifier::new(20)
        .with_max_depth(Some(3))
        .with_seed(42);
    let fitted = Fit::fit(&model, &x, &y).unwrap();
    let preds = fitted.predict(&x).unwrap();
    let our_acc = accuracy(&preds, &y);

    assert!(
        our_acc >= sklearn_acc - 0.2,
        "ExtraTreesClassifier accuracy {} vs sklearn {}",
        our_acc,
        sklearn_acc
    );
}

#[test]
fn test_golden_extra_trees_regressor_r2() {
    let cases = load_golden_data("ensemble_batch.json");
    let case = cases
        .iter()
        .find(|c| c["name"] == "extra_trees_regressor")
        .unwrap();

    let x = json_to_array2(&case["X"]);
    let y = json_to_array1(&case["y"]);
    let sklearn_r2 = case["r2"].as_f64().unwrap();

    let model = ExtraTreesRegressor::new(20)
        .with_max_depth(Some(3))
        .with_seed(42);
    let fitted = Fit::fit(&model, &x, &y).unwrap();
    let preds = fitted.predict(&x).unwrap();
    let our_r2 = r2(&preds, &y);

    assert!(
        our_r2 >= sklearn_r2 - 0.3,
        "ExtraTreesRegressor R² {} vs sklearn {}",
        our_r2,
        sklearn_r2
    );
}

#[test]
fn test_golden_adaboost_classifier_accuracy() {
    let cases = load_golden_data("ensemble_batch.json");
    let case = cases
        .iter()
        .find(|c| c["name"] == "adaboost_classifier")
        .unwrap();

    let x = json_to_array2(&case["X"]);
    let y = json_to_array1(&case["y"]);
    let sklearn_acc = case["accuracy"].as_f64().unwrap();

    let model = AdaBoostClassifier::new()
        .with_n_estimators(20)
        .with_seed(42);
    let fitted = Fit::fit(&model, &x, &y).unwrap();
    let preds = fitted.predict(&x).unwrap();
    let our_acc = accuracy(&preds, &y);

    assert!(
        our_acc >= sklearn_acc - 0.2,
        "AdaBoostClassifier accuracy {} vs sklearn {}",
        our_acc,
        sklearn_acc
    );
}

#[test]
fn test_golden_adaboost_regressor_r2() {
    let cases = load_golden_data("ensemble_batch.json");
    let case = cases
        .iter()
        .find(|c| c["name"] == "adaboost_regressor")
        .unwrap();

    let x = json_to_array2(&case["X"]);
    let y = json_to_array1(&case["y"]);
    let sklearn_r2 = case["r2"].as_f64().unwrap();

    let model = AdaBoostRegressor::new().with_n_estimators(20).with_seed(42);
    let fitted = Fit::fit(&model, &x, &y).unwrap();
    let preds = fitted.predict(&x).unwrap();
    let our_r2 = r2(&preds, &y);

    // Now using the same default base estimator as sklearn (max_depth=3),
    // we should reach near-perfect R² on this tiny linear dataset.
    assert!(
        our_r2 >= sklearn_r2 - 0.05,
        "AdaBoostRegressor R² {:.4} vs sklearn {:.4}",
        our_r2,
        sklearn_r2
    );
}

#[test]
fn test_golden_bagging_classifier_accuracy() {
    let cases = load_golden_data("ensemble_batch.json");
    let case = cases
        .iter()
        .find(|c| c["name"] == "bagging_classifier")
        .unwrap();

    let x = json_to_array2(&case["X"]);
    let y = json_to_array1(&case["y"]);
    let sklearn_acc = case["accuracy"].as_f64().unwrap();

    let model = BaggingClassifier::new(20).with_seed(42);
    let fitted = Fit::fit(&model, &x, &y).unwrap();
    let preds = fitted.predict(&x).unwrap();
    let our_acc = accuracy(&preds, &y);

    assert!(
        our_acc >= sklearn_acc - 0.2,
        "BaggingClassifier accuracy {} vs sklearn {}",
        our_acc,
        sklearn_acc
    );
}

#[test]
fn test_golden_bagging_regressor_r2() {
    let cases = load_golden_data("ensemble_batch.json");
    let case = cases
        .iter()
        .find(|c| c["name"] == "bagging_regressor")
        .unwrap();

    let x = json_to_array2(&case["X"]);
    let y = json_to_array1(&case["y"]);
    let sklearn_r2 = case["r2"].as_f64().unwrap();

    let model = BaggingRegressor::new(20).with_seed(42);
    let fitted = Fit::fit(&model, &x, &y).unwrap();
    let preds = fitted.predict(&x).unwrap();
    let our_r2 = r2(&preds, &y);

    assert!(
        our_r2 >= sklearn_r2 - 0.3,
        "BaggingRegressor R² {} vs sklearn {}",
        our_r2,
        sklearn_r2
    );
}

#[test]
fn test_golden_voting_classifier_accuracy() {
    let cases = load_golden_data("ensemble_batch.json");
    let case = cases
        .iter()
        .find(|c| c["name"] == "voting_classifier")
        .unwrap();

    let x = json_to_array2(&case["X"]);
    let y = json_to_array1(&case["y"]);
    let sklearn_acc = case["accuracy"].as_f64().unwrap();

    let vc = VotingClassifier::new()
        .push(
            "t1",
            DecisionTreeClassifier {
                max_depth: Some(2),
                ..Default::default()
            },
        )
        .push(
            "t2",
            DecisionTreeClassifier {
                max_depth: Some(3),
                ..Default::default()
            },
        )
        .push(
            "t3",
            DecisionTreeClassifier {
                max_depth: Some(5),
                ..Default::default()
            },
        );

    let fitted = Fit::<f64>::fit(&vc, &x, &y).unwrap();
    let preds = fitted.predict(&x).unwrap();
    let our_acc = accuracy(&preds, &y);

    assert!(
        our_acc >= sklearn_acc - 0.2,
        "VotingClassifier accuracy {} vs sklearn {}",
        our_acc,
        sklearn_acc
    );
}

#[test]
fn test_golden_voting_regressor_r2() {
    let cases = load_golden_data("ensemble_batch.json");
    let case = cases
        .iter()
        .find(|c| c["name"] == "voting_regressor")
        .unwrap();

    let x = json_to_array2(&case["X"]);
    let y = json_to_array1(&case["y"]);
    let sklearn_r2 = case["r2"].as_f64().unwrap();

    let vr = VotingRegressor::new()
        .push(
            "t1",
            DecisionTreeRegressor {
                max_depth: Some(2),
                ..Default::default()
            },
        )
        .push(
            "t2",
            DecisionTreeRegressor {
                max_depth: Some(3),
                ..Default::default()
            },
        );

    let fitted = Fit::<f64>::fit(&vr, &x, &y).unwrap();
    let preds = fitted.predict(&x).unwrap();
    let our_r2 = r2(&preds, &y);

    assert!(
        our_r2 >= sklearn_r2 - 0.3,
        "VotingRegressor R² {} vs sklearn {}",
        our_r2,
        sklearn_r2
    );
}

#[test]
fn test_golden_stacking_regressor_r2() {
    let cases = load_golden_data("ensemble_batch.json");
    let case = cases
        .iter()
        .find(|c| c["name"] == "stacking_regressor")
        .unwrap();

    let x = json_to_array2(&case["X"]);
    let y = json_to_array1(&case["y"]);
    let sklearn_r2 = case["r2"].as_f64().unwrap();

    let sr = StackingRegressor::new(DecisionTreeRegressor {
        max_depth: Some(2),
        ..Default::default()
    })
    .push(
        "t1",
        DecisionTreeRegressor {
            max_depth: Some(2),
            ..Default::default()
        },
    )
    .push(
        "t2",
        DecisionTreeRegressor {
            max_depth: Some(3),
            ..Default::default()
        },
    )
    .with_cv_folds(2);

    let fitted = Fit::<f64>::fit(&sr, &x, &y).unwrap();
    let preds = fitted.predict(&x).unwrap();
    let our_r2 = r2(&preds, &y);

    assert!(
        our_r2 >= sklearn_r2 - 0.5,
        "StackingRegressor R² {} vs sklearn {}",
        our_r2,
        sklearn_r2
    );
}

#[test]
fn test_golden_random_forest_oob_score() {
    let cases = load_golden_data("ensemble_batch.json");
    let case = cases
        .iter()
        .find(|c| c["name"] == "random_forest_oob")
        .unwrap();

    let x = json_to_array2(&case["X"]);
    let y = json_to_array1(&case["y"]);
    let sklearn_oob = case["oob_score"].as_f64().unwrap();

    let model = RandomForestClassifier::new(50)
        .with_max_depth(Some(3))
        .with_oob_score(true)
        .with_bootstrap(true)
        .with_seed(42);

    let fitted = Fit::<f64>::fit(&model, &x, &y).unwrap();

    // OOB score should be defined and reasonable
    let our_oob = fitted.oob_score().expect("oob_score should be set");
    assert!(
        (our_oob - sklearn_oob).abs() < 0.3,
        "OOB score {} vs sklearn {}",
        our_oob,
        sklearn_oob
    );
}
