//! Decision Tree Tests - TDD Red Phase
//!
//! These tests define the expected behavior of the decision tree implementation.
//! They will fail initially and drive the implementation.

use approx::assert_abs_diff_eq;
use rtx_ml_classic::trees::{DecisionTree, SplitCriterion};
use rtx_ml_classic::{MLError, Result};
use rtx_tensor::{DType, Device, Tensor};

/// Test fixture for decision tree tests
struct DecisionTreeTestFixture {
    device: Device,
    x_binary: Tensor,     // Simple binary classification data
    y_binary: Tensor,     // Binary labels
    x_multiclass: Tensor, // Multiclass classification data
    y_multiclass: Tensor, // Multiclass labels
    x_regression: Tensor, // Regression data
    y_regression: Tensor, // Regression targets
}

impl DecisionTreeTestFixture {
    fn new() -> Result<Self> {
        let device = Device::cpu();

        // Create simple linearly separable binary classification data
        // Features: [x1, x2] where x1 + x2 > 0.5 -> class 1, else class 0
        let x_binary = Tensor::from_slice(
            &[
                0.1f32, 0.1, 0.9, 0.1, 0.1, 0.9, 0.9, 0.9, 0.2, 0.2, 0.8, 0.8, 0.3, 0.1, 0.1, 0.8,
            ],
            &[8, 2],
            &device,
        )?;
        let y_binary =
            Tensor::from_slice(&[0.0f32, 1.0, 1.0, 1.0, 0.0, 1.0, 0.0, 1.0], &[8], &device)?;

        // Create 3-class classification data
        let x_multiclass = Tensor::from_slice(
            &[
                0.1f32, 0.1, 0.9, 0.9, 0.1, 0.9, 0.9, 0.1, 0.5, 0.1, 0.5, 0.9, 0.1, 0.5, 0.9, 0.5,
            ],
            &[8, 2],
            &device,
        )?;
        let y_multiclass =
            Tensor::from_slice(&[0.0f32, 2.0, 1.0, 1.0, 0.0, 2.0, 1.0, 2.0], &[8], &device)?;

        // Create simple regression data: y = x1 + 2*x2 + noise
        let x_regression = Tensor::from_slice(
            &[
                1.0f32, 2.0, 2.0, 1.0, 3.0, 3.0, 1.0, 1.0, 2.0, 2.0, 3.0, 1.0, 1.0, 3.0, 4.0, 2.0,
            ],
            &[8, 2],
            &device,
        )?;
        let y_regression =
            Tensor::from_slice(&[5.1f32, 4.1, 9.0, 3.0, 6.1, 5.0, 7.1, 8.0], &[8], &device)?;

        Ok(Self {
            device,
            x_binary,
            y_binary,
            x_multiclass,
            y_multiclass,
            x_regression,
            y_regression,
        })
    }
}

#[test]
fn test_decision_tree_creation() {
    let tree = DecisionTree::new();

    // Test default parameters
    assert_eq!(tree.get_max_depth(), None);
    assert_eq!(tree.get_min_samples_split(), 2);
    assert_eq!(tree.get_min_samples_leaf(), 1);
    assert_eq!(tree.get_criterion(), &SplitCriterion::Gini);
    assert!(tree.is_classifier());
    assert!(!tree.is_fitted());
}

#[test]
fn test_decision_tree_builder_pattern() {
    let tree = DecisionTree::new()
        .max_depth(5)
        .min_samples_split(4)
        .min_samples_leaf(2)
        .criterion("entropy")
        .task("regression");

    assert_eq!(tree.get_max_depth(), Some(5));
    assert_eq!(tree.get_min_samples_split(), 4);
    assert_eq!(tree.get_min_samples_leaf(), 2);
    assert_eq!(tree.get_criterion(), &SplitCriterion::Entropy);
    assert!(!tree.is_classifier()); // Should be regressor
}

#[test]
fn test_split_criterion_parsing() {
    let tree_gini = DecisionTree::new().criterion("gini");
    assert_eq!(tree_gini.get_criterion(), &SplitCriterion::Gini);

    let tree_entropy = DecisionTree::new().criterion("ENTROPY");
    assert_eq!(tree_entropy.get_criterion(), &SplitCriterion::Entropy);

    let tree_log_loss = DecisionTree::new().criterion("log_loss");
    assert_eq!(tree_log_loss.get_criterion(), &SplitCriterion::LogLoss);

    // Invalid criterion should default to Gini
    let tree_invalid = DecisionTree::new().criterion("invalid");
    assert_eq!(tree_invalid.get_criterion(), &SplitCriterion::Gini);
}

#[test]
fn test_prediction_before_fitting_fails() {
    let fixture = DecisionTreeTestFixture::new().unwrap();
    let tree = DecisionTree::new();

    let result = tree.predict(&fixture.x_binary);
    assert!(matches!(result, Err(MLError::ModelNotFitted)));

    let result = tree.predict_proba(&fixture.x_binary);
    assert!(matches!(result, Err(MLError::ModelNotFitted)));

    let result = tree.feature_importances();
    assert!(matches!(result, Err(MLError::ModelNotFitted)));
}

#[test]
fn test_binary_classification_gini() {
    let fixture = DecisionTreeTestFixture::new().unwrap();
    let mut tree = DecisionTree::new().max_depth(3).criterion("gini");

    // Fit should succeed
    tree.fit(&fixture.x_binary, &fixture.y_binary).unwrap();
    assert!(tree.is_fitted());

    // Predictions should have correct shape
    let predictions = tree.predict(&fixture.x_binary).unwrap();
    assert_eq!(predictions.shape().dims(), &[8]); // 8 samples

    // All predictions should be 0 or 1
    let pred_data = predictions.data().unwrap();
    for pred in pred_data {
        let pred_int = pred.round() as i32;
        assert!(pred_int == 0 || pred_int == 1);
    }

    // Probability predictions should sum to 1
    let proba = tree.predict_proba(&fixture.x_binary).unwrap();
    assert_eq!(proba.shape().dims(), &[8, 2]); // 8 samples, 2 classes

    let proba_data = proba.data().unwrap();
    for i in 0..8 {
        let row_sum = proba_data[i * 2] + proba_data[i * 2 + 1];
        assert_abs_diff_eq!(row_sum, 1.0, epsilon = 1e-6);
    }
}

#[test]
fn test_binary_classification_entropy() {
    let fixture = DecisionTreeTestFixture::new().unwrap();
    let mut tree = DecisionTree::new().max_depth(3).criterion("entropy");

    tree.fit(&fixture.x_binary, &fixture.y_binary).unwrap();
    let predictions = tree.predict(&fixture.x_binary).unwrap();

    // Should achieve perfect accuracy on this simple dataset
    let pred_data = predictions.data().unwrap();
    let y_data = fixture.y_binary.data().unwrap();

    let mut correct = 0;
    for (pred, true_val) in pred_data.iter().zip(y_data.iter()) {
        if (pred.round() - true_val).abs() < 0.1 {
            correct += 1;
        }
    }

    let accuracy = correct as f32 / y_data.len() as f32;
    assert!(
        accuracy >= 0.75,
        "Accuracy {} should be at least 75%",
        accuracy
    );
}

#[test]
fn test_multiclass_classification() {
    let fixture = DecisionTreeTestFixture::new().unwrap();
    let mut tree = DecisionTree::new().max_depth(5).criterion("gini");

    tree.fit(&fixture.x_multiclass, &fixture.y_multiclass)
        .unwrap();

    let predictions = tree.predict(&fixture.x_multiclass).unwrap();
    let pred_data = predictions.data().unwrap();

    // All predictions should be valid class indices (0, 1, or 2)
    for pred in pred_data {
        let pred_int = pred.round() as i32;
        assert!(pred_int >= 0 && pred_int <= 2);
    }

    // Probability predictions for 3 classes
    let proba = tree.predict_proba(&fixture.x_multiclass).unwrap();
    assert_eq!(proba.shape().dims(), &[8, 3]);

    let proba_data = proba.data().unwrap();
    for i in 0..8 {
        let row_sum = proba_data[i * 3] + proba_data[i * 3 + 1] + proba_data[i * 3 + 2];
        assert_abs_diff_eq!(row_sum, 1.0, epsilon = 1e-6);
    }
}

#[test]
fn test_regression() {
    let fixture = DecisionTreeTestFixture::new().unwrap();
    let mut tree = DecisionTree::new().task("regression").max_depth(4);

    tree.fit(&fixture.x_regression, &fixture.y_regression)
        .unwrap();

    let predictions = tree.predict(&fixture.x_regression).unwrap();
    assert_eq!(predictions.shape().dims(), &[8]);

    // predict_proba should fail for regression
    let result = tree.predict_proba(&fixture.x_regression);
    assert!(matches!(result, Err(MLError::InvalidParameter { .. })));

    // Calculate MSE - should be reasonable for this simple dataset
    let pred_data = predictions.data().unwrap();
    let y_data = fixture.y_regression.data().unwrap();

    let mse: f32 = pred_data
        .iter()
        .zip(y_data.iter())
        .map(|(pred, actual)| (pred - actual).powi(2))
        .sum::<f32>()
        / y_data.len() as f32;

    assert!(mse < 4.0, "MSE {} should be reasonable", mse);
}

#[test]
fn test_feature_importances() {
    let fixture = DecisionTreeTestFixture::new().unwrap();
    let mut tree = DecisionTree::new().max_depth(3);

    tree.fit(&fixture.x_binary, &fixture.y_binary).unwrap();

    let importances = tree.feature_importances().unwrap();
    assert_eq!(importances.shape().dims(), &[2]); // 2 features

    let imp_data = importances.data().unwrap();

    // All importances should be non-negative
    for imp in &imp_data {
        assert!(*imp >= 0.0);
    }

    // Importances should sum to 1.0
    let total: f32 = imp_data.iter().sum();
    assert_abs_diff_eq!(total, 1.0, epsilon = 1e-6);
}

#[test]
fn test_min_samples_constraints() {
    let fixture = DecisionTreeTestFixture::new().unwrap();

    // Tree with high min_samples_split should create simpler trees
    let mut tree_high_split = DecisionTree::new()
        .min_samples_split(10) // More than our dataset size
        .max_depth(10);

    // Should fail because min_samples_split > n_samples
    let result = tree_high_split.fit(&fixture.x_binary, &fixture.y_binary);
    assert!(result.is_err());

    // Let's test a valid case where min_samples_split causes simpler trees but doesn't fail
    let mut tree_simple = DecisionTree::new()
        .min_samples_split(4) // Reasonable value for 8 samples
        .max_depth(10);

    tree_simple
        .fit(&fixture.x_binary, &fixture.y_binary)
        .unwrap();
    let predictions = tree_simple.predict(&fixture.x_binary).unwrap();
    assert_eq!(predictions.shape().dims(), &[8]);
}

#[test]
fn test_max_depth_constraint() {
    let fixture = DecisionTreeTestFixture::new().unwrap();

    // Very shallow tree
    let mut tree_shallow = DecisionTree::new().max_depth(1);

    tree_shallow
        .fit(&fixture.x_binary, &fixture.y_binary)
        .unwrap();
    let predictions_shallow = tree_shallow.predict(&fixture.x_binary).unwrap();

    // Deeper tree
    let mut tree_deep = DecisionTree::new().max_depth(5);

    tree_deep.fit(&fixture.x_binary, &fixture.y_binary).unwrap();
    let predictions_deep = tree_deep.predict(&fixture.x_binary).unwrap();

    // Both should produce valid predictions
    assert_eq!(predictions_shallow.shape().dims(), &[8]);
    assert_eq!(predictions_deep.shape().dims(), &[8]);
}

#[test]
fn test_invalid_input_dimensions() {
    let device = Device::cpu();
    let mut tree = DecisionTree::new();

    // Mismatched X and y dimensions
    let x = Tensor::zeros([10, 3], &device).unwrap();
    let y = Tensor::zeros([5], &device).unwrap(); // Wrong size

    let result = tree.fit(&x, &y);
    assert!(result.is_err());

    // Empty data - use 1 sample of zeros instead since tensor doesn't allow zero dims
    let x_single = Tensor::zeros([1, 3], &device).unwrap();
    let y_single = Tensor::zeros([1], &device).unwrap();

    let result = tree.fit(&x_single, &y_single);
    // Should fail because insufficient data
    assert!(result.is_err());
}